pandora_llm.features.MinK

Module Contents

class pandora_llm.features.MinK.MinK(*args, **kwargs)

Bases: pandora_llm.features.base.FeatureComputer, pandora_llm.features.base.LLMHandler

Min-K computes token-level log probabilities for sequences and extracts the mean of the k% smallest log probabilities in each sequence. Introduced by Shi et al. 2024 (https://arxiv.org/pdf/2310.16789).

model

The pre-trained language model used for log probability computation.

Type:

AutoModelForCausalLM

compute_features(dataloader, accelerator)

Computes token-level log probabilities for a given dataloader.

Parameters:

• dataloader (torch.utils.data.DataLoader) – A DataLoader object containing the input sequences and optional attention masks.

• accelerator (accelerate.Accelerator) – An Accelerator object to handle distributed or mixed-precision computation.

Returns:

A tensor of shape (n, seq-1) containing log probabilities for each token in each sequence.

Raises:

Exception – If the model is not loaded prior to calling this method.

Return type:

jaxtyping.Float[torch.Tensor, n seq-1]

static reduce(token_log_probs, k)

Computes the average negative log probability of the k% rarest tokens in each sequence.

Parameters:

• token_log_probs (jaxtyping.Float[torch.Tensor, n seq-1]) – A tensor of log probabilities for each token in the sequence, of shape (n, seq-1), where: - n: number of sequences in the batch - seq-1: number of tokens in each sequence (excluding the first token)

• k (float) – A float between 0 and 1 indicating the proportion of the smallest values to select

• example (in each sequence. For)

• values. (k=0.2 selects the smallest 20% of)

Returns:

A tensor of shape (n,) containing the mean of the k% rarest negative log probabilities for each sequence in the batch.

Raises:

ValueError – If k is not in the range (0, 1].

Return type:

jaxtyping.Float[torch.Tensor, n]