# PyPilot Inference Script
from modeling_pypilot import PyPilotModel, PyPilotConfig

def test_inference():
    print("🧪 Testing PyPilot Inference...")
    
    config = PyPilotConfig()
    model = PyPilotModel(config)
    
    # Mock input (tokenized code)
    dummy_input = torch.tensor([[1, 2, 3, 4, 5]])
    
    with torch.no_grad():
        output = model(dummy_input)
    
    print(f"✅ Inference test passed! Output shape: {output.shape}")
    return output

if __name__ == "__main__":
    test_inference()