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

@@ -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",

astrai/inference/api/openai.py

@@ -21,7 +21,6 @@ logger = logging.getLogger(__name__)
 _UNSUPPORTED_PARAMS = (
     "n",
     "presence_penalty",
-    "frequency_penalty",
     "logit_bias",
     "user",
 )

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:

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()

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:

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:

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)

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

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)