import json import os import pathlib import re import sys from collections import Counter from enum import StrEnum, auto import yaml from dotenv import load_dotenv load_dotenv() from datasets import ( # noqa: E402 Dataset, DatasetDict, Features, Value, concatenate_datasets, load_dataset, ) # Pin the output schema so messages serialize as a clean list>; # without this, datasets carries the source arrow.json extension and writes each # message as a JSON string (forcing json.loads downstream). _SFT_FEATURES = Features({"messages": [{"role": Value("string"), "content": Value("string")}]}) class DataSplit(StrEnum): TRAIN = auto() TEST = auto() _SOURCE_REPO = "acon96/Home-Assistant-Requests-V2" _ENGLISH_DATA = { DataSplit.TRAIN: "home_assistant_train_english.jsonl", DataSplit.TEST: "home_assistant_test_english.jsonl", } _REPO_ROOT = pathlib.Path(__file__).resolve().parent.parent _ACORN96_CACHE_DIR = _REPO_ROOT / "data" / "acorn96_v2" _PROCESSED_DIR = _REPO_ROOT / "data" / "processed" _HA_ACTIONS = "ha_actions" _HA_ACTIONS_DIR = _PROCESSED_DIR / _HA_ACTIONS _TASKGEN_SFT = _PROCESSED_DIR / "taskgen" / "sft.parquet" _ROUTING = "routing" _ROUTING_DIR = _PROCESSED_DIR / _ROUTING _REPLAY_N = 2000 # ha_actions train rows mixed in for anti-forgetting (run-1's 300 forgot badly) _TEST_N = 1000 # total test rows (trainer's eval_ds selects 1000) _EVAL_AUTO_N = 100 # route_automation rows HELD OUT into test = new-skill eval gold _SEED = 3407 # match train_sft.py seed _CFG = yaml.safe_load((pathlib.Path(__file__).parent / "config.yaml").read_text()) _TIME = re.compile(r"The current time and date is (.+)") _DEVICES = re.compile(r"Devices:\n(.*?)(?:\nUser instruction:|\Z)", re.S) def pull_acon96() -> DatasetDict: return load_dataset( _SOURCE_REPO, data_files=_ENGLISH_DATA, cache_dir=str(_ACORN96_CACHE_DIR), ) def _add_intent(batch: dict) -> dict: return { "intent": [ "action" if any(turn.get("tool_calls") for turn in messages) else "query" for messages in batch["messages"] ] } def label_split(ds: DatasetDict, split: DataSplit) -> Dataset: return ds[split].map(_add_intent, batched=True) def _to_sft(row: dict) -> dict: msgs = row["messages"] original = msgs[0]["content"][0]["text"] system = ( f"{_CFG['system_prompt']}\n\nActions:\n{_CFG['intents']}" f"\n\nDevices:\n{_DEVICES.search(original).group(1).rstrip()}" f"\nCurrent time: {_TIME.search(original).group(1).strip()}" ) assistant = next(t for t in msgs if t["role"] == "assistant") if assistant.get("tool_calls"): payload = { "intents": [ {"name": c["function"]["name"], "slots": json.loads(c["function"]["arguments"])} for c in assistant["tool_calls"] ] } else: payload = {"response": assistant["content"][0]["text"]} return { "messages": [ {"role": "system", "content": system}, {"role": "user", "content": msgs[1]["content"][0]["text"]}, {"role": "assistant", "content": json.dumps(payload, ensure_ascii=False, separators=(",", ":"))}, ] } def label() -> None: _PROCESSED_DIR.mkdir(parents=True, exist_ok=True) raw_data = None for split in DataSplit: out = _PROCESSED_DIR / f"{split}.parquet" if out.exists(): print(f"[label] {split}: already labelled -> {out}") continue raw_data = raw_data or pull_acon96() labelled = label_split(raw_data, split) labelled.to_parquet(str(out)) print(f"[label] {split}: {dict(Counter(labelled['intent']))} -> {out}") def reform() -> None: _HA_ACTIONS_DIR.mkdir(parents=True, exist_ok=True) for split in DataSplit: ds = load_dataset("parquet", data_files=str(_PROCESSED_DIR / f"{split}.parquet"), split="train") sft = ds.map(_to_sft, remove_columns=ds.column_names, features=_SFT_FEATURES) out = _HA_ACTIONS_DIR / f"{split}.parquet" sft.to_parquet(str(out)) print(f"[reform] {split}: {len(sft)} rows -> {out}") def push() -> None: repo = f"{os.environ['HF_ORG']}/{_CFG['hf_dataset_name']}" dsd = DatasetDict( { str(s): load_dataset("parquet", data_files=str(_HA_ACTIONS_DIR / f"{s}.parquet"), split="train") for s in DataSplit } ) dsd.push_to_hub(repo, config_name=_HA_ACTIONS) print(f"[push] {repo} ({_HA_ACTIONS}): {dict((k, len(v)) for k, v in dsd.items())}") def routing() -> None: """Assemble the SFT-2 'routing' dataset, schema-pinned to match ha_actions. train = taskgen (minus held-out automation) + ha_actions replay (anti-forgetting). test = held-out route_automation gold (new-skill eval) + ha_actions (retention) -> eval.py scores BOTH the automation route AND forgetting from one set.""" _ROUTING_DIR.mkdir(parents=True, exist_ok=True) clean = lambda ds: ds.remove_columns( [c for c in ("category", "home") if c in ds.column_names] ).cast(_SFT_FEATURES) tg = load_dataset("parquet", data_files=str(_TASKGEN_SFT), split="train").shuffle(seed=_SEED) auto = tg.filter(lambda r: r["category"] == "route_automation") rest = tg.filter(lambda r: r["category"] != "route_automation") eval_auto = auto.select(range(_EVAL_AUTO_N)) # never seen in training taskgen_train = clean(concatenate_datasets([auto.select(range(_EVAL_AUTO_N, len(auto))), rest])) ha_train = load_dataset("parquet", data_files=str(_HA_ACTIONS_DIR / "train.parquet"), split="train") replay = clean(ha_train.shuffle(seed=_SEED).select(range(_REPLAY_N))) train = concatenate_datasets([taskgen_train, replay]).shuffle(seed=_SEED) train.to_parquet(str(_ROUTING_DIR / "train.parquet")) print(f"[routing] train: {len(taskgen_train)} taskgen + {len(replay)} replay = {len(train)} rows") ha_test = load_dataset("parquet", data_files=str(_HA_ACTIONS_DIR / "test.parquet"), split="train") ha_test = ha_test.shuffle(seed=_SEED).select(range(_TEST_N - _EVAL_AUTO_N)) test = concatenate_datasets([clean(eval_auto), ha_test.cast(_SFT_FEATURES)]).shuffle(seed=_SEED) test.to_parquet(str(_ROUTING_DIR / "test.parquet")) print(f"[routing] test: {len(eval_auto)} held-out automation + {len(ha_test)} ha_actions = {len(test)} rows") _CHATML = ( "{% for m in messages %}" "{{ '<|im_start|>' + m['role'] + '\n' + m['content'] + '<|im_end|>\n' }}" "{% endfor %}" ) def tokens() -> None: from transformers import AutoTokenizer tok = AutoTokenizer.from_pretrained(_CFG["base_models"]["brain1"], trust_remote_code=True) tok.chat_template = _CHATML overall = 0 for split in DataSplit: ds = load_dataset("parquet", data_files=str(_HA_ACTIONS_DIR / f"{split}.parquet"), split="train") # MiniCPM's custom tokenizer mis-tokenizes apply_chat_template(tokenize=True); # render to a string first, then tokenize that. lens = sorted( len(tok(tok.apply_chat_template(r["messages"], tokenize=False))["input_ids"]) for r in ds ) n = len(lens) pct = {p: lens[min(n - 1, p * n // 100)] for p in (50, 90, 99)} overall = max(overall, lens[-1]) print(f"[tokens] {split}: n={n} p50={pct[50]} p90={pct[90]} p99={pct[99]} max={lens[-1]}") print(f"[tokens] max over all splits = {overall} (set MAX_SEQ >= this)") def verify(n: int = 5) -> None: for split in DataSplit: ds = load_dataset("parquet", data_files=str(_HA_ACTIONS_DIR / f"{split}.parquet"), split="train") rows = [(r["messages"][1]["content"], r["messages"][2]["content"]) for r in ds] intents = [r for r in rows if '"intents"' in r[1]] responses = [r for r in rows if '"intents"' not in r[1]] print(f"\n[{split}] {len(ds)} rows | intents: {len(intents)} response: {len(responses)}") for label, group in (("intents", intents), ("response", responses)): for user, completion in group[:: max(1, len(group) // n)][:n]: print(f" [{label}]") print(f" in : {user!r}") print(f" out: {completion}") if __name__ == "__main__": { "label": label, "reform": reform, "push": push, "verify": verify, "tokens": tokens, "routing": routing, }[sys.argv[1] if len(sys.argv) > 1 else "label"]()