#!/usr/bin/env python3
"""
Compare logits from original E4B model vs our Swift implementation.
Uses transformers library since MLX doesn't support Gemma-4 yet.
"""

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import json
import numpy as np

# Load original E4B model from HuggingFace
model_path = "google/gemma-4-4b-it"

print("Loading original E4B model from HuggingFace...")
print("This may take a moment to download...")

try:
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        device_map="auto",
        torch_dtype=torch.bfloat16
    )
    print("Model loaded successfully!")
except Exception as e:
    print(f"Error loading model: {e}")
    print("\nFalling back to checking config from converted model...")
    
    # Check config from converted model
    conv_path = "/Users/accusys/MarkBase12B/models/E4B-MarkBase/config.json"
    with open(conv_path) as f:
        cfg = json.load(f)
    
    print("\nConverted model config:")
    for key in ["hidden_size", "num_hidden_layers", "vocab_size", 
                "num_attention_heads", "num_key_value_heads",
                "hidden_size_per_layer_input"]:
        print(f"  {key}: {cfg.get('text_config', {}).get(key, cfg.get(key))}")
    
    exit(1)

# Get BOS token
bos_token_id = tokenizer.bos_token_id if tokenizer.bos_token_id else 2
print(f"\nBOS token ID: {bos_token_id}")

# Run forward pass for position 0
input_ids = torch.tensor([[bos_token_id]])
print(f"Input IDs: {input_ids}")

with torch.no_grad():
    outputs = model(input_ids)
    logits = outputs.logits[0, 0].float().numpy()  # [vocab_size]

print(f"\nLogits shape: {logits.shape}")
print(f"Logits range: min={logits.min():.2f}, max={logits.max():.2f}")

# Get top 10 tokens
top_indices = np.argsort(logits)[-10:][::-1]
print("\nTop 10 tokens (position 0, BOS):")
for i, idx in enumerate(top_indices):
    token = tokenizer.decode([idx])
    print(f"  {i+1}. token {idx} '{token}': {logits[idx]:.2f}")

# Test with actual prompt
prompt = "<start_of_turn>user\nThe capital of France is<end_of_turn>\n<start_of_turn>model\n"
print(f"\nTesting prompt: '{prompt[:50]}...'")

# Tokenize
inputs = tokenizer(prompt, return_tensors="pt")
print(f"Tokens: {inputs['input_ids'][0][:20].tolist()}...")

# Generate
print("\nGenerating response...")
with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=30,
        do_sample=True,
        temperature=1.0,
        top_k=40
    )
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(f"Response: '{response}'")