|
|
import pytest |
|
|
|
|
|
from ether0.chat import ChatArguments |
|
|
from ether0.model_prompts import ProblemPrompt, SysPrompt |
|
|
|
|
|
|
|
|
class TestChatArguments: |
|
|
@pytest.mark.parametrize( |
|
|
("args", "row", "expected"), |
|
|
[ |
|
|
( |
|
|
ChatArguments(problem_prompt=ProblemPrompt.NONE), |
|
|
{"problem": "stub problem"}, |
|
|
{"prompt": [{"content": "stub problem", "role": "user"}]}, |
|
|
), |
|
|
( |
|
|
ChatArguments(problem_prompt=ProblemPrompt.NONE), |
|
|
{"problem": ["stub problem", "stub problem 2"]}, |
|
|
{ |
|
|
"prompt": [ |
|
|
[{"content": "stub problem", "role": "user"}], |
|
|
[{"content": "stub problem 2", "role": "user"}], |
|
|
] |
|
|
}, |
|
|
), |
|
|
( |
|
|
ChatArguments( |
|
|
sys_prompt=SysPrompt.SCIENTIFIC_AI, |
|
|
problem_prompt=ProblemPrompt.THINK_ANSWER, |
|
|
), |
|
|
{"problem": "stub problem"}, |
|
|
{ |
|
|
"prompt": [ |
|
|
{ |
|
|
"role": "system", |
|
|
"content": "You are a scientific reasoning AI assistant.", |
|
|
}, |
|
|
{ |
|
|
"role": "user", |
|
|
"content": ( |
|
|
"A conversation between User and Assistant. The user" |
|
|
" asks a question, and the Assistant solves it. The" |
|
|
" assistant first thinks about the reasoning process in" |
|
|
" the mind and then provides the user with the answer." |
|
|
" The reasoning process and answer are enclosed within" |
|
|
" <|think_start|> <|think_end|>" |
|
|
" and <|answer_start|> <|answer_end|> tags," |
|
|
" respectively, i.e., <|think_start|> reasoning process here" |
|
|
" <|think_end|><|answer_start|> answer here <|answer_end|>" |
|
|
"\n\nstub problem" |
|
|
), |
|
|
}, |
|
|
] |
|
|
}, |
|
|
), |
|
|
], |
|
|
) |
|
|
def test_rl_conversation( |
|
|
self, args: ChatArguments, row: dict, expected: dict |
|
|
) -> None: |
|
|
assert args.make_rl_conversation(row) == expected |
|
|
|
