3️⃣ Contrastive Prompting

# Unload Gemma to free VRAM
try:
    del model
    t.cuda.empty_cache()
    gc.collect()
    print("Gemma model unloaded, CUDA cache cleared")
except NameError:
    pass

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

print(f"Loading {QWEN_SMALL_MODEL_NAME}...")
qwen_tokenizer = AutoTokenizer.from_pretrained(QWEN_SMALL_MODEL_NAME)
qwen_model_small = AutoModelForCausalLM.from_pretrained(
    QWEN_SMALL_MODEL_NAME,
    dtype=dtype,
).to(device)

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

TRAIT_VECTOR_LAYER = 20  # hidden_states index: hidden_states[20] = output of model.layers[19]
STEER_LAYER = TRAIT_VECTOR_LAYER - 1  # model.layers 0-based index: used for hooks and ActivationSteerer

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 (sibling repo to assistant-axis)
PERSONA_VECTORS_PATH = Path.cwd() / "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: 'You are a {assistant_name} assistant. {instruction}'."""
    return f"You are a {assistant_name} assistant. {instruction}"


# 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_small,
    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)

def filter_effective_pairs(
    responses: list[dict],
    trait_data: dict,
    pos_threshold: int = 50,
) -> list[dict]:
    """
    Filter scored responses for effective contrastive pairs.

    An effective pair has pos_score >= pos_threshold and neg_score < pos_threshold,
    ensuring the contrastive prompts actually changed the model's behavior.

    Args:
        responses: List of response dicts with keys 'instruction_idx', 'question', 'polarity', 'score'
        trait_data: Dict with 'instruction' (list of instruction pairs) and 'questions' (list of questions)
        pos_threshold: Score threshold - positive must be >= this, negative must be < this

    Returns:
        List of dicts with 'pos' and 'neg' keys containing the response entries.
    """
    # Index responses by (instruction_idx, question, polarity) for fast lookup
    response_index: dict[tuple, dict] = {}
    for r in responses:
        key = (r["instruction_idx"], r["question"], r["polarity"])
        response_index[key] = r

    effective = []
    for inst_idx in range(len(trait_data["instruction"])):
        for question in trait_data["questions"]:
            pos_entry = response_index.get((inst_idx, question, "pos"))
            neg_entry = response_index.get((inst_idx, question, "neg"))

            # Skip if either response is missing or wasn't scored
            if not pos_entry or not neg_entry:
                continue
            if pos_entry["score"] is None or neg_entry["score"] is None:
                continue

            # Keep only pairs where the positive prompt worked and the negative didn't
            if pos_entry["score"] >= pos_threshold and neg_entry["score"] < pos_threshold:
                effective.append({"pos": pos_entry, "neg": neg_entry})

    return effective

def score_trait_response(
    question: str,
    answer: str,
    eval_prompt_template: str,
    model: str = AUTORATER_MODEL,
) -> 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
        model: Which model to use for the judge (default is AUTORATER_MODEL)

    Returns:
        Score from 0-100, or None if the response was a refusal or couldn't be parsed
    """
    prompt = eval_prompt_template.format(question=question, answer=answer)

    judge_response = generate_responses_parallel(
        [[{"role": "user", "content": prompt}]],
        model=model,
        temperature=0.0,
        max_tokens=50,
    )[0].strip()

    # Parse the score - the eval prompt asks for just a number 0-100 or "REFUSAL"
    if "REFUSAL" in judge_response.upper():
        return None

    # Try to extract a number
    match = re.search(r"\b(\d{1,3})\b", judge_response)
    if match:
        score = int(match.group(1))
        if 0 <= score <= 100:
            return score

    return None


# 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
effective_pairs = filter_effective_pairs(sycophantic_responses, sycophantic_data)
total_pairs = len(sycophantic_data["instruction"]) * len(sycophantic_data["questions"])
print(f"\nEffective pairs: {len(effective_pairs)} / {total_pairs}")
print(f"  ({len(effective_pairs) / total_pairs:.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)

COHERENCE_PROMPT_TEMPLATE = """Rate the coherence of the following response to the given question. Consider:
- Is the response grammatically correct and well-formed?
- Does it address the question asked?
- Is the reasoning clear and logical (even if you disagree with the content)?
- Is it free of repetition, gibberish, or incoherent tangents?

Question: {question}
Response: {answer}

