# nnsight Tutorials

## Tutorial 1: Basic Activation Analysis

### Goal
Load a model, access internal activations, and analyze them.

### Step-by-Step

```python
from nnsight import LanguageModel
import torch

# 1. Load model
model = LanguageModel("openai-community/gpt2", device_map="auto")

# 2. Trace and collect activations
prompt = "The capital of France is"

with model.trace(prompt) as tracer:
    # Collect from multiple layers
    activations = {}
    for i in range(12):  # GPT-2 has 12 layers
        activations[i] = model.transformer.h[i].output[0].save()

    # Get final logits
    logits = model.output.save()

# 3. Analyze (outside context)
print("Layer-wise activation norms:")
for layer, act in activations.items():
    print(f"  Layer {layer}: {act.norm().item():.2f}")

# 4. Check predictions
probs = torch.softmax(logits[0, -1], dim=-1)
top_tokens = probs.topk(5)
print("\nTop predictions:")
for token_id, prob in zip(top_tokens.indices, top_tokens.values):
    token_str = model.tokenizer.decode(token_id)
    print(f"  {token_str!r}: {prob.item():.3f}")
```

---

## Tutorial 2: Activation Patching

### Goal
Patch activations from one prompt into another to test causal relationships.

### Step-by-Step

```python
from nnsight import LanguageModel
import torch

model = LanguageModel("gpt2", device_map="auto")

clean_prompt = "The Eiffel Tower is in the city of"
corrupted_prompt = "The Colosseum is in the city of"

# 1. Get clean activations
with model.trace(clean_prompt) as tracer:
    clean_hidden = model.transformer.h[8].output[0].save()
    clean_logits = model.output.save()

# 2. Define metric
paris_token = model.tokenizer.encode(" Paris")[0]
rome_token = model.tokenizer.encode(" Rome")[0]

def logit_diff(logits):
    return (logits[0, -1, paris_token] - logits[0, -1, rome_token]).item()

print(f"Clean logit diff: {logit_diff(clean_logits):.3f}")

# 3. Patch clean into corrupted
with model.trace(corrupted_prompt) as tracer:
    # Replace layer 8 output with clean activations
    model.transformer.h[8].output[0][:] = clean_hidden
    patched_logits = model.output.save()

print(f"Patched logit diff: {logit_diff(patched_logits):.3f}")

# 4. Systematic patching sweep
results = torch.zeros(12)  # 12 layers

for layer in range(12):
    # Get clean activation for this layer
    with model.trace(clean_prompt) as tracer:
        clean_act = model.transformer.h[layer].output[0].save()

    # Patch into corrupted
    with model.trace(corrupted_prompt) as tracer:
        model.transformer.h[layer].output[0][:] = clean_act
        logits = model.output.save()

    results[layer] = logit_diff(logits)
    print(f"Layer {layer}: {results[layer]:.3f}")

print(f"\nMost important layer: {results.argmax().item()}")
```

---

## Tutorial 3: Cross-Prompt Activation Sharing

### Goal
Transfer activations between different prompts in a single trace.

### Step-by-Step

```python
from nnsight import LanguageModel

model = LanguageModel("gpt2", device_map="auto")

with model.trace() as tracer:
    # First prompt - get "cat" representations
    with tracer.invoke("The cat sat on the mat"):
        cat_hidden = model.transformer.h[6].output[0].save()

    # Second prompt - inject "cat" into "dog"
    with tracer.invoke("The dog ran through the park"):
        # Replace with cat's activations
        model.transformer.h[6].output[0][:] = cat_hidden
        modified_logits = model.output.save()

# The dog prompt now has cat's internal representations
print(f"Modified logits shape: {modified_logits.shape}")
```

---

## Tutorial 4: Remote Execution with NDIF

### Goal
Run the same interpretability code on massive models (70B+).

### Step-by-Step

