train/data_prep.py

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<struct<role,content>>;
# 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"]()