pandora_llm.features.LOSS

Module Contents

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

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

Computes the negative log-likelihood (NLL) for a given dataset using a pre-trained language model. Under strong assumptions, thresholding this is approximately optimal by the Neyman-Pearson lemma: MALT from Sablayrolles et al. 2019 (https://arxiv.org/pdf/1908.11229).

model

The pre-trained language model to compute the NLL.

Type:

AutoModelForCausalLM

compute_features(dataloader, accelerator)

Computes the negative log-likelihood (NLL) feature for the given dataloader.

Parameters:

• dataloader (torch.utils.data.DataLoader) – The dataloader providing input sequences.

• accelerator (accelerate.Accelerator) – The Accelerator object for distributed or mixed-precision training.

Returns:

The NLL feature for each sequence in the dataloader.

Raises:

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

Return type:

jaxtyping.Float[torch.Tensor, n]

pandora_llm.features.LOSS.compute_log_probs_dl(model, dataloader, accelerator, mode='mean')

Computes log probabilities for sequences in a dataloader.

This function processes a dataloader using a language model and an Accelerator for efficiency. It can return various forms of log probabilities based on the specified mode.

Parameters:

• model (transformers.AutoModelForCausalLM) – The pre-trained causal language model.

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

• accelerator (accelerate.Accelerator) – The Accelerator object to manage model and data parallelism.

• mode (str) – Specifies the type of output: - “mean”: Returns the mean of log probabilities (log-likelihood). - “tokens”: Returns log probabilities for each token in the sequence. - “all”: Returns log probabilities for the entire vocabulary for each token. - “tokens+all”: Returns target log probs stacked on the 0th index of the vocab dimension - “tokens+z”: Returns target log probs stacked with mean and std of all vocab tokens as the last dimension Default is “mean”.

Returns:

The computed log probabilities based on the selected mode.

Return type:

Union[jaxtyping.Float[torch.Tensor, n], jaxtyping.Float[torch.Tensor, n seq-1], jaxtyping.Float[torch.Tensor, n seq-1 vocab], jaxtyping.Float[torch.Tensor, n seq-1 vocab+1], jaxtyping.Float[torch.Tensor, n seq-1 3]]

pandora_llm.features.LOSS.compute_log_probs(model, input_ids, attention_mask=None, mode='mean')

Computes log probabilities for a batch of input sequences.

This function calculates the log probabilities of target tokens in a batch using a causal language model. It supports multiple output modes for different use cases.

Parameters:

• model (transformers.AutoModelForCausalLM) – The pre-trained causal language model.

• input_ids (jaxtyping.Integer[torch.Tensor, batch seq]) – Input token IDs of shape (batch, sequence length).

• attention_mask (jaxtyping.Bool[torch.Tensor, batch seq]) – Attention mask to ignore padding tokens. If None, a default mask will be generated from input_ids>0.

• mode (str) – Specifies the type of output: - “mean”: Returns the mean of log probabilities (log-likelihood) per sequence. - “tokens”: Returns log probabilities for each token in the sequence. - “all”: Returns log probabilities for all vocabulary tokens at each sequence position. - “tokens+all”: Returns target log probs stacked on the 0th index of the vocab dimension - “tokens+z”: Returns target log probs stacked with mean and std of all vocab tokens as the last dimension Default is “mean”.

Returns:

The computed log probabilities based on the selected mode.

Raises:

ValueError – If the mode is not one of “mean”, “tokens”, or “all”.

Return type:

Union[jaxtyping.Float[torch.Tensor, batch], jaxtyping.Float[torch.Tensor, batch seq-1], jaxtyping.Float[torch.Tensor, batch seq-1 vocab], jaxtyping.Float[torch.Tensor, batch seq-1 vocab+1], jaxtyping.Float[torch.Tensor, batch seq-1 3]]