From d08a92c7bd9e6d3ec6608f6a7095944c6d8dc639 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Fri, 17 Jul 2026 20:43:32 +0800 Subject: [PATCH] feat: add frequency penalty to inference sampling pipeline - Add FrequencyPenaltyStrategy (logit -= penalty * count) - Per-task rep_window for penalty history lookup - Wire through engine, task, executor, API layer - Add --frequency_penalty and --rep_window to stream_chat.py - 9 unit tests for frequency penalty strategy --- astrai/inference/__init__.py | 4 +- astrai/inference/api/openai.py | 1 - astrai/inference/api/protocol.py | 1 + astrai/inference/core/executor.py | 30 ++++++ astrai/inference/core/task.py | 8 ++ astrai/inference/engine.py | 68 ++++++++++++- astrai/inference/sample.py | 157 +++++++++++++++++++++++++++--- scripts/demo/stream_chat.py | 15 +++ tests/inference/test_sample.py | 106 ++++++++++++++++++++ 9 files changed, 370 insertions(+), 20 deletions(-) diff --git a/astrai/inference/__init__.py b/astrai/inference/__init__.py index e3b4b9d..d112b89 100644 --- a/astrai/inference/__init__.py +++ b/astrai/inference/__init__.py @@ -6,7 +6,7 @@ Layers: - protocols/: Response builders (OpenAI, Anthropic) - transport/: SSE transport utilities - engine.py: Facade (InferenceEngine), Value Object (GenerationRequest) - - sample.py: Strategy pattern (TemperatureStrategy, TopKStrategy, TopPStrategy) + - sample.py: Strategy pattern (TemperatureStrategy, TopKStrategy, TopPStrategy, FrequencyPenaltyStrategy) """ from astrai.inference.api import ( @@ -50,6 +50,7 @@ from astrai.inference.core import ( from astrai.inference.engine import GenerationRequest, InferenceEngine from astrai.inference.sample import ( BaseSamplingStrategy, + FrequencyPenaltyStrategy, SamplingPipeline, TemperatureStrategy, TopKStrategy, @@ -83,6 +84,7 @@ __all__ = [ "TemperatureStrategy", "TopKStrategy", "TopPStrategy", + "FrequencyPenaltyStrategy", "SamplingPipeline", "ProtocolHandler", "StopChecker", diff --git a/astrai/inference/api/openai.py b/astrai/inference/api/openai.py index 2078797..948ffb9 100644 --- a/astrai/inference/api/openai.py +++ b/astrai/inference/api/openai.py @@ -21,7 +21,6 @@ logger = logging.getLogger(__name__) _UNSUPPORTED_PARAMS = ( "n", "presence_penalty", - "frequency_penalty", "logit_bias", "user", ) diff --git a/astrai/inference/api/protocol.py b/astrai/inference/api/protocol.py index 2fb7726..b00bf2d 100644 --- a/astrai/inference/api/protocol.py +++ b/astrai/inference/api/protocol.py @@ -125,6 +125,7 @@ class ProtocolHandler: temperature=self.request.temperature, top_p=self.request.top_p, top_k=self.request.top_k, + frequency_penalty=getattr(self.request, "frequency_penalty", 0.0), ) if self.request.stream: diff --git a/astrai/inference/core/executor.py b/astrai/inference/core/executor.py index 6d0bd55..abaa164 100644 --- a/astrai/inference/core/executor.py +++ b/astrai/inference/core/executor.py @@ -75,6 +75,33 @@ class Executor: temperatures = torch.tensor([t.temperature for t in tasks], device=self.device) top_ks = torch.tensor([t.top_k for t in tasks], device=self.device) top_ps = torch.tensor([t.top_p for t in tasks], device=self.device) + freq_penalties = torch.tensor( + [t.frequency_penalty for t in tasks], device=self.device + ) + + history_lists = [] + mask_lists = [] + for t in tasks: + window = t.rep_window + prompt_part = t.prompt_ids[-window:] + ids = prompt_part + t.output_ids + history_lists.append(ids) + mask_lists.append([True] * len(ids)) + + max_len = max(len(h) for h in history_lists) + padded_ids = torch.zeros( + len(tasks), max_len, dtype=torch.long, device=self.device + ) + padded_mask = torch.zeros( + len(tasks), max_len, dtype=torch.bool, device=self.device + ) + for i, (h, m) in enumerate(zip(history_lists, mask_lists)): + padded_ids[i, : len(h)] = torch.tensor( + h, dtype=torch.long, device=self.device + ) + padded_mask[i, : len(m)] = torch.tensor( + m, dtype=torch.bool, device=self.device + ) with torch.inference_mode(): outputs = self.model( @@ -89,4 +116,7 @@ class Executor: temperature=temperatures, top_k=top_ks, top_p=top_ps, + frequency_penalty=freq_penalties, + input_ids=padded_ids, + input_mask=padded_mask, ).tolist() diff --git a/astrai/inference/core/task.py b/astrai/inference/core/task.py index b58af57..8006567 100644 --- a/astrai/inference/core/task.py +++ b/astrai/inference/core/task.py @@ -33,6 +33,8 @@ class Task: temperature: float = 1.0, top_p: float = 1.0, top_k: int = 50, + frequency_penalty: float = 0.0, + rep_window: int = 64, ): self.task_id = task_id self.prompt_ids = prompt_ids @@ -40,6 +42,8 @@ class Task: self.temperature = temperature self.top_p = top_p self.top_k = top_k + self.frequency_penalty = frequency_penalty + self.rep_window = rep_window self.status = TaskStatus.PENDING self.output_ids: List[int] = [] @@ -92,6 +96,8 @@ class TaskManager: temperature: float = 1.0, top_p: float = 1.0, top_k: int = 50, + frequency_penalty: float = 0.0, + rep_window: int = 64, stream_callback: Optional[Callable[[str], None]] = None, ) -> str: task_id = f"task_{int(time.time())}_{uuid.uuid4().hex[:8]}" @@ -116,6 +122,8 @@ class TaskManager: temperature=temperature, top_p=top_p, top_k=top_k, + frequency_penalty=frequency_penalty, + rep_window=rep_window, ) with self._lock: diff --git a/astrai/inference/engine.py b/astrai/inference/engine.py index bbcf71e..19ea9f8 100644 --- a/astrai/inference/engine.py +++ b/astrai/inference/engine.py @@ -74,6 +74,8 @@ class GenerationRequest: top_p: float = 1.0, temperature: float = 1.0, max_tokens: Optional[int] = None, + frequency_penalty: float = 0.0, + rep_window: int = 64, stream: bool = False, ): if not (isinstance(top_k, int) and top_k >= 0): @@ -82,12 +84,21 @@ class GenerationRequest: raise ValueError("top_p must be a float between 0.0 and 1.0") if not (isinstance(temperature, (int, float)) and temperature > 0): raise ValueError("temperature must be a positive number") + if not ( + isinstance(frequency_penalty, (int, float)) + and -2.0 <= frequency_penalty <= 2.0 + ): + raise ValueError("frequency_penalty must be between -2.0 and 2.0") + if not (isinstance(rep_window, int) and rep_window > 0): + raise ValueError("rep_window must be a positive integer") self.messages = messages self.top_k = top_k self.top_p = top_p self.temperature = temperature self.max_tokens = max_tokens + self.frequency_penalty = frequency_penalty + self.rep_window = rep_window self.stream = stream @@ -132,17 +143,33 @@ class InferenceEngine: temperature: float = 1.0, top_p: float = 1.0, top_k: int = 50, + frequency_penalty: float = 0.0, + rep_window: int = 64, ) -> Union[Generator, str, List[str]]: is_batch = isinstance(prompt, list) prompts = prompt if is_batch else [prompt] if stream: return self._generate_streaming( - prompts, is_batch, max_tokens, temperature, top_p, top_k + prompts, + is_batch, + max_tokens, + temperature, + top_p, + top_k, + frequency_penalty, + rep_window, ) else: return self._generate_non_streaming( - prompts, is_batch, max_tokens, temperature, top_p, top_k + prompts, + is_batch, + max_tokens, + temperature, + top_p, + top_k, + frequency_penalty, + rep_window, ) def generate_async( @@ -152,9 +179,18 @@ class InferenceEngine: temperature: float = 1.0, top_p: float = 1.0, top_k: int = 50, + frequency_penalty: float = 0.0, + rep_window: int = 64, ) -> AsyncGenerator[str, None]: sync_gen = self._generate_streaming( - [prompt], False, max_tokens, temperature, top_p, top_k + [prompt], + False, + max_tokens, + temperature, + top_p, + top_k, + frequency_penalty, + rep_window, ) async def _agen(): @@ -185,6 +221,8 @@ class InferenceEngine: temperature=request.temperature, top_p=request.top_p, top_k=request.top_k, + frequency_penalty=request.frequency_penalty, + rep_window=request.rep_window, ) def _submit_tasks( @@ -194,6 +232,8 @@ class InferenceEngine: temperature: float, top_p: float, top_k: int, + frequency_penalty: float, + rep_window: int, ) -> Tuple[GenerateResult, List[str]]: n = len(prompts) result = GenerateResult(count=n) @@ -206,6 +246,8 @@ class InferenceEngine: temperature=temperature, top_p=top_p, top_k=top_k, + frequency_penalty=frequency_penalty, + rep_window=rep_window, stream_callback=cb, ) task_ids.append(task_id) @@ -226,9 +268,17 @@ class InferenceEngine: temperature: float, top_p: float, top_k: int, + frequency_penalty: float, + rep_window: int, ) -> Generator: result, task_ids = self._submit_tasks( - prompts, max_tokens, temperature, top_p, top_k + prompts, + max_tokens, + temperature, + top_p, + top_k, + frequency_penalty, + rep_window, ) n = len(prompts) remaining = n @@ -262,9 +312,17 @@ class InferenceEngine: temperature: float, top_p: float, top_k: int, + frequency_penalty: float, + rep_window: int, ) -> Union[str, List[str]]: result, task_ids = self._submit_tasks( - prompts, max_tokens, temperature, top_p, top_k + prompts, + max_tokens, + temperature, + top_p, + top_k, + frequency_penalty, + rep_window, ) try: diff --git a/astrai/inference/sample.py b/astrai/inference/sample.py index b66099f..5ac0104 100644 --- a/astrai/inference/sample.py +++ b/astrai/inference/sample.py @@ -1,15 +1,15 @@ """Composable sampling strategies for logit transformation. Implements the Strategy pattern: each sampling technique -(temperature, top-k, top-p) is a pluggable strategy that -can be composed into a pipeline. +(temperature, top-k, top-p, frequency penalty) is a pluggable +strategy that can be composed into a pipeline. All strategies accept both scalar and per-sample tensor parameters, so a single pipeline works for any batch size. """ from abc import ABC, abstractmethod -from typing import List, Union +from typing import List, Optional, Union import torch from torch import Tensor @@ -19,12 +19,23 @@ class BaseSamplingStrategy(ABC): """Abstract base for a logit transformation strategy.""" @abstractmethod - def apply(self, logits: Tensor, filter_value: float = -float("inf")) -> Tensor: + def apply( + self, + logits: Tensor, + filter_value: float = -float("inf"), + input_ids: Optional[Tensor] = None, + input_mask: Optional[Tensor] = None, + ) -> Tensor: """Applies the strategy to logits. Args: logits: Raw logits tensor (batch, vocab_size). filter_value: Value assigned to filtered-out positions. + input_ids: Previously generated token IDs ``[batch, seq_len]``, + padded with 0. Used by frequency penalty. + input_mask: Boolean mask ``[batch, seq_len]``, True for real + tokens, False for padding. Used to exclude padding from + penalty computation. Returns: Transformed logits tensor. @@ -42,7 +53,13 @@ class TemperatureStrategy(BaseSamplingStrategy): def __init__(self, temperature: Union[float, Tensor] = 1.0): self.temperature = temperature - def apply(self, logits: Tensor, filter_value: float = -float("inf")) -> Tensor: + def apply( + self, + logits: Tensor, + filter_value: float = -float("inf"), + input_ids: Optional[Tensor] = None, + input_mask: Optional[Tensor] = None, + ) -> Tensor: t = self.temperature if isinstance(t, Tensor): t = t.to(logits.device, non_blocking=True).view(-1, 1) @@ -64,7 +81,13 @@ class TopKStrategy(BaseSamplingStrategy): def __init__(self, top_k: Union[int, Tensor] = 0): self.top_k = top_k - def apply(self, logits: Tensor, filter_value: float = -float("inf")) -> Tensor: + def apply( + self, + logits: Tensor, + filter_value: float = -float("inf"), + input_ids: Optional[Tensor] = None, + input_mask: Optional[Tensor] = None, + ) -> Tensor: tk = self.top_k if isinstance(tk, Tensor): tk = tk.to(logits.device, non_blocking=True).long().clamp(min=0) @@ -114,7 +137,13 @@ class TopPStrategy(BaseSamplingStrategy): logits[mask] = filter_value return logits - def apply(self, logits: Tensor, filter_value: float = -float("inf")) -> Tensor: + def apply( + self, + logits: Tensor, + filter_value: float = -float("inf"), + input_ids: Optional[Tensor] = None, + input_mask: Optional[Tensor] = None, + ) -> Tensor: tp = self.top_p if isinstance(tp, Tensor): tp = tp.to(logits.device, non_blocking=True) @@ -125,6 +154,84 @@ class TopPStrategy(BaseSamplingStrategy): return logits +class FrequencyPenaltyStrategy(BaseSamplingStrategy): + """Penalizes tokens based on how many times they appeared in history. + + Subtracts ``penalty * count(token)`` from each token's logit, where + ``count(token)`` is the number of occurrences in the generation history + (prompt + output). A penalty of ``0.0`` disables the strategy. + + Unlike repetition penalty (which only checks *presence*), frequency + penalty scales linearly with occurrence count: the first use is + penalized once, the third use three times. This allows natural + repetition of common words while suppressing degenerate loops. + + Reference: OpenAI API ``frequency_penalty`` parameter. + + Args: + penalty: Scalar or ``[batch]`` tensor (0.0 disables, range -2.0~2.0). + """ + + def __init__(self, penalty: Union[float, Tensor] = 0.0): + self.penalty = penalty + + def apply( + self, + logits: Tensor, + filter_value: float = -float("inf"), + input_ids: Optional[Tensor] = None, + input_mask: Optional[Tensor] = None, + ) -> Tensor: + if input_ids is None: + return logits + + p = self.penalty + if isinstance(p, Tensor): + p = p.to(logits.device, non_blocking=True).view(-1, 1) + if (p == 0.0).all(): + return logits + elif p == 0.0: + return logits + + input_ids = input_ids.to(logits.device, non_blocking=True) + + if input_mask is not None: + input_mask = input_mask.to(logits.device, non_blocking=True) + masked_ids = input_ids.clone() + masked_ids[~input_mask] = -1 + else: + masked_ids = input_ids + + batch_sz, seq_len = masked_ids.shape + vocab_size = logits.size(-1) + + if isinstance(p, Tensor): + penalty_per_row = p.expand(batch_sz, 1) + else: + penalty_per_row = torch.full( + (batch_sz, 1), float(p), device=logits.device, dtype=logits.dtype + ) + + counts = torch.zeros( + batch_sz, vocab_size, device=logits.device, dtype=logits.dtype + ) + valid_mask = masked_ids >= 0 + if valid_mask.any(): + valid_ids = masked_ids[valid_mask] + row_indices = ( + torch.arange(batch_sz, device=logits.device) + .unsqueeze(1) + .expand_as(masked_ids)[valid_mask] + ) + counts.index_put_( + (row_indices, valid_ids), + torch.ones_like(valid_ids, dtype=logits.dtype), + accumulate=True, + ) + + return logits - penalty_per_row * counts + + class SamplingPipeline(BaseSamplingStrategy): """Composes multiple sampling strategies into a single transformation. @@ -145,23 +252,39 @@ class SamplingPipeline(BaseSamplingStrategy): def __init__(self, strategies: List[BaseSamplingStrategy]): self.strategies = strategies - def apply(self, logits: Tensor, filter_value: float = -float("inf")) -> Tensor: + def apply( + self, + logits: Tensor, + filter_value: float = -float("inf"), + input_ids: Optional[Tensor] = None, + input_mask: Optional[Tensor] = None, + ) -> Tensor: for strategy in self.strategies: - logits = strategy.apply(logits, filter_value) + logits = strategy.apply(logits, filter_value, input_ids, input_mask) return logits - @torch.no_grad() - def sample(self, logits: Tensor, filter_value: float = -float("inf")) -> Tensor: + @torch.inference_mode() + def sample( + self, + logits: Tensor, + filter_value: float = -float("inf"), + input_ids: Optional[Tensor] = None, + input_mask: Optional[Tensor] = None, + ) -> Tensor: """Apply strategies then sample (softmax + multinomial). Args: logits: Raw logits ``[batch, vocab_size]``. + input_ids: Previously generated token IDs ``[batch, seq_len]``. + input_mask: Boolean mask for ``input_ids`` padding. Returns: Sampled token IDs ``[batch]``. """ return torch.multinomial( - torch.softmax(self.apply(logits, filter_value), dim=-1), + torch.softmax( + self.apply(logits, filter_value, input_ids, input_mask), dim=-1 + ), num_samples=1, ).squeeze(-1) @@ -172,6 +295,9 @@ def sample( temperature: Union[float, Tensor] = 1.0, top_k: Union[int, Tensor] = 0, top_p: Union[float, Tensor] = 1.0, + frequency_penalty: Union[float, Tensor] = 0.0, + input_ids: Optional[Tensor] = None, + input_mask: Optional[Tensor] = None, filter_value: float = -float("inf"), ) -> Tensor: """Apply sampling strategies then sample (softmax + multinomial). @@ -180,6 +306,10 @@ def sample( Args: logits: Raw logits ``[batch, vocab_size]``. + frequency_penalty: Penalty per occurrence for repeated tokens + (0.0 disables, range -2.0~2.0). + input_ids: Previously generated token IDs ``[batch, seq_len]``. + input_mask: Boolean mask for ``input_ids`` padding. Returns: Sampled token IDs ``[batch]``. @@ -189,5 +319,6 @@ def sample( TemperatureStrategy(temperature), TopKStrategy(top_k), TopPStrategy(top_p), + FrequencyPenaltyStrategy(frequency_penalty), ] - ).sample(logits, filter_value) + ).sample(logits, filter_value, input_ids, input_mask) diff --git a/scripts/demo/stream_chat.py b/scripts/demo/stream_chat.py index 3578564..7ebd395 100644 --- a/scripts/demo/stream_chat.py +++ b/scripts/demo/stream_chat.py @@ -42,6 +42,19 @@ def parse_args(): default=2048, help="Maximum tokens to generate", ) + parser.add_argument( + "--frequency_penalty", + type=float, + default=0.5, + help="Penalty per occurrence for repeated tokens (0.0 disables, " + "range -2.0~2.0, typical 0.3-1.0)", + ) + parser.add_argument( + "--rep_window", + type=int, + default=64, + help="Number of recent prompt tokens to include in penalty history", + ) parser.add_argument( "--system_prompt", type=str, @@ -79,6 +92,8 @@ def chat(): temperature=args.temperature, top_p=args.top_p, top_k=args.top_k, + frequency_penalty=args.frequency_penalty, + rep_window=args.rep_window, ): print(token, end="", flush=True) full_response += token diff --git a/tests/inference/test_sample.py b/tests/inference/test_sample.py index 9942a26..863b770 100644 --- a/tests/inference/test_sample.py +++ b/tests/inference/test_sample.py @@ -3,6 +3,7 @@ import torch from astrai.inference.sample import ( + FrequencyPenaltyStrategy, SamplingPipeline, TemperatureStrategy, TopKStrategy, @@ -125,3 +126,108 @@ def test_module_sample_batch(): assert tokens.shape == (2,) for t in tokens: assert 0 <= t < logits.size(-1) + + +def test_frequency_penalty_noop_when_zero(): + logits = torch.tensor([[1.0, 2.0, 3.0]]) + input_ids = torch.tensor([[0, 2]]) + s = FrequencyPenaltyStrategy(penalty=0.0) + result = s.apply(logits.clone(), input_ids=input_ids) + assert torch.equal(result, logits) + + +def test_frequency_penalty_noop_when_no_input_ids(): + logits = torch.tensor([[1.0, 2.0, 3.0]]) + s = FrequencyPenaltyStrategy(penalty=0.5) + result = s.apply(logits.clone()) + assert torch.equal(result, logits) + + +def test_frequency_penalty_single_occurrence(): + logits = torch.tensor([[4.0, 1.0, 2.0]]) + input_ids = torch.tensor([[0, 2]]) + input_mask = torch.tensor([[True, True]]) + s = FrequencyPenaltyStrategy(penalty=0.5) + result = s.apply(logits.clone(), input_ids=input_ids, input_mask=input_mask) + assert result[0, 0] == 3.5 + assert result[0, 1] == 1.0 + assert result[0, 2] == 1.5 + + +def test_frequency_penalty_multiple_occurrences(): + logits = torch.tensor([[4.0, 1.0, 2.0]]) + input_ids = torch.tensor([[0, 2, 0]]) + input_mask = torch.tensor([[True, True, True]]) + s = FrequencyPenaltyStrategy(penalty=0.5) + result = s.apply(logits.clone(), input_ids=input_ids, input_mask=input_mask) + assert result[0, 0] == 3.0 + assert result[0, 1] == 1.0 + assert result[0, 2] == 1.5 + + +def test_frequency_penalty_respects_padding_mask(): + logits = torch.tensor([[4.0, 1.0, 2.0]]) + input_ids = torch.tensor([[0, 2, 0]]) + input_mask = torch.tensor([[True, True, False]]) + s = FrequencyPenaltyStrategy(penalty=0.5) + result = s.apply(logits.clone(), input_ids=input_ids, input_mask=input_mask) + assert result[0, 0] == 3.5 + assert result[0, 1] == 1.0 + assert result[0, 2] == 1.5 + + +def test_frequency_penalty_batch_tensor(): + logits = torch.tensor( + [ + [4.0, 1.0, 2.0], + [3.0, 5.0, 1.0], + ] + ) + input_ids = torch.tensor([[0, 2, 0], [1, 1, 0]]) + input_mask = torch.tensor([[True, True, True], [True, True, False]]) + s = FrequencyPenaltyStrategy(penalty=torch.tensor([0.5, 1.0])) + result = s.apply(logits.clone(), input_ids=input_ids, input_mask=input_mask) + assert result[0, 0] == 3.0 + assert result[0, 2] == 1.5 + assert result[1, 1] == 3.0 + + +def test_frequency_penalty_negative_penalty_boosts_repeats(): + logits = torch.tensor([[4.0, 1.0, 2.0]]) + input_ids = torch.tensor([[0, 0]]) + input_mask = torch.tensor([[True, True]]) + s = FrequencyPenaltyStrategy(penalty=-0.5) + result = s.apply(logits.clone(), input_ids=input_ids, input_mask=input_mask) + assert result[0, 0] == 5.0 + + +def test_frequency_penalty_in_pipeline(): + logits = torch.tensor([[5.0, 1.0, 2.0, 3.0]]) + input_ids = torch.tensor([[0, 2, 0]]) + input_mask = torch.tensor([[True, True, True]]) + pipeline = SamplingPipeline( + [ + TemperatureStrategy(1.0), + FrequencyPenaltyStrategy(0.5), + ] + ) + result = pipeline.apply(logits.clone(), input_ids=input_ids, input_mask=input_mask) + assert result[0, 0] == 4.0 + assert result[0, 2] == 1.5 + + +def test_sample_with_frequency_penalty(): + logits = torch.tensor([[5.0, 1.0, 2.0, 3.0]]) + input_ids = torch.tensor([[0, 2, 0]]) + input_mask = torch.tensor([[True, True, True]]) + tokens = sample( + logits, + temperature=1.0, + top_k=0, + top_p=1.0, + frequency_penalty=0.5, + input_ids=input_ids, + input_mask=input_mask, + ) + assert tokens.shape == (1,) + assert 0 <= tokens[0] < logits.size(-1)