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

def compute_angular_distance(a, b):
    """Compute the angular distance between two tensors."""
    a = a.detach().cpu().numpy().flatten()
    b = b.detach().cpu().numpy().flatten()
    cosine_similarity = np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
    angular_distance = np.arccos(np.clip(cosine_similarity, -1.0, 1.0)) / np.pi
    return angular_distance

def get_model_hidden_states(model, tokenizer, input_text):
    """Pass input through the model and get hidden states."""
    inputs = tokenizer(input_text, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs, output_hidden_states=True)
    hidden_states = outputs.hidden_states
    return hidden_states

def calculate_layer_distances(hidden_states):
    """Calculate and return angular distances between consecutive hidden states."""
    distances = []
    for i in range(len(hidden_states) - 1):
        distance = compute_angular_distance(hidden_states[i], hidden_states[i+1])
        distances.append(distance)
    return distances

def find_optimal_pruning_layer(distances):
    """Find the optimal layer to prune based on minimum angular distance."""
    min_distance = min(distances)
    optimal_layer = distances.index(min_distance)
    return optimal_layer, min_distance

# Load the model and tokenizer
model_name_or_path = "/home/joe/Downloads/checkpoint-2000-math/checkpoint-2004"  # Replace with your model path.
print(model_name_or_path)
model = AutoModelForCausalLM.from_pretrained(model_name_or_path)
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)

# Define your sample input text
input_text = "I have a reasoning math problem for you. What is 4+4*7=?"

# Get model hidden states for the input text
hidden_states = get_model_hidden_states(model, tokenizer, input_text)

# Calculate angular distances for all consecutive layers
distances = calculate_layer_distances(hidden_states)

# Display angular distances
for i, distance in enumerate(distances, start=1):
    print(f"Layer {i} to {i+1}: Angular Distance = {distance:.4f}")

# Find the optimal layer for pruning
optimal_layer, min_distance = find_optimal_pruning_layer(distances)

print(f"\nLayer Pair with Minimal Angular Distance Suggesting Redundancy: {optimal_layer + 1} to {optimal_layer + 2}")
print(f"Minimum Angular Distance (Suggesting Redundancy): {min_distance:.4f}")