"""Utilities for scoring sequences using Language Models."""

from typing import (

Iterable,

Union,

List,

Collection,

Optional,

Callable,

Tuple,

Any,

cast,

)

import re

import torch

import warnings

from collections import defaultdict

from itertools import chain

import numpy as np

from transformers import (

AutoModelForCausalLM,

AutoModelForMaskedLM,

AutoModelForSeq2SeqLM,

AutoTokenizer,

BatchEncoding,

AutoProcessor,

AutoModelForImageTextToText

)

from PIL import Image

try:

from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel

except:

pass

from transformers.utils.logging import set_verbosity_error

from .utils import batch_wise_logprobs, all_equal

set_verbosity_error()

class LMScorer:

"""

Base LM scorer class intended to store models and tokenizers along

with methods to facilitate the analysis of language model output scores.

"""

def __init__(

self,

model: Union[str, torch.nn.Module],

device: Optional[str] = "cpu",

tokenizer=None,

**kwargs,

) -> None:

"""

:param model: should be path to a model (.pt or .bin file) stored

locally, or name of a pretrained model stored on the Huggingface

Model Hub, or a model (torch.nn.Module) that have the same

signature as the corresponding Huggingface model (see the subclass

for details).

:param device: device type that the model should be loaded on,

options: `cpu or cuda:{0, 1, ...}`

:type device: str, optional

:param tokenizer: if provided, use this tokenizer.

"""

if tokenizer is not None:

if isinstance(tokenizer, str):

self.tokenizer = AutoTokenizer.from_pretrained(tokenizer, **kwargs)

else:

self.tokenizer = tokenizer

elif isinstance(model, str):

self.tokenizer = AutoTokenizer.from_pretrained(model, use_fast=True)

else:

raise Exception("Must provide either model name or tokenizer.")

self.device = device

self.vocab = defaultdict(list)

# {self.vocab[x.strip()].append(i) for x, i in [(self.tokenizer.decode([i]), i) for i in range(self.tokenizer.vocab_size)]}

for i in range(self.tokenizer.vocab_size):

decoded = [(self.tokenizer.decode(i), i)]

for x, j in decoded:

self.vocab[x.strip()].append(j)

def add_special_tokens(self, text: Union[str, Iterable[str]]) -> List[str]:

raise NotImplementedError

def distribution(self, batch: Iterable) -> torch.Tensor:

raise NotImplementedError

def topk(self, distribution: torch.Tensor, k: int = 1) -> Tuple:

top_k = distribution.topk(k)

probs = top_k.values.squeeze(1).exp().tolist()

if k == 1:

tokens = self.decode(top_k.indices.squeeze(1))

else:

tokens = [self.decode(x) for x in top_k.indices.squeeze(1)]

return tokens, probs

# def query(self, distribution: torch.Tensor, queries: List[str]) -> Tuple:

# # this will be self.vocab tho

# query_ids = [self.vocab[a] for a in queries]

# maxlen = max(map(len, query_ids))

# query_ids = [

# (

# q + [self.tokenizer.pad_token_id] * (maxlen - len(q))

# if len(q) < maxlen

# else q

# )

# for q in query_ids

# ]

# current_batch_size = distribution.shape[0]

# probs = (

# distribution[torch.arange(current_batch_size)[:, None], query_ids]

# .max(1)

# .values.exp()

# .tolist()

# )

# inv_ranks = distribution.argsort().argsort() + 1

# ranks = distribution.shape[1] - inv_ranks + 1

# token_ranks = (

# ranks[torch.arange(current_batch_size)[:, None], query_ids]

# .min(1)

# .values.tolist()

# )

# return probs, token_ranks

def query(

self,

distribution: torch.Tensor,

queries: List[List[str]],

prob: bool = True,

tolist: bool = True,

):

"""Queries distributions for (log)probabilities of target tokens."""

if isinstance(queries[0], str):

queries = [queries]

