From fb77f8c9f9563c1eaea93fdf9b16bef0e5ec94ef Mon Sep 17 00:00:00 2001
From: zhulin1 <zhulin1@pjlab.org.cn>
Date: Mon, 1 Jul 2024 13:37:13 +0800
Subject: [PATCH] update

---
 tests/test_hf_model.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/test_hf_model.py b/tests/test_hf_model.py
index c810f1e..debef03 100644
--- a/tests/test_hf_model.py
+++ b/tests/test_hf_model.py
@@ -213,12 +213,12 @@ class TestReward:
         score2 = model.get_score(tokenizer, chat_2)
         print('score1: ', score1)
         print('score2: ', score2)
-        assert score1 > 0 & score2 < 0
+        assert score1.startswith('0.') & score2.startswith('-')
 
         # batch inference, get multiple scores at once
         scores = model.get_scores(tokenizer, [chat_1, chat_2])
         print('scores: ', scores)
-        assert scores[0] > 0 & scores[1] < 0
+        assert scores[0].startswith('0.') & scores[1].startswith('-')
 
         # compare whether chat_1 is better than chat_2
         compare_res = model.compare(tokenizer, chat_1, chat_2)
@@ -269,7 +269,7 @@ class TestReward:
             model_name, trust_remote_code=True)
 
         # prepare the chat prompt
-        prompt = 'Write an article about artificial intelligence revolution.'
+        prompt = 'Write an short bedtime story.'
         messages = [{
             'role': 'system',
             'content': 'You are a helpful assistant.'