"""Prepare a bounded, reproducible text-only dataset on the cloud machine.""" import argparse from collections import Counter import hashlib import json from pathlib import Path import random import re from .runtime import CONFIG def normalize_messages(row): messages = row.get("messages") if not isinstance(messages, list) or not messages: raise ValueError("Expected a nonempty messages list.") clean = [] for index, item in enumerate(messages): role, content = item.get("role"), item.get("content") if role not in {"system", "user", "assistant"} or not isinstance(content, str) or not content.strip(): raise ValueError("Only nonempty text messages are supported.") if role == "system" and index != 0: raise ValueError("System messages must come first.") clean.append({"role": role, "content": content.strip()}) turns = clean[1:] if clean[0]["role"] == "system" else clean if len(turns) < 2 or len(turns) % 2: raise ValueError("Expected complete conversation turns.") if any(m["role"] != ("user" if i % 2 == 0 else "assistant") for i, m in enumerate(turns)): raise ValueError("Expected alternating user and assistant turns.") if clean[0]["role"] != "system": clean.insert(0, {"role": "system", "content": CONFIG["system_prompt"]}) return clean def prompt_key(messages): # Group identical opening questions, even when answers differ. first = next(m["content"] for m in messages if m["role"] == "user") normalized = re.sub(r"\s+", " ", first).strip().casefold() return hashlib.sha256(normalized.encode()).hexdigest() def encode_example(tokenizer, messages, max_length): """Supervise assistant content and end-of-turn tokens; mask every prompt.""" ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False) if len(ids) > max_length: raise ValueError("Overlength example; skip rather than truncate code or its answer.") labels = [-100] * len(ids) for index, message in enumerate(messages): if message["role"] != "assistant": continue prefix = tokenizer.apply_chat_template(messages[:index], tokenize=True, add_generation_prompt=True) end = tokenizer.apply_chat_template(messages[:index + 1], tokenize=True, add_generation_prompt=False) if ids[:len(prefix)] != prefix or ids[:len(end)] != end: raise ValueError("Chat template does not preserve assistant boundaries.") labels[len(prefix):len(end)] = ids[len(prefix):len(end)] if not any(x != -100 for x in labels[1:]): raise ValueError("No assistant target tokens.") return {"input_ids": ids, "attention_mask": [1] * len(ids), "labels": labels} def split_for(key): bucket = int(key[:8], 16) % 100 return "validation" if bucket < 2 else "test" if bucket < 4 else "train" def prepare(args): from datasets import load_dataset from transformers import AutoTokenizer output = Path(args.output) if output.exists() and any(output.iterdir()): raise FileExistsError("Use an empty data output directory.") tokenizer = AutoTokenizer.from_pretrained( CONFIG["base_model"], revision=CONFIG["base_revision"], trust_remote_code=False) result = {name: [] for name in ["train", "validation", "test"]} seen = set() stats = {} for source in CONFIG["sources"]: targets = { "train": round(args.train * source["weight"]), "validation": args.holdout_per_source, "test": args.holdout_per_source, } counts, skipped = Counter(), Counter() stream = load_dataset( CONFIG["dataset"], source["config"], split="train", revision=CONFIG["dataset_revision"], streaming=True) stream = stream.shuffle(seed=CONFIG["seed"], buffer_size=4096) for row in stream: try: messages = normalize_messages(row) key = prompt_key(messages) if key in seen: skipped["duplicate_prompt"] += 1 continue encode_example(tokenizer, messages, CONFIG["max_sequence_length"]) except (ValueError, TypeError, KeyError): skipped["invalid_or_overlength"] += 1 continue seen.add(key) split = split_for(key) if counts[split] >= targets[split]: continue content_hash = hashlib.sha256(json.dumps(messages, ensure_ascii=False, sort_keys=True).encode()).hexdigest() result[split].append({"id": content_hash, "prompt_hash": key, "source": source["config"], "messages": messages}) counts[split] += 1 if all(counts[s] >= targets[s] for s in targets): break if any(counts[s] < targets[s] for s in targets): raise RuntimeError(f"Not enough usable data for {source['config']}: {dict(counts)}") stats[source["config"]] = {"selected": dict(counts), "skipped": dict(skipped)} print(source["config"], dict(counts), flush=True) output.mkdir(parents=True, exist_ok=True) hashes = {} for split, rows in result.items(): random.Random(CONFIG["seed"]).shuffle(rows) content = "".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows) (output / f"{split}.jsonl").write_bytes(content.encode("utf-8")) hashes[split] = hashlib.sha256(content.encode()).hexdigest() manifest = {"config": CONFIG, "counts": {k: len(v) for k, v in result.items()}, "sources": stats, "sha256": hashes, "split_policy": "Hash of normalized first user question; official source test splits unused.", "limitations": ["Synthetic data may contain mistakes.", "Exact prompt deduplication is not semantic or repository-level decontamination.", "The pretrained base may already have seen these data.", "This pilot's code source focuses on Python; other languages need additional data and evaluation."]} (output / "manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8") if __name__ == "__main__": parser = argparse.ArgumentParser(description="Prepare Chat & Code data on the cloud machine") parser.add_argument("--output", default="chatcode-data") parser.add_argument("--train", type=int, default=CONFIG["train_examples"]) parser.add_argument("--holdout-per-source", type=int, default=128) args = parser.parse_args() if args.train < 10 or args.train % 10 or args.holdout_per_source < 1: parser.error("--train must be a positive multiple of 10; holdouts must be positive.") prepare(args)