File size: 5,878 Bytes
fa82a9d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | #!/usr/bin/env python3
"""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 # noqa: E402
from vbench_prompts_compile.training import TASK_INSTRUCTIONS # noqa: E402
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())
|