context_cite/utils.py

import nltk
import numpy as np
import pandas as pd
import torch as ch
from numpy.typing import NDArray
from spacy.lang.en import English
from tqdm.auto import tqdm
from typing import Any, List, Optional, Tuple
from datasets import Dataset
from torch.utils.data import DataLoader
from transformers import DataCollatorForSeq2Seq



def split_text(text: str, split_by: str) -> Tuple[List[str], List[str], List[str]]:
    """Split response into parts and return the parts, start indices, and separators."""
    parts = []
    separators = []
    start_indices = []

    for line in text.splitlines():
        if split_by == "sentence":
            parts.extend(nltk.sent_tokenize(line))
        elif split_by == "word":
            tokenizer = English().tokenizer
            parts = [token.text for token in tokenizer(text)]
        else:
            raise ValueError(f"Cannot split response by '{split_by}'")

    cur_start = 0
    for part in parts:
        cur_end = text.find(part, cur_start)
        separator = text[cur_start:cur_end]
        separators.append(separator)
        start_indices.append(cur_end)
        cur_start = cur_end + len(part)

    return parts, separators, start_indices


def highlight_word_indices(words, indices, separators, color: bool):
    formatted_words = []

    # ANSI escape code for red color
    if color:
        RED = "\033[36m"  # ANSI escape code for light gray
        RESET = "\033[0m"  # Reset color to default
    else:
        RED = ""
        RESET = ""

    for word, idx in zip(words, indices):
        # Wrap index with red color
        formatted_words.append(f"{RED}[{idx}]{RESET}{word}")

    result = "".join(sep + word for sep, word in zip(separators, formatted_words))
    return result


def _create_mask(num_sources, alpha, seed):
    random = np.random.RandomState(seed)
    p = [1 - alpha, alpha]
    return random.choice([False, True], size=num_sources, p=p)


def _create_regression_dataset(
    num_masks, num_sources, get_prompt_ids, response_ids, alpha, base_seed=0
):
    masks = np.zeros((num_masks, num_sources), dtype=bool)
    data_dict = {
        "input_ids": [],
        "attention_mask": [],
        "labels": [],
    }
    for seed in range(num_masks):
        mask = _create_mask(num_sources, alpha, seed + base_seed)
        masks[seed] = mask
        prompt_ids = get_prompt_ids(mask=mask)
        input_ids = prompt_ids + response_ids
        data_dict["input_ids"].append(input_ids)
        data_dict["attention_mask"].append([1] * len(input_ids))
        data_dict["labels"].append([-100] * len(prompt_ids) + response_ids)
    return masks, Dataset.from_dict(data_dict)


def _compute_logit_probs(logits, labels):
    batch_size, seq_length = labels.shape
    # [num_tokens x vocab_size]
    reshaped_logits = logits.reshape(batch_size * seq_length, -1)
    reshaped_labels = labels.reshape(batch_size * seq_length)
    correct_logits = reshaped_logits.gather(-1, reshaped_labels[:, None])[:, 0]
    cloned_logits = reshaped_logits.clone()
    cloned_logits.scatter_(-1, reshaped_labels[:, None], -ch.inf)
    other_logits = cloned_logits.logsumexp(dim=-1)
    reshaped_outputs = correct_logits - other_logits
    return reshaped_outputs.reshape(batch_size, seq_length)


def _make_loader(dataset, tokenizer, batch_size):
    collate_fn = DataCollatorForSeq2Seq(tokenizer=tokenizer, padding="longest")
    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        collate_fn=collate_fn,
    )
    return loader


def _get_response_logit_probs(dataset, model, tokenizer, response_length, batch_size):
    if batch_size > 1:
        assert tokenizer.padding_side == "left", "Tokenizer must use left padding"
    loader = _make_loader(dataset, tokenizer, batch_size)
    logit_probs = ch.zeros((len(dataset), response_length), device=model.device)

    start_index = 0
    for batch in tqdm(loader):
        batch = {key: value.to(model.device) for key, value in batch.items()}
        with ch.no_grad(), ch.cuda.amp.autocast():
            output = model(**batch)
        logits = output.logits[:, -(response_length + 1) : -1]
        labels = batch["labels"][:, -response_length:]
        batch_size, _ = labels.shape
        cur_logit_probs = _compute_logit_probs(logits, labels)
        logit_probs[start_index : start_index + batch_size] = cur_logit_probs
        start_index += batch_size

    return logit_probs.cpu().numpy()


def get_masks_and_logit_probs(
    model,
    tokenizer,
    num_masks,
    num_sources,
    get_prompt_ids,
    response_ids,
    ablation_keep_prob,
    batch_size,
    base_seed=0,
):
    masks, dataset = _create_regression_dataset(
        num_masks,
        num_sources,
        get_prompt_ids,
        response_ids,
        ablation_keep_prob,
        base_seed=base_seed,
    )
    logit_probs = _get_response_logit_probs(
        dataset, model, tokenizer, len(response_ids), batch_size
    )
    return masks, logit_probs.astype(np.float32)


def aggregate_logit_probs(logit_probs, output_type="logit_prob"):
    """Compute sequence-level outputs from token-level logit-probabilities."""
    logit_probs = ch.tensor(logit_probs)
    log_probs = ch.nn.functional.logsigmoid(logit_probs).sum(dim=1)
    if output_type == "log_prob":
        return log_probs.numpy()
    elif output_type == "logit_prob":
        log_1mprobs = ch.log1p(-ch.exp(log_probs))
        return (log_probs - log_1mprobs).numpy()
    elif output_type == "total_token_logit_prob":
        return logit_probs.mean(dim=1).numpy()
    else:
        raise ValueError(f"Cannot aggregate log probs for output type '{output_type}'")


def _color_scale(val, max_val):
    start_color = (255, 255, 255)
    end_color = (80, 180, 80)
    if val == 0:
        return f"background-color: rgb{start_color}"
    elif val == max_val:
        return f"background-color: rgb{end_color}"
    else:
        fraction = val / max_val
        interpolated_color = tuple(
            start_color[i] + (end_color[i] - start_color[i]) * fraction
            for i in range(3)
        )
        return f"background-color: rgb{interpolated_color}"


def _apply_color_scale(df):
    # A score of np.log(10) means that the ablating this sources causes the
    # logit-probability to drop by np.log(10), which (roughly) corresponds to
    # a decrease in probability of 10x.
    max_val = max([df["Score"].max(), np.log(10)])
    return df.style.applymap(lambda val: _color_scale(val, max_val), subset=["Score"])


def get_attributions_df(
    attributions: NDArray[Any],
    context_partitioner,
    top_k: Optional[int] = None,
) -> Any:
    order = attributions.argsort()[::-1]
    selected_attributions = []
    selected_sources = []

    if top_k is not None:
        order = order[:top_k]

    for i in order:
        selected_attributions.append(attributions[i])
        selected_sources.append(context_partitioner.get_source(i))

    df = pd.DataFrame.from_dict(
        {"Score": selected_attributions, "Source": selected_sources}
    )
    df = _apply_color_scale(df).format(precision=3)
    return df


# The Llama 3 char_to_token is buggy (start and end chars for a given token
# are often the same), so we implement our own
def char_to_token(output_tokens, char_index):
    for i in range(len(output_tokens["input_ids"]) - 1):
        if char_index < output_tokens.token_to_chars(i + 1).start:
            return i
    return i + 1