```python
from nnsight import LanguageModel
import os

# 1. Setup API key
os.environ["NDIF_API_KEY"] = "your_key_here"

# 2. Load large model (runs remotely)
model = LanguageModel("meta-llama/Llama-3.1-70B")

# 3. Same code, just remote=True
prompt = "The meaning of life is"

with model.trace(prompt, remote=True) as tracer:
    # Access layer 40 of 70B model!
    hidden = model.model.layers[40].output[0].save()
    logits = model.output.save()

# 4. Results returned from NDIF
print(f"Hidden shape: {hidden.shape}")
print(f"Logits shape: {logits.shape}")

# 5. Check predictions
import torch
probs = torch.softmax(logits[0, -1], dim=-1)
top_tokens = probs.topk(5)
print("\nTop predictions from Llama-70B:")
for token_id, prob in zip(top_tokens.indices, top_tokens.values):
    print(f"  {model.tokenizer.decode(token_id)!r}: {prob.item():.3f}")
```

### Batching with Sessions

```python
# Run multiple experiments in one NDIF request
with model.session(remote=True) as session:
    with model.trace("What is 2+2?"):
        math_hidden = model.model.layers[30].output[0].save()

    with model.trace("The capital of France is"):
        fact_hidden = model.model.layers[30].output[0].save()

# Compare representations
similarity = torch.cosine_similarity(
    math_hidden.mean(dim=1),
    fact_hidden.mean(dim=1),
    dim=-1
)
print(f"Similarity: {similarity.item():.3f}")
```

---

## Tutorial 5: Steering with Activation Addition

### Goal
Add a steering vector to change model behavior.

### Step-by-Step

```python
from nnsight import LanguageModel
import torch

model = LanguageModel("gpt2", device_map="auto")

# 1. Get contrasting activations
with model.trace("I love this movie, it's wonderful") as tracer:
    positive_hidden = model.transformer.h[6].output[0].save()

with model.trace("I hate this movie, it's terrible") as tracer:
    negative_hidden = model.transformer.h[6].output[0].save()

# 2. Compute steering direction
steering_vector = positive_hidden.mean(dim=1) - negative_hidden.mean(dim=1)

# 3. Generate without steering
test_prompt = "This restaurant is"
with model.trace(test_prompt) as tracer:
    normal_logits = model.output.save()

# 4. Generate with steering
with model.trace(test_prompt) as tracer:
    # Add steering at layer 6
    model.transformer.h[6].output[0][:] += 3.0 * steering_vector
    steered_logits = model.output.save()

# 5. Compare predictions
def top_prediction(logits):
    token = logits[0, -1].argmax()
    return model.tokenizer.decode(token)

print(f"Normal: {top_prediction(normal_logits)}")
print(f"Steered (positive): {top_prediction(steered_logits)}")
```

---

## Tutorial 6: Logit Lens

### Goal
See what the model "believes" at each layer.

### Step-by-Step

```python
from nnsight import LanguageModel
import torch

model = LanguageModel("gpt2", device_map="auto")

prompt = "The quick brown fox jumps over the lazy"

with model.trace(prompt) as tracer:
    # Collect residual stream at each layer
    residuals = []
    for i in range(12):
        resid = model.transformer.h[i].output[0].save()
        residuals.append(resid)

# Access model's unembedding and final layernorm
W_U = model._model.lm_head.weight.T  # [d_model, vocab]
ln_f = model._model.transformer.ln_f

print("Layer-by-layer predictions for final token:")
for i, resid in enumerate(residuals):
    # Apply final layernorm
    normed = ln_f(resid)

    # Project to vocabulary
    layer_logits = normed @ W_U

    # Get prediction
    probs = torch.softmax(layer_logits[0, -1], dim=-1)
    top_token = probs.argmax()
    top_prob = probs[top_token].item()

    print(f"Layer {i}: {model.tokenizer.decode(top_token)!r} ({top_prob:.3f})")
```

---

## External Resources

### Official Resources
- [Getting Started](https://nnsight.net/start/)
- [Features Overview](https://nnsight.net/features/)
- [Documentation](https://nnsight.net/documentation/)
- [Tutorials](https://nnsight.net/tutorials/)

### NDIF Resources
- [NDIF Homepage](https://ndif.us/)
- [Available Models](https://ndif.us/models)
- [API Key Signup](https://login.ndif.us/)

### Paper
- [NNsight and NDIF](https://arxiv.org/abs/2407.14561) - ICLR 2025

### Community
- [Discussion Forum](https://discuss.ndif.us/)
- [GitHub Issues](https://github.com/ndif-team/nnsight/issues)