"""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="