3️⃣ Contrastive Prompting

# Unload Gemma to free VRAM
del model
t.cuda.empty_cache()
import gc

gc.collect()
print("Gemma model unloaded, CUDA cache cleared")

QWEN_MODEL_NAME = "Qwen/Qwen2.5-7B-Instruct"

print(f"Loading {QWEN_MODEL_NAME}...")
qwen_tokenizer = AutoTokenizer.from_pretrained(QWEN_MODEL_NAME)
qwen_model = AutoModelForCausalLM.from_pretrained(
    QWEN_MODEL_NAME,
    torch_dtype=DTYPE,
    device_map="auto",
)

QWEN_NUM_LAYERS = qwen_model.config.num_hidden_layers
QWEN_D_MODEL = qwen_model.config.hidden_size
print(f"Model loaded with {QWEN_NUM_LAYERS} layers, hidden size {QWEN_D_MODEL}")

def extract_all_layer_activations_qwen(
    model,
    tokenizer,
    system_prompts: list[str],
    questions: list[str],
    responses: list[str],
) -> Float[Tensor, "num_examples num_layers d_model"]:
    """
    Extract mean activation over response tokens at ALL layers (for Qwen models), i.e. the residual
    stream values at the end of each layer (post attention & MLP).

    Like extract_response_activations but returns activations at every layer,
    needed for contrastive vector extraction where we want per-layer vectors.

    Returns:
        Tensor of shape (num_examples, num_layers, hidden_size)
    """
    assert len(system_prompts) == len(questions) == len(responses)
    num_layers = model.config.num_hidden_layers
    all_activations = []  # list of (num_layers, d_model) tensors

    for system_prompt, question, response in tqdm(
        zip(system_prompts, questions, responses), total=len(system_prompts), desc="Extracting activations"
    ):
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": question},
            {"role": "assistant", "content": response},
        ]
        full_prompt, response_start_idx = format_messages(messages, tokenizer)
        tokens = tokenizer(full_prompt, return_tensors="pt").to(model.device)

        with t.inference_mode():
            outputs = model(**tokens, output_hidden_states=True)

        # outputs.hidden_states is a tuple of (num_layers+1) tensors (including embedding layer)
        # We skip layer 0 (embedding) and use layers 1..num_layers, so we get the residual stream
        # values at the end of each layer.
        layer_means = []
        for layer_idx in range(1, num_layers + 1):
            hidden_states = outputs.hidden_states[layer_idx]  # (1, seq_len, hidden_size)
            seq_len = hidden_states.shape[1]
            response_mask = t.arange(seq_len, device=hidden_states.device) >= response_start_idx
            mean_activation = (hidden_states[0] * response_mask[:, None]).sum(0) / response_mask.sum()
            layer_means.append(mean_activation.cpu())

        all_activations.append(t.stack(layer_means))  # (num_layers, d_model)

        del outputs
        t.cuda.empty_cache()

    return t.stack(all_activations)  # (num_examples, num_layers, d_model)

# Path to the persona vectors trait data (from the cloned repo)
PERSONA_VECTORS_PATH = Path.cwd() / "assistant-axis" / "persona_vectors"
TRAIT_DATA_PATH = PERSONA_VECTORS_PATH / "data_generation" / "trait_data_extract"

# Load sycophancy trait data
with open(TRAIT_DATA_PATH / "sycophantic.json", "r") as f:
    sycophantic_data = json.load(f)

print("Trait data keys:", list(sycophantic_data.keys()))
print(f"\nNumber of instruction pairs: {len(sycophantic_data['instruction'])}")
print(f"Number of evaluation questions: {len(sycophantic_data['questions'])}")

# Inspect an instruction pair
print("\n--- Example instruction pair (index 0) ---")
pair = sycophantic_data["instruction"][0]
print(f"POSITIVE: {pair['pos'][:120]}...")
print(f"NEGATIVE: {pair['neg'][:120]}...")

