| """Small offline checks; no pretrained weights, dataset downloads, or GPU jobs.""" |
| import hashlib |
| import json |
| from pathlib import Path |
| import tempfile |
| import unittest |
| from unittest.mock import patch |
|
|
| from .data import encode_example, normalize_messages, prompt_key |
| from .runtime import CONFIG, conversation, load_engine |
| from .train import collate, read_data |
|
|
|
|
| class PipelineTests(unittest.TestCase): |
| @classmethod |
| def setUpClass(cls): |
| from tokenizers import Tokenizer |
| from tokenizers.models import WordLevel |
| from tokenizers.pre_tokenizers import Whitespace |
| from transformers import PreTrainedTokenizerFast |
| words = ["[UNK]", "[PAD]", "<|im_start|>", "<|im_end|>", |
| "system", "user", "assistant", "first", "second", "answer", "other"] |
| tokenizer = Tokenizer(WordLevel({word: i for i, word in enumerate(words)}, unk_token="[UNK]")) |
| tokenizer.pre_tokenizer = Whitespace() |
| cls.tokenizer = PreTrainedTokenizerFast( |
| tokenizer_object=tokenizer, unk_token="[UNK]", pad_token="[PAD]", eos_token="<|im_end|>", |
| additional_special_tokens=["<|im_start|>", "<|im_end|>"]) |
| cls.tokenizer.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 test_multiturn_assistant_targets_and_padding(self): |
| messages = normalize_messages({"messages": [ |
| {"role": "user", "content": "first"}, {"role": "assistant", "content": "answer"}, |
| {"role": "user", "content": "second"}, {"role": "assistant", "content": "other"}]}) |
| encoded = encode_example(self.tokenizer, messages, 2048) |
| targets = [x for x in encoded["labels"] if x != -100] |
| self.assertEqual(self.tokenizer.decode(targets), "answer <|im_end|> other <|im_end|>") |
| short = encode_example(self.tokenizer, messages[:3], 2048) |
| batch = collate([encoded, short], self.tokenizer.pad_token_id) |
| self.assertTrue((batch["labels"][1, len(short["input_ids"]):] == -100).all()) |
| self.assertTrue((batch["attention_mask"][1, len(short["input_ids"]):] == 0).all()) |
| with self.assertRaises(ValueError): |
| encode_example(self.tokenizer, messages, 4) |
|
|
| def test_no_image_records_and_no_partial_turns(self): |
| for messages in [ |
| [{"role": "user", "content": [{"type": "image"}]}], |
| [{"role": "user", "content": "first"}], |
| [{"role": "assistant", "content": "answer"}, {"role": "user", "content": "first"}], |
| ]: |
| with self.assertRaises(ValueError): |
| normalize_messages({"messages": messages}) |
| self.assertEqual(conversation("first")[-1], {"role": "user", "content": "first"}) |
| with self.assertRaises(ValueError): |
| conversation("first", [{"role": "user", "content": "second"}]) |
|
|
| def test_cpu_training_and_inference_fail_before_download(self): |
| from argparse import Namespace |
| from .train import train |
| with patch("torch.cuda.is_available", return_value=False): |
| with self.assertRaisesRegex(RuntimeError, "cloud model"): |
| load_engine() |
| with self.assertRaisesRegex(RuntimeError, "No training or model download"): |
| train(Namespace()) |
|
|
| def test_manifest_integrity_and_prompt_separation(self): |
| with tempfile.TemporaryDirectory() as temp: |
| folder = Path(temp) |
| manifest = {"config": CONFIG, "sha256": {}, "counts": {}} |
| for split in ["train", "validation", "test"]: |
| messages = normalize_messages({"messages": [ |
| {"role": "user", "content": f"question {split}"}, |
| {"role": "assistant", "content": "answer"}]}) |
| row = {"messages": messages, "prompt_hash": prompt_key(messages)} |
| raw = (json.dumps(row) + "\n").encode() |
| (folder / f"{split}.jsonl").write_bytes(raw) |
| manifest["sha256"][split] = hashlib.sha256(raw).hexdigest() |
| manifest["counts"][split] = 1 |
| (folder / "manifest.json").write_text(json.dumps(manifest), encoding="utf-8") |
| self.assertEqual(set(read_data(folder)), {"train", "validation", "test"}) |
| (folder / "train.jsonl").write_text("changed", encoding="utf-8") |
| with self.assertRaisesRegex(ValueError, "Data changed"): |
| read_data(folder) |
|
|
| def test_small_qwen_lora_forward_and_gradients(self): |
| import torch |
| from peft import LoraConfig, TaskType, get_peft_model |
| from transformers import Qwen3Config, Qwen3ForCausalLM |
| torch.set_num_threads(2) |
| model = get_peft_model(Qwen3ForCausalLM(Qwen3Config( |
| vocab_size=len(self.tokenizer), hidden_size=32, intermediate_size=64, |
| num_hidden_layers=1, num_attention_heads=2, num_key_value_heads=1, |
| head_dim=16, max_position_embeddings=512)), |
| LoraConfig(r=2, lora_alpha=4, target_modules="all-linear", task_type=TaskType.CAUSAL_LM)) |
| messages = normalize_messages({"messages": [ |
| {"role": "user", "content": "first"}, {"role": "assistant", "content": "answer"}]}) |
| batch = collate([encode_example(self.tokenizer, messages, 512)], self.tokenizer.pad_token_id) |
| loss = model(**batch).loss |
| self.assertTrue(torch.isfinite(loss)) |
| loss.backward() |
| self.assertTrue(any(p.grad is not None and p.grad.abs().sum() > 0 |
| for name, p in model.named_parameters() if "lora_B" in name)) |
| self.assertFalse(any(p.requires_grad for name, p in model.named_parameters() if "lora_" not in name)) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
