utils.py

import json
import random
from pathlib import Path

import torch


def prepare_generation_prompts(raw_prompts, tokenizer, model, device, task):
    """
    Tokenize chat prompts for generation and activation extraction.

    The prompts are wrapped as single-turn user chat messages, formatted with
    the tokenizer's chat template, left-padded, and moved to `device`. Explicit
    `position_ids` are added so left padding does not shift token positions.

    For the `numbers` task, the function first greedily generates one token and
    appends it to the prompt tensors. This makes downstream probing extract the
    residual stream after the model has emitted the answer-prefix token expected
    by the arithmetic prompt format.

    Args:
        raw_prompts: Sequence of raw prompt strings.
        tokenizer: HuggingFace-compatible tokenizer with `apply_chat_template`.
        model: Model object exposing `generate`, used only for `task == "numbers"`.
        device: Torch device or device string for returned tensors.
        task: Task name. Currently `numbers` triggers one-token pre-generation;
            other tasks only tokenize the prompts.

    Returns:
        A tokenizer-output dictionary containing `input_ids`, `attention_mask`,
        and `position_ids`, with tensors on `device`.
    """
    prompts = [[{"role": "user", "content": prompt}] for prompt in raw_prompts]
    tokenizer.padding_side = "left"
    prompts = tokenizer.apply_chat_template(
        prompts,
        add_generation_prompt=True,
        return_tensors="pt",
        tokenize=True,
        padding=True,
    )
    prompts = {k: v.to(device) for k, v in prompts.items()}
    prompts["position_ids"] = prompts["attention_mask"].long().cumsum(-1) - 1
    prompts["position_ids"].masked_fill_(prompts["attention_mask"] == 0, 1)

    if task == "numbers":
        pre_generated = model.generate(
            prompts["input_ids"],
            attention_mask=prompts["attention_mask"],
            position_ids=prompts["position_ids"],
            max_new_tokens=1,
            temperature=0.0,
            do_sample=False,
            return_type="tokens",
            verbose=False,
        )
        prompts["input_ids"] = pre_generated
        prompts["attention_mask"] = torch.cat(
            [prompts["attention_mask"], torch.ones_like(pre_generated[:, -1:])], dim=1
        )
        prompts["position_ids"] = prompts["attention_mask"].long().cumsum(-1) - 1
        prompts["position_ids"].masked_fill_(prompts["attention_mask"] == 0, 1)

    return prompts


def load_rows(path: Path, limit: int | None = None) -> list[dict]:
    """Load rows, optionally taking a deterministic random subset."""

    with path.open("r", encoding="utf-8") as handle:
        rows = [json.loads(line) for line in handle if line.strip()]

    if limit is not None:
        rng = random.Random(0)
        rows = rng.sample(rows, k=min(limit, len(rows)))

    return rows


def n_layers_from_hook_dict(model) -> int:
    """Infer the number of transformer layers from available hook names."""

    layer_ids = set()
    for name in model.hook_dict.keys():
        parts = name.split(".")
        if len(parts) >= 2 and parts[0] == "blocks" and parts[1].isdigit():
            layer_ids.add(int(parts[1]))

    if not layer_ids:
        raise RuntimeError("Could not infer the number of layers from model.hook_dict.")

    return max(layer_ids) + 1


def prompt_tokens(model, prompt: str) -> torch.Tensor:
    """Tokenize one month prompt with the repository chat-template convention."""

    device = "cuda" if torch.cuda.is_available() else "cpu"
    with torch.no_grad():
        return prepare_generation_prompts(
            [prompt], model.tokenizer, model, device, task="months"
        )["input_ids"]


def tokenize_continuation(model, text: str) -> torch.Tensor:
    """Tokenize a continuation without adding special tokens."""

    with torch.no_grad():
        try:
            return model.to_tokens(text, prepend_bos=False)[0]
        except TypeError:
            encoded = model.tokenizer(
                text,
                return_tensors="pt",
                add_special_tokens=False,
            )
            return encoded["input_ids"][0]


def rendered_prompt_text(tokenizer, prompt: str) -> str:
    """Render a single user prompt with the tokenizer chat template."""

    return tokenizer.apply_chat_template(
        [[{"role": "user", "content": prompt}]],
        add_generation_prompt=True,
        tokenize=False,
    )[0]


def find_text_token_position(model, prompt: str, needle: str) -> int:
    """Find the first token position whose offset overlaps `needle` in a prompt."""

    tokenizer = model.tokenizer
    rendered = rendered_prompt_text(tokenizer, prompt)
    char_start = rendered.find(needle)
    if char_start == -1:
        raise ValueError(f"Could not find {needle!r} in rendered prompt: {rendered!r}")
    char_end = char_start + len(needle)

    try:
        encoded = tokenizer(
            rendered,
            return_tensors="pt",
            add_special_tokens=False,
            return_offsets_mapping=True,
        )
    except Exception:
        return find_token_subsequence_position(model, prompt, needle)

    for idx, (start, end) in enumerate(encoded["offset_mapping"][0].tolist()):
        if start < char_end and end > char_start:
            return idx

    raise ValueError(f"Could not map {needle!r} to a token position.")


def find_token_subsequence_position(model, prompt: str, needle: str) -> int:
    """Fallback token-subsequence search when tokenizer offsets are unavailable."""

    full = prompt_tokens(model, prompt)[0].tolist()
    candidates = [needle, " " + needle]
    for candidate in candidates:
        sub = tokenize_continuation(model, candidate).tolist()
        if not sub:
            continue
        for pos in range(0, len(full) - len(sub) + 1):
            if full[pos : pos + len(sub)] == sub:
                return pos

    raise ValueError(f"Could not find token subsequence for {needle!r} in {prompt!r}")


def rotate_existing_output(path: Path) -> None:
    """Move an existing output file to the first available numbered suffix."""

    if not path.exists():
        return

    suffix_idx = 1
    while True:
        rotated_path = path.with_name(f"{path.name}.{suffix_idx}")
        if not rotated_path.exists():
            path.rename(rotated_path)
            return
        suffix_idx += 1