MambaScorer¶

class minicons.scorer.MambaScorer(model: str | torch.nn.Module, device: str | None = 'cuda', tokenizer=None, dtype=torch.float16, **kwargs)¶

Bases: LMScorer

encode(text: str | List[str], bos_token: bool = False, eos_token: bool = False) → transformers.BatchEncoding¶

Encode a batch of sentences using the model’s tokenizer. Equivalent of calling model.tokenizer(input)

Parameters:

text (Union[str, List[str]]) – Input batch/sentence to be encoded.
manual_special (bool) – Specification of whether special tokens will be manually encoded.
return_tensors (str) – returned tensor format. Default ‘pt’

Returns:

Encoded batch

Return type:

BatchEncoding

word_spans_tokenized(batch: Iterable, tokenize_function: Callable)¶

Aligns the spans of a string tokenized by a function to the tokenizer of the LM.

Parameters:

batch – batch of sentences to be prepared for scoring.
tokenize_function – the tokenizer function – we recommend nltk.TweetTokenizer

Returns:

Batch of index spans and the tokenized words.

prepare_text(text: str | List[str] | transformers.BatchEncoding, bos_token: bool = False, eos_token: bool = False) → Tuple¶

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

Parameters:: text – batch of sentences to be prepared for scoring.
Returns:: Batch of formatted input that can be passed to compute_stats

prime_text(preamble: str | List[str], stimuli: str | List[str], separator: str | List[str] = ' ', bos_token=False, eos_token=False) → Tuple¶

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

Parameters:

preamble (Union[str, List[str]]) – Batch of prefixes/prime/preambles on which the LM is conditioned.
stimuli (Union[str, List[str]]) – Batch of continuations that are scored based on the conditioned text (provided in the preamble). The positions of the elements match their counterparts in the preamble.

Returns:

Batch of formatted input that can be passed to compute_stats

distribution(batch: Iterable) → torch.Tensor¶

Returns a distribution over the vocabulary of the model.

Parameters:: batch (Iterable) – A batch of inputs fit to pass to a transformer LM.
Returns:: Tensor consisting of log probabilies over vocab items.

next_word_distribution(queries: List, bos_token=False, eos_token=False, surprisal: bool = False)¶: Returns the log probability distribution of the next word.

compute_stats(batch: Iterable, rank: bool = False, prob: bool = False, base_two: bool = False, return_tensors: bool = False, bow_correction: bool = False) → Tuple[List[float], List[float]] | List[float]¶

Primary computational method that processes a batch of prepared sentences and returns per-token scores for each sentence. By default, returns log-probabilities.

Parameters:

batch (Iterable) – batched input as processed by prepare_text or prime_text.
rank (bool) – whether the model should also return ranks per word (based on the conditional log-probability of the word in context).
prob (bool) – whether the model should return probabilities instead of log-probabilities. Can only be True when base_two is False.
base_two (bool) – whether the base of the log should be 2 (usually preferred when reporting results in bits). Can only be True when prob is False.
return_tensors (bool) – whether the model should return scores as a list of tensors instead of a list of lists. This is important in some other convenient methods used in the package.
bow_correction (bool) – whether to apply the beginning of word correction, as pointed out in Pimentel and Meister (2024) and Oh and Schuler (2024).

Returns:

Either a tuple of lists, each containing probabilities and ranks per token in each sentence passed in the input.

Return type:

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

word_score_tokenized(batch: Iterable, tokenize_function: Callable, bow_correction: bool = False, bos_token: bool = False, eos_token: bool = False, surprisal: bool = False, prob: bool = False, base_two: bool = False, return_tensors: bool = False)¶

Returns the logprobs per word, as tokenized by the tokenize_function.

Parameters:

batch – batch of sentences to be prepared for scoring.
tokenize_function – tokenize function to maps strings to tokenized words.
bow_correction – whether to apply Oh and Schuler’s correction.
bos_token – if a beginning of sentence token should be added (specify this as True for GPT2 and Pythia, for other they do it by default).
eos_token – if an end of sentence token should be added
prob – if the scores should be probabilities instead of logprobabilities
base_two – if the logprobabilities should

Returns:

Batch of words and their corresponding log-probs (summed log-probs of the tokens)

sequence_score(batch, reduction=<function MambaScorer.<lambda>>, base_two=False, bow_correction=False, **kwargs)¶: TODO: reduction should be a string, if it’s a function, specify what kind of function. –> how to ensure it is always that type?

token_score(batch: str | List[str], surprisal: bool = False, prob: bool = False, base_two: bool = False, rank: bool = False, decode: bool = False, bow_correction: bool = False, **kwargs) → 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.

Parameters:

batch (Union[str, List[str]]) – a single sentence or a batch of sentences.
surprisal (bool) – If True, returns per-word surprisals instead of log-probabilities.
prob (bool) – If True, returns per-word probabilities instead of log-probabilities.
base_two (bool) – If True, uses log base 2 instead of natural-log (returns bits of values in case of surprisals)
rank (bool) – If True, also returns the rank of each word in context (based on the log-probability value)
bow_correction (bool) – whether to apply the beginning of word correction, as pointed out in Pimentel and Meister (2024) and Oh and Schuler (2024).

Returns:

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

Return type:

Union[List[Tuple[str, float]], List[Tuple[str, float, int]]]