2️⃣ Implementing oracle components

def collect_activations_multiple_layers(
    model: AutoModelForCausalLM,
    submodules: dict[int, torch.nn.Module],
    inputs_BL: dict[str, Int[Tensor, "batch seq"]],
    start_offset: int | None,
    end_offset: int | None,
) -> dict[int, Float[Tensor, "batch seq d_model"]]:
    """
    Collect activations from multiple layers using forward hooks.

    Args:
        model: The target model
        submodules: Dict mapping layer number to submodule to hook
        inputs_BL: Tokenized inputs (input_ids, attention_mask)
        start_offset: Start of the token slice (negative index from end). Only used when `end_offset`
            is also non-None; if `end_offset` is None, this must also be None (no slicing is applied).
        end_offset: End of the token slice (negative index from end, exclusive). Set both `start_offset`
            and `end_offset` to non-None values to enable token-position slicing; if both are None,
            the full sequence activations are returned.

    Returns:
        Dict mapping layer → activations tensor [batch, length, d_model]
    """
    if end_offset is not None:
        assert start_offset is not None
        assert start_offset < end_offset
        assert end_offset < 0
        assert start_offset < 0
    else:
        assert start_offset is None

    activations_BLD_by_layer = {}
    module_to_layer = {submodule: layer for layer, submodule in submodules.items()}
    max_layer = max(submodules.keys())

    def gather_target_act_hook(module, inputs, outputs):
        layer = module_to_layer[module]
        # Handle different output formats
        if isinstance(outputs, tuple):
            activations_BLD_by_layer[layer] = outputs[0]
        else:
            activations_BLD_by_layer[layer] = outputs

        # Slice if requested
        if end_offset is not None:
            activations_BLD_by_layer[layer] = activations_BLD_by_layer[layer][:, start_offset:end_offset, :]

        # Early stop after max layer
        if layer == max_layer:
            raise EarlyStopException("Early stopping after capturing activations")

    # Register hooks
    handles = []
    for layer, submodule in submodules.items():
        handles.append(submodule.register_forward_hook(gather_target_act_hook))

    try:
        with torch.no_grad():
            _ = model(**inputs_BL)
    except EarlyStopException:
        pass  # Expected
    except Exception as e:
        print(f"Unexpected error during forward pass: {str(e)}")
        raise
    finally:
        # Clean up hooks
        for handle in handles:
            handle.remove()

    return activations_BLD_by_layer


# Test the function
test_prompt = "The capital of France is"
test_inputs = tokenizer(test_prompt, return_tensors="pt", add_special_tokens=False).to(device)

# Extract from layer 18 (50% of 36 layers)
layer = layer_fraction_to_layer(MODEL_NAME, 0.5)
submodules = {layer: get_hf_submodule(model, layer)}

activations = collect_activations_multiple_layers(
    model=model,
    submodules=submodules,
    inputs_BL=test_inputs,
    start_offset=None,
    end_offset=None,
)

print(f"Extracted activations from layer {layer}")
print(f"Shape: {activations[layer].shape}")  # Should be [1, seq_len, d_model]

tests.test_collect_activations_multiple_layers(collect_activations_multiple_layers, model, tokenizer, device)

def find_pattern_in_tokens(
    token_ids: list[int],
    special_token_str: str,
    num_positions: int,
    tokenizer: AutoTokenizer,
) -> list[int]:
    """
    Find positions of special token in tokenized sequence.

    Args:
        token_ids: List of token IDs
        special_token_str: The special token string (e.g., " ?")
        num_positions: Expected number of occurrences
        tokenizer: Tokenizer to encode special token

    Returns:
        List of positions where special token appears
    """
    special_token_id = tokenizer.encode(special_token_str, add_special_tokens=False)
    assert len(special_token_id) == 1, f"Expected single token, got {len(special_token_id)}"
    special_token_id = special_token_id[0]

    # Find ALL positions where this token appears (don't stop early)
    positions = [i for i, tid in enumerate(token_ids) if tid == special_token_id]

    if len(positions) != num_positions:
        raise ValueError(
            f"Expected {num_positions} occurrences of special token, but found {len(positions)}. "
            f"This can happen if your oracle prompt contains a literal '?' character."
        )
    assert positions[-1] - positions[0] == num_positions - 1, f"Positions are not consecutive: {positions}"

    return positions


# Test the function
test_text = "Layer: 18\n ? ? ? \nWhat is this?"
test_tokens = tokenizer.encode(test_text, add_special_tokens=False)
positions = find_pattern_in_tokens(test_tokens, SPECIAL_TOKEN, 3, tokenizer)
print(f"Found ? tokens at positions: {positions}")

