training/prepare_home.py

"""Prepare the audited home-command snapshot without loading millions of rows into RAM."""
from collections import Counter, defaultdict
import hashlib
import json
import math
from pathlib import Path
import re
import sqlite3
import time

import numpy as np
from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders
from transformers import PreTrainedTokenizerFast

from tiny_json_llm.data import SPECIAL, SEED, canonical, sha256, write_json, choose_architecture

OUT = Path("prepared/home_commands_v1")
PREFIXES = sorted(["can you please", "could you please", "would you please", "can you", "could you",
    "would you", "go ahead and", "ek kaam karo", "please", "kindly", "quickly", "hey", "hi", "listen",
    "boss", "yaar", "suno", "sun", "acha", "achha", "accha", "arre", "are", "bhai", "jara", "zara"], key=len, reverse=True)
SUFFIXES = sorted(["right away", "right now", "for me", "please", "na", "yaar", "abhi", "zara", "jaldi",
    "turant", "fatafat", "now", "immediately", "asap", "quickly"], key=len, reverse=True)


def phrase_group(key):
    text = re.sub(r"[^\w\s]", " ", key)
    text = re.sub(r"\s+", " ", text).strip()
    for _ in range(12):
        before = text
        for prefix in PREFIXES:
            if text.startswith(prefix + " "):
                text = text[len(prefix):].strip()
                break
        for suffix in SUFFIXES:
            if text.endswith(" " + suffix):
                text = text[:-len(suffix)].strip()
                break
        if before == text:
            break
    text = re.sub(r"\bthe\b", "", text)
    return re.sub(r"\s+", " ", text).strip()


