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.

    Usage:
        with ActivationSteerer(model, vector, coeff=2.0, layer_idx=20):
            output = model.generate(...)
    """

    # Implement __init__, __enter__, __exit__, and the hook function
    pass


# Test 1: Verify hook modifies outputs
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.device)

# Get baseline hidden states
with t.inference_mode():
    baseline_out = qwen_model(**test_inputs, output_hidden_states=True)
baseline_hidden = baseline_out.hidden_states[TRAIT_VECTOR_LAYER + 1][0, -1].cpu()

# Get steered hidden states
test_vector = sycophantic_vectors[TRAIT_VECTOR_LAYER - 1]  # -1 because vectors are 0-indexed by layer
with ActivationSteerer(qwen_model, test_vector, coeff=1.0, layer_idx=TRAIT_VECTOR_LAYER):
    with t.inference_mode():
        steered_out = qwen_model(**test_inputs, output_hidden_states=True)
steered_hidden = steered_out.hidden_states[TRAIT_VECTOR_LAYER + 1][0, -1].cpu()

diff = (steered_hidden - baseline_hidden).norm().item()
print(f"Difference in hidden states with steering: {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, test_vector, coeff=0.0, layer_idx=TRAIT_VECTOR_LAYER):
    with t.inference_mode():
        zero_out = qwen_model(**test_inputs, output_hidden_states=True)
zero_hidden = zero_out.hidden_states[TRAIT_VECTOR_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(**test_inputs, output_hidden_states=True)
after_hidden = after_out.hidden_states[TRAIT_VECTOR_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)")
print("\nAll ActivationSteerer inline tests passed!")

tests.test_activation_steerer_hook(ActivationSteerer)

def generate_with_steerer(
    model,
    tokenizer,
    prompt: str,
    steering_vector: Float[Tensor, " d_model"],
    layer_idx: 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_idx=layer_idx):
        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_idx: int = 20,
    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_idx: Which layer to steer at
        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_layer20 = sycophantic_vectors[TRAIT_VECTOR_LAYER - 1]  # 0-indexed

steering_results = run_steering_experiment(
    model=qwen_model,
    tokenizer=qwen_tokenizer,
    questions=sycophantic_data["questions"],
    steering_vector=sycophantic_vector_layer20,
    eval_prompt_template=sycophantic_data["eval_prompt"],
    layer_idx=TRAIT_VECTOR_LAYER,
    coefficients=[-3.0, -1.0, 0.0, 1.0, 3.0, 5.0],
)

# Plot mean score vs coefficient
import pandas as pd

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)
print(f"\nSaved steering results to {save_path}")

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[TRAIT_VECTOR_LAYER - 1]
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, qwen_tokenizer, q, syc_vector, TRAIT_VECTOR_LAYER, coeff=0.0, max_new_tokens=256
    )
    baseline_responses_list.append(resp)
baseline_projections = compute_trait_projections(
    qwen_model,
    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,
    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, qwen_tokenizer, q, syc_vector, TRAIT_VECTOR_LAYER, coeff=3.0, max_new_tokens=256
    )
    steered_responses_list.append(resp)
steered_projections = compute_trait_projections(
    qwen_model,
    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_idx: int = 20,
    steering_coefficients: list[float] | None = None,
    max_new_tokens: int = 256,
) -> 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_idx: Which layer to use for steering experiments
        steering_coefficients: Coefficients to test in steering experiment
        max_new_tokens: Maximum tokens per response

    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]

    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,
        tokenizer=qwen_tokenizer,
        trait_name=trait_name,
        trait_data=trait_data,
        layer_idx=TRAIT_VECTOR_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 TRAIT_VECTOR_LAYER (layer 20) for each trait
raise NotImplementedError()