| |
| """Build a local random micro model + tokenizer for CPU smoke tests. |
| |
| Nothing is downloaded: the tokenizer is trained on the project's own fixture and |
| contract text, and the causal LM is randomly initialised. The result validates |
| the training/saving/routing path only; it has no semantic ability and its |
| metrics must never be reported as model quality. |
| |
| Usage:: |
| |
| uv run --no-sync --extra train python scripts/make_tiny_model.py --output models/tiny-qwen-smoke |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
| import sys |
| from typing import Any, Sequence |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT / "src")) |
|
|
| from vbench_prompts_compile import records as R |
| from vbench_prompts_compile.training import TASK_INSTRUCTIONS |
|
|
| SPECIAL_TOKENS = ["<unk>", "<pad>", "<s>", "</s>", "<|im_start|>", "<|im_end|>"] |
| CHAT_TEMPLATE = ( |
| "{% for message in messages %}" |
| "{{ '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n' }}" |
| "{% endfor %}" |
| "{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}" |
| ) |
|
|
|
|
| def corpus_lines() -> list[str]: |
| lines: list[str] = list(TASK_INSTRUCTIONS.values()) |
| lines += ['{"relationships":[{"subject":"cat","relation":"left","object":"dog"}]}', |
| '{"relationships":[]}', '{"actions":["riding a bike"]}', '{"actions":[]}', |
| '{"entities":["cat","dog","bicycle"]}', 'supported', 'contradicted', 'insufficient'] |
| fixtures = ROOT / "tests" / "fixtures" |
| for path in sorted(fixtures.glob("*.jsonl")): |
| for line in path.read_text(encoding="utf-8").splitlines(): |
| if not line.strip(): |
| continue |
| row = json.loads(line) |
| if "prompt" in row: |
| lines.append(row["prompt"]) |
| if "caption" in row: |
| lines.append(row["caption"]) |
| if "target" in row: |
| lines.append(row["target"] if isinstance(row["target"], str) else R.canonical_json(row["target"])) |
| small = ROOT / "data" / "smoke" |
| if small.exists(): |
| for path in sorted(small.glob("*/*.jsonl")): |
| for line in path.read_text(encoding="utf-8").splitlines(): |
| if not line.strip(): |
| continue |
| row = json.loads(line) |
| lines.append(row["input"]["prompt"]) |
| if "caption" in row["input"]: |
| lines.append(row["input"]["caption"]) |
| lines.append(row["target"] if isinstance(row["target"], str) else R.canonical_json(row["target"])) |
| return lines |
|
|
|
|
| def build(output: Path, *, vocab_size: int, hidden_size: int, layers: int, heads: int, seed: int) -> dict[str, Any]: |
| from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers |
| from transformers import PreTrainedTokenizerFast, Qwen2Config, Qwen2ForCausalLM |
|
|
| tokenizer = Tokenizer(models.BPE(unk_token="<unk>")) |
| tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) |
| tokenizer.decoder = decoders.ByteLevel() |
| trainer = trainers.BpeTrainer(vocab_size=vocab_size, special_tokens=SPECIAL_TOKENS, initial_alphabet=pre_tokenizers.ByteLevel.alphabet()) |
| lines = corpus_lines() |
| tokenizer.train_from_iterator(lines, trainer=trainer, length=len(lines)) |
|
|
| output.mkdir(parents=True, exist_ok=True) |
| fast = PreTrainedTokenizerFast( |
| tokenizer_object=tokenizer, |
| unk_token="<unk>", |
| pad_token="<pad>", |
| bos_token="<s>", |
| eos_token="</s>", |
| model_max_length=1024, |
| model_input_names=["input_ids", "attention_mask"], |
| ) |
| fast.chat_template = CHAT_TEMPLATE |
| fast.save_pretrained(output) |
|
|
| config = Qwen2Config( |
| vocab_size=len(tokenizer.get_vocab()), |
| hidden_size=hidden_size, |
| intermediate_size=hidden_size * 2, |
| num_hidden_layers=layers, |
| num_attention_heads=heads, |
| num_key_value_heads=max(1, heads // 2), |
| max_position_embeddings=1024, |
| tie_word_embeddings=True, |
| bos_token_id=fast.bos_token_id, |
| eos_token_id=fast.eos_token_id, |
| pad_token_id=fast.pad_token_id, |
| ) |
| model = Qwen2ForCausalLM(config) |
| model.save_pretrained(output) |
|
|
| info = { |
| "output": str(output), |
| "vocab_size": len(tokenizer.get_vocab()), |
| "corpus_lines": len(lines), |
| "hidden_size": hidden_size, |
| "layers": layers, |
| "parameters": sum(p.numel() for p in model.parameters()), |
| "note": "randomly initialised micro model for engineering smoke only; no semantic ability, no downloads", |
| } |
| (output / "tiny_model_info.json").write_text(json.dumps(info, indent=2) + "\n", encoding="utf-8") |
| return info |
|
|
|
|
| def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| parser.add_argument("--output", default=str(ROOT / "models" / "tiny-qwen-smoke")) |
| parser.add_argument("--vocab-size", type=int, default=2048) |
| parser.add_argument("--hidden-size", type=int, default=64) |
| parser.add_argument("--layers", type=int, default=2) |
| parser.add_argument("--heads", type=int, default=4) |
| parser.add_argument("--seed", type=int, default=20260919) |
| return parser.parse_args(argv) |
|
|
|
|
| def main(argv: Sequence[str] | None = None) -> int: |
| args = parse_args(argv) |
| import random |
|
|
| import torch |
|
|
| random.seed(args.seed) |
| torch.manual_seed(args.seed) |
| info = build(Path(args.output), vocab_size=args.vocab_size, hidden_size=args.hidden_size, layers=args.layers, heads=args.heads, seed=args.seed) |
| print(json.dumps(info, ensure_ascii=False, indent=2)) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
