mirror of https://github.com/InternLM/InternLM
Update test_hf_model.py
zhulinJulia24
GitHub
parent
091774d928
commit
ed26d1d76c
|
|
@ -161,164 +161,6 @@ class TestMath:
|
||||||
assert '2' in response
|
assert '2' in response
|
||||||
|
|
||||||
|
|
||||||
class TestReward:
|
|
||||||
"""Test cases for base model."""
|
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
|
||||||
'model_name',
|
|
||||||
[
|
|
||||||
'internlm/internlm-reward-1_8b', 'internlm/internlm-reward-7b',
|
|
||||||
'internlm/internlm-reward-20b'
|
|
||||||
],
|
|
||||||
)
|
|
||||||
@pytest.mark.parametrize(
|
|
||||||
'usefast',
|
|
||||||
[
|
|
||||||
True,
|
|
||||||
False,
|
|
||||||
],
|
|
||||||
)
|
|
||||||
def test_demo_default(self, model_name, usefast):
|
|
||||||
tokenizer = AutoTokenizer.from_pretrained(model_name,
|
|
||||||
trust_remote_code=True,
|
|
||||||
use_fast=usefast)
|
|
||||||
model = AutoModel.from_pretrained(
|
|
||||||
model_name,
|
|
||||||
device_map='cuda',
|
|
||||||
torch_dtype=torch.float16,
|
|
||||||
trust_remote_code=True,
|
|
||||||
)
|
|
||||||
tokenizer = AutoTokenizer.from_pretrained(model_name,
|
|
||||||
trust_remote_code=True)
|
|
||||||
|
|
||||||
chat_1 = [{
|
|
||||||
'role': 'user',
|
|
||||||
'content': "Hello! What's your name?"
|
|
||||||
}, {
|
|
||||||
'role':
|
|
||||||
'assistant',
|
|
||||||
'content':
|
|
||||||
'I am InternLM2! A helpful AI assistant. What can I do for you?'
|
|
||||||
}]
|
|
||||||
chat_2 = [{
|
|
||||||
'role': 'user',
|
|
||||||
'content': "Hello! What's your name?"
|
|
||||||
}, {
|
|
||||||
'role': 'assistant',
|
|
||||||
'content': 'I have no idea.'
|
|
||||||
}]
|
|
||||||
|
|
||||||
# get reward score for a single chat
|
|
||||||
score1 = model.get_score(tokenizer, chat_1)
|
|
||||||
score2 = model.get_score(tokenizer, chat_2)
|
|
||||||
print('score1: ', score1)
|
|
||||||
print('score2: ', score2)
|
|
||||||
assert score1 > 0
|
|
||||||
assert score2 < 0
|
|
||||||
|
|
||||||
# batch inference, get multiple scores at once
|
|
||||||
scores = model.get_scores(tokenizer, [chat_1, chat_2])
|
|
||||||
print('scores: ', scores)
|
|
||||||
assert scores[0] > 0
|
|
||||||
assert scores[1] < 0
|
|
||||||
|
|
||||||
# compare whether chat_1 is better than chat_2
|
|
||||||
compare_res = model.compare(tokenizer, chat_1, chat_2)
|
|
||||||
print('compare_res: ', compare_res)
|
|
||||||
assert compare_res
|
|
||||||
# >>> compare_res: True
|
|
||||||
|
|
||||||
# rank multiple chats, it will return the ranking index of each chat
|
|
||||||
# the chat with the highest score will have ranking index as 0
|
|
||||||
rank_res = model.rank(tokenizer, [chat_1, chat_2])
|
|
||||||
print('rank_res: ', rank_res) # lower index means higher score
|
|
||||||
# >>> rank_res: [0, 1]
|
|
||||||
assert rank_res[0] == 0
|
|
||||||
assert rank_res[1] == 1
|
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
|
||||||
'model_name',
|
|
||||||
[
|
|
||||||
'internlm/internlm-reward-1_8b', 'internlm/internlm-reward-7b',
|
|
||||||
'internlm/internlm-reward-20b'
|
|
||||||
],
|
|
||||||
)
|
|
||||||
@pytest.mark.parametrize(
|
|
||||||
'usefast',
|
|
||||||
[
|
|
||||||
True,
|
|
||||||
False,
|
|
||||||
],
|
|
||||||
)
|
|
||||||
def test_demo_topn(self, model_name, usefast):
|
|
||||||
# prepare the llm model and tokenizer
|
|
||||||
llm = AutoModel.from_pretrained(
|
|
||||||
'internlm/internlm2-chat-7b',
|
|
||||||
device_map='cuda',
|
|
||||||
torch_dtype=torch.float16,
|
|
||||||
trust_remote_code=True,
|
|
||||||
)
|
|
||||||
llm_tokenizer = AutoTokenizer.from_pretrained(
|
|
||||||
'internlm/internlm2-chat-7b', trust_remote_code=True)
|
|
||||||
|
|
||||||
# prepare the reward model and tokenizer
|
|
||||||
reward = AutoModel.from_pretrained(
|
|
||||||
model_name,
|
|
||||||
device_map='cuda',
|
|
||||||
torch_dtype=torch.float16,
|
|
||||||
trust_remote_code=True,
|
|
||||||
)
|
|
||||||
reward_tokenizer = AutoTokenizer.from_pretrained(
|
|
||||||
model_name, trust_remote_code=True)
|
|
||||||
|
|
||||||
# prepare the chat prompt
|
|
||||||
prompt = 'Write an short bedtime story.'
|
|
||||||
messages = [{
|
|
||||||
'role': 'system',
|
|
||||||
'content': 'You are a helpful assistant.'
|
|
||||||
}, {
|
|
||||||
'role': 'user',
|
|
||||||
'content': prompt
|
|
||||||
}]
|
|
||||||
text = llm_tokenizer.apply_chat_template(messages,
|
|
||||||
tokenize=False,
|
|
||||||
add_generation_prompt=True)
|
|
||||||
model_inputs = llm_tokenizer([text], return_tensors='pt').to('cuda')
|
|
||||||
|
|
||||||
# generate best of N candidates
|
|
||||||
num_candidates = 3 # N=3
|
|
||||||
candidates = []
|
|
||||||
|
|
||||||
outputs = llm.generate(
|
|
||||||
**model_inputs,
|
|
||||||
max_new_tokens=512,
|
|
||||||
num_return_sequences=num_candidates,
|
|
||||||
pad_token_id=llm_tokenizer.eos_token_id,
|
|
||||||
do_sample=True,
|
|
||||||
top_k=50,
|
|
||||||
top_p=0.95,
|
|
||||||
temperature=0.8,
|
|
||||||
)
|
|
||||||
outputs = outputs[:, model_inputs['input_ids'].shape[1]:]
|
|
||||||
for i in range(num_candidates):
|
|
||||||
candidate = llm_tokenizer.decode(outputs[i],
|
|
||||||
skip_special_tokens=True)
|
|
||||||
candidates.append(messages + [{
|
|
||||||
'role': 'assistant',
|
|
||||||
'content': candidate
|
|
||||||
}])
|
|
||||||
|
|
||||||
rank_indices = reward.rank(reward_tokenizer, candidates)
|
|
||||||
sorted_candidates = sorted(zip(rank_indices, candidates),
|
|
||||||
key=lambda x: x[0])
|
|
||||||
|
|
||||||
# print the best response
|
|
||||||
best_response = sorted_candidates[0][1][-1]['content']
|
|
||||||
print(sorted_candidates)
|
|
||||||
print(best_response)
|
|
||||||
assert len(sorted_candidates) == 3
|
|
||||||
|
|
||||||
|
|
||||||
class TestMMModel:
|
class TestMMModel:
|
||||||
"""Test cases for base model."""
|
"""Test cases for base model."""
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue