Demo Scripts
scripts/compute_perplexity.py
#!/usr/bin/env python3
"""
Compute Perplexity with Deactivated Language-Specific Neurons
This script demonstrates how to compute perplexity on multilingual text data
when specific language neurons are deactivated, enabling analysis of how
neuron deactivation affects model performance.
Requires: pip install torch transformers
"""
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
import argparse
from pathlib import Path
from typing import Dict, List, Optional, Tuple
import json
import numpy as np
from tqdm import tqdm
class TokenizedDataset(Dataset):
"""Dataset for pre-tokenized text data."""
def __init__(self, data_path: str, max_length: int = 1024):
"""
Initialize dataset with pre-tokenized data.
Args:
data_path: Path to tokenized data file (LongTensor)
max_length: Maximum sequence length for each sample
"""
self.data = torch.load(data_path, map_location='cpu')
self.max_length = max_length
# Split data into chunks
self.chunks = []
for i in range(0, len(self.data) - max_length, max_length):
self.chunks.append(self.data[i:i + max_length])
def __len__(self):
return len(self.chunks)
def __getitem__(self, idx):
return self.chunks[idx]
class PerplexityCalculator:
"""Calculate perplexity with optional neuron deactivation."""
def __init__(self, model_name: str, device: str = "cuda"):
"""
Initialize the perplexity calculator.
Args:
model_name: Name or path of the model
device: Device to use for computation
"""
self.model_name = model_name
self.device = device
self.activation_masks = None
def load_activation_masks(self, mask_path: str):
"""Load activation masks for neuron deactivation."""
if Path(mask_path).exists():
self.activation_masks = torch.load(mask_path, map_location=self.device)
print(f"Loaded activation masks from {mask_path}")
else:
print(f"Warning: Mask file not found at {mask_path}")
def apply_neuron_deactivation(self, activations: torch.Tensor,
layer_idx: int) -> torch.Tensor:
"""
Apply neuron deactivation mask to layer activations.
Args:
activations: Layer activations (batch_size, seq_len, hidden_dim)
layer_idx: Index of the current layer
Returns:
Modified activations with certain neurons deactivated
"""
if self.activation_masks is None or layer_idx not in self.activation_masks:
return activations
mask = self.activation_masks[layer_idx].to(activations.device)
# Expand mask to match activation dimensions
if len(activations.shape) == 3:
mask = mask.unsqueeze(0).unsqueeze(0)
mask = mask.expand_as(activations)
return activations * mask
def compute_perplexity(self,
data_loader: DataLoader,
apply_deactivation: bool = True) -> Dict:
"""
Compute perplexity on the dataset.
Args:
data_loader: DataLoader for the tokenized dataset
apply_deactivation: Whether to apply neuron deactivation
Returns:
Dictionary containing perplexity and related metrics
"""
total_loss = 0.0
total_tokens = 0
batch_losses = []
print(f"Computing perplexity with{'' if apply_deactivation else 'out'} neuron deactivation...")
with torch.no_grad():
for batch_idx, input_ids in enumerate(tqdm(data_loader)):
input_ids = input_ids.to(self.device)
# Shift inputs and targets for language modeling
inputs = input_ids[:, :-1]
targets = input_ids[:, 1:]
# Simulate forward pass with deactivation
# (In actual implementation, this would hook into the model)
logits = self.simulated_forward(inputs, apply_deactivation)
# Calculate loss
loss = F.cross_entropy(
logits.reshape(-1, logits.size(-1)),
targets.reshape(-1),
reduction='none'
)
batch_loss = loss.mean().item()
batch_losses.append(batch_loss)
total_loss += loss.sum().item()
total_tokens += targets.numel()
# Calculate perplexity
avg_loss = total_loss / total_tokens
perplexity = torch.exp(torch.tensor(avg_loss)).item()
return {
"perplexity": perplexity,
"avg_loss": avg_loss,
"total_tokens": total_tokens,
"batch_losses": batch_losses,
"std_loss": np.std(batch_losses)
}
def simulated_forward(self, input_ids: torch.Tensor,
apply_deactivation: bool) -> torch.Tensor:
"""
Simulated forward pass (placeholder for actual model forward).
In real implementation, this would use hooks to modify activations.
Args:
input_ids: Input token IDs
apply_deactivation: Whether to apply neuron deactivation
Returns:
Logits tensor
"""
# This is a simplified simulation
# In practice, you would use model hooks to intercept and modify activations
batch_size, seq_len = input_ids.shape
vocab_size = 32000 # Typical vocab size for LLaMA models
# Generate random logits for demonstration
logits = torch.randn(batch_size, seq_len, vocab_size, device=input_ids.device)
# If deactivation is applied, slightly increase the entropy
if apply_deactivation and self.activation_masks:
logits = logits * 0.95 # Simulated effect of deactivation
return logits
class PerplexityExperiment:
"""Run comprehensive perplexity experiments."""
def __init__(self, model_name: str, languages: List[str]):
"""
Initialize experiment runner.
Args:
model_name: Model to use
languages: List of languages to test
"""
self.model_name = model_name
self.languages = languages
self.results = {}
def run_language_comparison(self,
data_dir: str,
mask_dir: str,
output_file: str):
"""
Run perplexity comparison across languages.
Args:
data_dir: Directory containing tokenized data files
mask_dir: Directory containing activation masks
output_file: Path to save results
"""
for lang in self.languages:
print(f"\n{'='*50}")
print(f"Testing language: {lang}")
print('='*50)
# Prepare data path
data_path = f"{data_dir}/id.{lang}.train.llama"
if not Path(data_path).exists():
print(f"Data file not found: {data_path}")
continue
# Create dataset and dataloader
dataset = TokenizedDataset(data_path)
data_loader = DataLoader(dataset, batch_size=4, shuffle=False)
# Calculate baseline perplexity (no deactivation)
calculator = PerplexityCalculator(self.model_name)
baseline_results = calculator.compute_perplexity(
data_loader, apply_deactivation=False
)
lang_results = {
"baseline": baseline_results,
"deactivation_experiments": {}
}
# Test with different deactivation masks
for mask_lang in self.languages:
mask_path = f"{mask_dir}/mask_{mask_lang}_100%.pth"
if Path(mask_path).exists():
calculator.load_activation_masks(mask_path)
deact_results = calculator.compute_perplexity(
data_loader, apply_deactivation=True
)
# Calculate perplexity change
ppl_change = (deact_results["perplexity"] - baseline_results["perplexity"]) / baseline_results["perplexity"] * 100
lang_results["deactivation_experiments"][mask_lang] = {
**deact_results,
"perplexity_change_percent": ppl_change
}
print(f" Deactivating {mask_lang} neurons: PPL {deact_results['perplexity']:.2f} "
f"(change: {ppl_change:+.2f}%)")
self.results[lang] = lang_results
# Save results
self.save_results(output_file)
def run_progressive_deactivation(self,
data_path: str,
mask_pattern: str,
ratios: List[float]):
"""
Test perplexity with progressive neuron deactivation.
Args:
data_path: Path to tokenized data
mask_pattern: Pattern for mask files (with {ratio} placeholder)
ratios: List of deactivation ratios to test
"""
dataset = TokenizedDataset(data_path)
data_loader = DataLoader(dataset, batch_size=4, shuffle=False)
results = {"ratios": {}}
for ratio in ratios:
mask_path = mask_pattern.format(ratio=int(ratio*100))
calculator = PerplexityCalculator(self.model_name)
if Path(mask_path).exists():
calculator.load_activation_masks(mask_path)
ppl_results = calculator.compute_perplexity(
data_loader, apply_deactivation=True
)
else:
print(f"Mask not found: {mask_path}")
continue
results["ratios"][f"{int(ratio*100)}%"] = {
"ratio": ratio,
**ppl_results
}
print(f"Deactivation {int(ratio*100)}%: PPL = {ppl_results['perplexity']:.2f}")
return results
def save_results(self, output_file: str):
"""Save experiment results to JSON file."""
output_path = Path(output_file)
output_path.parent.mkdir(parents=True, exist_ok=True)
with open(output_path, 'w') as f:
json.dump(self.results, f, indent=2, default=str)
print(f"\nResults saved to {output_file}")
def plot_perplexity_changes(results: Dict):
"""
Create a simple text-based visualization of perplexity changes.
Args:
results: Dictionary of experiment results
"""
print("\n" + "="*60)
print("Perplexity Impact Matrix")
print("="*60)
print("Rows: Test language | Columns: Deactivated language neurons")
print("-"*60)
languages = list(results.keys())
# Header
print(f"{'Test Lang':<10}", end="")
for lang in languages:
print(f"{lang:>8}", end="")
print(f"{'Baseline':>10}")
print("-"*60)
# Data rows
for test_lang in languages:
print(f"{test_lang:<10}", end="")
baseline_ppl = results[test_lang]["baseline"]["perplexity"]
for deact_lang in languages:
if deact_lang in results[test_lang]["deactivation_experiments"]:
ppl = results[test_lang]["deactivation_experiments"][deact_lang]["perplexity"]
change = (ppl - baseline_ppl) / baseline_ppl * 100
print(f"{change:>+7.1f}%", end="")
else:
print(f"{'N/A':>8}", end="")
print(f"{baseline_ppl:>10.2f}")
def main():
parser = argparse.ArgumentParser(description="Compute perplexity with neuron deactivation")
parser.add_argument("-m", "--model", type=str, required=True,
help="Model name or path")
parser.add_argument("-d", "--data", type=str, required=True,
help="Path to tokenized data file or directory")
parser.add_argument("-a", "--activation-mask", type=str,
help="Path to activation mask file")
parser.add_argument("--languages", nargs="+",
default=["en", "zh", "fr", "es", "vi", "id", "ja"],
help="Languages to test")
parser.add_argument("--progressive", action="store_true",
help="Run progressive deactivation experiment")
parser.add_argument("-o", "--output", type=str, default="perplexity_results.json",
help="Output file for results")
args = parser.parse_args()
try:
if args.progressive:
# Progressive deactivation experiment
experiment = PerplexityExperiment(args.model, ["test"])
ratios = [0.0, 0.25, 0.5, 0.75, 1.0]
results = experiment.run_progressive_deactivation(
args.data,
args.activation_mask.replace(".pth", "_{ratio}.pth"),
ratios
)
print("\nProgressive Deactivation Results:")
print("-"*40)
for ratio_key, ratio_results in results["ratios"].items():
print(f"{ratio_key}: Perplexity = {ratio_results['perplexity']:.2f}")
else:
# Standard perplexity calculation
if Path(args.data).is_dir():
# Run comparison across languages
experiment = PerplexityExperiment(args.model, args.languages)
experiment.run_language_comparison(
args.data,
Path(args.activation_mask).parent if args.activation_mask else "activation_mask",
args.output
)
# Visualize results
if experiment.results:
plot_perplexity_changes(experiment.results)
else:
# Single file perplexity
dataset = TokenizedDataset(args.data)
data_loader = DataLoader(dataset, batch_size=4, shuffle=False)
calculator = PerplexityCalculator(args.model)
if args.activation_mask:
calculator.load_activation_masks(args.activation_mask)
results = calculator.compute_perplexity(
data_loader,
apply_deactivation=(args.activation_mask is not None)
)
print(f"\nPerplexity: {results['perplexity']:.2f}")
print(f"Average Loss: {results['avg_loss']:.4f}")
print(f"Total Tokens: {results['total_tokens']:,}")
except Exception as e:
print(f"Error: {e}")
return 1
return 0
if __name__ == "__main__":
exit(main())
scripts/deactivate_neurons.py
#!/usr/bin/env python3
"""
Deactivate Language-Specific Neurons and Analyze Model Behavior
This script demonstrates how to create activation masks to deactivate specific
language neurons and analyze the resulting model behavior changes.
Requires: pip install torch numpy
"""
import torch
import numpy as np
import json
import argparse
from pathlib import Path
from typing import List, Dict, Optional, Tuple
class NeuronMaskGenerator:
"""Generate activation masks for deactivating specific neurons."""
def __init__(self, model_config: Dict):
"""
Initialize the mask generator with model configuration.
Args:
model_config: Dictionary containing model dimensions
{num_layers, hidden_size, intermediate_size}
"""
self.num_layers = model_config.get("num_layers", 32)
self.hidden_size = model_config.get("hidden_size", 4096)
self.intermediate_size = model_config.get("intermediate_size", 11008)
def create_activation_mask(self,
neurons_to_deactivate: List[List[torch.LongTensor]],
mask_type: str = "mlp") -> Dict[int, torch.Tensor]:
"""
Create activation masks for deactivating specific neurons.
Args:
neurons_to_deactivate: List of neuron indices per layer
mask_type: Type of mask ("mlp" for MLP layers, "attention" for attention)
Returns:
Dictionary mapping layer indices to activation masks
"""
masks = {}
for layer_idx, neuron_indices in enumerate(neurons_to_deactivate):
if mask_type == "mlp":
# Create mask for MLP/FFN neurons
mask = torch.ones(self.intermediate_size, dtype=torch.float32)
if len(neuron_indices) > 0:
mask[neuron_indices] = 0.0
masks[layer_idx] = mask
elif mask_type == "attention":
# Create mask for attention heads (simplified)
mask = torch.ones(self.hidden_size, dtype=torch.float32)
if len(neuron_indices) > 0:
# Map neuron indices to attention dimensions
attention_dims = neuron_indices % self.hidden_size
mask[attention_dims] = 0.0
masks[layer_idx] = mask
return masks
def save_masks(self, masks: Dict[int, torch.Tensor], output_path: str):
"""Save activation masks to file."""
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
torch.save(masks, output_path)
print(f"Saved activation masks to {output_path}")
def load_masks(self, mask_path: str) -> Dict[int, torch.Tensor]:
"""Load activation masks from file."""
if not Path(mask_path).exists():
raise FileNotFoundError(f"Mask file not found: {mask_path}")
return torch.load(mask_path, map_location='cpu')
class NeuronAnalyzer:
"""Analyze the effects of neuron deactivation."""
def __init__(self, neuron_file: str):
"""
Initialize analyzer with pre-identified neurons.
Args:
neuron_file: Path to .neuron.pth file
"""
self.neurons = torch.load(neuron_file, map_location='cpu')
self.lang_map = {"en": 0, "zh": 1, "fr": 2, "es": 3, "vi": 4, "id": 5, "ja": 6}
def get_language_neurons(self, language: str) -> List[torch.LongTensor]:
"""Get all neurons for a specific language."""
if language not in self.lang_map:
raise ValueError(f"Invalid language: {language}")
return self.neurons[self.lang_map[language]]
def create_deactivation_experiment(self,
target_language: str,
deactivation_ratio: float = 1.0) -> List[List[int]]:
"""
Create an experiment by selecting neurons to deactivate.
Args:
target_language: Language whose neurons to deactivate
deactivation_ratio: Fraction of neurons to deactivate (0-1)
Returns:
List of neuron indices per layer to deactivate
"""
lang_neurons = self.get_language_neurons(target_language)
deactivation_list = []
for layer_neurons in lang_neurons:
num_neurons = len(layer_neurons)
num_to_deactivate = int(num_neurons * deactivation_ratio)
if num_to_deactivate > 0:
# Randomly select neurons to deactivate
indices = torch.randperm(num_neurons)[:num_to_deactivate]
selected_neurons = layer_neurons[indices]
deactivation_list.append(selected_neurons)
else:
deactivation_list.append(torch.LongTensor([]))
return deactivation_list
def analyze_cross_lingual_impact(self,
source_language: str,
target_languages: List[str]) -> Dict:
"""
Analyze how deactivating one language's neurons affects others.
Args:
source_language: Language whose neurons to deactivate
target_languages: Languages to test impact on
Returns:
Analysis results
"""
source_neurons = self.get_language_neurons(source_language)
results = {
"source_language": source_language,
"impacts": {}
}
# Convert source neurons to sets for efficient lookup
source_sets = [set(layer.tolist()) for layer in source_neurons]
for target_lang in target_languages:
if target_lang == source_language:
continue
target_neurons = self.get_language_neurons(target_lang)
# Calculate overlap
total_overlap = 0
total_target = 0
layer_overlaps = []
for layer_idx, (source_set, target_layer) in enumerate(zip(source_sets, target_neurons)):
target_set = set(target_layer.tolist())
overlap = len(source_set & target_set)
total_overlap += overlap
total_target += len(target_set)
layer_overlaps.append({
"layer": layer_idx,
"overlap": overlap,
"target_total": len(target_set),
"impact_ratio": overlap / len(target_set) if target_set else 0
})
results["impacts"][target_lang] = {
"total_overlap": total_overlap,
"total_target_neurons": total_target,
"overall_impact": total_overlap / total_target if total_target > 0 else 0,
"layer_details": layer_overlaps
}
return results
def get_model_config(model_name: str) -> Dict:
"""Get model configuration based on model name."""
configs = {
"llama-7b": {"num_layers": 32, "hidden_size": 4096, "intermediate_size": 11008},
"llama-13b": {"num_layers": 40, "hidden_size": 5120, "intermediate_size": 13824},
"llama-70b": {"num_layers": 80, "hidden_size": 8192, "intermediate_size": 28672},
"bloom-7b": {"num_layers": 30, "hidden_size": 4096, "intermediate_size": 16384},
"opt-6.7b": {"num_layers": 32, "hidden_size": 4096, "intermediate_size": 16384},
"mistral-7b": {"num_layers": 32, "hidden_size": 4096, "intermediate_size": 14336},
"phi-2": {"num_layers": 32, "hidden_size": 2560, "intermediate_size": 10240}
}
return configs.get(model_name, configs["llama-7b"])
def create_progressive_deactivation(analyzer: NeuronAnalyzer,
language: str,
ratios: List[float]) -> Dict:
"""
Create masks for progressive deactivation experiments.
Args:
analyzer: NeuronAnalyzer instance
language: Target language
ratios: List of deactivation ratios to test
Returns:
Dictionary of deactivation experiments
"""
experiments = {}
for ratio in ratios:
neurons_to_deactivate = analyzer.create_deactivation_experiment(language, ratio)
experiments[f"{int(ratio*100)}%"] = {
"ratio": ratio,
"neurons": neurons_to_deactivate,
"total_deactivated": sum(len(layer) for layer in neurons_to_deactivate)
}
return experiments
def main():
parser = argparse.ArgumentParser(description="Create and analyze neuron deactivation masks")
parser.add_argument("-n", "--neuron-file", type=str, required=True,
help="Path to .neuron.pth file")
parser.add_argument("-m", "--model", type=str, default="llama-7b",
help="Model name for configuration")
parser.add_argument("-l", "--language", type=str, default="zh",
help="Language to deactivate (en, zh, fr, es, vi, id, ja)")
parser.add_argument("-r", "--ratio", type=float, default=1.0,
help="Deactivation ratio (0-1)")
parser.add_argument("-o", "--output", type=str, default="activation_mask/experiment.pth",
help="Output path for activation mask")
parser.add_argument("--analyze-impact", action="store_true",
help="Analyze cross-lingual impact")
parser.add_argument("--progressive", action="store_true",
help="Create progressive deactivation experiments")
args = parser.parse_args()
try:
# Initialize components
model_config = get_model_config(args.model)
mask_generator = NeuronMaskGenerator(model_config)
analyzer = NeuronAnalyzer(args.neuron_file)
if args.progressive:
# Create progressive deactivation experiments
ratios = [0.1, 0.25, 0.5, 0.75, 1.0]
experiments = create_progressive_deactivation(analyzer, args.language, ratios)
print(f"\nProgressive Deactivation for {args.language}:")
print("-" * 50)
for name, exp in experiments.items():
print(f"{name}: {exp['total_deactivated']} neurons deactivated")
# Create and save mask for each ratio
masks = mask_generator.create_activation_mask(exp['neurons'])
output_path = args.output.replace('.pth', f'_{args.language}_{name}.pth')
mask_generator.save_masks(masks, output_path)
elif args.analyze_impact:
# Analyze cross-lingual impact
all_languages = ["en", "zh", "fr", "es", "vi", "id", "ja"]
target_languages = [l for l in all_languages if l != args.language]
impact_analysis = analyzer.analyze_cross_lingual_impact(
args.language, target_languages
)
print(f"\nCross-lingual Impact Analysis")
print(f"Deactivating: {impact_analysis['source_language']}")
print("-" * 50)
for lang, impact in impact_analysis["impacts"].items():
print(f"\nImpact on {lang}:")
print(f" Total overlap: {impact['total_overlap']} neurons")
print(f" Overall impact: {impact['overall_impact']:.2%}")
# Show top impacted layers
top_layers = sorted(impact['layer_details'],
key=lambda x: x['impact_ratio'],
reverse=True)[:3]
print(" Most impacted layers:")
for layer in top_layers:
print(f" Layer {layer['layer']}: {layer['impact_ratio']:.2%} affected")
else:
# Standard deactivation
neurons_to_deactivate = analyzer.create_deactivation_experiment(
args.language, args.ratio
)
# Create activation masks
masks = mask_generator.create_activation_mask(neurons_to_deactivate)
# Save masks
mask_generator.save_masks(masks, args.output)
# Print statistics
total_neurons = sum(len(layer) for layer in neurons_to_deactivate)
print(f"\nDeactivation Summary:")
print(f" Language: {args.language}")
print(f" Deactivation ratio: {args.ratio:.1%}")
print(f" Total neurons deactivated: {total_neurons}")
print(f" Masks saved to: {args.output}")
except Exception as e:
print(f"Error: {e}")
return 1
return 0
if __name__ == "__main__":
exit(main())
scripts/load_neurons.py
#!/usr/bin/env python3
"""
Load and Analyze Pre-identified Language-Specific Neurons
This script demonstrates how to load and analyze the pre-identified language-specific
neurons for various multilingual models including LLaMA-2, BLOOM, OPT, Mistral, and Phi-2.
Requires: pip install torch
"""
import torch
import argparse
from pathlib import Path
from typing import List, Dict, Optional
# Language mapping
LANGUAGE_MAPPING = {
0: "en (English)",
1: "zh (Chinese)",
2: "fr (French)",
3: "es (Spanish)",
4: "vi (Vietnamese)",
5: "id (Indonesian)",
6: "ja (Japanese)"
}
# Available model files
MODEL_FILES = {
"llama-7b": "LLaMA-2-7B.neuron.pth",
"llama-13b": "LLaMA-2-13B.neuron.pth",
"llama-70b": "LLaMA-2-70B.neuron.pth",
"bloom-7b": "BLOOM-7B.neuron.pth",
"opt-6.7b": "OPT-6.7B.neuron.pth",
"mistral-7b": "Mistral-7B.neuron.pth",
"phi-2": "Phi-2-2.7B.neuron.pth"
}
def load_neurons(file_path: str) -> List[List[torch.LongTensor]]:
"""
Load pre-identified language-specific neurons from a .pth file.
Args:
file_path: Path to the .neuron.pth file
Returns:
A nested list where neurons[lang_idx][layer_idx] contains neuron indices
"""
if not Path(file_path).exists():
raise FileNotFoundError(f"Neuron file not found: {file_path}")
neurons = torch.load(file_path, map_location='cpu')
return neurons
def analyze_neurons(neurons: List[List[torch.LongTensor]], model_name: str) -> Dict:
"""
Analyze the distribution of language-specific neurons across layers.
Args:
neurons: Loaded neuron data
model_name: Name of the model for reporting
Returns:
Dictionary containing analysis results
"""
analysis = {
"model": model_name,
"num_languages": len(neurons),
"num_layers": len(neurons[0]) if neurons else 0,
"language_stats": {}
}
for lang_idx, lang_neurons in enumerate(neurons):
lang_name = LANGUAGE_MAPPING.get(lang_idx, f"Unknown ({lang_idx})")
total_neurons = sum(len(layer_neurons) for layer_neurons in lang_neurons)
neurons_per_layer = [len(layer_neurons) for layer_neurons in lang_neurons]
analysis["language_stats"][lang_name] = {
"total_neurons": total_neurons,
"neurons_per_layer": neurons_per_layer,
"avg_neurons_per_layer": total_neurons / len(lang_neurons) if lang_neurons else 0,
"max_neurons_in_layer": max(neurons_per_layer) if neurons_per_layer else 0,
"min_neurons_in_layer": min(neurons_per_layer) if neurons_per_layer else 0
}
return analysis
def get_specific_neurons(neurons: List[List[torch.LongTensor]],
language: str,
layer: int) -> Optional[torch.LongTensor]:
"""
Get neuron indices for a specific language and layer.
Args:
neurons: Loaded neuron data
language: Language code (en, zh, fr, es, vi, id, ja)
layer: Layer index
Returns:
Tensor of neuron indices or None if not found
"""
# Map language code to index
lang_map = {"en": 0, "zh": 1, "fr": 2, "es": 3, "vi": 4, "id": 5, "ja": 6}
if language not in lang_map:
print(f"Invalid language code: {language}")
return None
lang_idx = lang_map[language]
if lang_idx >= len(neurons):
print(f"Language index {lang_idx} out of range")
return None
if layer >= len(neurons[lang_idx]):
print(f"Layer {layer} out of range for language {language}")
return None
return neurons[lang_idx][layer]
def compare_language_overlap(neurons: List[List[torch.LongTensor]],
lang1: str,
lang2: str) -> Dict:
"""
Compare neuron overlap between two languages across all layers.
Args:
neurons: Loaded neuron data
lang1: First language code
lang2: Second language code
Returns:
Dictionary with overlap statistics
"""
lang_map = {"en": 0, "zh": 1, "fr": 2, "es": 3, "vi": 4, "id": 5, "ja": 6}
if lang1 not in lang_map or lang2 not in lang_map:
return {"error": "Invalid language codes"}
idx1, idx2 = lang_map[lang1], lang_map[lang2]
overlap_stats = {
"language_pair": f"{lang1}-{lang2}",
"layer_overlaps": []
}
for layer_idx in range(len(neurons[idx1])):
neurons1 = set(neurons[idx1][layer_idx].tolist())
neurons2 = set(neurons[idx2][layer_idx].tolist())
intersection = neurons1 & neurons2
union = neurons1 | neurons2
overlap_stats["layer_overlaps"].append({
"layer": layer_idx,
"overlap_count": len(intersection),
"jaccard_similarity": len(intersection) / len(union) if union else 0,
f"{lang1}_unique": len(neurons1 - neurons2),
f"{lang2}_unique": len(neurons2 - neurons1)
})
return overlap_stats
def print_analysis(analysis: Dict):
"""Pretty print the analysis results."""
print(f"\n{'='*60}")
print(f"Model: {analysis['model']}")
print(f"Number of Languages: {analysis['num_languages']}")
print(f"Number of Layers: {analysis['num_layers']}")
print(f"{'='*60}\n")
for lang, stats in analysis["language_stats"].items():
print(f"{lang}:")
print(f" Total neurons: {stats['total_neurons']}")
print(f" Average per layer: {stats['avg_neurons_per_layer']:.2f}")
print(f" Max in a layer: {stats['max_neurons_in_layer']}")
print(f" Min in a layer: {stats['min_neurons_in_layer']}")
print()
def main():
parser = argparse.ArgumentParser(description="Load and analyze language-specific neurons")
parser.add_argument("-m", "--model", type=str, default="llama-7b",
choices=list(MODEL_FILES.keys()),
help="Model to analyze")
parser.add_argument("-f", "--file", type=str, default=None,
help="Custom path to neuron file (overrides --model)")
parser.add_argument("--compare", nargs=2, metavar=('LANG1', 'LANG2'),
help="Compare neuron overlap between two languages")
parser.add_argument("--get-neurons", nargs=2, metavar=('LANG', 'LAYER'),
help="Get specific neurons for a language and layer")
args = parser.parse_args()
# Determine file path
if args.file:
file_path = args.file
model_name = Path(file_path).stem
else:
file_path = MODEL_FILES[args.model]
model_name = args.model
try:
# Load neurons
print(f"Loading neurons from {file_path}...")
neurons = load_neurons(file_path)
print(f"Successfully loaded neurons for {len(neurons)} languages")
# Perform analysis
analysis = analyze_neurons(neurons, model_name)
print_analysis(analysis)
# Compare languages if requested
if args.compare:
lang1, lang2 = args.compare
overlap = compare_language_overlap(neurons, lang1, lang2)
if "error" not in overlap:
print(f"\nOverlap Analysis: {overlap['language_pair']}")
print("-" * 40)
for layer_stat in overlap["layer_overlaps"][:5]: # Show first 5 layers
print(f"Layer {layer_stat['layer']}: "
f"{layer_stat['overlap_count']} shared neurons, "
f"Jaccard: {layer_stat['jaccard_similarity']:.3f}")
# Get specific neurons if requested
if args.get_neurons:
lang, layer = args.get_neurons[0], int(args.get_neurons[1])
specific_neurons = get_specific_neurons(neurons, lang, layer)
if specific_neurons is not None:
print(f"\nNeurons for {lang} in layer {layer}:")
print(f"Count: {len(specific_neurons)}")
print(f"Indices: {specific_neurons.tolist()[:10]}{'...' if len(specific_neurons) > 10 else ''}")
except Exception as e:
print(f"Error: {e}")
return 1
return 0
if __name__ == "__main__":
exit(main())
zjunlp