scores = []

token_ranks = []

for i, querylist in enumerate(queries):

query_ids = [self.vocab[a] for a in querylist]

maxlen = max(map(len, query_ids))

query_ids = [

(q + [self.pad_token_id] * (maxlen - len(q)) if len(q) < maxlen else q)

for q in query_ids

]

query_logprobs = distribution[i, query_ids].max(1).values

inv_ranks = distribution.argsort().argsort() + 1

ranks = distribution.shape[1] - inv_ranks + 1

query_token_ranks = ranks[i, query_ids].min(1).values.tolist()

scores.append(query_logprobs)

token_ranks.append(query_token_ranks)

if prob:

scores = [s.exp() for s in scores]

if tolist:

scores = [s.tolist() for s in scores]

return scores, token_ranks

def logprobs(

self, batch: Iterable, rank: bool = False

) -> Union[float, List[float]]:

"""

.. deprecated::

Use :meth:`compute_stats` instead.

"""

warnings.warn(

"logprobs is deprecated, use compute_stats instead", DeprecationWarning

)

raise NotImplementedError

def compute_stats(self, batch: Iterable, rank: bool = False) -> Union[

Tuple[List[float], List[int]],

List[float],

torch.Tensor,

Tuple[torch.Tensor, torch.Tensor],

]:

raise NotImplementedError

def prepare_text(self, text: Union[str, List[str], BatchEncoding]) -> Any:

raise NotImplementedError

def prime_text(

self,

prefix: Union[str, List[str]],

stimuli: Union[str, List[str]],

separator: str = " ",

) -> Tuple:

raise NotImplementedError

def token_score(

self,

batch: Union[str, List[str]],

surprisal: bool = False,

prob: bool = False,

base_two: bool = False,

rank: bool = False,

) -> Union[List[Tuple[str, float]], List[Tuple[str, float, int]]]:

"""

For every input sentence, returns a list of tuples in the following format:

`(token, score)`,

where score represents the log-probability (by default) of the token given context. Can also return ranks along with scores.

:param ``Union[str, List[str]]`` batch: a single sentence or a batch of sentences.

:param ``bool`` surprisal: If `True`, returns per-word surprisals instead of log-probabilities.

:param ``bool`` prob: If `True`, returns per-word probabilities instead of log-probabilities.

:param ``bool`` base_two: If `True`, uses log base 2 instead of natural-log (returns bits of values in case of surprisals)

:param ``bool`` rank: If `True`, also returns the rank of each word in context (based on the log-probability value)

:return: A `List` containing a `Tuple` consisting of the word, its associated score, and optionally, its rank.

:rtype: ``Union[List[Tuple[str, float]], List[Tuple[str, float, int]]]``

"""

raise NotImplementedError

def word_score(

self,

batch: Union[str, List[str]],

surprisal: bool = False,

prob: bool = False,

base_two: bool = False,

rank: bool = False,

bow_correction: bool = False,

) -> Union[List[Tuple[str, float]], List[Tuple[str, float, int]]]:

"""

Wraps token_score's outputs into word-level metrics:

`(word, score)`,

where score represents the log-probability (by default) of the word given context.

Token probabilities are summed across the whole word. Words are currently split on spaces and punctuation.

Args are the same as token_score except for `agg_method`

:param ``Union[str, List[str]]`` batch: a single sentence or a batch of sentences.

:param ``bool`` surprisal: If `True`, returns per-word surprisals instead of log-probabilities.

:param ``bool`` prob: If `True`, returns per-word probabilities instead of log-probabilities.

:param ``bool`` base_two: If `True`, uses log base 2 instead of natural-log (returns bits of values in case of surprisals)

:param ``bool`` rank: If `True`, also returns the rank of each word in context (based on the log-probability value)

:return: A `List` containing a `Tuple` consisting of the word, its associated score, and optionally, its rank.

:rtype: ``Union[List[Tuple[str, float]], List[Tuple[str, float, int]]]``

"""

