78 lines
2.5 KiB
Python
78 lines
2.5 KiB
Python
import json
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
|
|
MODULE_DIR = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(MODULE_DIR))
|
|
|
|
import prepare_data
|
|
import purpose_data
|
|
|
|
|
|
def example(index: int):
|
|
return {
|
|
"prompt": f"Implement sample endpoint number {index} with stable pagination",
|
|
"purpose": purpose_data.LABELS[index % len(purpose_data.LABELS)],
|
|
"secondary": None,
|
|
"mixed": False,
|
|
"difficulty": 0.4,
|
|
"slice": "vague-eval" if index < 5 else "core",
|
|
"lang": "en",
|
|
}
|
|
|
|
|
|
class PrepareIntegrationTests(unittest.TestCase):
|
|
def test_refresh_then_verify_frozen_split(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
source = root / "source.jsonl"
|
|
purpose_data.write_jsonl(source, (example(index) for index in range(80)))
|
|
fixtures = root / "fixtures.json"
|
|
fixtures.write_text(
|
|
json.dumps(
|
|
[
|
|
{"prompt": "Plan the cache migration", "purpose": "planning"},
|
|
{"prompt": "Anything else?", "purpose": "general"},
|
|
]
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
output = root / "output"
|
|
frozen = root / "frozen.jsonl"
|
|
manifest = root / "manifest.json"
|
|
|
|
first = prepare_data.prepare(
|
|
sources=[source],
|
|
fixtures_path=fixtures,
|
|
output_dir=output,
|
|
frozen_test_path=frozen,
|
|
manifest_path=manifest,
|
|
refresh_frozen_test=True,
|
|
seed=23,
|
|
near_duplicate_threshold=0.92,
|
|
)
|
|
second = prepare_data.prepare(
|
|
sources=[source],
|
|
fixtures_path=fixtures,
|
|
output_dir=output,
|
|
frozen_test_path=frozen,
|
|
manifest_path=manifest,
|
|
refresh_frozen_test=False,
|
|
seed=23,
|
|
near_duplicate_threshold=0.92,
|
|
)
|
|
|
|
self.assertEqual(first, second)
|
|
self.assertEqual(65, first["splits"]["train"]["records"])
|
|
self.assertEqual(8, first["splits"]["validation"]["records"])
|
|
self.assertEqual(8, first["splits"]["test"]["logicalRecords"])
|
|
train = purpose_data.load_jsonl(output / "train.jsonl")
|
|
self.assertFalse(any(row["slice"] == "vague-eval" for row in train))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|