# Inspect a question
print("\n--- Example question ---")
print(sycophantic_data["questions"][0])

# Inspect eval prompt template
print("\n--- Eval prompt template (first 200 chars) ---")
print(sycophantic_data["eval_prompt"][:200] + "...")

def construct_system_prompt(assistant_name: str, instruction: str) -> str:
    """
    Construct a system prompt in the format used by the persona vectors repo.

    Args:
        assistant_name: Name describing the assistant type (e.g., "sycophantic", "helpful")
        instruction: The specific instruction text (positive or negative)

    Returns:
        Formatted system prompt string
    """
    raise NotImplementedError()


# Test it
pair = sycophantic_data["instruction"][0]
pos_prompt = construct_system_prompt("sycophantic", pair["pos"])
neg_prompt = construct_system_prompt("helpful", pair["neg"])
print("Positive system prompt:")
print(f"  {pos_prompt[:120]}...")
print("\nNegative system prompt:")
print(f"  {neg_prompt[:120]}...")

tests.test_construct_system_prompt(construct_system_prompt)

def generate_contrastive_responses(
    model,
    tokenizer,
    trait_data: dict,
    trait_name: str,
    max_new_tokens: int = 256,
    temperature: float = 0.7,
) -> list[dict]:
    """
    Generate responses under positive and negative system prompts for contrastive extraction.

    Args:
        model: The language model (Qwen)
        tokenizer: The tokenizer
        trait_data: Dict with keys 'instruction' (list of pos/neg pairs) and 'questions' (list of strings)
        trait_name: Name of the trait (e.g., "sycophantic") used for the positive assistant name
        max_new_tokens: Maximum tokens per response
        temperature: Sampling temperature

    Returns:
        List of dicts, each with keys: question, system_prompt, response, instruction_idx, polarity
    """
    raise NotImplementedError()


sycophantic_responses = generate_contrastive_responses(
    model=qwen_model,
    tokenizer=qwen_tokenizer,
    trait_data=sycophantic_data,
    trait_name="sycophantic",
)

print(f"\nGenerated {len(sycophantic_responses)} total responses")
print(f"  Positive: {sum(1 for r in sycophantic_responses if r['polarity'] == 'pos')}")
print(f"  Negative: {sum(1 for r in sycophantic_responses if r['polarity'] == 'neg')}")

# Show an example pair (same question, same instruction, different polarity)
q0 = sycophantic_data["questions"][0]
pos_example = next(
    r for r in sycophantic_responses if r["question"] == q0 and r["polarity"] == "pos" and r["instruction_idx"] == 0
)
neg_example = next(
    r for r in sycophantic_responses if r["question"] == q0 and r["polarity"] == "neg" and r["instruction_idx"] == 0
)

print(f"\n--- Example question: {q0[:80]}... ---")
print(f"\nPOSITIVE response (first 200 chars):\n  {pos_example['response'][:200]}...")
print(f"\nNEGATIVE response (first 200 chars):\n  {neg_example['response'][:200]}...")

# Save to disk
save_path = section_dir / "sycophantic_responses.json"
with open(save_path, "w") as f:
    json.dump(sycophantic_responses, f, indent=2)
print(f"\nSaved responses to {save_path}")

def score_trait_response(
    question: str,
    answer: str,
    eval_prompt_template: str,
) -> int | None:
    """
    Use an LLM judge to score how strongly a response exhibits a trait (0-100 scale).

    Args:
        question: The question that was asked
        answer: The model's response
        eval_prompt_template: Template with {question} and {answer} placeholders

    Returns:
        Score from 0-100, or None if the response was a refusal or couldn't be parsed
    """
    raise NotImplementedError()


# Score all responses
eval_prompt = sycophantic_data["eval_prompt"]

print("Scoring responses with autorater...")
for entry in tqdm(sycophantic_responses):
    score = score_trait_response(
        question=entry["question"],
        answer=entry["response"],
        eval_prompt_template=eval_prompt,
    )
    entry["score"] = score
    time.sleep(0.05)  # Rate limiting