tests.test_find_pattern_in_tokens(find_pattern_in_tokens, tokenizer)

@contextlib.contextmanager
def add_hook(module: torch.nn.Module, hook: Callable):
    """Temporarily adds a forward hook to a model module."""
    handle = module.register_forward_hook(hook)
    try:
        yield
    finally:
        handle.remove()


def get_hf_activation_steering_hook(
    vectors: Float[Tensor, "num_pos d_model"],
    positions: list[int],
    steering_coefficient: float,
    device: torch.device,
    dtype: torch.dtype,
) -> Callable:
    """
    Create hook that injects activations at specified positions (assumes batch_size=1).

    Args:
        vectors: Steering vectors [K, d_model] where K is number of positions
        positions: List of positions to inject at
        steering_coefficient: Multiplier for steering strength
        device: Device for tensors
        dtype: Data type for steering

    Returns:
        Hook function that modifies activations during forward pass
    """
    # Normalize vectors to unit norm
    normed_vectors = torch.nn.functional.normalize(vectors, dim=-1).detach()
    positions_tensor = torch.tensor(positions, dtype=torch.long, device=device)

    def hook_fn(module, _input, output):
        # Extract residual stream tensor
        if isinstance(output, tuple):
            resid_BLD, *rest = output
            output_is_tuple = True
        else:
            resid_BLD = output
            output_is_tuple = False

        B, L, d_model = resid_BLD.shape

        if B != 1:
            raise ValueError(f"Expected batch_size=1, got B={B}")

        if L <= 1:
            return (resid_BLD, *rest) if output_is_tuple else resid_BLD

        # Check positions are valid
        assert positions_tensor.min() >= 0
        assert positions_tensor.max() < L, f"Position {positions_tensor.max()} >= sequence length {L}"

        # Get original activations at steering positions
        orig_KD = resid_BLD[0, positions_tensor, :]  # [K, d_model]
        norms_K1 = orig_KD.norm(dim=-1, keepdim=True)  # [K, 1]

        # Scale normalized steering vectors by original magnitudes
        steered_KD = (normed_vectors * norms_K1 * steering_coefficient).to(dtype)

        # Inject (add to original)
        resid_BLD[0, positions_tensor, :] = steered_KD.detach() + orig_KD

        return (resid_BLD, *rest) if output_is_tuple else resid_BLD

    return hook_fn


# Test the function
# Create dummy data (batch_size=1)
test_positions = [5, 6, 7]  # Inject at positions 5, 6, 7
test_vectors = torch.randn(len(test_positions), model.config.hidden_size, device=device)

hook_fn = get_hf_activation_steering_hook(
    vectors=test_vectors,
    positions=test_positions,
    steering_coefficient=1.0,
    device=device,
    dtype=dtype,
)

# Create dummy activations
dummy_resid = torch.randn(1, 20, model.config.hidden_size, device=device)
orig_values = dummy_resid[0, test_positions, :].clone()

# Apply hook
modified_resid = hook_fn(None, None, dummy_resid)

# Check modifications occurred
new_values = modified_resid[0, test_positions[0], :]
assert not torch.allclose(orig_values, new_values), "Hook should modify activations"
print("Steering hook test passed!")

tests.test_get_hf_activation_steering_hook(get_hf_activation_steering_hook, device, model.config.hidden_size)
tests.test_get_hf_activation_steering_hook_matches_reference(
    get_hf_activation_steering_hook, device, model.config.hidden_size
)

def create_oracle_input(
    prompt: str,
    layer: int,
    num_positions: int,
    tokenizer: AutoTokenizer,
    acts_BD: Float[Tensor, "num_pos d_model"],
) -> OracleInput:
    """
    Create an oracle input for inference.

    Args:
        prompt: Question to ask the oracle
        layer: Layer the activations came from
        num_positions: Number of ? tokens (equals length of acts_BD)
        tokenizer: Tokenizer
        acts_BD: Activation vectors [num_positions, d_model]

    Returns:
        OracleInput ready for generation
    """
    # Add introspection prefix with ? tokens
    prefix = get_introspection_prefix(layer, num_positions)
    prompt = prefix + prompt
    input_messages = [{"role": "user", "content": prompt}]

    # Create prompt with generation template (ends with <|im_start|>assistant\n)
    input_prompt_ids = tokenizer.apply_chat_template(
        input_messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors=None,
        padding=False,
        enable_thinking=False,
    )

    # Find ? token positions in the prompt
    positions = find_pattern_in_tokens(input_prompt_ids, SPECIAL_TOKEN, num_positions, tokenizer)

    # Ensure activations are on CPU and detached
    acts_BD = acts_BD.cpu().clone().detach()

    return OracleInput(
        input_ids=input_prompt_ids,
        layer=layer,
        steering_vectors=acts_BD,
        positions=positions,
    )


