179 lines
7.3 KiB
Python
179 lines
7.3 KiB
Python
import argparse
|
|
import json
|
|
import shutil
|
|
import subprocess
|
|
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"
|
|
self.dictionary_dir.mkdir()
|
|
shutil.copy2(self.repo / "dictionary" / "static-base.json", self.dictionary_dir / "static-base.json")
|
|
subprocess.run(
|
|
[
|
|
"cargo",
|
|
"run",
|
|
"--",
|
|
"dictionary",
|
|
"generate",
|
|
"--out-dir",
|
|
str(self.dictionary_dir),
|
|
],
|
|
cwd=self.repo,
|
|
check=True,
|
|
stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.DEVNULL,
|
|
)
|
|
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,
|
|
"sdg_model_alias": "nvidia-text",
|
|
"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())
|
|
|
|
def test_training_data_relative_path_prefers_custom_root_shape(self):
|
|
path = Path("/Volumes/Zoe/custom-local-llm/training-data/module_controller_intents_123")
|
|
self.assertEqual(
|
|
bridge.relative_data_path(path),
|
|
"training-data/module_controller_intents_123",
|
|
)
|
|
|
|
def test_data_designer_config_declares_structured_expansion(self):
|
|
config = bridge.data_designer_config(
|
|
dataset_key="module_controller_intents_123",
|
|
generated_at=123,
|
|
remote_seed="/tmp/seed.jsonl",
|
|
action_ids=["logger"],
|
|
model_alias="nvidia-text",
|
|
)
|
|
self.assertIn("load_config_builder", config)
|
|
self.assertIn('column_type="llm-structured"', config)
|
|
self.assertIn('model_alias="nvidia-text"', config)
|
|
self.assertIn("expanded_record", config)
|
|
|
|
def test_remote_data_designer_command_sources_mimir_env(self):
|
|
command = bridge.remote_data_designer_create_command(
|
|
remote_config="/tmp/controller.py",
|
|
num_records=42,
|
|
dataset_key="module_controller_intents_123",
|
|
artifact_path="/tmp/artifacts",
|
|
)
|
|
self.assertIn(bridge.MIMIR_DATA_DESIGNER_ENV, command)
|
|
self.assertIn(bridge.MIMIR_DATA_DESIGNER_BIN, command)
|
|
self.assertIn("--num-records 42", command)
|
|
self.assertIn("--no-tui", command)
|
|
|
|
def test_remote_bash_command_quotes_inner_command(self):
|
|
command = bridge.remote_bash_command("set -a; echo ready")
|
|
self.assertTrue(command.startswith("bash -lc "))
|
|
self.assertIn("'set -a; echo ready'", command)
|
|
|
|
def test_normalize_expanded_jsonl_extracts_expanded_record(self):
|
|
raw = self.root / "expanded.raw.jsonl"
|
|
normalized = self.root / "expanded.jsonl"
|
|
source_row = {
|
|
"expanded_record": {
|
|
"instruction": bridge.INSTRUCTION,
|
|
"input": "record this message",
|
|
"output": json.dumps({"action_id": "logger", "payload": {"message": "hello"}}),
|
|
"metadata": {
|
|
"action_id": "logger",
|
|
"source": "data_designer",
|
|
"split": "train",
|
|
"generated_at": 123,
|
|
},
|
|
}
|
|
}
|
|
bridge.write_jsonl(raw, [source_row])
|
|
bridge.normalize_expanded_jsonl(raw, normalized)
|
|
rows = bridge.read_jsonl(normalized)
|
|
self.assertEqual(len(rows), 1)
|
|
self.assertEqual(rows[0]["metadata"]["action_id"], "logger")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|