all_token_scores = self.token_score(

batch=batch,

surprisal=surprisal,

prob=prob,

base_two=base_two,

rank=rank,

bow_correction=bow_correction,

)

all_word_scores = []

for i in range(len(all_token_scores)):

if type(batch) == str:

sentence = batch

else:

sentence = batch[i]

words = re.findall(r"[\w']+|[.,!?;]", sentence)

token_scores = all_token_scores[i]

# if token_score pads the beginning with a token (i.e. like llama)

if token_scores[0][0] == self.tokenizer.special_tokens_map["bos_token"]:

token_scores = token_scores[1:]

token_index = 0

word_index = 0

word_scores = [] # list of word, surprisal tuples

try:

while token_index < len(token_scores):

current_word = words[word_index]

current_token, current_surprisal = token_scores[token_index]

# token does not match, alignment must be adjusted

mismatch = current_token != current_word

while mismatch:

token_index += 1

current_token += token_scores[token_index][0]

current_surprisal += token_scores[token_index][1]

mismatch = current_token != current_word

word_scores.append((current_word, current_surprisal))

token_index += 1

word_index += 1

except Exception:

warning_message = f"Failed to aggregate word-level scores for {sentence}, returning token-level scores"

warnings.warn(warning_message)

word_scores = token_scores

all_word_scores.append(word_scores)

return all_word_scores

def score(

self, batch: Union[str, List[str]], pool: Callable = torch.mean, *args

) -> Union[float, List[float]]:

"""

.. deprecated:: 0.1.18

Use :meth:`sequence_score` or :meth:`token_score` instead.

Pooled estimates of sentence log probabilities, computed by the

language model. Pooling is usually done using a function that

is passed to the method.

:param batch: a list of sentences that will be passed to the

language model to score.

:type batch: Union[str, List[str]]

:param pool: Pooling function, is selected to be

`torch.mean()` by default.

:type pool: Callable

:return: Float or list of floats specifying the log

probabilities of the input sentence(s).

:rtype: Union[float, List[float]]

"""

warnings.warn(

"score is deprecated, use sequence_score or token_score instead",

DeprecationWarning,

)

def adapt_score(

self,

preamble: Union[str, List[str]],

stimuli: Union[str, List[str]],

pool: Callable = torch.mean,

*args,

) -> None:

"""

.. deprecated:: 0.2.10

Use :meth:`conditional_score` or :meth:`token_score` instead.

"""

warnings.warn(

"adapt_score is deprecated, use conditional_score or token_score instead",

DeprecationWarning,

)

def partial_score(

self,

preamble: Union[str, List[str]],

stimuli: Union[str, List[str]],

separator: str = " ",

reduction: Callable = lambda x: x.mean(0).item(),

**kwargs,

) -> List[float]:

"""

.. deprecated::

Use :meth:`conditional_score` instead.

"""

warnings.warn(

"partial_score is deprecated, use conditional_score instead",

DeprecationWarning,

)

def conditional_score(

self,

prefix: Union[str, List[str]],

stimuli: Union[str, List[str]],

separator: str = " ",

reduction: Callable = lambda x: x.mean(0).item(),

prob: bool = False,

base_two: bool = False,

bow_correction: bool = False,

**kw,

) -> List[float]:

"""

Pooled estimates of sequence log probabilities (or some modification of it), given a prefix. Pooling is usually done using a function that is passed to the method.

:param prefix: a batch of prefixes or primes passed to the

language model. This is what the sequence is conditioned on, and the model ignores the word probabilities of this part of the input in estimating the overall score.

:type prefix: ``Union[str, List[str]]``

:param stimuli: a batch of sequences (same length as prefix)

that form the main input consisting of the sequence whose

score you want to calculate.

:type stimuli: ``Union[str, List[str]]``

:param reduction: Reduction function, is selected to be

``lambda x: x.mean(0).item()`` by default, which stands for the avg. log-probability per token for each sequence in the batch.

:type reduction: Callable

:param kw: model-specific keyword arguments to pass to the `prepare_text` function

:return: List of floats specifying the desired score for the stimuli part of the input, e.g., P(stimuli | preamble).

:rtype: ``List[float]``

"""