# Test the function
test_activations = torch.randn(3, model.config.hidden_size)
datapoint = create_oracle_input(
    prompt="What is the model thinking about?",
    layer=18,
    num_positions=3,
    tokenizer=tokenizer,
    acts_BD=test_activations,
)

print(f"Created datapoint with {len(datapoint.input_ids)} tokens")
print(f"? tokens at positions: {datapoint.positions}")

tests.test_create_oracle_input(create_oracle_input, tokenizer, model.config.hidden_size)

def run_oracle(
    model: AutoModelForCausalLM,
    tokenizer: AutoTokenizer,
    target_prompt: str,
    oracle_prompt: str,
    layer_fraction: float = 0.5,
    device: torch.device = device,
) -> str:
    """
    Run oracle query from scratch using components we built.

    Args:
        model: Model with oracle LoRA loaded
        tokenizer: Tokenizer
        target_prompt: Prompt to analyze (already formatted with chat template)
        oracle_prompt: Question to ask about activations
        layer_fraction: Which layer to extract from (as fraction of total, 0.0-1.0)
        device: Device

    Returns:
        Oracle's response as string
    """
    # For oracle sampling
    generation_kwargs = {"do_sample": False, "temperature": 0.0, "max_new_tokens": 50}

    # Tokenize target prompt and extract non-padding positions
    inputs_BL = tokenizer(target_prompt, return_tensors="pt", add_special_tokens=False).to(device)
    model_name = model.config._name_or_path
    act_layer = layer_fraction_to_layer(model_name, layer_fraction)
    submodules = {act_layer: get_hf_submodule(model, act_layer)}

    seq_len = inputs_BL["input_ids"].shape[1]
    attn_mask = inputs_BL["attention_mask"][0]
    real_len = int(attn_mask.sum().item())
    left_pad = seq_len - real_len

    # Step 1: Collect activations from the target model
    model.set_adapter("default")
    acts_by_layer = collect_activations_multiple_layers(
        model=model,
        submodules=submodules,
        inputs_BL=inputs_BL,
        start_offset=None,
        end_offset=None,
    )

    # Extract activations for all non-padding positions
    num_positions = real_len
    acts_BD = acts_by_layer[act_layer][0, left_pad:, :]  # [num_positions, d_model]

    # Step 2: Create oracle input
    datapoint = create_oracle_input(
        prompt=oracle_prompt,
        layer=act_layer,
        num_positions=num_positions,
        tokenizer=tokenizer,
        acts_BD=acts_BD,
    )

    # Step 3: Build steering hook, switch to oracle, generate
    input_ids = torch.tensor([datapoint.input_ids], dtype=torch.long, device=device)
    attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
    steering_vectors = datapoint.steering_vectors.to(device)
    positions = datapoint.positions

    injection_layer = 1  # Inject at layer 1 (gives oracle max processing depth)
    injection_submodule = get_hf_submodule(model, injection_layer)

    hook_fn = get_hf_activation_steering_hook(
        vectors=steering_vectors,
        positions=positions,
        steering_coefficient=1.0,
        device=device,
        dtype=dtype,
    )

    model.set_adapter("oracle")

    with add_hook(injection_submodule, hook_fn):
        output_ids = model.generate(input_ids=input_ids, attention_mask=attention_mask, **generation_kwargs)

    # Decode response
    generated_tokens = output_ids[:, input_ids.shape[1] :]
    response = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]

    return response


# Test our implementation
target_prompt_dict = [{"role": "user", "content": "The capital of France is"}]
target_prompt = tokenizer.apply_chat_template(target_prompt_dict, tokenize=False, add_generation_prompt=True)
oracle_prompt = "What answer will the model give, as a single token?"

our_response = run_oracle(
    model=model,
    tokenizer=tokenizer,
    target_prompt=target_prompt,
    oracle_prompt=oracle_prompt,
    layer_fraction=0.5,
    device=device,
)

print(f"Our implementation response: {our_response!r}")

# Compare to library version
library_results = utils.run_oracle(
    model=model,
    tokenizer=tokenizer,
    device=device,
    target_prompt=target_prompt,
    target_lora_path=None,
    oracle_prompt=oracle_prompt,
    oracle_lora_path="oracle",
    oracle_input_type="full_seq",
)
library_response = library_results.full_sequence_responses[0]

print(f"Library response: {library_response!r}")
assert our_response.strip().lower() == library_response.strip().lower()