File size: 2,988 Bytes
948a05a | 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 | """Test model: forward pass, generation, save/load."""
import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
import numpy as np
from splitbit_llm.config import get_model_config, HardwareTier
from splitbit_llm.model.model import SplitBitLLM
from splitbit_llm.model.tokenizer import BPETokenizer
def test_model_forward():
"""Test model forward pass."""
cfg = get_model_config(HardwareTier.MOBILE)
cfg.vocab_size = 256
model = SplitBitLLM(config=cfg)
token_ids = np.array([[1, 5, 10, 15, 20]], dtype=np.int64)
logits = model.forward(token_ids)
assert logits.shape == (1, 5, 256), f"Wrong shape: {logits.shape}"
print(f" Logits shape: {logits.shape}")
print(f" Param count: {model.param_count:,}")
def test_model_generate():
"""Test text generation."""
cfg = get_model_config(HardwareTier.MOBILE)
cfg.vocab_size = 256
model = SplitBitLLM(config=cfg)
output = model.generate("Hello", max_tokens=10, temperature=0.7)
assert isinstance(output, str), f"Expected str, got {type(output)}"
assert len(output) > 0, "Empty output"
print(f" Generated: {repr(output[:50])}")
def test_model_generate_stream():
"""Test streaming generation."""
cfg = get_model_config(HardwareTier.MOBILE)
cfg.vocab_size = 256
model = SplitBitLLM(config=cfg)
chunks = list(model.generate_stream("Hello", max_tokens=10, temperature=0.7))
assert len(chunks) > 0, "No chunks generated"
print(f" Chunks: {len(chunks)}")
def test_model_with_tokenizer():
"""Test model with trained tokenizer."""
tok = BPETokenizer(vocab_size=256)
tok.train("Hello world! This is a test. Hello world again. The quick brown fox jumps.")
cfg = get_model_config(HardwareTier.MOBILE)
cfg.vocab_size = 256
model = SplitBitLLM(config=cfg, tokenizer=tok)
output = model.generate("Hello", max_tokens=10, temperature=0.7)
assert isinstance(output, str)
print(f" Generated with tokenizer: {repr(output[:50])}")
def test_model_truncation():
"""Test that long inputs are truncated to max_seq_len."""
cfg = get_model_config(HardwareTier.MOBILE)
cfg.vocab_size = 256
cfg.max_seq_len = 32
model = SplitBitLLM(config=cfg)
# Input longer than max_seq_len
long_input = np.array([[i for i in range(100)]], dtype=np.int64)
logits = model.forward(long_input)
assert logits.shape[1] == 32, f"Should truncate to 32, got {logits.shape[1]}"
print(f" Truncated to: {logits.shape[1]}")
if __name__ == "__main__":
print("Running model tests...")
test_model_forward()
print(" ✓ test_model_forward")
test_model_generate()
print(" ✓ test_model_generate")
test_model_generate_stream()
print(" ✓ test_model_generate_stream")
test_model_with_tokenizer()
print(" ✓ test_model_with_tokenizer")
test_model_truncation()
print(" ✓ test_model_truncation")
print("\nAll model tests passed!")
|