primed = self.prime_text(prefix, stimuli, separator, **kw)

result = self.compute_stats(

primed,

rank=False,

base_two=base_two,

prob=prob,

bow_correction=bow_correction,

return_tensors=True,

)

logprob = result

reduced = list(map(reduction, logprob))

return reduced

def sequence_score(

self,

batch,

reduction=lambda x: x.mean(0).item(),

prob: bool = False,

base_two: bool = False,

bow_correction: bool = False,

**kw,

):

"""

Pooled estimates of sequence log probabilities (or some modification of it).

:param batch: a batch of sequences whose score you want to calculate.

:type batch: ``Union[str, List[str]]``

:param reduction: Reduction function, is selected to be

``lambda x: x.mean(0).item()`` by default, which stands for the avg. log-probability per token for each sequence in the batch.

:type reduction: Callable

:param kw: model-specific keyword arguments to pass to the `prepare_text` function

:return: List of floats specifying the desired score for the stimuli part of the input, e.g., P(stimuli | preamble).

:rtype: ``List[float]``

TODO: reduction should be a string, if it's a function, specify what kind of function. --> how to ensure it is always that type?

"""

tokenized = self.prepare_text(batch, **kw)

scores = self.compute_stats(

tokenized,

rank=False,

base_two=base_two,

prob=prob,

bow_correction=bow_correction,

return_tensors=True,

)

reduced = list(map(reduction, scores))

return reduced

def encode(

self,

text: Union[str, List[str]],

manual_special: bool = True,

return_tensors: Optional[str] = "pt",

) -> BatchEncoding:

"""

Encode a batch of sentences using the model's tokenizer.

Equivalent of calling `model.tokenizer(input)`

:param ``Union[str, List[str]]`` text: Input batch/sentence to

be encoded.

:param manual_special: Specification of whether special tokens

will be manually encoded.

:type manual_special: bool

:param return_tensors: returned tensor format. Default `'pt'`

:type return_tensors: str

:return: Encoded batch

:rtype: ``BatchEncoding``

"""

sentences = [text] if isinstance(text, str) else text

if manual_special:

# manually add special tokens

sentences = self.add_special_tokens(sentences)

if return_tensors:

tokens = self.tokenizer(

sentences,

add_special_tokens=False,

padding="longest",

return_attention_mask=True,

return_tensors=return_tensors,

)

else:

# mostly for masked LMs

tokens = self.tokenizer(

sentences, padding="longest", return_attention_mask=True

)

return tokens

def decode(self, idx: List[int]):

"""

Decode input ids using the model's tokenizer.

:param ``List[int]`` idx: List of ids.

:return: Decoded strings

:rtype: List[str]

"""

return [

self.tokenizer.decode([x]).strip()

for x in self.tokenizer.convert_tokens_to_ids(

self.tokenizer.convert_ids_to_tokens(idx)

)

]

class MaskedLMScorer(LMScorer):

"""

Class for Masked Langauge Models such as BERT, RoBERTa, etc.

:param model: should be path to a model (.pt or .bin file) stored locally,

or name of a pretrained model stored on the Huggingface Model Hub, or

a model (torch.nn.Module) that have the same signature as a

Huggingface model obtained from `AutoModelForMaskedLM`. In the last

case, a corresponding tokenizer must also be provided.

:param device: device type that the model should be loaded on,

options: `cpu or cuda:{0, 1, ...} or auto`

:type device: str, optional

:param tokenizer: if provided, use this tokenizer.

"""

