Demo Scripts
scripts/basic_analysis.py
#!/usr/bin/env python3
"""
Basic LogitLens Analysis Script
This script demonstrates how to use LogitLens4LLMs to analyze layer-wise predictions
in large language models like Llama-3.1-8B and Qwen-2.5-7B.
Requirements:
pip install torch transformers matplotlib seaborn numpy
"""
import json
import os
from typing import List, Dict, Any
from enum import Enum
# Mock implementation of the LogitLens components
# In actual usage, import from the installed package
class ModelType(Enum):
"""Enumeration of supported model types"""
LLAMA_3_1_8B = "llama_3_1_8b"
QWEN_2_5_7B = "qwen_2_5_7b"
LLAMA_2_7B = "llama_2_7b"
class LogitLensAnalyzer:
"""Main analyzer class for LogitLens analysis"""
def __init__(self, model_type: ModelType, use_local: bool = False):
"""
Initialize the LogitLens analyzer.
Args:
model_type: Type of model to analyze
use_local: Whether to use locally cached model
"""
self.model_type = model_type
self.use_local = use_local
self.model_name = self._get_model_name()
def _get_model_name(self) -> str:
"""Get the Hugging Face model identifier"""
model_map = {
ModelType.LLAMA_3_1_8B: "meta-llama/Meta-Llama-3.1-8B",
ModelType.QWEN_2_5_7B: "Qwen/Qwen2.5-7B",
ModelType.LLAMA_2_7B: "meta-llama/Llama-2-7b-hf"
}
return model_map.get(self.model_type, "")
def analyze_prompt(self,
prompt: str,
max_new_tokens: int = 10,
temperature: float = 0.7,
print_details: bool = True) -> List[Dict[str, Any]]:
"""
Perform LogitLens analysis on a given prompt.
Args:
prompt: Input text to analyze
max_new_tokens: Number of tokens to generate
temperature: Sampling temperature
print_details: Whether to print detailed layer information
Returns:
List of prediction steps with layer-wise analysis
"""
prediction_steps = []
print(f"Running LogitLens analysis on {self.model_type.value}")
print(f"Prompt: {prompt}")
print("-" * 50)
# Simulate generation steps
generated_tokens = self._simulate_generation(prompt, max_new_tokens)
for step_idx, token in enumerate(generated_tokens):
step_data = self._analyze_step(prompt, token, step_idx, print_details)
prediction_steps.append(step_data)
if print_details:
self._print_step_analysis(step_data)
return prediction_steps
def _simulate_generation(self, prompt: str, max_tokens: int) -> List[str]:
"""Simulate token generation for demonstration"""
# In real implementation, this would use the actual model
example_completions = {
"The cat sat on the": ["mat", "and", "looked", "at", "the", "mouse"],
"Once upon a time": ["there", "was", "a", "brave", "knight"],
"The weather today is": ["sunny", "with", "clear", "blue", "skies"]
}
for key in example_completions:
if key in prompt:
return example_completions[key][:max_tokens]
# Default tokens
return ["token"] * min(max_tokens, 5)
def _analyze_step(self,
prompt: str,
token: str,
step_idx: int,
include_all_layers: bool = True) -> Dict[str, Any]:
"""
Analyze a single generation step.
Args:
prompt: Current prompt text
token: Generated token
step_idx: Step index
include_all_layers: Whether to include all layer predictions
Returns:
Dictionary containing step analysis data
"""
# Simulate layer-wise predictions
num_layers = 32 if "8b" in self.model_type.value.lower() else 28
layer_predictions = []
important_layers = []
for layer_idx in range(num_layers):
# Simulate confidence scores decreasing with layer depth
confidence = 100 - (layer_idx * 2) + (step_idx * 5)
confidence = max(10, min(100, confidence))
layer_data = {
"layer_idx": layer_idx,
"predicted_token": token,
"confidence": confidence,
"top_k_predictions": self._get_top_k_predictions(token, layer_idx)
}
layer_predictions.append(layer_data)
# Mark important layers (high confidence)
if confidence > 70:
important_layers.append(layer_idx)
return {
"step_idx": step_idx,
"predicted_token": token,
"current_text": f"{prompt} {token}",
"layer_predictions": layer_predictions if include_all_layers else None,
"important_layers": important_layers,
"attention_weights": self._simulate_attention_weights(num_layers),
"mlp_contributions": self._simulate_mlp_contributions(num_layers)
}
def _get_top_k_predictions(self, token: str, layer_idx: int) -> List[tuple]:
"""Get top-k token predictions for a layer"""
# Simulate alternative predictions
alternatives = {
"mat": ["floor", "carpet", "rug", "ground"],
"sunny": ["cloudy", "rainy", "warm", "bright"],
"there": ["once", "lived", "existed", "came"]
}
alt_tokens = alternatives.get(token, ["alt1", "alt2", "alt3"])
predictions = [(token, 85 - layer_idx)]
for i, alt in enumerate(alt_tokens[:3]):
score = max(5, 70 - layer_idx - (i * 20))
predictions.append((alt, score))
return predictions
def _simulate_attention_weights(self, num_layers: int) -> Dict[int, float]:
"""Simulate attention weights across layers"""
weights = {}
for i in range(num_layers):
# Higher attention in middle layers
if i < num_layers // 3:
weights[i] = 0.3 + (i * 0.02)
elif i < 2 * num_layers // 3:
weights[i] = 0.6 + (i * 0.01)
else:
weights[i] = 0.8 + (i * 0.005)
return weights
def _simulate_mlp_contributions(self, num_layers: int) -> Dict[int, float]:
"""Simulate MLP contributions across layers"""
contributions = {}
for i in range(num_layers):
# MLP contribution increases with depth
contributions[i] = 0.2 + (i / num_layers) * 0.6
return contributions
def _print_step_analysis(self, step_data: Dict[str, Any]):
"""Print formatted analysis for a generation step"""
print(f"\nStep {step_data['step_idx'] + 1}: Generated '{step_data['predicted_token']}'")
print(f"Current text: {step_data['current_text']}")
if step_data['important_layers']:
print(f"Important layers: {step_data['important_layers'][:5]}")
print("Top predictions from final layers:")
if step_data['layer_predictions']:
for layer in step_data['layer_predictions'][-3:]:
preds = layer['top_k_predictions'][:3]
pred_str = ", ".join([f"{t}({s}%)" for t, s in preds])
print(f" Layer {layer['layer_idx']}: {pred_str}")
def save_analysis(self,
prediction_steps: List[Dict[str, Any]],
output_path: str = "output/analysis_results.json"):
"""
Save analysis results to JSON file.
Args:
prediction_steps: List of prediction step data
output_path: Path to save the JSON file
"""
os.makedirs(os.path.dirname(output_path), exist_ok=True)
# Prepare data for JSON serialization
output_data = {
"model_type": self.model_type.value,
"num_steps": len(prediction_steps),
"predictions": []
}
for step in prediction_steps:
step_summary = {
"step_idx": step["step_idx"],
"token": step["predicted_token"],
"text": step["current_text"],
"important_layers": step["important_layers"][:5] if step["important_layers"] else []
}
output_data["predictions"].append(step_summary)
with open(output_path, 'w') as f:
json.dump(output_data, f, indent=2)
print(f"\nAnalysis saved to: {output_path}")
def main():
"""Main function demonstrating LogitLens analysis workflow"""
# Example 1: Analyze Llama-3.1-8B
print("=" * 60)
print("Example 1: Llama-3.1-8B Analysis")
print("=" * 60)
analyzer = LogitLensAnalyzer(ModelType.LLAMA_3_1_8B, use_local=False)
prompt = "Complete this sentence: The cat sat on the"
results = analyzer.analyze_prompt(
prompt=prompt,
max_new_tokens=5,
temperature=0.7,
print_details=True
)
# Save results
analyzer.save_analysis(results, "output/llama_analysis.json")
# Example 2: Analyze Qwen-2.5-7B
print("\n" + "=" * 60)
print("Example 2: Qwen-2.5-7B Analysis")
print("=" * 60)
qwen_analyzer = LogitLensAnalyzer(ModelType.QWEN_2_5_7B, use_local=False)
prompt = "The weather today is"
qwen_results = qwen_analyzer.analyze_prompt(
prompt=prompt,
max_new_tokens=5,
temperature=0.8,
print_details=True
)
# Save results
qwen_analyzer.save_analysis(qwen_results, "output/qwen_analysis.json")
# Example 3: Batch analysis
print("\n" + "=" * 60)
print("Example 3: Batch Analysis")
print("=" * 60)
prompts = [
"Once upon a time",
"The secret to happiness is",
"In the future, AI will"
]
batch_analyzer = LogitLensAnalyzer(ModelType.LLAMA_2_7B, use_local=False)
all_results = []
for prompt in prompts:
print(f"\nAnalyzing: {prompt}")
results = batch_analyzer.analyze_prompt(
prompt=prompt,
max_new_tokens=3,
print_details=False
)
all_results.extend(results)
print(f"\nCompleted batch analysis of {len(prompts)} prompts")
print(f"Total prediction steps analyzed: {len(all_results)}")
if __name__ == "__main__":
main()
scripts/visualization_generator.py
#!/usr/bin/env python3
"""
Visualization Generator for LogitLens Analysis
This script generates heatmap visualizations from LogitLens analysis results,
showing layer-wise predictions and confidence scores.
Requirements:
pip install matplotlib seaborn numpy pandas
"""
import json
import os
from typing import List, Dict, Any, Optional, Tuple
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import pandas as pd
from datetime import datetime
class LogitLensVisualizer:
"""Generate visualizations for LogitLens analysis results"""
def __init__(self, output_dir: str = "output/visualizations"):
"""
Initialize the visualizer.
Args:
output_dir: Directory to save visualization outputs
"""
self.output_dir = output_dir
os.makedirs(output_dir, exist_ok=True)
# Set visualization style
sns.set_style("whitegrid")
plt.rcParams['figure.figsize'] = (12, 8)
plt.rcParams['font.size'] = 10
def create_layer_heatmap(self,
layer_predictions: List[Dict[str, Any]],
step_idx: int,
title: str = "Layer-wise Token Predictions",
show_all_layers: bool = True) -> str:
"""
Create a heatmap showing token predictions across layers.
Args:
layer_predictions: List of layer prediction data
step_idx: Generation step index
title: Title for the heatmap
show_all_layers: Whether to show all layers or just important ones
Returns:
Path to saved visualization file
"""
# Prepare data for heatmap
num_layers = len(layer_predictions)
top_k = 5 # Show top 5 predictions per layer
# Extract unique tokens and create matrix
all_tokens = set()
for layer_data in layer_predictions:
for token, _ in layer_data.get('top_k_predictions', [])[:top_k]:
all_tokens.add(token)
tokens = sorted(list(all_tokens))
# Create confidence matrix
confidence_matrix = np.zeros((num_layers, len(tokens)))
for i, layer_data in enumerate(layer_predictions):
for token, score in layer_data.get('top_k_predictions', []):
if token in tokens:
j = tokens.index(token)
confidence_matrix[i, j] = score
# Create heatmap
fig, ax = plt.subplots(figsize=(14, 10))
# Use custom colormap
cmap = sns.color_palette("RdYlBu_r", as_cmap=True)
# Create heatmap with annotations
sns.heatmap(confidence_matrix,
xticklabels=tokens,
yticklabels=[f"Layer {i}" for i in range(num_layers)],
cmap=cmap,
cbar_kws={'label': 'Confidence (%)'},
fmt='.0f',
linewidths=0.5,
linecolor='gray',
ax=ax)
ax.set_title(f"{title} - Step {step_idx + 1}", fontsize=14, fontweight='bold')
ax.set_xlabel("Predicted Tokens", fontsize=12)
ax.set_ylabel("Model Layers", fontsize=12)
# Rotate x-axis labels for better readability
plt.xticks(rotation=45, ha='right')
# Tight layout
plt.tight_layout()
# Save figure
filename_prefix = "all_layers" if show_all_layers else "important_layers"
filepath = os.path.join(self.output_dir, f"{filename_prefix}_step_{step_idx}.png")
plt.savefig(filepath, dpi=150, bbox_inches='tight')
plt.close()
return filepath
def create_confidence_progression(self,
prediction_steps: List[Dict[str, Any]],
model_name: str = "Model") -> str:
"""
Create a line plot showing confidence progression across generation steps.
Args:
prediction_steps: List of prediction step data
model_name: Name of the model for title
Returns:
Path to saved visualization file
"""
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 10))
# Plot 1: Average confidence per step
steps = []
avg_confidences = []
max_confidences = []
min_confidences = []
for step in prediction_steps:
step_idx = step['step_idx']
steps.append(step_idx)
# Calculate statistics from layer predictions
confidences = []
if step.get('layer_predictions'):
for layer in step['layer_predictions']:
if layer['top_k_predictions']:
# Get confidence of top prediction
confidences.append(layer['top_k_predictions'][0][1])
if confidences:
avg_confidences.append(np.mean(confidences))
max_confidences.append(np.max(confidences))
min_confidences.append(np.min(confidences))
else:
avg_confidences.append(0)
max_confidences.append(0)
min_confidences.append(0)
# Plot confidence bands
ax1.plot(steps, avg_confidences, 'b-', linewidth=2, label='Average')
ax1.fill_between(steps, min_confidences, max_confidences, alpha=0.3, color='blue')
ax1.plot(steps, max_confidences, 'g--', linewidth=1, alpha=0.7, label='Maximum')
ax1.plot(steps, min_confidences, 'r--', linewidth=1, alpha=0.7, label='Minimum')
ax1.set_xlabel("Generation Step", fontsize=12)
ax1.set_ylabel("Confidence (%)", fontsize=12)
ax1.set_title(f"Confidence Progression - {model_name}", fontsize=14, fontweight='bold')
ax1.legend(loc='best')
ax1.grid(True, alpha=0.3)
# Plot 2: Number of high-confidence layers per step
high_conf_counts = []
threshold = 70 # Confidence threshold
for step in prediction_steps:
count = len([l for l in step.get('important_layers', [])])
high_conf_counts.append(count)
ax2.bar(steps, high_conf_counts, color='steelblue', alpha=0.7)
ax2.set_xlabel("Generation Step", fontsize=12)
ax2.set_ylabel(f"Layers with >{threshold}% Confidence", fontsize=12)
ax2.set_title("High-Confidence Layer Count", fontsize=14, fontweight='bold')
ax2.grid(True, alpha=0.3, axis='y')
plt.tight_layout()
# Save figure
filepath = os.path.join(self.output_dir, "confidence_progression.png")
plt.savefig(filepath, dpi=150, bbox_inches='tight')
plt.close()
return filepath
def create_attention_mlp_comparison(self,
step_data: Dict[str, Any],
step_idx: int) -> str:
"""
Create visualization comparing attention and MLP contributions.
Args:
step_data: Single step analysis data
step_idx: Step index
Returns:
Path to saved visualization file
"""
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6))
# Get attention and MLP data
attention_weights = step_data.get('attention_weights', {})
mlp_contributions = step_data.get('mlp_contributions', {})
if not attention_weights or not mlp_contributions:
# Create dummy data for demonstration
num_layers = 32
attention_weights = {i: 0.5 + i * 0.015 for i in range(num_layers)}
mlp_contributions = {i: 0.3 + i * 0.02 for i in range(num_layers)}
layers = list(range(len(attention_weights)))
att_values = [attention_weights[i] for i in layers]
mlp_values = [mlp_contributions[i] for i in layers]
# Plot attention weights
ax1.plot(layers, att_values, 'b-', linewidth=2, marker='o', markersize=4)
ax1.fill_between(layers, 0, att_values, alpha=0.3, color='blue')
ax1.set_xlabel("Layer Index", fontsize=12)
ax1.set_ylabel("Attention Weight", fontsize=12)
ax1.set_title("Attention Mechanism Contribution", fontsize=13, fontweight='bold')
ax1.grid(True, alpha=0.3)
ax1.set_xlim([0, len(layers) - 1])
# Plot MLP contributions
ax2.plot(layers, mlp_values, 'g-', linewidth=2, marker='s', markersize=4)
ax2.fill_between(layers, 0, mlp_values, alpha=0.3, color='green')
ax2.set_xlabel("Layer Index", fontsize=12)
ax2.set_ylabel("MLP Contribution", fontsize=12)
ax2.set_title("MLP Output Contribution", fontsize=13, fontweight='bold')
ax2.grid(True, alpha=0.3)
ax2.set_xlim([0, len(layers) - 1])
# Add step information
fig.suptitle(f"Component Analysis - Step {step_idx + 1}: '{step_data.get('predicted_token', '')}'",
fontsize=14, fontweight='bold', y=1.02)
plt.tight_layout()
# Save figure
filepath = os.path.join(self.output_dir, f"component_analysis_step_{step_idx}.png")
plt.savefig(filepath, dpi=150, bbox_inches='tight')
plt.close()
return filepath
def create_summary_dashboard(self,
analysis_results: Dict[str, Any],
model_name: str = "Model") -> str:
"""
Create a comprehensive dashboard summarizing the analysis.
Args:
analysis_results: Complete analysis results dictionary
model_name: Name of the model
Returns:
Path to saved dashboard file
"""
fig = plt.figure(figsize=(16, 12))
# Create grid for subplots
gs = fig.add_gridspec(3, 3, hspace=0.3, wspace=0.3)
# Title
fig.suptitle(f"LogitLens Analysis Dashboard - {model_name}",
fontsize=16, fontweight='bold')
# Subplot 1: Token sequence
ax1 = fig.add_subplot(gs[0, :])
predictions = analysis_results.get('predictions', [])
tokens = [p['token'] for p in predictions]
token_text = ' '.join(tokens)
ax1.text(0.5, 0.5, f"Generated Sequence:\n\n{token_text}",
ha='center', va='center', fontsize=14,
bbox=dict(boxstyle='round', facecolor='lightblue', alpha=0.5))
ax1.set_xlim([0, 1])
ax1.set_ylim([0, 1])
ax1.axis('off')
# Subplot 2: Layer importance distribution
ax2 = fig.add_subplot(gs[1, 0])
all_important_layers = []
for pred in predictions:
all_important_layers.extend(pred.get('important_layers', []))
if all_important_layers:
ax2.hist(all_important_layers, bins=20, color='steelblue', alpha=0.7, edgecolor='black')
ax2.set_xlabel("Layer Index")
ax2.set_ylabel("Frequency")
ax2.set_title("Important Layer Distribution")
ax2.grid(True, alpha=0.3)
# Subplot 3: Step-wise token confidence
ax3 = fig.add_subplot(gs[1, 1:])
step_indices = list(range(len(predictions)))
# Simulate confidence scores
confidences = [85 - i * 5 + np.random.randn() * 5 for i in step_indices]
ax3.bar(step_indices, confidences, color='green', alpha=0.6)
ax3.set_xlabel("Step Index")
ax3.set_ylabel("Confidence (%)")
ax3.set_title("Token Generation Confidence")
ax3.set_xticks(step_indices)
ax3.set_xticklabels([f"Step {i+1}" for i in step_indices], rotation=45)
ax3.grid(True, alpha=0.3, axis='y')
# Subplot 4: Model statistics
ax4 = fig.add_subplot(gs[2, :])
stats_text = f"""
Analysis Statistics:
• Model Type: {analysis_results.get('model_type', 'Unknown')}
• Total Steps: {analysis_results.get('num_steps', 0)}
• Average Important Layers per Step: {len(all_important_layers) / max(len(predictions), 1):.1f}
• Timestamp: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
"""
ax4.text(0.1, 0.5, stats_text, fontsize=11, verticalalignment='center',
family='monospace', bbox=dict(boxstyle='round', facecolor='lightyellow', alpha=0.5))
ax4.axis('off')
plt.tight_layout()
# Save dashboard
filepath = os.path.join(self.output_dir, "analysis_dashboard.png")
plt.savefig(filepath, dpi=150, bbox_inches='tight')
plt.close()
return filepath
def generate_all_visualizations(self,
analysis_json_path: str,
model_name: Optional[str] = None) -> Dict[str, str]:
"""
Generate all visualizations from a JSON analysis file.
Args:
analysis_json_path: Path to the JSON file with analysis results
model_name: Optional model name override
Returns:
Dictionary mapping visualization type to file path
"""
# Load analysis results
with open(analysis_json_path, 'r') as f:
analysis_data = json.load(f)
if not model_name:
model_name = analysis_data.get('model_type', 'Model')
generated_files = {}
# Generate dashboard
dashboard_path = self.create_summary_dashboard(analysis_data, model_name)
generated_files['dashboard'] = dashboard_path
print(f"Generated dashboard: {dashboard_path}")
# For demonstration, create sample visualizations
# In real implementation, these would use actual layer prediction data
# Sample layer predictions for heatmap
sample_layer_predictions = []
for i in range(32): # 32 layers
predictions = [
("token1", 90 - i * 2),
("token2", 70 - i * 1.5),
("token3", 50 - i),
("token4", 30 - i * 0.5),
]
sample_layer_predictions.append({
'layer_idx': i,
'top_k_predictions': predictions
})
# Generate heatmaps for first 3 steps
for step_idx in range(min(3, analysis_data.get('num_steps', 1))):
heatmap_path = self.create_layer_heatmap(
sample_layer_predictions,
step_idx,
title="Layer-wise Token Predictions",
show_all_layers=True
)
generated_files[f'heatmap_step_{step_idx}'] = heatmap_path
print(f"Generated heatmap for step {step_idx}: {heatmap_path}")
return generated_files
def main():
"""Main function demonstrating visualization generation"""
# Create visualizer
visualizer = LogitLensVisualizer(output_dir="output/visualizations")
# Example: Generate visualizations from mock analysis data
mock_analysis_data = {
"model_type": "llama_3_1_8b",
"num_steps": 5,
"predictions": [
{"step_idx": 0, "token": "The", "important_layers": [15, 20, 25]},
{"step_idx": 1, "token": "cat", "important_layers": [16, 21, 26]},
{"step_idx": 2, "token": "sat", "important_layers": [17, 22, 27]},
{"step_idx": 3, "token": "on", "important_layers": [18, 23, 28]},
{"step_idx": 4, "token": "the", "important_layers": [19, 24, 29]},
]
}
# Save mock data to JSON
os.makedirs("output", exist_ok=True)
mock_json_path = "output/mock_analysis.json"
with open(mock_json_path, 'w') as f:
json.dump(mock_analysis_data, f, indent=2)
print("Generating visualizations from analysis data...")
print("=" * 60)
# Generate all visualizations
generated_files = visualizer.generate_all_visualizations(
mock_json_path,
model_name="Llama-3.1-8B"
)
print("\n" + "=" * 60)
print("Visualization generation complete!")
print(f"Generated {len(generated_files)} visualization files:")
for viz_type, filepath in generated_files.items():
print(f" - {viz_type}: {filepath}")
if __name__ == "__main__":
main()
zjunlp