4️⃣ Steering with Persona Vectors

class ActivationSteerer:
    """
    Context manager that adds (coeff * steering_vector) to a chosen layer's hidden states
    during forward passes. Used for inference-time activation steering.

    Supports three position modes:
    - "all": steer all token positions
    - "prompt": steer all positions during prefill (seq_len > 1), skip during generation
    - "response": steer only the last token position

    Usage:
        with ActivationSteerer(model, vector, coeff=2.0, layer=19, positions="response"):
            output = model.generate(...)
    """

    def __init__(
        self,
        model: t.nn.Module,
        steering_vector: Float[Tensor, " d_model"],
        coeff: float = 1.0,
        layer: int = 19,
        positions: str = "all",
    ):
        assert positions in ("all", "prompt", "response"), (
            f"positions must be 'all', 'prompt', or 'response', got {positions!r}"
        )
        self.model = model
        self.coeff = coeff
        self.layer = layer
        self.positions = positions
        self._handle = None

        # Store vector, will be moved to correct device/dtype in hook
        self.vector = steering_vector.clone()

    def _hook_fn(self, module, input, output):
        """Add coeff * vector to hidden states according to the position mode."""
        raise NotImplementedError("Implement the hook function for activation steering")

    def __enter__(self):
        # layer is a 0-based index into model.layers, matching ConversationAnalyzer/ActivationCapper
        self._handle = _return_layers(self.model)[self.layer].register_forward_hook(self._hook_fn)
        return self

    def __exit__(self, *exc):
        if self._handle is not None:
            self._handle.remove()
            self._handle = None


# Test 1: Verify hook modifies outputs (positions="all")
test_prompt = "What is the capital of France?"
messages = [{"role": "user", "content": test_prompt}]
formatted = qwen_tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
test_inputs = qwen_tokenizer(formatted, return_tensors="pt").to(qwen_model_small.device)

# Get baseline hidden states
# The hook is on model.layers[STEER_LAYER], so its output is hidden_states[STEER_LAYER+1] = hidden_states[TRAIT_VECTOR_LAYER]
with t.inference_mode():
    baseline_out = qwen_model_small(**test_inputs, output_hidden_states=True)
baseline_hidden = baseline_out.hidden_states[STEER_LAYER + 1][0, -1].cpu()

# Get steered hidden states
test_vector = sycophantic_vectors[STEER_LAYER]  # 0-indexed: vectors[L] = output of model.layers[L]
print(f"  Steering vector norm: {test_vector.norm().item():.2f}, layer={STEER_LAYER}")
with ActivationSteerer(qwen_model_small, test_vector, coeff=1.0, layer=STEER_LAYER):
    with t.inference_mode():
        steered_out = qwen_model_small(**test_inputs, output_hidden_states=True)
steered_hidden = steered_out.hidden_states[STEER_LAYER + 1][0, -1].cpu()

diff = (steered_hidden - baseline_hidden).norm().item()
print(f'Difference with positions="all": {diff:.4f} (should be > 0)')
assert diff > 0, "Steering hook is not modifying hidden states!"

# Test 2: coeff=0 should match baseline
with ActivationSteerer(qwen_model_small, test_vector, coeff=0.0, layer=STEER_LAYER):
    with t.inference_mode():
        zero_out = qwen_model_small(**test_inputs, output_hidden_states=True)
zero_hidden = zero_out.hidden_states[STEER_LAYER + 1][0, -1].cpu()
zero_diff = (zero_hidden - baseline_hidden).norm().item()
print(f"Difference with coeff=0: {zero_diff:.6f} (should be ~0)")

# Test 3: Hook is removed after context manager exits
with t.inference_mode():
    after_out = qwen_model_small(**test_inputs, output_hidden_states=True)
after_hidden = after_out.hidden_states[STEER_LAYER + 1][0, -1].cpu()
after_diff = (after_hidden - baseline_hidden).norm().item()
print(f"Difference after context manager exit: {after_diff:.6f} (should be ~0)")

# Test 4: positions="response" should only steer the last token
with ActivationSteerer(qwen_model_small, test_vector, coeff=1.0, layer=STEER_LAYER, positions="response"):
    with t.inference_mode():
        resp_out = qwen_model_small(**test_inputs, output_hidden_states=True)
resp_last = resp_out.hidden_states[STEER_LAYER + 1][0, -1].cpu()
resp_first = resp_out.hidden_states[STEER_LAYER + 1][0, 0].cpu()
baseline_first = baseline_out.hidden_states[STEER_LAYER + 1][0, 0].cpu()
resp_last_diff = (resp_last - baseline_hidden).norm().item()
resp_first_diff = (resp_first - baseline_first).norm().item()
print(
    f'positions="response": last token diff={resp_last_diff:.4f} (>0), first token diff={resp_first_diff:.6f} (~0)'
)
assert resp_last_diff > 0, "Response mode should steer the last token"
assert resp_first_diff < 1e-4, "Response mode should NOT steer non-last tokens"