def __init__(

self,

model: Union[str, torch.nn.Module],

device: Optional[str] = "cpu",

tokenizer=None,

PLL_metric: str = "original",

**kwargs,

) -> None:

"""

:param model: should be path to a model (.pt or .bin file) stored

locally, or name of a pretrained model stored on the Huggingface

Model Hub, or a model (torch.nn.Module) that have the same

signature as a Huggingface model obtained from

`AutoModelForMaskedLM`. In the last case, a corresponding tokenizer

must also be provided.

:param device: device type that the model should be loaded on,

options: `cpu or cuda:{0, 1, ...}`

:type device: str, optional

:param tokenizer: if provided, use this tokenizer.

"""

super(MaskedLMScorer, self).__init__(model, device=device, tokenizer=tokenizer)

if isinstance(model, str):

if self.device == "auto":

self.model = AutoModelForMaskedLM.from_pretrained(

model, device_map=self.device, return_dict=True, **kwargs

)

else:

self.model = AutoModelForMaskedLM.from_pretrained(

model, return_dict=True, **kwargs

)

# self.model.to(self.device)

else:

self.model = model

if self.device != "auto":

self.model.to(self.device)

self.model.eval()

self.PLL_metric: str = PLL_metric

# define CLS and SEP tokens

self.bos_token_id = self.tokenizer.cls_token_id

self.eos_token_id = self.tokenizer.sep_token_id

self.cls_token_id = self.tokenizer.cls_token_id

self.sep_token_id = self.tokenizer.sep_token_id

self.mask_token_id = self.tokenizer.mask_token_id

self.pad_token_id = self.tokenizer.pad_token_id

def add_special_tokens(self, text: Union[str, Iterable[str]]) -> List[str]:

"""

Reformats input text to add special model-dependent tokens.

:param text: single string or batch of strings to be

modified.

:type text: ``Union[str, List[str]]``

:return: Modified input, containing special tokens as per

tokenizer specification

:rtype: ``List[str]``

"""

sentences = [text] if isinstance(text, str) else text

sentences = [

self.tokenizer.cls_token + " " + sentence + " " + self.tokenizer.sep_token

for sentence in sentences

]

return sentences

def mask(

self, sentence_words: Union[Tuple[str, str], List[Tuple[str, str]]]

) -> Tuple[List[str], List[str]]:

"""

Processes a list of (sentence, word) into input that has the

word masked out of the sentence.

Note: only works for masked LMs.

:param ``Union[Tuple[str], List[Tuple[str]]]`` sentence_words:

Input consisting of `[(sentence, word)]`, where sentence

is an input sentence, and word is a word present in the

sentence that will be masked out.

:return: Tuple `(sentences, words)`

"""

sentence_words = (

[sentence_words] if isinstance(sentence_words, tuple) else sentence_words

)

sentences: List[str] = []

words: List[str] = []

for sentence, word in sentence_words:

words.append(word)

sentences.append(

re.sub(

rf"(?<![\w\/-])({word})(?=[^\w\/-])",

self.tokenizer.mask_token,

sentence,

)

)

return (sentences, words)

def cloze(

self,

sentence_words: Union[Tuple[str, str], List[Tuple[str, str]]],

PLL_metric: Optional[str] = None,

probs: Optional[bool] = False,

) -> List[float]:

"""

Runs inference on masked input.

Note: only works for masked LMs.

:param ``Union[Tuple[str], List[Tuple[str]]]`` sentence_words:

Input consisting of `[(sentence, word)]`, where sentence

is an input sentence, and word is a word present in the

sentence that will be masked out and inferred.

:param PLL_metric: PLL scoring strategy to be used.

Options: `original` or `within_word_l2r`. Default: `original`

For motivation as to why to use `within_word_l2r` PLL scoring, see Kauf & Ivanova (2023):

https://arxiv.org/abs/2305.10588

:param probs: whether to return probabilities (if True) or log probabilities (if False)

:return: A list of tensors corresponding to (log) probabilities for the desired word

in context

"""

