From 6e6927270cfdc26d0f69e624de00f50c41175ad9 Mon Sep 17 00:00:00 2001 From: Andrew Lapp Date: Mon, 9 Sep 2024 17:29:42 -0400 Subject: [PATCH] WIP logging logits processor --- docs/reference/logits_processors.md | 51 ++++ outlines/generate/base.py | 22 ++ outlines/processors/__init__.py | 29 +- outlines/processors/base_logits_processor.py | 36 ++- outlines/processors/logging.py | 64 ++++ outlines/processors/sampling.py | 276 +++++++++++++++++ tests/processors/test_sampling.py | 302 +++++++++++++++++++ 7 files changed, 778 insertions(+), 2 deletions(-) create mode 100644 docs/reference/logits_processors.md create mode 100644 outlines/generate/base.py create mode 100644 outlines/processors/logging.py create mode 100644 outlines/processors/sampling.py create mode 100644 tests/processors/test_sampling.py diff --git a/docs/reference/logits_processors.md b/docs/reference/logits_processors.md new file mode 100644 index 0000000000..e6f765cce5 --- /dev/null +++ b/docs/reference/logits_processors.md @@ -0,0 +1,51 @@ +# Logits Processors + +TODO: Explanation of what logits processors do + +TODO: List of `OutlinesLogitsProcessor` + +TODO: Example of how to implement TemperatureLogitsProcessor + +TODO: using logits processors with models directly vs using with outlines + +TODO: Explanation of pipelines + +TODO: Link to log logits + +## Using Logits Processors in Outlines + +TODO Explanation + +``` +import outlines +``` + + +## Chaining Logits Processors + +``` +import outlines +import outlines.processors as processors + +model = outlines.models.llamacpp( + repo_id="M4-ai/TinyMistral-248M-v2-Instruct-GGUF", + filename="TinyMistral-248M-v2-Instruct.Q4_K_M.gguf" +) + +# Create a chained logits processor +logits_processor = ( + processors.sequence_logging(model.tokenizer) | # Log the generated sequence + processors.logits_logging(model.tokenizer) | # Log the raw logits + processors.regex(r"[0-9]*", model.tokenizer) | # Restrict the logits to match the pattern + processors.temperature(0.5) | # Set temperature to 0.5 + processors.logits_logging(model.tokenizer) # Log the restricted, temperature-augmentent, sampled logits +) + +generator = outlines.generate.base(model, logits_process) +generator("What is your favorite number? ") +``` + +Output: +``` +TODO +``` diff --git a/outlines/generate/base.py b/outlines/generate/base.py new file mode 100644 index 0000000000..9364b46a92 --- /dev/null +++ b/outlines/generate/base.py @@ -0,0 +1,22 @@ +from functools import singledispatch + +from outlines.generate.api import SequenceGeneratorAdapter +from outlines.models import MLXLM, OpenAI +from outlines.processors import OutlinesLogitsProcessor +from outlines.samplers import Sampler, multinomial + + +@singledispatch +def base( + model, logits_processor: OutlinesLogitsProcessor, sampler: Sampler = multinomial() +): + return SequenceGeneratorAdapter(model, logits_processor, sampler) + + +@base.register(OpenAI) +def base_openai( + model: OpenAI, + logits_processor: OutlinesLogitsProcessor, + sampler: Sampler = multinomial(), +) -> Exception: + raise NotImplementedError("The OpenAI API does not support logits processing.") diff --git a/outlines/processors/__init__.py b/outlines/processors/__init__.py index f0f0f829b5..8de066d61e 100644 --- a/outlines/processors/__init__.py +++ b/outlines/processors/__init__.py @@ -1,7 +1,34 @@ +from .base_logits_processor import ChainedLogitsProcessor, OutlinesLogitsProcessor +from .logging import LogitsLoggingLogitsProcessor, SequenceLoggingLogitsProcessor +from .sampling import ( + FrequencyPenaltyLogitsProcessor, + MinPLogitsProcessor, + NoRepeatNGramLogitsProcessor, + PresencePenaltyLogitsProcessor, + QuadraticSmoothingLogitsProcessor, + RepetitionPenaltyLogitsProcessor, + TemperatureLogitsProcessor, + TFSLogitsProcessor, + TopKLogitsProcessor, + TopPLogitsProcessor, +) from .structured import ( CFGLogitsProcessor, GuideLogitsProcessor, JSONLogitsProcessor, - OutlinesLogitsProcessor, RegexLogitsProcessor, ) + +# aliases for convenience +chained = ChainedLogitsProcessor + +cfg = CFGLogitsProcessor +guide = GuideLogitsProcessor +json = JSONLogitsProcessor +regex = RegexLogitsProcessor + +temperature = TemperatureLogitsProcessor +min_p = MinPLogitsProcessor + +sequence_logging = SequenceLoggingLogitsProcessor +logits_logging = LogitsLoggingLogitsProcessor diff --git a/outlines/processors/base_logits_processor.py b/outlines/processors/base_logits_processor.py index feedf52535..e207faed37 100644 --- a/outlines/processors/base_logits_processor.py +++ b/outlines/processors/base_logits_processor.py @@ -1,5 +1,6 @@ +import reprlib from abc import abstractmethod -from typing import TYPE_CHECKING, List, Protocol, Type, Union +from typing import TYPE_CHECKING, List, Protocol, Type, Union, runtime_checkable import numpy as np import torch @@ -20,6 +21,7 @@ def is_mlx_array_type(array_type): return issubclass(array_type, mx.array) +@runtime_checkable class OutlinesLogitsProcessor(Protocol): """ Base class for logits processors which normalizes types of logits: @@ -133,3 +135,35 @@ def _from_torch(tensor: torch.Tensor, target_type: Type) -> Array: raise TypeError( f"Failed to convert torch tensors to target_type `{target_type}`" ) + + def __repr__(self): + return f"{self.__class__.__name__}({reprlib.repr(self.__dict__)})" + + def __or__(self, other): + if not isinstance(other, OutlinesLogitsProcessor): + raise ValueError( + "Can only chain with another OutlinesLogitsProcessor instance." + ) + return ChainedLogitsProcessor([self, other]) + + +class ChainedLogitsProcessor(OutlinesLogitsProcessor): + """Handle chaining two logits processors to process logits sequentially""" + + processors: List[OutlinesLogitsProcessor] + + def __init__(self, processors: List[OutlinesLogitsProcessor]): + self.processors = [] + for processor in processors: + if isinstance(processor, ChainedLogitsProcessor): + self.processors.extend(processor.processors) + else: + self.processors.append(processor) + + def process_logits( + self, input_ids: List[List[int]], logits: torch.Tensor + ) -> torch.Tensor: + result = logits + for processor in self.processors: + result = processor.process_logits(input_ids, result) + return result diff --git a/outlines/processors/logging.py b/outlines/processors/logging.py new file mode 100644 index 0000000000..8d95eb89cf --- /dev/null +++ b/outlines/processors/logging.py @@ -0,0 +1,64 @@ +import logging +import sys +import warnings +from typing import List + +import torch + +from .base_logits_processor import OutlinesLogitsProcessor + + +class LogitsLoggingLogitsProcessor(OutlinesLogitsProcessor): + """Handle chaining two logits processors to process logits sequentially""" + + def __init__(self, tokenizer, top_n=8, logger=None, warn=True): + self.tokenizer = tokenizer + self.top_n = top_n + if logger is not None: + self.logger = logger + else: + self.logger = logging.getLogger("logits_logger") + self.logger.setLevel(logging.info) + self.logger.addHandler(logging.StreamHandler(sys.stderr)) + if warn: + warnings.warn( + "Do not use LoggingLogitsProcessor in production, it slows down generation." + ) + + def process_logits( + self, input_ids: List[List[int]], logits: torch.Tensor + ) -> torch.Tensor: + # all token probs for the current batch + probs = torch.nn.functional.softmax(logits, dim=-1) + # top candidate tokens probs + top_indices = torch.topk(probs, self.top_n).indices + top_indices = [ + set(row.tolist()) | {self.tokenizer.eos_token_id} for row in top_indices + ] + batch_top_probs = [ + {token_idx: probs[batch_num, token_idx] for token_idx in row_indices} + for batch_num, row_indices in enumerate(top_indices) + ] + self.logger.info(batch_top_probs) + return logits + + +class SequenceLoggingLogitsProcessor(OutlinesLogitsProcessor): + def __init__(self, tokenizer, top_n=8, logger=None, warn=True): + self.tokenizer = tokenizer + if logger is not None: + self.logger = logger + else: + self.logger = logging.getLogger("sequence_logger") + self.logger.setLevel(logging.info) + self.logger.addHandler(logging.StreamHandler(sys.stderr)) + if warn: + warnings.warn( + "Do not use SequenceLoggingProcessor in production, it slows down generation." + ) + + def process_logits( + self, input_ids: List[List[int]], logits: torch.Tensor + ) -> torch.Tensor: + self.logger.info(self.tokenizer.decode(input_ids)) + return logits diff --git a/outlines/processors/sampling.py b/outlines/processors/sampling.py new file mode 100644 index 0000000000..3e91b3301d --- /dev/null +++ b/outlines/processors/sampling.py @@ -0,0 +1,276 @@ +from typing import List + +import torch +from torch.nn import functional as F + +from .base_logits_processor import OutlinesLogitsProcessor + + +class TemperatureLogitsProcessor(OutlinesLogitsProcessor): + """Processor to apply temperature scaling to logits. + + Args: + temperature (float): The temperature value to scale logits. Must be > 0. + A value of 0 will result in greedy sampling. + + Raises: + ValueError: If temperature is less than 0. + """ + + def __init__(self, temperature=0): + if temperature < 0: + raise ValueError("Temperature must be > 0.") + self.temperature = temperature + + def process_logits( + self, input_ids: List[List[int]], logits: torch.Tensor + ) -> torch.Tensor: + """ + Apply temperature scaling to logits. + For temperature 0, set the highest logit to inf (equivalent to greedy sampling) + """ + if self.temperature == 0: + max_indices = logits.argmax(dim=-1, keepdim=True) + mask = torch.full_like(logits, -torch.inf).scatter_(-1, max_indices, 0) + return logits + mask + return logits / self.temperature + + +class MinPLogitsProcessor(OutlinesLogitsProcessor): + """Processor to ensure a minimum probability for each element in the logits. + + Args: + min_p (float): The minimum probability value. Must be between 0 and 1. + + Raises: + ValueError: If min_p is not between 0 and 1. + """ + + def __init__(self, min_p=0.01): + if min_p <= 0 or min_p >= 1: + raise ValueError("min_p must be between 0 and 1") + self.min_p = min_p + + def process_logits( + self, input_ids: List[List[int]], logits: torch.Tensor + ) -> torch.Tensor: + """ + Ensure each logit's probability is at least min_p by + setting logits with lower probabilities to -inf. + """ + return logits.masked_fill(F.softmax(logits, dim=-1) < self.min_p, -torch.inf) + + +class TopPLogitsProcessor(OutlinesLogitsProcessor): + """Processor to apply top-p (nucleus) filtering to logits. + + Args: + p (float): The cumulative probability threshold. Must be between 0 and 1. + + Raises: + ValueError: If p is not between 0 and 1. + """ + + def __init__(self, p: float = 0.9): + if p <= 0 or p > 1: + raise ValueError("p must be between 0 and 1") + self.p = p + + def process_logits( + self, input_ids: List[List[int]], logits: torch.Tensor + ) -> torch.Tensor: + """ + Apply top-p filtering by setting lowest logits + with cumulative probability above p to -inf. + """ + sorted_logits, sorted_indices = logits.sort(descending=True, dim=-1) + cumulative_probs = F.softmax(sorted_logits, dim=-1).cumsum(dim=-1) + mask = cumulative_probs > self.p + mask[..., 1:] = mask[..., :-1].clone() + mask[..., 0] = 0 + return logits.masked_fill(mask.scatter(1, sorted_indices, mask), -float("inf")) + + +class TopKLogitsProcessor(OutlinesLogitsProcessor): + """Processor to apply top-k filtering to logits. + + Args: + k (int): The number of highest probability tokens to keep. + + Raises: + ValueError: If k is not a positive integer. + """ + + def __init__(self, k: int = 50): + if k <= 0: + raise ValueError("k must be a positive integer.") + self.k = k + + def process_logits( + self, input_ids: List[List[int]], logits: torch.Tensor + ) -> torch.Tensor: + """ + Apply top-k filtering by keeping the top k logits and setting the rest to -inf. + """ + if self.k > logits.size(-1): + raise ValueError("k cannot be greater than the number of logits") + top_k = torch.topk(logits, self.k, dim=-1).indices + mask = torch.ones_like(logits, dtype=torch.bool).scatter_(-1, top_k, False) + return logits.masked_fill(mask, -torch.inf) + + +class TFSLogitsProcessor(OutlinesLogitsProcessor): + """Processor to apply Tail Free Sampling (TFS) to logits.""" + + def __init__(self, threshold: float = 0.9): + if threshold <= 0 or threshold > 1: + raise ValueError("threshold must be between 0 and 1") + self.threshold = threshold + + def process_logits( + self, input_ids: List[List[int]], logits: torch.Tensor + ) -> torch.Tensor: + """ + Apply TFS by keeping top logits with cumulative probability + below the threshold and setting others to -inf. + """ + if logits.numel() == 0: + return logits + sorted_logits, sorted_indices = logits.sort(descending=True, dim=-1) + cumulative_probs = F.softmax(sorted_logits, dim=-1).cumsum(dim=-1) + sorted_indices_to_keep = cumulative_probs <= self.threshold + + last_valid_index = sorted_indices_to_keep.sum(dim=-1, keepdim=True) - 1 + sorted_indices_to_keep.scatter_(1, last_valid_index, 1) + + mask = torch.zeros_like(logits, dtype=torch.bool).scatter( + 1, sorted_indices, sorted_indices_to_keep + ) + return logits.masked_fill(~mask, -torch.inf) + + +class QuadraticSmoothingLogitsProcessor(OutlinesLogitsProcessor): + """Processor to apply quadratic smoothing to logits.""" + + def __init__(self, alpha: float = 0.5): + if alpha < 0 or alpha > 1: + raise ValueError("alpha must be between 0 and 1") + self.alpha = alpha + + def process_logits( + self, input_ids: List[List[int]], logits: torch.Tensor + ) -> torch.Tensor: + """ + Apply quadratic smoothing by mixing logits with their square, controlled by alpha. + """ + return logits * (1 - self.alpha) + logits.pow(2) * self.alpha + + +class RepetitionPenaltyLogitsProcessor(OutlinesLogitsProcessor): + """Processor to apply a repetition penalty to logits. + + Args: + penalty (float): The penalty to apply to repeated tokens. Must be > 0. + """ + + def __init__(self, penalty: float = 1.2): + if penalty < 0: + raise ValueError("penalty must be <= 1") + self.penalty = penalty + + def process_logits( + self, input_ids: List[List[int]], logits: torch.Tensor + ) -> torch.Tensor: + """ + Apply a penalty to logits of tokens that appear + in the input sequence to reduce repetition. + """ + input_ids_tensor = torch.tensor(input_ids, device=logits.device) + token_counts = torch.bincount( + input_ids_tensor.view(-1), minlength=logits.size(-1) + ).float() + return logits / (1 + token_counts * (self.penalty - 1)).unsqueeze(0) + + +class PresencePenaltyLogitsProcessor(OutlinesLogitsProcessor): + """Processor to apply a presence penalty to logits. + + Args: + penalty (float): The penalty to apply to tokens present in the sequence. Must be > 0. + """ + + def __init__(self, penalty: float = 0.1): + if penalty < 0: + raise ValueError("penalty must be greater than or equal to 0") + self.penalty = penalty + + def process_logits( + self, input_ids: List[List[int]], logits: torch.Tensor + ) -> torch.Tensor: + """ + Apply a contsant penalty to logits of tokens that are present in the input sequence. + """ + input_ids_tensor = torch.tensor(input_ids, device=logits.device) + unique_tokens = torch.unique(input_ids_tensor, sorted=False) + logits[:, unique_tokens] -= self.penalty + return logits + + +class FrequencyPenaltyLogitsProcessor(OutlinesLogitsProcessor): + """Processor to apply a frequency penalty to logits based on token frequency. + + Args: + penalty (float): The penalty to apply to tokens based on frequency. Must be > 0. + """ + + def __init__(self, penalty: float = 0.1): + if penalty < 0: + raise ValueError("penalty must be greater than or equal to 0") + self.penalty = penalty + + def process_logits( + self, input_ids: List[List[int]], logits: torch.Tensor + ) -> torch.Tensor: + """ + Apply a penalty to logits based on a multiple of the frequenccy of tokens in the input sequence. + """ + input_ids_tensor = torch.tensor(input_ids, device=logits.device) + token_counts = torch.bincount( + input_ids_tensor.view(-1), minlength=logits.size(-1) + ).float() + return logits - token_counts * self.penalty + + +class NoRepeatNGramLogitsProcessor(OutlinesLogitsProcessor): + """Processor to ensure no repeated n-grams in the sequence. + + Args: + n (int): The n-gram size to check for repetitions. + """ + + def __init__(self, n: int = 3): + if n <= 0: + raise ValueError("n must be greater than 0") + self.n = n + + def process_logits( + self, input_ids: List[List[int]], logits: torch.Tensor + ) -> torch.Tensor: + """ + Set logits which would result in n-gram repetition exceeding `n` to -inf + """ + if len(input_ids[0]) < self.n: + return logits + input_ids_tensor = torch.tensor(input_ids, device=logits.device) + ngrams = [ + tuple(input_ids_tensor[0, i : i + self.n]) + for i in range(len(input_ids_tensor[0]) - self.n + 1) + ] + last_ngram = tuple(input_ids_tensor[0, -self.n + 1 :]) + mask = torch.zeros(logits.shape[-1], dtype=torch.bool, device=logits.device) + # TODO: don't iterate + for token in range(logits.shape[-1]): + if tuple(list(last_ngram) + [token]) in ngrams: + mask[token] = True + logits.masked_fill_(mask, -float("inf")) + return logits diff --git a/tests/processors/test_sampling.py b/tests/processors/test_sampling.py new file mode 100644 index 0000000000..1d1935bb83 --- /dev/null +++ b/tests/processors/test_sampling.py @@ -0,0 +1,302 @@ +import pytest +import torch + +from outlines.processors import ( + FrequencyPenaltyLogitsProcessor, + MinPLogitsProcessor, + NoRepeatNGramLogitsProcessor, + PresencePenaltyLogitsProcessor, + QuadraticSmoothingLogitsProcessor, + RepetitionPenaltyLogitsProcessor, + TemperatureLogitsProcessor, + TFSLogitsProcessor, + TopKLogitsProcessor, + TopPLogitsProcessor, +) + + +@pytest.fixture() +def input_logits(): + return torch.tensor([[1.0, 2.0, 3.0]]) + + +@pytest.fixture() +def input_logits_2d(): + return torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + + +# TemperatureLogitsProcessor tests +def test_normal_temperature(input_logits): + processor = TemperatureLogitsProcessor(temperature=0.5) + expected = torch.tensor([[2.0, 4.0, 6.0]]) + torch.testing.assert_close(processor.process_logits([], input_logits), expected) + + +def test_temperature_one(input_logits): + processor = TemperatureLogitsProcessor(temperature=1.0) + expected = input_logits.clone() + torch.testing.assert_close(processor.process_logits([], input_logits), expected) + + +def test_extreme_temperature_high(input_logits): + processor = TemperatureLogitsProcessor(temperature=100.0) + expected = input_logits / 100.0 + torch.testing.assert_close(processor.process_logits([], input_logits), expected) + + +def test_extreme_temperature_low(input_logits): + processor = TemperatureLogitsProcessor(temperature=0.01) + expected = input_logits / 0.01 + torch.testing.assert_close(processor.process_logits([], input_logits), expected) + + +def test_temperature_zero(input_logits): + processor = TemperatureLogitsProcessor(temperature=0) + result = processor.process_logits([], input_logits) + max_idx = input_logits.argmax(dim=-1, keepdim=True).item() + assert torch.softmax(result, dim=-1)[0][max_idx] == 1.0 + + +def test_temperature_negative(input_logits): + with pytest.raises(ValueError): + TemperatureLogitsProcessor(temperature=-1.0) + + +# MinPLogitsProcessor tests +def test_valid_min_p(input_logits): + processor = MinPLogitsProcessor(min_p=0.1) + processed_logits = processor.process_logits([], input_logits) + probs = torch.softmax(processed_logits, dim=-1) + assert not torch.isinf(probs).any() + assert (probs >= 0.1).any() # Ensure at least one probability meets the minimum + + +def test_all_below_min_p(input_logits): + processor = MinPLogitsProcessor(min_p=0.99) + processed_logits = processor.process_logits([], input_logits) + assert torch.isinf(processed_logits).all() + + +def test_some_below_min_p(input_logits): + processor = MinPLogitsProcessor(min_p=0.5) + probs = torch.softmax(input_logits, dim=-1) + mask = probs < 0.5 + expected = input_logits.masked_fill(mask, -torch.inf) + torch.testing.assert_close(processor.process_logits([], input_logits), expected) + + +def test_no_below_min_p(input_logits): + processor = MinPLogitsProcessor(min_p=0.01) + processed_logits = processor.process_logits([], input_logits) + torch.testing.assert_close(processed_logits, input_logits) + + +def test_min_p_boundary_values(): + with pytest.raises(ValueError): + MinPLogitsProcessor(min_p=0) + + with pytest.raises(ValueError): + MinPLogitsProcessor(min_p=1) + + +def test_high_dimensional_logits(): + processor = MinPLogitsProcessor(min_p=0.1) + input_logits = torch.randn(3, 4, 5) + processed_logits = processor.process_logits([], input_logits) + assert processed_logits.shape == input_logits.shape + + +# TopPLogitsProcessor tests +def test_valid_top_p(input_logits): + processor = TopPLogitsProcessor(p=0.9) + processed_logits = processor.process_logits([], input_logits) + probs = torch.softmax(processed_logits, dim=-1) + cumulative_probs = torch.cumsum(probs, dim=-1) + assert (cumulative_probs <= 0.9).any() + + +def test_top_p_zero(input_logits): + with pytest.raises(ValueError): + TopPLogitsProcessor(p=0) + + +def test_top_p_one(input_logits): + processor = TopPLogitsProcessor(p=1.0) + torch.testing.assert_close(processor.process_logits([], input_logits), input_logits) + + +def test_top_p_greater_than_one(input_logits): + with pytest.raises(ValueError): + TopPLogitsProcessor(p=1.1) + + +# TopKLogitsProcessor tests +def test_valid_top_k(input_logits): + processor = TopKLogitsProcessor(k=2) + processed_logits = processor.process_logits([], input_logits) + assert (processed_logits == float("-inf")).sum() == 1 + + +def test_top_k_zero(input_logits): + with pytest.raises(ValueError): + TopKLogitsProcessor(k=0) + + +def test_top_k_equal_to_logits(input_logits): + processor = TopKLogitsProcessor(k=3) + torch.testing.assert_close(processor.process_logits([], input_logits), input_logits) + + +def test_top_k_greater_than_logits(input_logits): + processor = TopKLogitsProcessor(k=4) + with pytest.raises(ValueError): + processor.process_logits([], input_logits) + + +# TFSLogitsProcessor tests +def test_valid_tfs(input_logits): + processor = TFSLogitsProcessor(threshold=0.9) + processed_logits = processor.process_logits([], input_logits) + probs = torch.softmax(processed_logits, dim=-1) + cumulative_probs = torch.cumsum(probs, dim=-1) + assert (cumulative_probs <= 0.9).all() or (cumulative_probs <= 1).all() + + +def test_tfs_zero(input_logits): + with pytest.raises(ValueError): + TFSLogitsProcessor(threshold=0) + + +def test_tfs_one(input_logits): + processor = TFSLogitsProcessor(threshold=1.0) + torch.testing.assert_close(processor.process_logits([], input_logits), input_logits) + + +def test_tfs_greater_than_one(input_logits): + with pytest.raises(ValueError): + TFSLogitsProcessor(threshold=1.1) + + +# QuadraticSmoothingLogitsProcessor tests +def test_valid_quadratic_smoothing(input_logits): + processor = QuadraticSmoothingLogitsProcessor(alpha=0.5) + processed_logits = processor.process_logits([], input_logits) + expected = input_logits**2 * 0.5 + input_logits * 0.5 + torch.testing.assert_close(processed_logits, expected) + + +def test_quadratic_smoothing_zero(input_logits): + processor = QuadraticSmoothingLogitsProcessor(alpha=0.0) + torch.testing.assert_close(processor.process_logits([], input_logits), input_logits) + + +def test_quadratic_smoothing_negative(input_logits): + with pytest.raises(ValueError): + QuadraticSmoothingLogitsProcessor(alpha=-0.1) + + +def test_quadratic_smoothing_greater_than_one(input_logits): + with pytest.raises(ValueError): + QuadraticSmoothingLogitsProcessor(alpha=1.1) + + +# RepetitionPenaltyLogitsProcessor tests +def test_valid_repetition_penalty(input_logits): + processor = RepetitionPenaltyLogitsProcessor(penalty=1.5) + input_ids = [[0, 1, 2]] + processed_logits = processor.process_logits(input_ids, torch.tensor(input_logits)) + expected = input_logits.clone() + for seq in input_ids: + for token in seq: + expected[..., token] /= 1.5 + torch.testing.assert_close(processed_logits, expected) + + +def test_repetition_penalty_one(input_logits): + processor = RepetitionPenaltyLogitsProcessor(penalty=1.0) + processed_logits = processor.process_logits([[0, 1, 2]], input_logits) + torch.testing.assert_close(processed_logits, input_logits) + + +def test_repetition_penalty_negative(input_logits): + with pytest.raises(ValueError): + RepetitionPenaltyLogitsProcessor(penalty=-0.5) + + +# PresencePenaltyLogitsProcessor tests +def test_valid_presence_penalty(input_logits): + processor = PresencePenaltyLogitsProcessor(penalty=0.5) + input_ids = [[0, 1, 2]] + processed_logits = processor.process_logits(input_ids, torch.tensor(input_logits)) + expected = torch.tensor(input_logits) + expected[0] -= 0.5 + torch.testing.assert_close(processed_logits, expected) + + +def test_presence_penalty_zero(input_logits): + processor = PresencePenaltyLogitsProcessor(penalty=0.0) + torch.testing.assert_close( + processor.process_logits([[0, 1, 2]], input_logits), input_logits + ) + + +def test_presence_penalty_negative(input_logits): + with pytest.raises(ValueError): + PresencePenaltyLogitsProcessor(penalty=-0.5) + + +# FrequencyPenaltyLogitsProcessor tests +def test_valid_frequency_penalty(input_logits): + processor = FrequencyPenaltyLogitsProcessor(penalty=0.5) + input_ids = [[0, 1, 2]] + processed_logits = processor.process_logits(input_ids, torch.tensor(input_logits)) + expected = input_logits.clone() + token_counts = torch.zeros( + input_logits.shape[-1], dtype=input_logits.dtype, device=input_logits.device + ) + for seq in input_ids: + for token in seq: + token_counts[token] += 1 + for token in range(input_logits.shape[-1]): + expected[..., token] -= token_counts[token] * 0.5 + torch.testing.assert_close(processed_logits, expected) + + +def test_frequency_penalty_zero(input_logits): + processor = FrequencyPenaltyLogitsProcessor(penalty=0.0) + torch.testing.assert_close( + processor.process_logits([[0, 1, 2]], input_logits), input_logits + ) + + +def test_frequency_penalty_negative(input_logits): + with pytest.raises(ValueError): + FrequencyPenaltyLogitsProcessor(penalty=-0.5) + + +# NoRepeatNGramLogitsProcessor tests +def test_valid_no_repeat_ngram(input_logits): + processor = NoRepeatNGramLogitsProcessor(n=2) + processed_logits = processor.process_logits( + [[0, 1, 0, 1]], torch.tensor(input_logits) + ) + assert processed_logits[0, 0] == float("-inf") + + +def test_no_repeat_ngram_zero(input_logits): + with pytest.raises(ValueError): + NoRepeatNGramLogitsProcessor(n=0) + + +def test_no_repeat_ngram_one(input_logits): + processor = NoRepeatNGramLogitsProcessor(n=1) + torch.testing.assert_close( + processor.process_logits([[0, 1, 2]], input_logits), input_logits + ) + + +def test_no_repeat_ngram_greater_than_sequence_length(input_logits): + processor = NoRepeatNGramLogitsProcessor(n=10) + torch.testing.assert_close( + processor.process_logits([[0, 1, 2]], input_logits), input_logits + )