| |
| """Prepare compact, assistant-only code SFT binaries for a roughly 2B model.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| from pathlib import Path |
| from typing import Any, Iterable |
|
|
| import numpy as np |
| from tokenizers import Tokenizer |
|
|
|
|
| def arguments() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description=__doc__) |
| source = parser.add_mutually_exclusive_group(required=True) |
| source.add_argument("--input", type=Path, help="Local JSON or JSONL records") |
| source.add_argument("--dataset", help="Hugging Face dataset ID, for example HuggingFaceH4/CodeAlpaca_20K") |
| parser.add_argument("--dataset-config") |
| parser.add_argument("--split", default="train") |
| parser.add_argument("--tokenizer", type=Path, required=True) |
| parser.add_argument("--out", type=Path, required=True) |
| parser.add_argument("--max-records", type=int, default=30000) |
| parser.add_argument("--max-length", type=int, default=1024) |
| parser.add_argument("--max-answer-tokens", type=int, default=768) |
| parser.add_argument("--validation-fraction", type=float, default=0.02) |
| return parser.parse_args() |
|
|
|
|
| def records(args: argparse.Namespace) -> Iterable[dict[str, Any]]: |
| if args.input: |
| text = args.input.read_text(encoding="utf-8") |
| value = json.loads(text) if args.input.suffix == ".json" else [json.loads(line) for line in text.splitlines() if line.strip()] |
| if isinstance(value, dict): |
| value = value.get("data", value.get("records")) |
| if not isinstance(value, list): |
| raise ValueError("local input must be a JSON array/JSONL, or contain data/records array") |
| yield from value |
| return |
| try: |
| from datasets import load_dataset |
| except ImportError as exc: |
| raise SystemExit("Hugging Face input requires: pip install datasets") from exc |
| dataset = load_dataset(args.dataset, args.dataset_config, split=args.split, streaming=True) |
| yield from dataset |
|
|
|
|
| def instruction_answer(row: dict[str, Any]) -> tuple[str, str] | None: |
| instruction = row.get("instruction") or row.get("prompt") or row.get("question") |
| answer = row.get("output") or row.get("completion") or row.get("response") or row.get("answer") |
| if not isinstance(instruction, str) or not isinstance(answer, str): |
| messages = row.get("messages") |
| if isinstance(messages, list): |
| users = [m.get("content") for m in messages if m.get("role") == "user"] |
| assistants = [m.get("content") for m in messages if m.get("role") == "assistant"] |
| if users and assistants: |
| instruction, answer = users[-1], assistants[-1] |
| if not isinstance(instruction, str) or not isinstance(answer, str): |
| return None |
| extra = row.get("input") |
| prompt = instruction.strip() |
| if isinstance(extra, str) and extra.strip(): |
| prompt += "\n\nInput:\n" + extra.strip() |
| answer = answer.strip() |
| if len(prompt) < 8 or len(answer) < 8 or "\x00" in prompt or "\x00" in answer: |
| return None |
| return prompt, answer |
|
|
|
|
| def write_split(path: Path, examples: list[tuple[np.ndarray, np.ndarray]]) -> dict[str, int]: |
| path.mkdir(parents=True, exist_ok=True) |
| lengths = np.asarray([len(tokens) for tokens, _ in examples], dtype="<i4") |
| tokens = np.concatenate([item[0] for item in examples]) if examples else np.asarray([], dtype="<u2") |
| masks = np.concatenate([item[1] for item in examples]) if examples else np.asarray([], dtype="u1") |
| tokens.astype("<u2", copy=False).tofile(path / "sft_tokens.u16") |
| masks.tofile(path / "sft_mask.u8") |
| lengths.tofile(path / "sft_lens.i32") |
| return {"records": len(examples), "tokens": int(tokens.size), "assistant_tokens": int(masks.sum())} |
|
|
|
|
| def main() -> None: |
| args = arguments() |
| if not 0 < args.validation_fraction < 0.5: |
| raise SystemExit("--validation-fraction must be between 0 and 0.5") |
| tokenizer = Tokenizer.from_file(str(args.tokenizer)) |
| if tokenizer.get_vocab_size() > 65536: |
| raise SystemExit("uint16 output requires tokenizer vocabulary <= 65,536") |
| special = {name: tokenizer.token_to_id(name) for name in ("<s>", "</s>", "<|user|>", "<|assistant|>")} |
| missing = [name for name, token_id in special.items() if token_id is None] |
| if missing: |
| raise SystemExit("tokenizer lacks required tokens: " + ", ".join(missing)) |
|
|
| train: list[tuple[np.ndarray, np.ndarray]] = [] |
| validation: list[tuple[np.ndarray, np.ndarray]] = [] |
| seen: set[str] = set() |
| rejected = 0 |
| for row in records(args): |
| pair = instruction_answer(row) |
| if pair is None: |
| rejected += 1 |
| continue |
| prompt, answer = pair |
| digest = hashlib.sha256((prompt + "\0" + answer).encode()).hexdigest() |
| if digest in seen: |
| rejected += 1 |
| continue |
| answer_ids = tokenizer.encode("\n" + answer).ids[: args.max_answer_tokens] |
| prefix = [special["<s>"], special["<|user|>"]] + tokenizer.encode("\n" + prompt).ids + [special["</s>"], special["<|assistant|>"]] |
| room = args.max_length - len(prefix) - 1 |
| if room < 16 or len(answer_ids) < 4: |
| rejected += 1 |
| continue |
| answer_ids = answer_ids[:room] |
| sequence = prefix + answer_ids + [special["</s>"]] |
| mask = [0] * len(prefix) + [1] * (len(answer_ids) + 1) |
| item = (np.asarray(sequence, dtype="<u2"), np.asarray(mask, dtype="u1")) |
| bucket = int(digest[:8], 16) / 0xFFFFFFFF |
| (validation if bucket < args.validation_fraction else train).append(item) |
| seen.add(digest) |
| if len(train) + len(validation) >= args.max_records: |
| break |
|
|
| if not train or not validation: |
| raise SystemExit("dataset filtering produced an empty train or validation split") |
| summary = { |
| "format": "v25-assistant-only-u16-v1", |
| "source": str(args.input) if args.input else args.dataset, |
| "dataset_config": args.dataset_config, |
| "source_split": args.split, |
| "tokenizer": str(args.tokenizer), |
| "max_length": args.max_length, |
| "max_answer_tokens": args.max_answer_tokens, |
| "split_method": "sha256(prompt + NUL + answer)", |
| "validation_fraction": args.validation_fraction, |
| "train": write_split(args.out / "train", train), |
| "validation": write_split(args.out / "validation", validation), |
| "rejected_or_duplicate_records": rejected, |
| } |
| (args.out / "meta.json").write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8") |
| print(json.dumps(summary, sort_keys=True)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