sentences = list(map(lambda x: x[0], sentence_words))

encoded = self.tokenizer(sentences, return_tensors="pt", padding=True)

targets_start = []

targets_end = []

# Iterating over sentence-target word pairs

for batch_index, (sentence, word) in enumerate(sentence_words):

desired_tokens = self.tokenizer(

word, return_tensors="pt", add_special_tokens=False

)["input_ids"][0]

if PLL_metric == "within_word_l2r":

start_idx = None

word_ids = encoded.word_ids(batch_index=batch_index)

# Iterating over all words in the sentence

for word_id in set(word_ids):

# Ignoring special tokens

if word_id is None:

continue

# Finding all tokens corresponding to the chosen word

indices = np.where(list(map(lambda x: x == word_id, word_ids)))[0]

tokens = encoded["input_ids"][batch_index][indices]

# Checking if the chosen word matches the target word

if torch.equal(tokens, desired_tokens):

start_idx = indices[0]

if start_idx:

targets_start.append(start_idx)

targets_end.append(start_idx + len(desired_tokens))

else:

raise ValueError(

f"Word ``{word}'' not found in sentence ``{sentence}''. PLL=within_word_l2r won't work if ``{word}'' is a subword or multiple words."

)

else:

for start_idx in range(

len(encoded["input_ids"][batch_index]) - len(desired_tokens)

):

# Checking if the chosen sequence of tokens matches the target sequence of tokens

if torch.equal(

encoded["input_ids"][batch_index][

start_idx : len(desired_tokens) + start_idx

],

desired_tokens,

):

targets_start.append(start_idx)

targets_end.append(start_idx + len(desired_tokens))

if self.device != "auto":

encoded = encoded.to(self.device)

masked_tensors = self.get_masked_tensors(

encoded,

PLL_metric=PLL_metric,

targets_start=targets_start,

targets_end=targets_end,

)

target_prob_list = []

with torch.no_grad():

for masked_tensor, attn_mask, token_ids, token_indices in masked_tensors:

masked_logits = (

self.model(input_ids=masked_tensor, attention_mask=attn_mask)

.logits[torch.arange(len(token_indices)), token_indices]

.squeeze()

.detach()

)

if len(token_indices) > 1:

logprobs = masked_logits - masked_logits.logsumexp(1).unsqueeze(1)

target_prob = (

logprobs[torch.arange(len(token_indices)), token_ids]

.squeeze()

.sum()

)

else:

logprobs = masked_logits - masked_logits.logsumexp(0)

target_prob = logprobs[token_ids].squeeze().sum()

target_prob_list.append(target_prob)

target_probs_tensor = torch.tensor(target_prob_list)

if probs:

target_probs_tensor = target_probs_tensor.exp()

return target_probs_tensor.tolist()

def get_masked_tensors(

self,

encoded: BatchEncoding,

PLL_metric: Optional[str] = None,

targets_start: Optional[List[int]] = None,

targets_end: Optional[List[int]] = None,

) -> Iterable[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]]:

"""Create batches of tokenized instances with each token in the sequence replaced by the mask token."""

if PLL_metric is None:

PLL_metric = self.PLL_metric

for batch_index, (token_ids, attention_mask) in enumerate(

zip(encoded["input_ids"], encoded["attention_mask"])

):

token_ids = torch.as_tensor(token_ids)

attention_mask = torch.as_tensor(attention_mask)

target_token_indices: List[int] = []

target_token_ids: List[torch.Tensor] = []

# select tokens (and their indices) that will be predicted

for token_index, token_id in enumerate(token_ids):

if (

token_id != self.pad_token_id

and token_id != self.cls_token_id

and token_id != self.sep_token_id

and (

targets_start is None

or token_index >= targets_start[batch_index]

)

and (targets_end is None or token_index < targets_end[batch_index])

):

target_token_ids.append(token_id)

