3d0eda94be
Model: Josephgflowers/Tinyllama-1.5B-Cinder-New-Test-1-with-gguf Source: Original Platform
62 lines
2.4 KiB
Python
62 lines
2.4 KiB
Python
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}")
|
|
|
|
|
|
|