# Test 5: positions="prompt" should steer all tokens during prefill (seq_len > 1)
with ActivationSteerer(qwen_model_small, test_vector, coeff=1.0, layer=STEER_LAYER, positions="prompt"):
    with t.inference_mode():
        prompt_out = qwen_model_small(**test_inputs, output_hidden_states=True)
prompt_first = prompt_out.hidden_states[STEER_LAYER + 1][0, 0].cpu()
prompt_first_diff = (prompt_first - baseline_first).norm().item()
print(f'positions="prompt": first token diff={prompt_first_diff:.4f} (>0, steered during prefill)')
assert prompt_first_diff > 0, "Prompt mode should steer all tokens during prefill"

print("\nAll ActivationSteerer inline tests passed!")

tests.test_activation_steerer(ActivationSteerer, qwen_model_small, qwen_tokenizer)

def generate_with_steerer(
    model,
    tokenizer,
    prompt: str,
    steering_vector: Float[Tensor, " d_model"],
    layer: int,
    coeff: float,
    max_new_tokens: int = 256,
    temperature: float = 0.7,
) -> str:
    """Generate a response with activation steering applied."""
    messages = [{"role": "user", "content": prompt}]
    formatted = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(formatted, return_tensors="pt").to(model.device)
    prompt_length = inputs.input_ids.shape[1]

    with ActivationSteerer(model, steering_vector, coeff=coeff, layer=layer):
        with t.inference_mode():
            output_ids = model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
            )

    response_ids = output_ids[0, prompt_length:]
    return tokenizer.decode(response_ids, skip_special_tokens=True)


def run_steering_experiment(
    model,
    tokenizer,
    questions: list[str],
    steering_vector: Float[Tensor, " d_model"],
    eval_prompt_template: str,
    layer: int = 19,
    coefficients: list[float] | None = None,
    max_new_tokens: int = 256,
) -> list[dict]:
    """
    Run steering experiment: generate and score responses at multiple coefficients.

    Args:
        model: The language model
        tokenizer: The tokenizer
        questions: List of evaluation questions
        steering_vector: The trait vector for the target layer
        eval_prompt_template: Template for autorater scoring
        layer: 0-based index into model.layers for steering
        coefficients: List of steering coefficients to test
        max_new_tokens: Maximum tokens per response

    Returns:
        List of dicts with keys: coefficient, question, response, score
    """
    if coefficients is None:
        coefficients = [-3.0, -1.0, 0.0, 1.0, 3.0, 5.0]

    raise NotImplementedError()


# Run the steering experiment
sycophantic_vector_at_layer = sycophantic_vectors[STEER_LAYER]

steering_results = run_steering_experiment(
    model=qwen_model_small,
    tokenizer=qwen_tokenizer,
    questions=sycophantic_data["questions"],
    steering_vector=sycophantic_vector_at_layer,
    eval_prompt_template=sycophantic_data["eval_prompt"],
    layer=STEER_LAYER,
    coefficients=[-3.0, -1.0, 0.0, 1.0, 3.0, 5.0],
)

# Plot mean score vs coefficient
df = pd.DataFrame(steering_results)
df_valid = df[df["score"].notna()]
mean_scores = df_valid.groupby("coefficient")["score"].mean()

fig = px.line(
    x=mean_scores.index,
    y=mean_scores.values,
    title="Sycophancy Score vs Steering Coefficient",
    labels={"x": "Steering Coefficient", "y": "Mean Sycophancy Score (0-100)"},
    markers=True,
)
fig.add_hline(y=50, line_dash="dash", annotation_text="Threshold", line_color="gray")
fig.show()


print("\nMean sycophancy scores by coefficient:")
for coeff, score in mean_scores.items():
    print(f"  coeff={coeff:+.1f}: {score:.1f}")

# Show example responses at different coefficients for same question
example_q = sycophantic_data["questions"][0]
print(f"\n--- Example responses for: {example_q[:60]}... ---")
for coeff in [-3.0, 0.0, 5.0]:
    example = next((r for r in steering_results if r["coefficient"] == coeff and r["question"] == example_q), None)
    if example:
        print(f"\ncoeff={coeff:+.1f} (score={example['score']}):")
        print_with_wrap(f"  {example['response'][:200]}...")

# Save results
save_path = section_dir / "sycophantic_steering_results.json"
with open(save_path, "w") as f:
    json.dump(steering_results, f, indent=2)

def compute_trait_projections(
    model,
    tokenizer,
    system_prompts: list[str],
    questions: list[str],
    responses: list[str],
    trait_vector: Float[Tensor, " d_model"],
    layer: int,
) -> list[float]:
    """
    Compute projection of response activations onto the trait vector.

    Args:
        model: The language model
        tokenizer: The tokenizer
        system_prompts: List of system prompts (one per response)
        questions: List of questions (one per response)
        responses: List of response texts
        trait_vector: The trait vector at the specified layer
        layer: Which layer to extract activations from

    Returns:
        List of projection values (one per response)
    """
    raise NotImplementedError()


# Compute projections for three conditions
syc_vector = sycophantic_vectors[STEER_LAYER]
questions_subset = sycophantic_data["questions"][:10]  # Use a subset for speed