target_token_indices.append(token_index)

target_token_indices = list(target_token_indices)

target_token_ids = list(target_token_ids)

# mask tokens based on the current token to be predicted

mask_indices: List[List[int]]

if PLL_metric == "within_word_l2r":

"""

Future tokens belonging to the same word as the target token are masked during token inference as well.

"""

word_ids = encoded.word_ids(

batch_index=batch_index

) # only used for this PLL_metric

mask_indices = [

# mask the target token and all following tokens which belong to the same word

(

[mask_pos]

+ [

j

for j in range(mask_pos + 1, target_token_indices[-1] + 1)

if word_ids[j] == word_ids[mask_pos]

]

if word_ids[mask_pos] is not None

else [mask_pos]

) # mask this token

for mask_pos in target_token_indices

]

elif PLL_metric == "original":

# Original PLL metric

mask_indices = [[target] for target in target_token_indices]

else:

raise ValueError(f"PLL metric '{PLL_metric}' not supported.")

# repeat the token ids and mask each set of tokens in a separate row

token_ids_masked = token_ids.repeat(len(target_token_indices), 1)

for i, mask_set in enumerate(mask_indices):

token_ids_masked[i, mask_set] = self.mask_token_id

yield (

# token ids with some replaced by the mask token (effective tokens are replaced, but potentially more)

token_ids_masked,

# the attention mask is identical for all masked sets

attention_mask.expand(len(target_token_indices), -1),

# ids of the tokens to be predicted

torch.tensor(target_token_ids),

# indices of the tokens to be predicted

torch.tensor(target_token_indices),

)

def prepare_text(

self,

text: Union[str, List[str], BatchEncoding],

PLL_metric: Optional[str] = "original",

) -> Iterable[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]]:

"""

Prepares a batch of input text into a format fit to run MLM

scoring on.

Borrows preprocessing algorithm from Salazar et al. (2020), and

modifies code from the following github repository by simonpri:

https://github.com/simonepri/lm-scorer

:param text: batch of sentences to be prepared for scoring.

:param PLL_metric: PLL scoring strategy to be used.

Options: `original` or `within_word_l2r`. Default: `original`

For motivation as to why to use `within_word_l2r` PLL scoring, see Kauf & Ivanova (2023):

https://arxiv.org/abs/2305.10588

:return: Batch of formatted input that can be passed to `compute_stats`

"""

# converts input text to batch of tensors with every position except the cls and sep token masked

if isinstance(text, BatchEncoding):

encoded = text

else:

sentences = [text] if isinstance(text, str) else text

encoded = self.encode(sentences, manual_special=False)

return self.get_masked_tensors(encoded, PLL_metric)

def prime_text(

self,

prefix: Union[str, List[str]],

stimuli: Union[str, List[str]],

suffix: Union[None, str, List[str]] = None,

separator: str = " ",

PLL_metric: Optional[str] = None,

) -> Iterable[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]]:

"""

Prepares a batch of input text into a format fit to run LM

scoring on.

Borrows preprocessing algorithm from Salazar et al. (2020), and

modifies code from the following github repository by simonpri:

https://github.com/simonepri/lm-scorer

:param ``Union[str, List[str]]`` prefix: Batch of prefixes/prime/preambles on which the LM is conditioned.

:param ``Union[str, List[str]]`` stimuli: Batch of continuations that are scored based on the conditioned text (provided in the ``prefix``). The positions of the elements match their counterparts in the ``prefix``.

:return: Batch of formatted input that can be passed to

``compute_stats``

"""

if isinstance(stimuli, str):

assert isinstance(prefix, str)

prefix = [prefix]

stimuli = [stimuli]

if suffix is not None:

assert isinstance(suffix, str)

assert isinstance(prefix, list)

assert suffix is None or isinstance(suffix, list)

assert len(prefix) == len(stimuli)

assert suffix is None or len(suffix) == len(stimuli)

