528 lines
19 KiB
Python
Executable File
528 lines
19 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Generate module-controller training data from dictionary artifacts."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
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
|
|
|
|
|
|
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="../workspace_Data/data/")
|
|
parser.add_argument("--manifest", default="../workspace_Data/manifest.llm.json")
|
|
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("--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),
|
|
}
|
|
|
|
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()),
|
|
)
|
|
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_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],
|
|
) -> Path:
|
|
require_sdg_args(args)
|
|
remote_base = "/mnt/storage/data-designer/managed-assets"
|
|
remote_seed = f"{remote_base}/module_controller_seed_{generated_at}.jsonl"
|
|
remote_config = f"{remote_base}/module_controller_{generated_at}.yaml"
|
|
expanded_path = dataset_dir / "expanded.jsonl"
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
tmp_path = Path(tmp)
|
|
local_seed = tmp_path / "seed.jsonl"
|
|
local_config = tmp_path / "module_controller.yaml"
|
|
write_jsonl(local_seed, seed_records)
|
|
local_config.write_text(data_designer_config(dataset_key, generated_at, remote_seed, action_ids))
|
|
|
|
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,
|
|
"/home/aaron-pressey/.venvs/data-designer/bin/data-designer",
|
|
"create",
|
|
remote_config,
|
|
"--num-records",
|
|
str(args.num_records),
|
|
"--dataset-name",
|
|
dataset_key,
|
|
],
|
|
check=True,
|
|
)
|
|
subprocess.run([*scp_base, f"{target}:{remote_base}/{dataset_key}.jsonl", str(expanded_path)], check=True)
|
|
return expanded_path
|
|
|
|
|
|
def data_designer_config(dataset_key: str, generated_at: int, remote_seed: str, action_ids: list[str]) -> str:
|
|
action_ids_json = json.dumps(action_ids)
|
|
return f"""# auto-generated by generate_training_data.py - do not edit by hand
|
|
# generated_at: {generated_at}
|
|
|
|
model_config_path: /mnt/storage/data-designer/model_configs.yaml
|
|
model_providers_path: /mnt/storage/data-designer/model_providers.yaml
|
|
|
|
dataset:
|
|
name: {dataset_key}
|
|
schema_profile: instruction
|
|
seed_file: {remote_seed}
|
|
|
|
columns:
|
|
- name: input
|
|
type: seed_passthrough
|
|
|
|
- name: output
|
|
type: llm_text
|
|
model_alias: nvidia-text
|
|
prompt: |
|
|
You are generating training data for a local action dispatcher.
|
|
Given the natural language request below, produce valid JSON with
|
|
exactly two fields: "action_id" (string) and "payload" (object).
|
|
The action_id must be one of: {action_ids_json}.
|
|
Vary the phrasing of the input naturally but keep the output schema strict.
|
|
|
|
Request: {{{{input}}}}
|
|
output_schema:
|
|
type: object
|
|
required: [action_id, payload]
|
|
additionalProperties: false
|
|
properties:
|
|
action_id:
|
|
type: string
|
|
enum: {action_ids_json}
|
|
payload:
|
|
type: object
|
|
"""
|
|
|
|
|
|
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_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": "nvidia/nemotron-3-nano-30b-a3b" 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 "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())
|