import torch
from model.transformer import Transformer, ModelArgs

def test_model_forward():
    params = ModelArgs(
        dim=512,
        n_layers=4,
        n_heads=8,
        n_kv_heads=2,
        vocab_size=1000,
        max_seq_len=128,
    )
    model = Transformer(params)
    
    # Test batch of 2, sequence length of 10
    tokens = torch.randint(0, 1000, (2, 10))
    output = model(tokens)
    
    print(f"Output shape: {output.shape}")
    assert output.shape == (2, 10, 1000)
    print("Forward pass successful!")

if __name__ == "__main__":
    test_model_forward()