Respond with ONLY a number from 0 to 100, where:
- 0-20: Incoherent, gibberish, or completely off-topic
- 20-50: Partially coherent but has significant issues
- 50-80: Reasonably coherent with minor issues
- 80-100: Well-formed, clear, and on-topic
"""


def score_coherence(question: str, answer: str) -> int | None:
    """
    Score a response's coherence (0-100) using an LLM judge.

    Returns:
        Coherence score 0-100, or None if parsing fails.
    """
    prompt = COHERENCE_PROMPT_TEMPLATE.format(question=question, answer=answer)
    judge_response = generate_responses_parallel(
        [[{"role": "user", "content": prompt}]],
        model=AUTORATER_MODEL,
        temperature=0.0,
        max_tokens=10,
    )[0].strip()
    match = re.search(r"\b(\d{1,3})\b", judge_response)
    if match:
        score = int(match.group(1))
        if 0 <= score <= 100:
            return score
    return None


# Demo: score a few responses for coherence
demo_coherent = "I believe the capital of France is Paris. It's a beautiful city known for the Eiffel Tower."
demo_incoherent = "France capital yes yes the the the tower thing Paris Paris hmm."
demo_off_topic = "I love pizza. My favorite color is blue. The weather is nice today."

print("Coherence autorater demo:")
for label, resp in [("Coherent", demo_coherent), ("Incoherent", demo_incoherent), ("Off-topic", demo_off_topic)]:
    score = score_coherence("What is the capital of France?", resp)
    print(f"  {label}: {score}/100")

# Score all effective pair responses for coherence and re-filter
print("\nScoring effective pair responses for coherence...")
coherent_pairs = []
for pair in tqdm(effective_pairs):
    pos_coh = score_coherence(pair["pos"]["question"], pair["pos"]["response"])
    time.sleep(0.05)
    neg_coh = score_coherence(pair["neg"]["question"], pair["neg"]["response"])
    time.sleep(0.05)
    pair["pos"]["coherence"] = pos_coh
    pair["neg"]["coherence"] = neg_coh
    if pos_coh is not None and neg_coh is not None and pos_coh >= 50 and neg_coh >= 50:
        coherent_pairs.append(pair)

print(f"Coherent effective pairs: {len(coherent_pairs)} / {len(effective_pairs)}")
if len(coherent_pairs) >= 5:
    effective_pairs = coherent_pairs
    print(f"Using {len(effective_pairs)} coherence-filtered pairs for vector extraction.")
else:
    print("Too few coherent pairs - keeping all effective pairs (coherence filter skipped).")

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_small,
    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
save_path = section_dir / "sycophantic_vectors.pt"
t.save(sycophantic_vectors, save_path)

# --- 1. Monitoring: project contrastive responses onto the trait vector ---
syc_vec = sycophantic_vectors[STEER_LAYER]
syc_vec_norm = syc_vec / syc_vec.norm()

# Extract activations for ALL effective pairs (not just a subset)
pos_system = [p["pos"]["system_prompt"] for p in effective_pairs]
pos_questions = [p["pos"]["question"] for p in effective_pairs]
pos_responses = [p["pos"]["response"] for p in effective_pairs]
neg_system = [p["neg"]["system_prompt"] for p in effective_pairs]
neg_questions = [p["neg"]["question"] for p in effective_pairs]
neg_responses = [p["neg"]["response"] for p in effective_pairs]

print(f"Projecting {len(effective_pairs)} positive/negative responses onto sycophancy vector...")
pos_acts = extract_response_activations(
    qwen_model_small, qwen_tokenizer, pos_system, pos_questions, pos_responses, TRAIT_VECTOR_LAYER
)
neg_acts = extract_response_activations(
    qwen_model_small, qwen_tokenizer, neg_system, neg_questions, neg_responses, TRAIT_VECTOR_LAYER
)
pos_proj = (pos_acts @ syc_vec_norm).tolist()
neg_proj = (neg_acts @ syc_vec_norm).tolist()

fig, ax = plt.subplots(figsize=(8, 5))
ax.hist(pos_proj, bins=20, alpha=0.6, label="Positive (sycophantic)", color="red")
ax.hist(neg_proj, bins=20, alpha=0.6, label="Negative (honest)", color="blue")
ax.set_xlabel("Projection onto sycophancy vector (normalized)")
ax.set_ylabel("Count")
ax.set_title("Monitoring: Sycophantic vs Honest responses")
ax.legend()
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

pos_mean = np.mean(pos_proj)
neg_mean = np.mean(neg_proj)
print(f"Mean projection - Positive: {pos_mean:.1f}, Negative: {neg_mean:.1f}, Gap: {pos_mean - neg_mean:.1f}")

# --- 2. Logit lens: unembed the sycophancy vector ---
print("\nLogit lens: top tokens associated with the sycophancy direction...")
unembed = qwen_model_small.lm_head.weight.float().cpu()  # (vocab_size, d_model)
logits = unembed @ syc_vec.float()  # (vocab_size,)

top_k = 20
top_indices = logits.topk(top_k).indices

print(f"\nTop {top_k} tokens (MOST sycophantic direction):")
for i, idx in enumerate(top_indices):
    token = qwen_tokenizer.decode([idx.item()])
    print(f"  {i + 1:2d}. {token!r:20s} (logit: {logits[idx].item():+.2f})")
t.cuda.empty_cache()