File size: 5,643 Bytes
2abcc30 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 | import json
import pyarrow as pa
import pyarrow.parquet as pq
from albedo_eval_service.shared.submit_protocol import CANONICAL_MARKER
from local_train.pack import (
candidate_turns,
family_of,
is_gold_submit,
pack,
phase_for,
)
from local_train.think import wrap_completion
def test_phase_and_family_helpers():
assert phase_for(1, 9) == "cold"
assert phase_for(7, 9) == "pre_edit"
assert phase_for(9, 9) == "at_edit"
assert phase_for(10, 9) == "post_edit"
assert family_of("repo.abc.pr_12") == "pr"
assert family_of("repo.abc.lm_9") == "lm"
assert family_of("repo.abc.combine_x") == "combine"
assert family_of("repo.abc.func_pm") == "mechanical"
def test_candidate_turns_skip_early_submit_and_include_edit():
golds = [
"```bash\nls\n```",
"```bash\ncat a.py\n```",
f"```bash\necho {CANONICAL_MARKER}\n```",
"```bash\nsed -i 's/a/b/' a.py\n```",
f"```bash\necho {CANONICAL_MARKER}\n```",
]
turns = candidate_turns(len(golds), first_edit=4, golds=golds)
idxs = {t for t, _ in turns}
assert 1 in idxs and 2 in idxs
assert 3 in idxs # first-edit turn (1-based 4)
assert 4 in idxs # later submit kept
# Early submit gold is not added as its own extra (turn 2 is only the official cold cut).
assert is_gold_submit(golds[4])
assert not is_gold_submit(golds[1])
def test_pack_from_tiny_parquet(tmp_path):
source = tmp_path / "mini-coder" / "data"
source.mkdir(parents=True)
native = (
"You are a coding agent.\n\n## Submission\n\n"
f"```bash\necho {CANONICAL_MARKER}\n```\n"
)
messages = [
{"role": "system", "content": native},
{"role": "user", "content": "Fix foo.py"},
{"role": "assistant", "content": "THOUGHT: look\n```bash\nls\n```"},
{"role": "user", "content": "<returncode>0</returncode>\n<output>\nok\n</output>"},
{"role": "assistant", "content": "THOUGHT: read\n```bash\ncat foo.py\n```"},
{"role": "user", "content": "<returncode>0</returncode>\n<output>\nx=1\n</output>"},
{
"role": "assistant",
"content": "THOUGHT: edit\n```bash\nsed -i 's/x=1/x=2/' foo.py\n```",
},
{"role": "user", "content": "<returncode>0</returncode>\n<output>\n\n</output>"},
{
"role": "assistant",
"content": f"THOUGHT: submit\n```bash\necho {CANONICAL_MARKER}\n```",
},
]
table = pa.table(
{
"instance_id": ["demo.repo.pr_1"],
"messages": [messages],
"first_edit": [3],
"family": ["pr"],
"repo": ["demo"],
"language": ["python"],
}
)
pq.write_table(table, source / "train-00000.parquet")
out = tmp_path / "pack"
path = pack(
dataset_root=tmp_path,
out_dir=out,
max_examples=16,
seed="unit",
n_salts=2,
)
rows = [json.loads(line) for line in path.read_text().splitlines() if line.strip()]
assert rows
kinds = {row["kind"] for row in rows}
assert "edit" in kinds or "explore" in kinds or "submit" in kinds
for row in rows:
assert row["completion"].count("```bash") == 1
assert "<think>" in row["completion"] and "</think>" in row["completion"]
assert row["prompt"]
if row["kind"] == "submit":
assert row["submit_command"]
body = row["completion"].split("</think>", 1)[-1]
assert row["submit_command"] in body
wrapped = wrap_completion(row["completion"])
assert wrapped is not None
meta = json.loads((out / "sft-16-unit.meta.json").read_text())
assert meta["n"] == len(rows)
def test_expand_submit_salts_covers_multiple_rewrites(tmp_path):
source = tmp_path / "mini-coder" / "data"
source.mkdir(parents=True)
native = (
"You are a coding agent.\n\n## Submission\n\n"
f"```bash\necho {CANONICAL_MARKER}\n```\n"
)
messages = [
{"role": "system", "content": native},
{"role": "user", "content": "Fix foo.py"},
{"role": "assistant", "content": "THOUGHT: look\n```bash\nls\n```"},
{"role": "user", "content": "<returncode>0</returncode>\n<output>\nok\n</output>"},
{"role": "assistant", "content": "THOUGHT: read\n```bash\ncat foo.py\n```"},
{"role": "user", "content": "<returncode>0</returncode>\n<output>\nx=1\n</output>"},
{
"role": "assistant",
"content": "THOUGHT: edit\n```bash\nsed -i 's/x=1/x=2/' foo.py\n```",
},
{"role": "user", "content": "<returncode>0</returncode>\n<output>\n\n</output>"},
{
"role": "assistant",
"content": f"THOUGHT: submit\n```bash\necho {CANONICAL_MARKER}\n```",
},
]
table = pa.table(
{
"instance_id": ["demo.repo.pr_1"],
"messages": [messages],
"first_edit": [3],
"family": ["pr"],
"repo": ["demo"],
"language": ["python"],
}
)
pq.write_table(table, source / "train-00000.parquet")
path = pack(
dataset_root=tmp_path,
out_dir=tmp_path / "pack",
max_examples=8,
seed="expand",
n_salts=4,
submit_frac=0.8,
edit_frac=0.2,
expand_submit_salts=True,
)
rows = [json.loads(line) for line in path.read_text().splitlines() if line.strip()]
submits = [row for row in rows if row["kind"] == "submit"]
assert len(submits) >= 2
assert len({row["submit_command"] for row in submits}) >= 2
|