32 lines
924 B
Python
32 lines
924 B
Python
import json
|
|
|
|
from code_agent import CodeRLDataset, collate_code_rl
|
|
|
|
|
|
class FakeTokenizer:
|
|
def apply_chat_template(self, messages, **kwargs):
|
|
return messages[0]["content"] + f"\nthinking={kwargs['open_thinking']}"
|
|
|
|
|
|
def test_code_rl_dataset_builds_prompt_and_preserves_tests(tmp_path):
|
|
path = tmp_path / "tasks.jsonl"
|
|
path.write_text(
|
|
json.dumps(
|
|
{
|
|
"task_id": "inc",
|
|
"prompt": "Implement inc(x).",
|
|
"entry_point": "inc",
|
|
"tests": [{"args": [1], "expected": 2}],
|
|
}
|
|
)
|
|
+ "\n",
|
|
encoding="utf-8",
|
|
)
|
|
dataset = CodeRLDataset(path, FakeTokenizer(), thinking_ratio=0.0)
|
|
sample = dataset[0]
|
|
|
|
assert "Required entry point: inc" in sample["prompt"]
|
|
assert sample["task"].tests[0].expected == 2
|
|
batch = collate_code_rl([sample])
|
|
assert batch["tasks"][0].task_id == "inc"
|