#!/usr/bin/env python3 """Generate module-controller training data from dictionary artifacts.""" from __future__ import annotations import argparse import json import shlex import subprocess import sys import tempfile import time from pathlib import Path from typing import Any INSTRUCTION = "Map the following natural language request to the correct action and payload." DATASET_PREFIX = "module_controller_intents" DEFAULT_MIN_RECORDS_PER_ACTION = 20 CUSTOM_LOCAL_LLM_ROOT = Path("/Volumes/Zoe/custom-local-llm") DEFAULT_OUTPUT_DIR = CUSTOM_LOCAL_LLM_ROOT / "training-data" DEFAULT_MANIFEST = CUSTOM_LOCAL_LLM_ROOT / "manifest.llm.json" MIMIR_DATA_DESIGNER_BIN = "/home/aaron-pressey/.venvs/data-designer/bin/data-designer" MIMIR_DATA_DESIGNER_ENV = "/home/aaron-pressey/.config/home-grown-llm-data/data-designer.env" def main() -> int: args = parse_args() try: summary = run_bridge(args) except Exception as exc: # noqa: BLE001 - CLI should report any bridge failure clearly. print(f"error: {exc}", file=sys.stderr) return 1 print(json.dumps(summary, indent=2, sort_keys=True)) return 0 def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--dictionary-dir", default="dictionary/") parser.add_argument("--output-dir", default=str(DEFAULT_OUTPUT_DIR)) parser.add_argument("--manifest", default=str(DEFAULT_MANIFEST)) parser.add_argument("--dry-run", action="store_true") parser.add_argument("--sdg-host") parser.add_argument("--sdg-user") parser.add_argument("--sdg-key") parser.add_argument("--sdg-model-alias", default="nvidia-text") parser.add_argument("--num-records", type=int, default=1000) parser.add_argument( "--timestamp", type=int, help="Override generation timestamp. Intended for deterministic tests.", ) return parser.parse_args() def run_bridge(args: argparse.Namespace) -> dict[str, Any]: dictionary_dir = Path(args.dictionary_dir) output_root = Path(args.output_dir) manifest_path = Path(args.manifest) generated_at = args.timestamp or int(time.time()) dataset_key = f"{DATASET_PREFIX}_{generated_at}" artifacts = load_dictionary_artifacts(dictionary_dir) records = build_seed_records(artifacts, generated_at) train_records, validation_records = split_records_by_action(records) summary = { "dataset_key": dataset_key, "generated_at": generated_at, "dry_run": bool(args.dry_run), "action_count": len(artifacts["actions"]), "train_records": len(train_records), "validation_records": len(validation_records), "output_dir": str(output_root / dataset_key), "manifest": str(manifest_path), "sdg_requested": has_sdg_args(args), "sdg_model_alias": args.sdg_model_alias, } if args.dry_run: return summary dataset_dir = output_root / dataset_key dataset_dir.mkdir(parents=True, exist_ok=False) write_jsonl(dataset_dir / "train.jsonl", train_records) write_jsonl(dataset_dir / "validation.jsonl", validation_records) write_dataset_readme(dataset_dir / "README.md", summary, artifacts) expanded_path = None sdg_error = None if has_sdg_args(args): try: expanded_path = run_mimir_expansion( args=args, dataset_dir=dataset_dir, dataset_key=dataset_key, generated_at=generated_at, seed_records=train_records + validation_records, action_ids=sorted(artifacts["actions"].keys()), model_alias=args.sdg_model_alias, ) except Exception as exc: # noqa: BLE001 - keep local seed output and report remote failure. sdg_error = str(exc) (dataset_dir / "SDG_ERROR.txt").write_text( f"Mimir/Data Designer expansion failed after local seed generation.\n\n{sdg_error}\n" ) patch_manifest( manifest_path=manifest_path, dataset_key=dataset_key, dataset_dir=dataset_dir, generated_at=generated_at, artifacts=artifacts, expanded_path=expanded_path, sdg_requested=has_sdg_args(args), sdg_model_alias=args.sdg_model_alias, sdg_error=sdg_error, ) summary["expanded_path"] = str(expanded_path) if expanded_path else None summary["sdg_error"] = sdg_error if sdg_error: raise RuntimeError(f"local seed dataset was written, but Mimir expansion failed: {sdg_error}") return summary def load_dictionary_artifacts(dictionary_dir: Path) -> dict[str, Any]: actions_index = read_json(dictionary_dir / "actions.index.json") codebook = read_json(dictionary_dir / "model.codebook.json") examples = read_jsonl(dictionary_dir / "model.examples.jsonl") actions = actions_index.get("actions", {}) if not actions: raise ValueError("actions.index.json contains no actions") return { "actions_index": actions_index, "codebook": codebook, "examples": examples, "actions": actions, "action_to_code": invert_unique(codebook["adaptive"]["actions"], "action codes"), "field_to_code": invert_unique(codebook["adaptive"]["fields"], "field codes"), } def build_seed_records(artifacts: dict[str, Any], generated_at: int) -> list[dict[str, Any]]: records: list[dict[str, Any]] = [] for action_id in sorted(artifacts["actions"]): action = artifacts["actions"][action_id] action_records: list[dict[str, Any]] = [] sample_payload = sample_payload_for_schema(action.get("payload_schema", {})) for phrase in action.get("intent_examples", []): action_records.append( build_record(artifacts, action_id, phrase, sample_payload, "seed_intent_example", generated_at) ) for alias in action.get("aliases", []): phrase = f"{alias} {payload_phrase(action_id, sample_payload)}".strip() action_records.append( build_record(artifacts, action_id, phrase, sample_payload, "alias_variant", generated_at) ) for phrase, payload in field_variants(action_id, action.get("payload_schema", {})): action_records.append( build_record(artifacts, action_id, phrase, payload, "field_variant", generated_at) ) for phrase, payload in template_variants(action_id, sample_payload, artifacts["examples"]): action_records.append( build_record(artifacts, action_id, phrase, payload, "template_expansion", generated_at) ) while len(action_records) < DEFAULT_MIN_RECORDS_PER_ACTION: ordinal = len(action_records) + 1 phrase = f"please run {action_id} example {ordinal} with {payload_phrase(action_id, sample_payload)}" action_records.append( build_record(artifacts, action_id, phrase, sample_payload, "template_expansion", generated_at) ) records.extend(dedupe_records(action_records)[:DEFAULT_MIN_RECORDS_PER_ACTION]) return records def build_record( artifacts: dict[str, Any], action_id: str, input_text: str, payload: dict[str, Any], source: str, generated_at: int, ) -> dict[str, Any]: verbose = {"action_id": action_id, "payload": payload} compact = compact_target(artifacts, action_id, payload) return { "instruction": INSTRUCTION, "input": input_text, "output": json.dumps(verbose, sort_keys=True), "metadata": { "action_id": action_id, "compact": json.dumps(compact, separators=(",", ":"), sort_keys=True), "source": source, "split": "", "generated_at": generated_at, }, } def compact_target(artifacts: dict[str, Any], action_id: str, payload: dict[str, Any]) -> dict[str, Any]: codebook = artifacts["codebook"] action_to_code = artifacts["action_to_code"] field_to_code = artifacts["field_to_code"] if action_id not in action_to_code: raise ValueError(f"missing compact action code for {action_id}") compact_payload = {} for field, value in sorted(payload.items()): if field not in field_to_code: raise ValueError(f"missing compact field code for {field}") compact_payload[field_to_code[field]] = value return { "v": codebook["dictionary_version"], "c": codebook["codebook_checksum"], "a": action_to_code[action_id], "p": compact_payload, } def split_records_by_action(records: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: grouped: dict[str, list[dict[str, Any]]] = {} for record in records: grouped.setdefault(record["metadata"]["action_id"], []).append(record) train: list[dict[str, Any]] = [] validation: list[dict[str, Any]] = [] for action_id in sorted(grouped): action_records = grouped[action_id] validation_count = max(1, round(len(action_records) * 0.1)) if len(action_records) > 1 else 0 split_at = len(action_records) - validation_count for record in action_records[:split_at]: train.append(with_split(record, "train")) for record in action_records[split_at:]: validation.append(with_split(record, "validation")) return train, validation def with_split(record: dict[str, Any], split: str) -> dict[str, Any]: cloned = json.loads(json.dumps(record)) cloned["metadata"]["split"] = split return cloned def sample_payload_for_schema(schema: dict[str, Any]) -> dict[str, Any]: properties = schema.get("properties", {}) return {field: sample_value(field, field_schema) for field, field_schema in sorted(properties.items())} def sample_value(field: str, schema: dict[str, Any]) -> Any: if "enum" in schema: return schema["enum"][0] if field == "expression": return "2 + 2 * 3" if field == "value": return 100 if field == "from": return "km" if field == "to": return "miles" if field == "min": return 1 if field == "max": return 100 if field == "max_chars": return 100 if field == "filename": return "hello.txt" if field == "content": return "hello from codex" if field in {"text", "message"}: return "hello from codex" if schema.get("type") in {"number", "integer"}: return 1 return f"{field} value" def field_variants(action_id: str, schema: dict[str, Any]) -> list[tuple[str, dict[str, Any]]]: properties = schema.get("properties", {}) if not properties: return [(f"{action_id} with empty payload", {})] variants: list[tuple[str, dict[str, Any]]] = [] base = sample_payload_for_schema(schema) for field, field_schema in sorted(properties.items()): if "enum" in field_schema: for enum_value in field_schema["enum"]: payload = dict(base) payload[field] = enum_value variants.append((f"{action_id} where {field} is {enum_value}", payload)) elif field_schema.get("type") in {"number", "integer"}: for value in [0, sample_value(field, field_schema), 999]: payload = dict(base) payload[field] = value variants.append((f"{action_id} with {field} {value}", payload)) else: payload = dict(base) payload[field] = sample_value(field, field_schema) variants.append((f"{action_id} using {field} {payload[field]}", payload)) return variants def template_variants( action_id: str, sample_payload: dict[str, Any], examples: list[dict[str, Any]], ) -> list[tuple[str, dict[str, Any]]]: template_count = max(1, sum(1 for example in examples if example.get("kind") == "synthetic_template")) variants = [] for index in range(1, max(5, template_count) + 1): variants.append((f"template request {index} for {action_id}: {payload_phrase(action_id, sample_payload)}", sample_payload)) return variants def payload_phrase(action_id: str, payload: dict[str, Any]) -> str: if not payload: return "with empty payload" parts = [f"{key} {value}" for key, value in sorted(payload.items())] return f"{action_id} with " + ", ".join(parts) def dedupe_records(records: list[dict[str, Any]]) -> list[dict[str, Any]]: seen = set() unique = [] for record in records: key = (record["input"], record["output"]) if key in seen: continue seen.add(key) unique.append(record) return unique def run_mimir_expansion( args: argparse.Namespace, dataset_dir: Path, dataset_key: str, generated_at: int, seed_records: list[dict[str, Any]], action_ids: list[str], model_alias: str, ) -> Path: require_sdg_args(args) remote_base = "/mnt/storage/data-designer/managed-assets" remote_artifact_path = f"{remote_base}/module_controller_artifacts_{generated_at}" remote_seed = f"{remote_base}/module_controller_seed_{generated_at}.jsonl" remote_config = f"{remote_base}/module_controller_{generated_at}.py" remote_expanded = f"{remote_artifact_path}/{dataset_key}.jsonl" expanded_path = dataset_dir / "expanded.jsonl" raw_expanded_path = dataset_dir / "expanded.raw.jsonl" with tempfile.TemporaryDirectory() as tmp: tmp_path = Path(tmp) local_seed = tmp_path / "seed.jsonl" local_config = tmp_path / "module_controller.py" write_jsonl(local_seed, seed_records) local_config.write_text(data_designer_config(dataset_key, generated_at, remote_seed, action_ids, model_alias)) target = f"{args.sdg_user}@{args.sdg_host}" scp_base = ["scp", "-i", str(Path(args.sdg_key).expanduser())] ssh_base = ["ssh", "-i", str(Path(args.sdg_key).expanduser()), target] subprocess.run([*scp_base, str(local_seed), f"{target}:{remote_seed}"], check=True) subprocess.run([*scp_base, str(local_config), f"{target}:{remote_config}"], check=True) subprocess.run( [ *ssh_base, remote_bash_command( remote_data_designer_create_command( remote_config=remote_config, num_records=args.num_records, dataset_key=dataset_key, artifact_path=remote_artifact_path, ) ), ], check=True, ) subprocess.run([*scp_base, f"{target}:{remote_expanded}", str(raw_expanded_path)], check=True) normalize_expanded_jsonl(raw_expanded_path, expanded_path) return expanded_path def remote_data_designer_create_command( remote_config: str, num_records: int, dataset_key: str, artifact_path: str, ) -> str: command = [ MIMIR_DATA_DESIGNER_BIN, "create", remote_config, "--num-records", str(num_records), "--dataset-name", dataset_key, "--artifact-path", artifact_path, "--output-format", "jsonl", "--no-tui", ] quoted_command = " ".join(shlex.quote(part) for part in command) quoted_env = shlex.quote(MIMIR_DATA_DESIGNER_ENV) return f"set -a; [ ! -f {quoted_env} ] || . {quoted_env}; set +a; {quoted_command}" def remote_bash_command(command: str) -> str: return f"bash -lc {shlex.quote(command)}" def data_designer_config( dataset_key: str, generated_at: int, remote_seed: str, action_ids: list[str], model_alias: str = "nvidia-text", ) -> str: action_ids_repr = repr(action_ids) action_ids_json = json.dumps(action_ids, separators=(",", ":")) return f'''# auto-generated by generate_training_data.py - do not edit by hand # generated_at: {generated_at} from data_designer.config import DataDesignerConfigBuilder, LocalFileSeedSource ACTION_IDS = {action_ids_repr} def load_config_builder(): builder = DataDesignerConfigBuilder( model_configs="/mnt/storage/data-designer/model_configs.yaml", ) builder.with_seed_dataset(LocalFileSeedSource(path="{remote_seed}")) builder.add_column( name="expanded_record", column_type="llm-structured", model_alias="{model_alias}", prompt="""You are generating additional instruction-tuning rows for a local action dispatcher. Given this seed natural-language request and target JSON, create one new varied request that maps to the same action_id. Keep the output schema strict. Allowed action_id values: {action_ids_json} Seed request: {{{{ input }}}} Seed output JSON: {{{{ output }}}} Return a complete instruction dataset row. The output field must be a JSON string with exactly action_id and payload. The metadata.action_id must match that output action_id.""", output_format={{ "type": "object", "required": ["instruction", "input", "output", "metadata"], "additionalProperties": False, "properties": {{ "instruction": {{"type": "string", "minLength": 1}}, "input": {{"type": "string", "minLength": 1}}, "output": {{"type": "string", "minLength": 1}}, "metadata": {{ "type": "object", "required": ["action_id", "source", "split", "generated_at"], "additionalProperties": True, "properties": {{ "action_id": {{"type": "string", "enum": ACTION_IDS}}, "source": {{"type": "string"}}, "split": {{"type": "string"}}, "generated_at": {{"type": "integer"}}, }}, }}, }}, }}, ) return builder ''' def normalize_expanded_jsonl(raw_path: Path, output_path: Path) -> None: normalized: list[dict[str, Any]] = [] for row in read_jsonl(raw_path): candidate = row.get("expanded_record") if isinstance(row, dict) else None if isinstance(candidate, dict): normalized.append(candidate) elif isinstance(row, dict) and {"instruction", "input", "output", "metadata"}.issubset(row): normalized.append(row) if not normalized: raise ValueError(f"no expanded instruction rows found in {raw_path}") write_jsonl(output_path, normalized) def has_sdg_args(args: argparse.Namespace) -> bool: return bool(args.sdg_host or args.sdg_user or args.sdg_key) def require_sdg_args(args: argparse.Namespace) -> None: missing = [name for name in ["sdg_host", "sdg_user", "sdg_key"] if not getattr(args, name)] if missing: raise ValueError(f"Mimir expansion requires all SDG flags; missing {', '.join('--' + name.replace('_', '-') for name in missing)}") def patch_manifest( manifest_path: Path, dataset_key: str, dataset_dir: Path, generated_at: int, artifacts: dict[str, Any], expanded_path: Path | None, sdg_requested: bool, sdg_model_alias: str | None, sdg_error: str | None = None, ) -> None: manifest = read_json(manifest_path) datasets = manifest.setdefault("datasets", {}) if dataset_key in datasets: raise ValueError(f"manifest already contains dataset {dataset_key}") codebook = artifacts["codebook"] actions_index = artifacts["actions_index"] datasets[dataset_key] = { "name": "Module Controller Intents", "category": "synthetic_sft", "schema_profile": "instruction", "source_url": None, "acquisition_method": "local_synthetic_generator", "license_spdx": "UNLICENSED", "redistribution_allowed": False, "estimated_tokens": None, "status": "generated_local_remote_failed" if sdg_error else "generated_local", "local_path": relative_data_path(dataset_dir), "local_format": "jsonl", "generator": "mp-ai-module-controller/scripts/generate_training_data.py", "sdg_engine": "data-designer" if sdg_requested else None, "sdg_host": "mimir (100.80.52.47)" if sdg_requested else None, "sdg_model": sdg_model_alias if sdg_requested else None, "sdg_model_alias": sdg_model_alias if sdg_requested else None, "expanded_path": str(expanded_path) if expanded_path else None, "sdg_error": sdg_error, "codebook_version": codebook["dictionary_version"], "codebook_checksum": codebook["codebook_checksum"], "registry_checksum": actions_index["registry_checksum"], "generated_at": generated_at, "doc": f"{relative_data_path(dataset_dir)}/README.md", } manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=False) + "\n") def relative_data_path(dataset_dir: Path) -> str: parts = dataset_dir.parts if "training-data" in parts: index = len(parts) - 1 - list(reversed(parts)).index("training-data") return "/".join(parts[index:]) if "data" in parts: index = len(parts) - 1 - list(reversed(parts)).index("data") return "/".join(parts[index:]) return str(dataset_dir) def write_dataset_readme(path: Path, summary: dict[str, Any], artifacts: dict[str, Any]) -> None: codebook = artifacts["codebook"] path.write_text( "\n".join( [ f"# {summary['dataset_key']}", "", "Instruction-format seed dataset for module-controller action dispatch.", "", f"- Generated at: `{summary['generated_at']}`", f"- Actions: `{summary['action_count']}`", f"- Train records: `{summary['train_records']}`", f"- Validation records: `{summary['validation_records']}`", f"- Codebook version: `{codebook['dictionary_version']}`", f"- Codebook checksum: `{codebook['codebook_checksum']}`", "", ] ) ) def read_json(path: Path) -> Any: return json.loads(path.read_text()) def read_jsonl(path: Path) -> list[dict[str, Any]]: rows = [] for line in path.read_text().splitlines(): if line.strip(): rows.append(json.loads(line)) return rows def write_jsonl(path: Path, records: list[dict[str, Any]]) -> None: with path.open("w") as file: for record in records: file.write(json.dumps(record, sort_keys=True) + "\n") def invert_unique(mapping: dict[str, str], label: str) -> dict[str, str]: inverted = {} for code, value in mapping.items(): if value in inverted: raise ValueError(f"duplicate {label} value {value}") inverted[value] = code return inverted if __name__ == "__main__": raise SystemExit(main())