| """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): |
| |
| 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) |
|
|
