import torch
from transformers import GPT2Model

# Load pretrained GPT-2
model = GPT2Model.from_pretrained("openai-community/gpt2")

ratios = []

# Loop through transformer blocks
for i, block in enumerate(model.h):  # model.h = list of transformer layers
    # GPT2 uses Conv1D for QKV, packed into one matrix
    W_qkv = block.attn.c_attn.weight.detach()  # shape [768, 2304]
    
    # Split into Q, K, V
    hidden_size = model.config.hidden_size  # 768
    W_q, W_k, W_v = W_qkv.split(hidden_size, dim=1)  # each [768, 768]
    
    # Compute norms
    l2_q = torch.norm(W_q, p=2)
    l2_k = torch.norm(W_k, p=2)
    
    ratio = l2_q / l2_k
    ratios.append(ratio.item())
    
    print(f"Layer {i+1}: Q/K L2 ratio = {ratio.item():.4f}")

# Convert to tensor for stats
ratios_tensor = torch.tensor(ratios)
mean_ratio = torch.mean(ratios_tensor).item()
std_ratio = torch.std(ratios_tensor).item()

print(f"\nMean Q/K ratio over {len(ratios)} layers: {mean_ratio:.4f}")
print(f"Std  Q/K ratio over {len(ratios)} layers: {std_ratio:.4f}")