73 lines
2.3 KiB
Python
73 lines
2.3 KiB
Python
"""Tests for publishable context-cutoff split materialization."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pyarrow as pa
|
|
import pyarrow.parquet as pq
|
|
|
|
from swe_data_processing.materialize import materialize_dataset
|
|
|
|
|
|
def _decision(
|
|
sample_id: str,
|
|
source: Path,
|
|
bucket: str,
|
|
resolved: int,
|
|
*,
|
|
review: bool = False,
|
|
) -> dict:
|
|
return {
|
|
"sample_id": sample_id,
|
|
"source_parquet": str(source),
|
|
"resolved": resolved,
|
|
"length_bucket": bucket,
|
|
"hard_reject_reasons": [],
|
|
"review_flags": ["EARLY_FAILURE_CLUSTER"] if review else [],
|
|
"metrics": {
|
|
"token_count": 50_000 if bucket == "LE_81920" else 100_000,
|
|
"failed_tool_call_count": 1,
|
|
"failed_tool_call_rate": 0.1,
|
|
"longest_consecutive_failure_run": 1,
|
|
},
|
|
}
|
|
|
|
|
|
def test_materialize_creates_nested_context_splits(tmp_path: Path) -> None:
|
|
source = tmp_path / "source" / "family" / "part.parquet"
|
|
source.parent.mkdir(parents=True)
|
|
pq.write_table(
|
|
pa.Table.from_pylist(
|
|
[
|
|
{"trajectory_id": "short", "resolved": 1},
|
|
{"trajectory_id": "medium", "resolved": 0},
|
|
{"trajectory_id": "review", "resolved": -1},
|
|
]
|
|
),
|
|
source,
|
|
)
|
|
decisions = tmp_path / "decisions.jsonl"
|
|
rows = [
|
|
_decision("short", source, "LE_81920", 1),
|
|
_decision("medium", source, "81921_TO_131072", 0),
|
|
_decision("review", source, "LE_81920", -1, review=True),
|
|
]
|
|
decisions.write_text(
|
|
"".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8"
|
|
)
|
|
|
|
output = tmp_path / "output"
|
|
manifest = materialize_dataset(decisions, output, workers=1, resume=False)
|
|
|
|
assert manifest["splits"]["context_131072"]["samples"] == 2
|
|
assert manifest["splits"]["context_81920"]["samples"] == 1
|
|
long_table = pq.read_table(output / "data/context_131072")
|
|
short_table = pq.read_table(output / "data/context_81920")
|
|
assert long_table.column("trajectory_id").to_pylist() == ["short", "medium"]
|
|
assert short_table.column("trajectory_id").to_pylist() == ["short"]
|
|
assert set(long_table.column("qc_filter_version").to_pylist()) == {
|
|
"deterministic-v3.1"
|
|
}
|