File size: 5,890 Bytes
5ba1bb4 | 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 | """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()
|