def main():
    started = time.time()
    audit = json.loads(Path("data/source_audit.json").read_text(encoding="utf-8"))
    if audit["counts"].get("invalid_rows", 0) or audit["counts"].get("conflicting_rows", 0):
        raise ValueError("Source audit has invalid or conflicting records; review before preparing")
    if OUT.exists() and any(OUT.iterdir()):
        raise ValueError(f"Snapshot already exists: {OUT}")
    OUT.mkdir(parents=True, exist_ok=True)
    db = sqlite3.connect("data/audit.sqlite", check_same_thread=False)
    db.execute("CREATE TABLE IF NOT EXISTS split_examples (key TEXT PRIMARY KEY, group_key TEXT, split TEXT)")
    groups, group_targets = Counter(), {}
    pending = []
    for key, target in db.execute("SELECT key,output FROM examples"):
        group = phrase_group(key)
        if not group:
            raise ValueError("Empty phrase group")
        if group in group_targets and group_targets[group] != target:
            raise ValueError(f"Phrase grouping produced conflicting targets: {group}")
        group_targets[group] = target
        groups[group] += 1
    # Stratified by output class, grouped by core wording; wrappers never cross splits.
    by_target = defaultdict(list)
    for group, target in group_targets.items():
        by_target[target].append(group)
    assignments = {}
    for target, group_list in sorted(by_target.items()):
        ordered = sorted(group_list, key=lambda value: hashlib.sha256(f"{SEED}:{value}".encode()).hexdigest())
        if len(ordered) < 10:
            raise ValueError(f"Fewer than 10 phrase families for class {target}; need a custom split")
        n = max(1, round(len(ordered) * .1))
        for i, group in enumerate(ordered):
            assignments[group] = "validation" if i < n else "test" if i < 2*n else "train"
    for key, in db.execute("SELECT key FROM examples"):
        group = phrase_group(key)
        pending.append((key, group, assignments[group]))
        if len(pending) >= 10000:
            db.executemany("INSERT OR IGNORE INTO split_examples VALUES (?,?,?)", pending)
            pending.clear()
    if pending:
        db.executemany("INSERT OR IGNORE INTO split_examples VALUES (?,?,?)", pending)
    db.commit()
    db.execute("CREATE INDEX IF NOT EXISTS split_index ON split_examples(split)")
    db.commit()
    print(json.dumps({"phrase_groups": len(groups), "classes": len(by_target),
                      "groups_per_split": dict(Counter(assignments.values()))}), flush=True)
    query = "SELECT e.instruction,e.output,e.line,s.group_key FROM examples e JOIN split_examples s ON e.key=s.key WHERE s.split=?"
    backend = Tokenizer(models.BPE())
    backend.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    backend.decoder = decoders.ByteLevel()
    trainer = trainers.BpeTrainer(vocab_size=1024, min_frequency=2, special_tokens=SPECIAL,
               initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), show_progress=False)
    def text_batches():
        cursor = db.execute(query, ("train",))
        while batch := cursor.fetchmany(10000):
            yield [text for instruction, target, _, _ in batch for text in (instruction, target)]
    backend.train_from_iterator(text_batches(), trainer=trainer)
    tokenizer = PreTrainedTokenizerFast(tokenizer_object=backend, bos_token=SPECIAL[1],
        eos_token=SPECIAL[2], pad_token=SPECIAL[0], additional_special_tokens=SPECIAL[3:],
        clean_up_tokenization_spaces=False, model_max_length=256)
    tokenizer.save_pretrained(OUT / "tokenizer")
    stats, max_length = {}, 0
    target_ids = {target: tokenizer.encode(target, add_special_tokens=False) + [tokenizer.eos_token_id]
                  for target in by_target}
    prefix = [tokenizer.bos_token_id, tokenizer.convert_tokens_to_ids("<|user|>")]
    middle = [tokenizer.convert_tokens_to_ids("<|turn|>"), tokenizer.convert_tokens_to_ids("<|assistant|>")]
    for split in ("train", "validation", "test"):
        cursor = db.execute(query, (split,))
        indices, position, predictions, responses, lengths, class_counts = [], 0, 0, 0, Counter(), Counter()
        with (OUT / f"{split}.ids.bin").open("wb") as ids_file, (OUT / f"{split}.jsonl").open("w", encoding="utf-8", newline="\n") as raw_file:
            while batch := cursor.fetchmany(10000):
                encoded = backend.encode_batch([r[0] for r in batch], add_special_tokens=False)
                blocks = []
                for (instruction, target, source_line, group), enc in zip(batch, encoded):
                    prompt = prefix + enc.ids + middle
                    ids = prompt + target_ids[target]
                    indices.append((position, len(ids), len(prompt)))
                    blocks.extend(ids)
                    position += len(ids)
                    predictions += len(ids) - 1
                    responses += len(target_ids[target])
                    lengths[len(ids)] += 1
                    class_counts[target] += 1
                    raw_file.write(canonical({"instruction": instruction, "output": json.loads(target),
                                             "group_id": group, "source_line": source_line}) + "\n")
                np.asarray(blocks, dtype="<u2").tofile(ids_file)
        np.save(OUT / f"{split}.index.npy", np.asarray(indices, dtype=np.uint64))
        max_length = max(max_length, max(lengths))
        stats[split] = {"rows": len(indices), "groups": sum(v == split for v in assignments.values()),
            "prediction_tokens": predictions, "sequence_tokens": position, "response_tokens": responses,
            "max_sequence_length": max(lengths), "max_response_tokens": max(map(len, target_ids.values())),
            "class_counts": dict(class_counts), "sha256": sha256(OUT / f"{split}.jsonl")}
        print(json.dumps({"prepared": split, "rows": len(indices), "tokens": predictions}), flush=True)
    architecture = choose_architecture(stats["train"]["prediction_tokens"], len(tokenizer))
    context = 2 ** math.ceil(math.log2(max(128, max_length + 32)))
    config = {k: architecture[k] for k in ("hidden_size", "intermediate_size", "num_hidden_layers", "num_attention_heads", "num_key_value_heads")}
    config.update(vocab_size=len(tokenizer), max_position_embeddings=context, tie_word_embeddings=True,
                  bos_token_id=tokenizer.bos_token_id, eos_token_id=tokenizer.eos_token_id,
                  pad_token_id=tokenizer.pad_token_id, attention_bias=False, mlp_bias=False,
                  attention_dropout=0., rms_norm_eps=1e-5)
    write_json(OUT / "model_config.json", config)
    allowed = sorted({target for group, target in group_targets.items() if assignments[group] == "train"})
    write_json(OUT / "allowed_outputs.json", [json.loads(value) for value in allowed])
    fields = {key: sorted({value[key] for value in map(json.loads, allowed)}) for key in ("activity", "subject", "action")}
    schema = {"type": "object", "properties": {k: {"type": "string", "enum": v} for k,v in fields.items()},
              "required": ["activity", "subject", "action"], "additionalProperties": False}
    write_json(OUT / "output_schema.json", schema)
    write_json(OUT / "phrase_groups.json", [{"core_phrase": key, "split": assignments[key], "rows": groups[key]} for key in sorted(groups)])
    hashes = {str(path.relative_to(OUT)).replace("\\", "/"): sha256(path) for path in sorted(OUT.rglob("*")) if path.is_file()}
    manifest = {"format_version": 2, "seed": SEED, "sources": [{k:audit[k] for k in ("source", "sha256", "bytes")}],
        "source_rows": audit["counts"]["rows"], "duplicates_removed": audit["counts"]["normalized_duplicate_rows"],
        "splits": stats, "architecture": architecture, "context_length": context,
        "tokenizer": {"type": "byte-level BPE", "vocabulary_size": len(tokenizer), "trained_on": "train only"},
        "objective": "causal next-token prediction on full input+JSON sequences; padding and first BOS excluded",
        "sizing_tokens": "one pass of nonpadding train prediction tokens, no epochs or held-out tokens added",
        "split_method": "80/10/10 phrase families stratified by JSON class; exact normalized input duplicates removed; casing, punctuation, common politeness prefixes/suffixes and the article 'the' grouped",
        "scaling_caveat": "D ~= 20N is a sizing heuristic, extrapolated below Chinchilla's 70M+ range. This corpus is highly templated; token count overstates independent linguistic information. No proof of compute optimality.",
        "supported_languages": ["English", "Romanized Hindi/Hinglish"],
        "limits": ["Closed set of observed training JSON commands", "No unknown/reject intent or negation guarantees", "No native-script Hindi claim", "Not a general chat model", "Automated grouping may miss other paraphrase families"],
        "data_rights": "Original user contributions: Copyright (c) 2026 sraivante, Apache-2.0. Third-party terms unchanged; source ownership/provenance not yet verified. Private snapshot, no publication authorized.",
        "artifact_sha256": hashes, "preparation_seconds": round(time.time()-started,2)}
    write_json(OUT / "manifest.json", manifest)
    print(json.dumps({"architecture": architecture, "context": context, "vocabulary": len(tokenizer), "seconds": manifest["preparation_seconds"]}, indent=2), flush=True)
    db.close()


if __name__ == "__main__":
    main()