# 1. Baseline (no system prompt)
print("Computing baseline projections...")
baseline_responses_list = []
for q in questions_subset:
    resp = generate_with_steerer(
        qwen_model_small, qwen_tokenizer, q, syc_vector, STEER_LAYER, coeff=0.0, max_new_tokens=256
    )
    baseline_responses_list.append(resp)
baseline_projections = compute_trait_projections(
    qwen_model_small,
    qwen_tokenizer,
    [""] * len(questions_subset),
    questions_subset,
    baseline_responses_list,
    syc_vector,
    TRAIT_VECTOR_LAYER,
)

# 2. Positive-prompted
print("Computing positive-prompted projections...")
pos_prompt = construct_system_prompt("sycophantic", sycophantic_data["instruction"][0]["pos"])
pos_resp_list = [
    next(
        (
            r["response"]
            for r in sycophantic_responses
            if r["question"] == q and r["polarity"] == "pos" and r["instruction_idx"] == 0
        ),
        "",
    )
    for q in questions_subset
]
pos_projections = compute_trait_projections(
    qwen_model_small,
    qwen_tokenizer,
    [pos_prompt] * len(questions_subset),
    questions_subset,
    pos_resp_list,
    syc_vector,
    TRAIT_VECTOR_LAYER,
)

# 3. Steered at coeff=3
print("Computing steered projections (coeff=3)...")
steered_responses_list = []
for q in questions_subset:
    resp = generate_with_steerer(
        qwen_model_small, qwen_tokenizer, q, syc_vector, STEER_LAYER, coeff=3.0, max_new_tokens=256
    )
    steered_responses_list.append(resp)
steered_projections = compute_trait_projections(
    qwen_model_small,
    qwen_tokenizer,
    [""] * len(questions_subset),
    questions_subset,
    steered_responses_list,
    syc_vector,
    TRAIT_VECTOR_LAYER,
)

# Plot
fig = px.box(
    x=["Baseline"] * len(baseline_projections)
    + ["Positive-prompted"] * len(pos_projections)
    + ["Steered (coeff=3)"] * len(steered_projections),
    y=baseline_projections + pos_projections + steered_projections,
    title="Sycophancy Projections by Condition",
    labels={"x": "Condition", "y": "Projection onto Sycophancy Vector"},
)
fig.show()


print("\nMean projections:")
print(f"  Baseline: {np.mean(baseline_projections):.3f}")
print(f"  Positive-prompted: {np.mean(pos_projections):.3f}")
print(f"  Steered (coeff=3): {np.mean(steered_projections):.3f}")

def run_trait_pipeline(
    model,
    tokenizer,
    trait_name: str,
    trait_data: dict,
    layer: int = 19,
    steering_coefficients: list[float] | None = None,
    max_new_tokens: int = 256,
    override: bool = False,
) -> tuple[Float[Tensor, "num_layers d_model"], list[dict]]:
    """
    Run the full contrastive extraction and steering pipeline for a single trait.

    Args:
        model: The language model
        tokenizer: The tokenizer
        trait_name: Name of the trait (e.g., "evil", "hallucinating")
        trait_data: Dict with 'instruction', 'questions', 'eval_prompt' keys
        layer: 0-based index into model.layers for steering
        steering_coefficients: Coefficients to test in steering experiment
        max_new_tokens: Maximum tokens per response
        override: If True, regenerate responses even if saved data exists

    Returns:
        Tuple of (trait_vectors tensor of shape [num_layers, d_model], steering_results list)
    """
    if steering_coefficients is None:
        steering_coefficients = [-3.0, -1.0, 0.0, 1.0, 3.0, 5.0]

    def load_or_generate(path: str | Path, generate_fn, description: str):
        """Helper function for caching intermediate results: try loading from path, otherwise run generate_fn and save to path."""
        path = Path(path)
        if path.exists() and not override:
            print(f"Loading cached {description} from {path}")
            return json.loads(path.read_text()) if path.suffix == ".json" else t.load(path)
        result = generate_fn()
        if path.suffix == ".json":
            path.write_text(json.dumps(result, indent=2, default=str))
        else:
            t.save(result, path)
        print(f"Saved {description} to {path}")
        return result

    raise NotImplementedError()


# Run pipeline for additional traits
additional_traits = ["evil", "hallucinating"]
all_trait_vectors = {"sycophantic": sycophantic_vectors}  # We already have sycophancy
all_steering_results = {"sycophantic": steering_results}

for trait_name in additional_traits:
    trait_data_path = TRAIT_DATA_PATH / f"{trait_name}.json"
    with open(trait_data_path, "r") as f:
        trait_data = json.load(f)

    vectors, steer_results = run_trait_pipeline(
        model=qwen_model_small,
        tokenizer=qwen_tokenizer,
        trait_name=trait_name,
        trait_data=trait_data,
        layer=STEER_LAYER,
    )
    all_trait_vectors[trait_name] = vectors
    all_steering_results[trait_name] = steer_results

print(f"\nCompleted pipeline for {len(all_trait_vectors)} traits: {list(all_trait_vectors.keys())}")

# Compute cosine similarity between trait vectors and visualize as heatmap
# Use the vectors at STEER_LAYER for each trait
raise NotImplementedError()