102 lines
4.3 KiB
Python
102 lines
4.3 KiB
Python
|
import argparse
|
||
|
|
import json
|
||
|
|
import shutil
|
||
|
|
import tempfile
|
||
|
|
import unittest
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
from scripts import generate_training_data as bridge
|
||
|
|
|
||
|
|
|
||
|
|
class GenerateTrainingDataTests(unittest.TestCase):
|
||
|
|
def setUp(self):
|
||
|
|
self.repo = Path(__file__).resolve().parents[1]
|
||
|
|
self.temp = tempfile.TemporaryDirectory()
|
||
|
|
self.root = Path(self.temp.name)
|
||
|
|
self.dictionary_dir = self.root / "dictionary"
|
||
|
|
shutil.copytree(self.repo / "dictionary", self.dictionary_dir)
|
||
|
|
self.output_dir = self.root / "workspace_Data" / "data"
|
||
|
|
self.output_dir.mkdir(parents=True)
|
||
|
|
self.manifest = self.root / "workspace_Data" / "manifest.llm.json"
|
||
|
|
self.manifest.write_text(json.dumps({"manifest_version": "0.0.0", "datasets": {}}))
|
||
|
|
|
||
|
|
def tearDown(self):
|
||
|
|
self.temp.cleanup()
|
||
|
|
|
||
|
|
def args(self, **overrides):
|
||
|
|
values = {
|
||
|
|
"dictionary_dir": str(self.dictionary_dir),
|
||
|
|
"output_dir": str(self.output_dir),
|
||
|
|
"manifest": str(self.manifest),
|
||
|
|
"dry_run": False,
|
||
|
|
"sdg_host": None,
|
||
|
|
"sdg_user": None,
|
||
|
|
"sdg_key": None,
|
||
|
|
"num_records": 1000,
|
||
|
|
"timestamp": 1234567890,
|
||
|
|
}
|
||
|
|
values.update(overrides)
|
||
|
|
return argparse.Namespace(**values)
|
||
|
|
|
||
|
|
def test_seed_generation_covers_current_dictionary(self):
|
||
|
|
artifacts = bridge.load_dictionary_artifacts(self.dictionary_dir)
|
||
|
|
records = bridge.build_seed_records(artifacts, generated_at=123)
|
||
|
|
actions = {record["metadata"]["action_id"] for record in records}
|
||
|
|
self.assertEqual(len(actions), 13)
|
||
|
|
self.assertIn("calculate", actions)
|
||
|
|
self.assertIn("delete_file", actions)
|
||
|
|
self.assertEqual(len(records), 260)
|
||
|
|
|
||
|
|
def test_compact_output_uses_codebook_codes(self):
|
||
|
|
artifacts = bridge.load_dictionary_artifacts(self.dictionary_dir)
|
||
|
|
compact = bridge.compact_target(
|
||
|
|
artifacts,
|
||
|
|
"hash_string",
|
||
|
|
{"text": "hello", "algorithm": "sha256"},
|
||
|
|
)
|
||
|
|
self.assertEqual(compact["a"], artifacts["action_to_code"]["hash_string"])
|
||
|
|
self.assertIn(artifacts["field_to_code"]["text"], compact["p"])
|
||
|
|
self.assertIn(artifacts["field_to_code"]["algorithm"], compact["p"])
|
||
|
|
|
||
|
|
def test_empty_payload_action(self):
|
||
|
|
artifacts = bridge.load_dictionary_artifacts(self.dictionary_dir)
|
||
|
|
records = bridge.build_seed_records(artifacts, generated_at=123)
|
||
|
|
uuid_records = [
|
||
|
|
record for record in records if record["metadata"]["action_id"] == "generate_uuid"
|
||
|
|
]
|
||
|
|
self.assertTrue(uuid_records)
|
||
|
|
output = json.loads(uuid_records[0]["output"])
|
||
|
|
compact = json.loads(uuid_records[0]["metadata"]["compact"])
|
||
|
|
self.assertEqual(output["payload"], {})
|
||
|
|
self.assertEqual(compact["p"], {})
|
||
|
|
|
||
|
|
def test_split_includes_every_action(self):
|
||
|
|
artifacts = bridge.load_dictionary_artifacts(self.dictionary_dir)
|
||
|
|
records = bridge.build_seed_records(artifacts, generated_at=123)
|
||
|
|
train, validation = bridge.split_records_by_action(records)
|
||
|
|
train_actions = {record["metadata"]["action_id"] for record in train}
|
||
|
|
validation_actions = {record["metadata"]["action_id"] for record in validation}
|
||
|
|
self.assertEqual(train_actions, validation_actions)
|
||
|
|
self.assertEqual(len(train_actions), 13)
|
||
|
|
|
||
|
|
def test_dry_run_writes_nothing(self):
|
||
|
|
summary = bridge.run_bridge(self.args(dry_run=True))
|
||
|
|
self.assertTrue(summary["dry_run"])
|
||
|
|
self.assertFalse((self.output_dir / "module_controller_intents_1234567890").exists())
|
||
|
|
manifest = json.loads(self.manifest.read_text())
|
||
|
|
self.assertEqual(manifest["datasets"], {})
|
||
|
|
|
||
|
|
def test_manifest_patch_adds_dataset_without_damaging_existing_keys(self):
|
||
|
|
bridge.run_bridge(self.args())
|
||
|
|
manifest = json.loads(self.manifest.read_text())
|
||
|
|
self.assertEqual(manifest["manifest_version"], "0.0.0")
|
||
|
|
dataset = manifest["datasets"]["module_controller_intents_1234567890"]
|
||
|
|
self.assertEqual(dataset["schema_profile"], "instruction")
|
||
|
|
self.assertEqual(dataset["codebook_version"], "1.0.0")
|
||
|
|
self.assertIn("codebook_checksum", dataset)
|
||
|
|
self.assertTrue((self.output_dir / "module_controller_intents_1234567890" / "train.jsonl").exists())
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|