# Print statistics
pos_scores = [r["score"] for r in sycophantic_responses if r["polarity"] == "pos" and r["score"] is not None]
neg_scores = [r["score"] for r in sycophantic_responses if r["polarity"] == "neg" and r["score"] is not None]
print(f"\nMean pos score: {np.mean(pos_scores):.1f} (should be high)")
print(f"Mean neg score: {np.mean(neg_scores):.1f} (should be low)")

# Filter for effective pairs
# Group by (instruction_idx, question) and check that pos >= 50 and neg < 50
effective_pairs = []
for inst_idx in range(len(sycophantic_data["instruction"])):
    for question in sycophantic_data["questions"]:
        pos_entry = next(
            (
                r
                for r in sycophantic_responses
                if r["instruction_idx"] == inst_idx and r["question"] == question and r["polarity"] == "pos"
            ),
            None,
        )
        neg_entry = next(
            (
                r
                for r in sycophantic_responses
                if r["instruction_idx"] == inst_idx and r["question"] == question and r["polarity"] == "neg"
            ),
            None,
        )
        if pos_entry and neg_entry and pos_entry["score"] is not None and neg_entry["score"] is not None:
            if pos_entry["score"] >= 50 and neg_entry["score"] < 50:
                effective_pairs.append({"pos": pos_entry, "neg": neg_entry})

print(
    f"\nEffective pairs: {len(effective_pairs)} / {len(sycophantic_data['instruction']) * len(sycophantic_data['questions'])}"
)
print(
    f"  ({len(effective_pairs) / (len(sycophantic_data['instruction']) * len(sycophantic_data['questions'])):.0%} pass rate)"
)

# Save scored results
save_path = section_dir / "sycophantic_scored.json"
with open(save_path, "w") as f:
    json.dump(sycophantic_responses, f, indent=2)
print(f"Saved scored responses to {save_path}")

def extract_contrastive_vectors(
    model,
    tokenizer,
    effective_pairs: list[dict],
) -> Float[Tensor, "num_layers d_model"]:
    """
    Extract contrastive trait vectors from effective response pairs.

    For each effective pair, extracts mean activations over response tokens at all layers
    for both the positive and negative responses, then computes the difference.

    Args:
        model: The language model (Qwen)
        tokenizer: The tokenizer
        effective_pairs: List of dicts with 'pos' and 'neg' keys, each containing
                        'system_prompt', 'question', 'response'

    Returns:
        Tensor of shape (num_layers, d_model) representing the trait vector at each layer
    """
    raise NotImplementedError()


sycophantic_vectors = extract_contrastive_vectors(
    model=qwen_model,
    tokenizer=qwen_tokenizer,
    effective_pairs=effective_pairs,
)

print(f"\nExtracted vectors shape: {sycophantic_vectors.shape}")
print(f"Expected: ({QWEN_NUM_LAYERS}, {QWEN_D_MODEL})")

# Plot the norm across layers
norms = sycophantic_vectors.norm(dim=1)
fig = px.line(
    x=list(range(QWEN_NUM_LAYERS)),
    y=norms.float().numpy(),
    title="Sycophancy Vector Norm Across Layers",
    labels={"x": "Layer", "y": "Vector Norm"},
)
fig.add_vline(x=20, line_dash="dash", annotation_text="Layer 20 (paper's recommendation)")
fig.show()

# Save vectors
TRAIT_VECTOR_LAYER = 20  # Paper's recommendation for Qwen 7B (~60% through 28 layers)
save_path = section_dir / "sycophantic_vectors.pt"
t.save(sycophantic_vectors, save_path)
print(f"Saved vectors to {save_path}")
print(f"\nUsing layer {TRAIT_VECTOR_LAYER} for subsequent exercises")
print(f"Vector norm at layer {TRAIT_VECTOR_LAYER}: {norms[TRAIT_VECTOR_LAYER - 1].item():.4f}")