Text Generation
Transformers
Safetensors
PyTorch
English
qtensorformer
q-tensorformer
tensor-train
tensor-networks
adaptive-computation
kv-cache
kv-cache-quantization
attention-sinks
pid-control
edge-ai
edge-inference
on-device-ai
model-compression
grouped-query-attention
rotary-position-embeddings
custom_code
File size: 2,219 Bytes
eaeea8f | 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 | """
Tests for Hugging Face Transformers Integration.
"""
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch
import tempfile
import pytest
from transformers import AutoConfig, AutoModelForCausalLM
from src.hf_model import QTensorFormerConfig, QTensorFormerForCausalLM
def test_hf_config_and_model_forward():
config = QTensorFormerConfig(
vocab_size=100,
d_model=32,
n_layers=2,
n_heads=2,
tt_rank=4,
max_seq_len=32,
use_quantum=True,
)
model = QTensorFormerForCausalLM(config)
input_ids = torch.randint(0, 100, (2, 8))
out = model(input_ids)
assert out.logits.shape == (2, 8, 100)
print("✓ test_hf_config_and_model_forward passed")
def test_hf_generation():
config = QTensorFormerConfig(
vocab_size=100,
d_model=32,
n_layers=2,
n_heads=2,
tt_rank=4,
max_seq_len=32,
)
model = QTensorFormerForCausalLM(config)
input_ids = torch.tensor([[1, 2, 3]])
generated = model.generate(input_ids, max_new_tokens=5, do_sample=False)
assert generated.shape == (1, 8), f"Expected shape (1, 8), got {generated.shape}"
print("✓ test_hf_generation passed")
def test_hf_save_and_load():
config = QTensorFormerConfig(
vocab_size=50,
d_model=32,
n_layers=2,
n_heads=2,
tt_rank=4,
)
model = QTensorFormerForCausalLM(config)
with tempfile.TemporaryDirectory() as tmpdir:
model.save_pretrained(tmpdir)
# Load back with AutoModelForCausalLM
loaded = AutoModelForCausalLM.from_pretrained(tmpdir)
assert isinstance(loaded, QTensorFormerForCausalLM)
model.eval()
loaded.eval()
inp = torch.randint(0, 50, (1, 4))
with torch.no_grad():
out1 = model(inp).logits
out2 = loaded(inp).logits
assert torch.allclose(out1, out2, atol=1e-4)
print("✓ test_hf_save_and_load passed")
if __name__ == "__main__":
test_hf_config_and_model_forward()
test_hf_generation()
test_hf_save_and_load()
print("All Hugging Face integration tests passed!")
|