Add extract_vectors.py

extract_vectors.py

@@ -0,0 +1,317 @@
+#!/usr/bin/env python3
+"""Extract emotion vectors from Gemma4 model using PyTorch hooks on residual stream."""
+
+import json
+import os
+import warnings
+import numpy as np
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel
+from collections import defaultdict
+
+warnings.filterwarnings("ignore")
+
+EXP_DIR = os.path.dirname(os.path.abspath(__file__))
+STORIES_FILE = os.path.join(EXP_DIR, "emotion_stories.jsonl")
+OUTPUT_DIR = os.path.join(EXP_DIR, "results")
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+# Model config - use Gemma4 E4B from HuggingFace
+MODEL_ID = "google/gemma-4-E4B-it"  # Gemma4 E4B (not gated, ~8GB float16)
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+START_TOKEN = 50  # Start averaging from token 50 (emotion should be apparent)
+
+
+def load_stories():
+    stories = defaultdict(list)
+    with open(STORIES_FILE, "r") as f:
+        for line in f:
+            d = json.loads(line)
+            stories[d["emotion"]].append(d["text"])
+    return dict(stories)
+
+
+def get_residual_stream_hooks(model):
+    """Attach hooks to capture residual stream after each layer."""
+    activations = {}
+
+    def make_hook(name):
+        def hook(module, input, output):
+            # output is typically (hidden_states, ...) or just hidden_states
+            if isinstance(output, tuple):
+                hidden = output[0]
+            else:
+                hidden = output
+            activations[name] = hidden.detach().cpu().float()
+        return hook
+
+    hooks = []
+    # Gemma4 multimodal: text layers at model.model.language_model.layers
+    # Gemma3/standard: text layers at model.model.layers
+    if hasattr(model.model, 'language_model'):
+        layers = model.model.language_model.layers
+    else:
+        layers = model.model.layers
+    for i, layer in enumerate(layers):
+        h = layer.register_forward_hook(make_hook(f"layer_{i}"))
+        hooks.append(h)
+
+    return activations, hooks
+
+
+def extract_activations(model, tokenizer, text, activations_dict, target_layer):
+    """Run text through model and return mean activation at target layer."""
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to(DEVICE)
+
+    with torch.no_grad():
+        model(**inputs)
+
+    key = f"layer_{target_layer}"
+    if key not in activations_dict:
+        return None
+
+    hidden = activations_dict[key]  # (1, seq_len, hidden_dim)
+    seq_len = hidden.shape[1]
+
+    if seq_len <= START_TOKEN:
+        # Short text, use all tokens
+        mean_act = hidden[0].mean(dim=0)
+    else:
+        mean_act = hidden[0, START_TOKEN:].mean(dim=0)
+
+    # Clear activations for next run
+    activations_dict.clear()
+
+    return mean_act.numpy()
+
+
+def compute_emotion_vectors(emotion_activations):
+    """
+    emotion_vector[e] = mean(activations[e]) - mean(activations[all])
+    """
+    # Global mean across all emotions
+    all_acts = []
+    for acts in emotion_activations.values():
+        all_acts.extend(acts)
+    global_mean = np.mean(all_acts, axis=0)
+
+    emotion_vectors = {}
+    for emotion, acts in emotion_activations.items():
+        emotion_mean = np.mean(acts, axis=0)
+        emotion_vectors[emotion] = emotion_mean - global_mean
+
+    return emotion_vectors, global_mean
+
+
+def denoise_vectors(emotion_vectors, neutral_activations, variance_threshold=0.5):
+    """Project out top PCA components from neutral text."""
+    if len(neutral_activations) == 0:
+        return emotion_vectors
+
+    neutral_matrix = np.stack(neutral_activations)
+    neutral_centered = neutral_matrix - neutral_matrix.mean(axis=0)
+
+    # SVD
+    U, S, Vt = np.linalg.svd(neutral_centered, full_matrices=False)
+
+    # Find components explaining 50% variance
+    total_var = (S ** 2).sum()
+    cumvar = np.cumsum(S ** 2) / total_var
+    n_components = np.searchsorted(cumvar, variance_threshold) + 1
+    print(f"Denoising: projecting out {n_components} components (explain {variance_threshold*100}% variance)")
+
+    # Projection matrix
+    V_noise = Vt[:n_components].T  # (hidden_dim, n_components)
+
+    denoised = {}
+    for emotion, vec in emotion_vectors.items():
+        # Project out noise components
+        projection = V_noise @ (V_noise.T @ vec)
+        denoised[emotion] = vec - projection
+
+    return denoised
+
+
+def logit_lens(model, tokenizer, emotion_vectors, top_k=10):
+    """Project emotion vectors through unembedding to see associated tokens."""
+    # Get the lm_head / embed_tokens weight
+    if hasattr(model, 'lm_head'):
+        W = model.lm_head.weight.detach().cpu().float().numpy()  # (vocab, hidden)
+    elif hasattr(model.model, 'language_model'):
+        W = model.model.language_model.embed_tokens.weight.detach().cpu().float().numpy()
+    else:
+        W = model.model.embed_tokens.weight.detach().cpu().float().numpy()
+
+    results = {}
+    for emotion, vec in emotion_vectors.items():
+        # Logits = W @ vec
+        logits = W @ vec
+        top_indices = np.argsort(logits)[-top_k:][::-1]
+        bottom_indices = np.argsort(logits)[:top_k]
+
+        top_tokens = [(tokenizer.decode([idx]), float(logits[idx])) for idx in top_indices]
+        bottom_tokens = [(tokenizer.decode([idx]), float(logits[idx])) for idx in bottom_indices]
+
+        results[emotion] = {"top": top_tokens, "bottom": bottom_tokens}
+
+    return results
+
+
+def pca_analysis(emotion_vectors):
+    """PCA on emotion vectors to find valence/arousal structure."""
+    emotions = list(emotion_vectors.keys())
+    matrix = np.stack([emotion_vectors[e] for e in emotions])
+
+    # Center
+    matrix_centered = matrix - matrix.mean(axis=0)
+
+    # SVD
+    U, S, Vt = np.linalg.svd(matrix_centered, full_matrices=False)
+
+    # Project onto first 2 PCs
+    projections = matrix_centered @ Vt[:2].T  # (n_emotions, 2)
+
+    explained_variance = (S[:2] ** 2) / (S ** 2).sum()
+
+    return {
+        "emotions": emotions,
+        "pc1": projections[:, 0].tolist(),
+        "pc2": projections[:, 1].tolist(),
+        "explained_variance_pc1": float(explained_variance[0]),
+        "explained_variance_pc2": float(explained_variance[1]),
+    }
+
+
+def main():
+    print("=== Emotion Vector Extraction Experiment ===\n")
+
+    # Load stories
+    stories = load_stories()
+    print(f"Loaded {sum(len(v) for v in stories.values())} stories across {len(stories)} emotions\n")
+
+    # Load model
+    print(f"Loading model {MODEL_ID}...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    # Gemma4 is multimodal — AutoModelForCausalLM maps to ForConditionalGeneration
+    # which is correct, but we need to handle the nested language_model layers
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.bfloat16,  # Gemma4 native dtype
+        device_map="auto",
+    )
+    model.eval()
+
+    if hasattr(model.model, 'language_model'):
+        num_layers = len(model.model.language_model.layers)
+    else:
+        num_layers = len(model.model.layers)
+    target_layer = int(num_layers * 2 / 3)  # 2/3 depth
+    print(f"Model loaded. {num_layers} layers, target layer: {target_layer}\n")
+
+    # Attach hooks
+    activations_dict, hooks = get_residual_stream_hooks(model)
+
+    # Extract activations for each emotion
+    print("Extracting activations...")
+    emotion_activations = defaultdict(list)
+    total = sum(len(v) for v in stories.values())
+    done = 0
+
+    for emotion, story_list in stories.items():
+        for story in story_list:
+            act = extract_activations(model, tokenizer, story, activations_dict, target_layer)
+            if act is not None:
+                emotion_activations[emotion].append(act)
+            done += 1
+            if done % 50 == 0:
+                print(f"  [{done}/{total}]")
+
+    print(f"  Extracted activations for {len(emotion_activations)} emotions\n")
+
+    # Extract neutral activations for denoising
+    print("Extracting neutral activations for denoising...")
+    neutral_texts = [
+        "The meeting is scheduled for 3pm tomorrow.",
+        "Please find the attached document.",
+        "The temperature today is 22 degrees Celsius.",
+        "The project deadline has been moved to next Friday.",
+        "The store is located on the corner of Main Street.",
+        "Chapter 3 discusses the economic implications.",
+        "The software update includes several bug fixes.",
+        "The report contains data from the past quarter.",
+        "The committee will review the proposal next week.",
+        "The library opens at 9am on weekdays.",
+    ] * 5  # 50 neutral texts
+
+    neutral_activations = []
+    for text in neutral_texts:
+        act = extract_activations(model, tokenizer, text, activations_dict, target_layer)
+        if act is not None:
+            neutral_activations.append(act)
+    print(f"  {len(neutral_activations)} neutral activations collected\n")
+
+    # Compute emotion vectors
+    print("Computing emotion vectors...")
+    raw_vectors, global_mean = compute_emotion_vectors(dict(emotion_activations))
+    print(f"  {len(raw_vectors)} raw emotion vectors computed")
+
+    # Denoise
+    print("Denoising...")
+    denoised_vectors = denoise_vectors(raw_vectors, neutral_activations)
+
+    # Logit Lens
+    print("\nRunning Logit Lens analysis...")
+    logit_results = logit_lens(model, tokenizer, denoised_vectors)
+
+    print("\n=== Logit Lens Results ===")
+    for emotion in sorted(logit_results.keys()):
+        top = logit_results[emotion]["top"][:5]
+        bottom = logit_results[emotion]["bottom"][:5]
+        top_str = ", ".join([f"{t[0].strip()}({t[1]:.1f})" for t in top])
+        bottom_str = ", ".join([f"{t[0].strip()}({t[1]:.1f})" for t in bottom])
+        print(f"  {emotion:12s} ↑ {top_str}")
+        print(f"  {' ':12s} ↓ {bottom_str}")
+
+    # PCA analysis
+    print("\nRunning PCA analysis...")
+    pca_results = pca_analysis(denoised_vectors)
+    print(f"  PC1 explains {pca_results['explained_variance_pc1']*100:.1f}% variance")
+    print(f"  PC2 explains {pca_results['explained_variance_pc2']*100:.1f}% variance")
+
+    print("\n=== Emotion Space (PC1 vs PC2) ===")
+    for i, emotion in enumerate(pca_results["emotions"]):
+        pc1 = pca_results["pc1"][i]
+        pc2 = pca_results["pc2"][i]
+        print(f"  {emotion:12s}  PC1={pc1:+.3f}  PC2={pc2:+.3f}")
+
+    # Save results
+    results = {
+        "model": MODEL_ID,
+        "target_layer": target_layer,
+        "num_layers": num_layers,
+        "num_emotions": len(denoised_vectors),
+        "stories_per_emotion": {e: len(v) for e, v in stories.items()},
+        "logit_lens": logit_results,
+        "pca": pca_results,
+    }
+
+    results_file = os.path.join(OUTPUT_DIR, "experiment_results.json")
+    with open(results_file, "w") as f:
+        json.dump(results, f, indent=2, ensure_ascii=False)
+
+    # Save raw vectors as numpy
+    vectors_file = os.path.join(OUTPUT_DIR, "emotion_vectors.npz")
+    np.savez(vectors_file, **{e: v for e, v in denoised_vectors.items()})
+
+    print(f"\nResults saved to {results_file}")
+    print(f"Vectors saved to {vectors_file}")
+
+    # Cleanup hooks
+    for h in hooks:
+        h.remove()
+
+    print("\n=== EXPERIMENT COMPLETE ===")
+
+
+if __name__ == "__main__":
+    main()