# compute the length of each preamble

prefix_encoded = self.encode(prefix, False)["input_ids"]

targets_start: List[int] = []

for prefix_tokens in prefix_encoded:

targets_start.append(

sum(

token != self.pad_token_id and token != self.sep_token_id

for token in prefix_tokens

)

)

targets_end: Optional[List[int]] = None

if suffix is None:

sentences = [p + separator + s for p, s in zip(prefix, stimuli)]

encoded = self.encode(sentences, manual_special=False)

else:

targets_end = []

for stimuli_tokens, prefix_len in zip(

self.encode(stimuli, False)["input_ids"], targets_start

):

targets_end.append(

prefix_len

+ sum(

token != self.cls_token_id

and token != self.pad_token_id

and token != self.sep_token_id

for token in stimuli_tokens

)

)

sentences = [

pre + separator + stim + separator + suff

for pre, stim, suff in zip(prefix, stimuli, suffix)

]

encoded = self.encode(sentences, manual_special=False)

return self.get_masked_tensors(

encoded,

PLL_metric=PLL_metric,

targets_start=targets_start,

targets_end=targets_end,

)

def distribution(self, batch: Iterable) -> torch.Tensor:

"""

Returns a distribution over the vocabulary of the model.

:param `Iterable` batch: A batch of inputs fit to pass to a

transformer LM.

:return: Tensor consisting of log probabilies over vocab items.

"""

# takes in prepared text and returns scores for each sentence in batch

token_ids, attention_masks, effective_token_ids, lengths, offsets = list(

zip(*batch)

)

token_ids = torch.cat(token_ids)

attention_masks = torch.cat(attention_masks)

if self.device != "auto":

token_ids = token_ids.to(self.device)

attention_masks = attention_masks.to(self.device)

effective_token_ids = torch.cat([torch.tensor(x) for x in effective_token_ids])

indices = list(

chain.from_iterable(

[list(range(o, o + n)) for n, o in zip(lengths, offsets)]

)

)

with torch.no_grad():

output = self.model(token_ids, attention_mask=attention_masks)

logits = output.logits[torch.arange(sum(lengths)), indices].detach()

logprob_distribution = logits - logits.logsumexp(1).unsqueeze(1)

return logprob_distribution

def cloze_distribution(

self, queries: Union[Collection[Tuple[str, str]], Tuple[str, str]]

) -> torch.Tensor:

"""

Accepts as input batch of [(s_i, bw_i)] where s_i is a prompt with an

abstract token (bw_i) representing a blank word and returns a distribution

over the vocabulary of the model.

:param `Iterable` queries: A batch of [(s_i, bw_i)] where s_i is a prompt with an abstract token (bw_i) representing a blank word

:return: Tensor contisting of log probabilities over vocab items.

"""

if len(queries) == 0:

return torch.tensor([])

if isinstance(next(iter(queries)), str):

queries = [cast(Tuple[str, str], queries)]

prompts, words = zip(*queries)

# modified_prompts = self.add_special_tokens(prompts)

modified_prompts = self.tokenizer.batch_decode(

self.tokenizer(prompts)["input_ids"]

)

# splits = [prompt.split(word) for prompt, word in zip(modified_prompts, words)]

splits = [

re.split(rf"{re.escape(word)}", prompt)

for prompt, word in zip(modified_prompts, words)

]

splits = [[x.strip() for x in s] for s in splits]

pre, post = list(zip(*splits))

pre_idx = self.tokenizer(list(pre), add_special_tokens=False, padding=False)[

"input_ids"

]

mask_idx = [len(item) for item in pre_idx]

# masked = [

# m.replace(w, self.tokenizer.mask_token)

# for m, w in zip(modified_prompts, words)

# ]

masked = [

re.sub(rf"\b{w}\b", self.tokenizer.mask_token, m)

for m, w in zip(prompts, words)

]

with torch.no_grad():