transformers 5.0.0rc2__py3-none-any.whl → 5.0.0rc3__py3-none-any.whl

transformers/generation/continuous_batching/requests.py

@@ -1,4 +1,3 @@
-# coding=utf-8
 # Copyright 2025 The HuggingFace Inc. team
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -120,7 +119,7 @@ class RequestState:
 
     # Required fields
     request_id: str
-    initial_tokens: list[int]  # Initial prompt tokens
+    initial_tokens: list[int]  # Initial prompt tokens # TODO: rename this as prefill tokens
     # Optional fields
     record_timestamps: bool = False  # Whether to record timestamps for the generated tokens
     num_children: int = 0  # Number of children requests
@@ -138,6 +137,8 @@ class RequestState:
     error: str | None = None  # Error message if the request failed
     lifespan: tuple[float, float] = (-1, -1)  # (time request was no longer pending, time request finished)
     _timestamps: list[float] = field(default_factory=list)  # Timestamps of the generated tokens
+    _true_initial_tokens: int = 0  # The true number of initial tokens, useful when soft resetting requests
+    # TODO: remove the attribute above to _num_initial_tokens once initial_tokens is renamed
 
     @property
     def status(self) -> RequestStatus:
@@ -221,6 +222,9 @@ class RequestState:
 
     def to_generation_output(self):
         """Convert the request state to a GenerationOutput object."""
+        if self._true_initial_tokens:
+            self.generated_tokens = self.initial_tokens[self._true_initial_tokens :] + self.generated_tokens
+            self.initial_tokens = self.initial_tokens[: self._true_initial_tokens]
         return GenerationOutput(
             request_id=self.request_id,
             prompt_ids=self.initial_tokens,
@@ -243,14 +247,31 @@ class RequestState:
             generated_tokens=self.generated_tokens[:],
             allocated_blocks=self.allocated_blocks,
             position_offset=self.position_offset,
-            status=self.status,
+            _status=self.status,
             max_new_tokens=self.max_new_tokens,
             eos_token_id=self.eos_token_id,
             streaming=self.streaming,
             created_time=t,
             lifespan=(t, -1),
-            timestamps=None if self.timestamps is None else self.timestamps[:],
+            _timestamps=None if self.timestamps is None else self.timestamps[:],
             error=self.error,
             record_timestamps=self.record_timestamps,
         )
         return new_request
+
+    def create_equivalent_initial_request(self) -> "RequestState":
+        """Creates an equivalent new request by removing the generated tokens and adding them to the initial prompt. The
+        created request has THE SAME request_id. Notably, we can retrieve the original request from the created one with
+        the _true_initial_tokens attribute."""
+        new_state = RequestState(
+            request_id=self.request_id,
+            initial_tokens=self.initial_tokens + self.generated_tokens,
+            num_children=self.num_children,
+            record_timestamps=self.record_timestamps,
+            tokens_to_process=self.initial_tokens + self.generated_tokens,
+            max_new_tokens=self.max_new_tokens - len(self.generated_tokens),
+            eos_token_id=self.eos_token_id,
+            streaming=self.streaming,
+        )
+        new_state._true_initial_tokens = self._true_initial_tokens + len(self.initial_tokens)
+        return new_state

transformers/generation/continuous_batching/scheduler.py

@@ -1,4 +1,3 @@
-# coding=utf-8
 # Copyright 2025 The HuggingFace Inc. team
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -18,7 +17,7 @@ from collections import deque
 
 from ...utils.metrics import attach_tracer, traced
 from .cache import PagedAttentionCache
-from .requests import RequestState, RequestStatus
+from .requests import RequestState, RequestStatus, logger
 
 
 class Scheduler(ABC):
@@ -37,6 +36,10 @@ class Scheduler(ABC):
         self._cancellation_lock = threading.Lock()
         self._requests_to_cancel: set[str] = set()
         self._requests_to_fork: list[RequestState] = []
+        # This state is used to avoid infinite loops when offloading requests
+        self.block_new_requests = False
+        # This is to compute the cache used by a new request being scheduled
+        self.cache_budget_module = None if cache.num_full_attention_groups else cache.config.sliding_window
 
     @traced
     def add_waiting_request(self, state: RequestState):
@@ -52,10 +55,11 @@ class Scheduler(ABC):
         self.waiting_requests_order.append(state.request_id)
 
     @abstractmethod
-    def schedule_batch(self, token_budget: int) -> list[RequestState]:
-        """Schedules requests for the next batch based on available token budget. This method selects which requests
-        should be processed in the current batch, considering the token budget and the scheduler's prioritization rules.
-        The token_budget is the maximum number of tokens that can be processed in this batch."""
+    def schedule_batch(self, token_budget: int, cache_budget: int) -> list[RequestState]:
+        """Schedules requests for the next batch based on available token and cache budgets. This method selects which
+        requests should be processed in the current batch, considering the budgets and the scheduler's prioritization
+        rules. The token_budget is the maximum number of tokens that can be processed in a batch, and the cache_budget
+        is the maximum number of KV cache entries that can be read in a batch."""
 
     @traced
     def has_pending_requests(self) -> bool:
@@ -63,14 +67,13 @@ class Scheduler(ABC):
         return len(self.active_requests) or len(self.waiting_requests)
 
     @traced
-    def finish_request(self, request_id: str, evict_from_cache: bool = True):
+    def finish_request(self, request_id: str, evict_from_cache: bool = True) -> None:
         """Completes processing of a request and optionally frees its allocated cache blocks. This method is called
         when a request has finished generation or encountered an error.
         """
         if evict_from_cache:
             self.cache.free_blocks(request_id)
-            if request_id in self.active_requests:
-                del self.active_requests[request_id]
+            self.active_requests.pop(request_id, None)
 
     @traced
     def get_active_request_static_outputs(self, request_id: str) -> list[int]:
@@ -90,10 +93,8 @@ class Scheduler(ABC):
         """Remove all cancelled requests from active and waiting queues."""
         with self._cancellation_lock:
             for request_id in self._requests_to_cancel:
-                if request_id in self.active_requests:
-                    del self.active_requests[request_id]
-                if request_id in self.waiting_requests:
-                    del self.waiting_requests[request_id]
+                self.active_requests.pop(request_id, None)
+                self.waiting_requests.pop(request_id, None)
                 if request_id in self.waiting_requests_order:
                     self.waiting_requests_order.remove(request_id)
                 self.cache.free_blocks(request_id)
@@ -107,7 +108,7 @@ class Scheduler(ABC):
         )
 
     @traced
-    def _allocate_blocks_if_needed(self, state: RequestState) -> bool:
+    def _allocate_blocks_if_needed(self, state: RequestState, len_next_tokens: int) -> bool:
         """Allocate additional cache blocks for a request if the currently allocated blocks are insufficient to
         accommodate the next tokens. It calculates how many blocks are needed based on the request's current
         cache occupancy and the number of tokens to be processed. The allocation itself is done by the CacheAllocator
@@ -116,20 +117,16 @@ class Scheduler(ABC):
         # 1. we check that the occupancy is less than the requested length
         # 2. we allocate enough blocks to cover the requested length
         current_len = state.current_len()
-        len_next_tokens = len(state.tokens_to_process)
         occupancy = state.allocated_blocks * self.cache.block_size - current_len
         if occupancy < len_next_tokens or state.allocated_blocks == 0:
             blocks_needed = ((len_next_tokens - occupancy + 1) // self.cache.block_size) + 1
-            allocated = self.cache.allocate_blocks(blocks_needed, state)
+            allocated = self.cache.allocate_blocks(blocks_needed, state.request_id, state.allocated_blocks)
             if allocated is None:
                 return False
             state.allocated_blocks += allocated
         return True
 
-    @traced(span_name="prepare_request")
-    def _prepare_request_for_processing(
-        self, state: RequestState, token_budget: int, request_ids_to_remove_from_waiting: set[str]
-    ) -> None:
+    def _infer_request_tokens(self, state: RequestState, request_ids_to_remove_from_waiting: set[str]) -> list[int]:
         """Prepares a request for processing in the current batch. If prefix sharing is enabled, and the request was
         pending, this is where we look for a prefix match and split the request if found."""
         # If prefix sharing is enabled, we look for a prefix match and split the request if found
@@ -139,6 +136,8 @@ class Scheduler(ABC):
                 self.active_requests[state.request_id] = state
                 request_ids_to_remove_from_waiting.add(state.request_id)
                 state.status = RequestStatus.SPLIT_PENDING_REMAINDER
+                # We keep track of the number of allocated blocks to avoid double allocation
+                state.allocated_blocks += prefill_length // self.cache.block_size
                 # Even if we match the whole request, we keep at least 1 token to start decoding
                 prefill_length = min(prefill_length, len(state.tokens_to_process) - 1)
                 state.remaining_prefill_tokens = state.tokens_to_process[prefill_length:]
@@ -151,8 +150,19 @@ class Scheduler(ABC):
         # Otherwise, the tokens to process are the prompt ids, which are the full prompt or the last predicted tokens
         else:
             request_tokens = state.tokens_to_process
-
-        # If the request has one or more children we make sure not to prefill it entrirely
+        return request_tokens
+
+    def _schedule_request(
+        self,
+        state: RequestState,
+        request_tokens: list[int],
+        token_budget: int,
+        request_ids_to_remove_from_waiting: set[str],
+    ) -> None:
+        """Schedules a request for the current batch, updating the request's status according to the token budget left.
+        If the request has children (for parallel decoding), it ensures at least one token remains before the request is
+        forked."""
+        # If the request has one or more children we make sure not to prefill it entirely
         if state.num_children > 0 and token_budget >= len(request_tokens) - 1:
             token_budget = len(request_tokens) - 1
             self._requests_to_fork.append(state)
@@ -196,7 +206,7 @@ class FIFOScheduler(Scheduler):
         self.safety_margin = safety_margin
 
     @traced
-    def schedule_batch(self, token_budget: int) -> list[RequestState]:
+    def schedule_batch(self, token_budget: int, cache_budget: int) -> list[RequestState] | None:
         priority_states: list[RequestState] = []
         second_priority_states: list[RequestState] = []
         scheduled_requests = []
@@ -208,33 +218,60 @@ class FIFOScheduler(Scheduler):
                 second_priority_states.append(state)
 
         # Add waiting requests to second priority
-        for req_id in self.waiting_requests_order:
-            second_priority_states.append(self.waiting_requests[req_id])
+        if not self.block_new_requests:
+            for req_id in self.waiting_requests_order:
+                second_priority_states.append(self.waiting_requests[req_id])
 
         candidates = priority_states + second_priority_states
         request_ids_to_remove_from_waiting = set()
         safety_margins = self.safety_margin * self.cache.num_blocks
 
+        one_allocation_failed = False
+
         for state in candidates:
             # If we are out the safety margin, we only accept decoding requests or the first prefill request
             num_free_blocks = self.cache.get_num_free_blocks()
             outside_safety_margin = num_free_blocks < safety_margins
             if outside_safety_margin and scheduled_requests and state.status != RequestStatus.DECODING:
+                logger.info(
+                    f"Outside safety margin, breaking out of scheduling loop. {num_free_blocks = } {safety_margins = }"
+                )
                 break
 
-            self._prepare_request_for_processing(state, token_budget, request_ids_to_remove_from_waiting)
-            request_len = len(state.tokens_to_process)
-            # If we can't allocate blocks, do not schedule the request and break if the cache is full
-            if not self._allocate_blocks_if_needed(state):
-                if self.cache.get_num_free_blocks() == 0:
+            # Check cache budget
+            cache_needed = state.current_len()
+            cache_needed = (
+                cache_needed if self.cache_budget_module is None else cache_needed % self.cache_budget_module
+            )
+            if cache_budget < cache_needed:
+                continue
+
+            # Infer the tokens that will be present in the batch if token budget is enough
+            request_tokens = self._infer_request_tokens(state, request_ids_to_remove_from_waiting)
+            # Account for token budget
+            request_len = min(len(request_tokens), token_budget)
+            # Check there will be enough cache for the new tokens
+            allocation_successful = self._allocate_blocks_if_needed(state, request_len)
+
+            # If the allocation would not be successful, we move on to the next request
+            if not allocation_successful:
+                one_allocation_failed = True
+                # If we have reached a request that was waiting, all subsequent requests are also waiting, and will need
+                # allocation as well. So if there is no more free blocks, we can safely break out of the loop.
+                if num_free_blocks == 0 and state.request_id in self.waiting_requests:
+                    logger.info(f"Breaking mid-loop for request {state.request_id} because the cache is full")
                     break
                 continue
 
-            # Add the request to the scheduled requests
+            # If this point is reached, it means we can safely schedule the request
+            self._schedule_request(state, request_tokens, token_budget, request_ids_to_remove_from_waiting)
+            request_len = len(state.tokens_to_process)  # it may change after scheduling
             scheduled_requests.append(state)
 
-            # Update the token budget
+            # Update the token and cache budgets
             token_budget -= request_len
+            cache_budget -= cache_needed
+
             # If using prefix sharing, we make note of the blocks that will be computed in the forward pass
             if self.cache.allow_block_sharing:
                 tokens_in_current_block = state.current_len() % self.cache.block_size
@@ -248,18 +285,24 @@ class FIFOScheduler(Scheduler):
             if was_waiting:
                 request_ids_to_remove_from_waiting.add(req_id)
 
-            # Early exit of the loop if we have no token budget left
-            if token_budget == 0:
+            # Early exit of the loop if we have no budget left
+            if token_budget == 0 or cache_budget == 0:
                 break
 
+        # We remove waiting requests before checking requests were scheduled, because there might have been prefill matches
         self.waiting_requests_order = deque(
             [req_id for req_id in self.waiting_requests_order if req_id not in request_ids_to_remove_from_waiting]
         )
 
+        # If no requests were scheduled and the cache is full, we signal it by returning None
+        if not scheduled_requests and one_allocation_failed:
+            return None
+
         return scheduled_requests
 
 
 # FIXME: prioritize adding from waiting reqs before scheduling `RequestStatus.DECODING` when cache space allows it
+# TODO: further consolidate the code by making more of it common. The reference Scheduler is FIFO, not this one.
 @attach_tracer()
 class PrefillFirstScheduler(Scheduler):
     """Scheduler that prioritizes split prefill requests over decoding requests. This scheduler ensures that split
@@ -267,7 +310,7 @@ class PrefillFirstScheduler(Scheduler):
     decoding requests."""
 
     @traced
-    def schedule_batch(self, token_budget: int) -> list[RequestState]:
+    def schedule_batch(self, token_budget: int, cache_budget: int) -> list[RequestState] | None:
         priority_states: list[RequestState] = []
         second_priority_states: list[RequestState] = []
         scheduled_requests = []
@@ -280,27 +323,47 @@ class PrefillFirstScheduler(Scheduler):
                 second_priority_states.append(state)
 
         # Add waiting requests to second priority
-        for req_id in self.waiting_requests_order:
-            second_priority_states.append(self.waiting_requests[req_id])
+        if not self.block_new_requests:
+            for req_id in self.waiting_requests_order:
+                second_priority_states.append(self.waiting_requests[req_id])
 
         candidates = priority_states + second_priority_states
-
         request_ids_to_remove_from_waiting = set()
+        one_allocation_failed = False
 
         for state in candidates:
-            self._prepare_request_for_processing(state, token_budget, request_ids_to_remove_from_waiting)
-            request_len = len(state.tokens_to_process)
-            # If we can't allocate blocks, do not schedule the request and break if the cache is full
-            if not self._allocate_blocks_if_needed(state):
-                if self.cache.get_num_free_blocks() == 0:
+            # Check cache budget
+            cache_needed = state.current_len()
+            cache_needed = (
+                cache_needed if self.cache_budget_module is None else cache_needed % self.cache_budget_module
+            )
+            if cache_budget < cache_needed:
+                continue
+
+            # Infer the tokens that will be present in the batch if token budget is enough
+            request_tokens = self._infer_request_tokens(state, request_ids_to_remove_from_waiting)
+            # Account for token budget
+            request_len = min(len(request_tokens), token_budget)
+            # Check there will be enough cache for the new tokens
+            allocation_successful = self._allocate_blocks_if_needed(state, request_len)
+
+            # If the allocation would not be successful, we move on to the next request
+            if not allocation_successful:
+                one_allocation_failed = True
+                # If the request was waiting, all requests afterwards will need allocation, so we break if the cache is full
+                if state.request_id in self.waiting_requests and self.cache.get_num_free_blocks() == 0:
                     break
                 continue
 
-            # Add the request to the scheduled requests
+            # If this point is reached, it means we can safely schedule the request
+            self._schedule_request(state, request_tokens, token_budget, request_ids_to_remove_from_waiting)
+            request_len = len(state.tokens_to_process)  # it may change after scheduling
             scheduled_requests.append(state)
 
-            # Update the token budget
+            # Update the token and cache budgets
             token_budget -= request_len
+            cache_budget -= cache_needed
+
             # If using prefix sharing, we make note of the blocks that will be computed in the forward pass
             if self.cache.allow_block_sharing:
                 tokens_in_current_block = state.current_len() % self.cache.block_size
@@ -310,18 +373,23 @@ class PrefillFirstScheduler(Scheduler):
 
             # Remove the request from the waiting queue and mark it as removed
             req_id = state.request_id
-            if req_id in self.waiting_requests:
-                del self.waiting_requests[req_id]
+            was_waiting = self.waiting_requests.pop(req_id, None) is not None
+            if was_waiting:
                 request_ids_to_remove_from_waiting.add(req_id)
 
-            # Early exit of the loop if we have no token budget left
-            if token_budget == 0:
+            # Early exit of the loop if we have no budget left
+            if token_budget == 0 or cache_budget == 0:
                 break
 
+        # We remove waiting requests before checking requests were scheduled, because there might have been prefill matches
         self.waiting_requests_order = deque(
             [req_id for req_id in self.waiting_requests_order if req_id not in request_ids_to_remove_from_waiting]
         )
 
+        # If no requests were scheduled and the cache is full, we signal it by returning None
+        if not scheduled_requests and one_allocation_failed:
+            return None
+
         return scheduled_requests
 
 

transformers/generation/logits_process.py

@@ -1,4 +1,3 @@
-# coding=utf-8
 # Copyright 2024 The HuggingFace Inc. team and Google DeepMind.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -1543,133 +1542,6 @@ class PrefixConstrainedLogitsProcessor(LogitsProcessor):
         return scores_processed
 
 
-class HammingDiversityLogitsProcessor(LogitsProcessor):
-    r"""
-    [`LogitsProcessor`] that enforces diverse beam search.
-    Note that this logits processor is only effective for [`PreTrainedModel.group_beam_search`]. See [Diverse Beam
-    Search: Decoding Diverse Solutions from Neural Sequence Models](https://huggingface.co/papers/1610.02424) for more
-    details.
-    Traditional beam search often generates very similar sequences across different beams.
-    `HammingDiversityLogitsProcessor` addresses this by penalizing beams that generate tokens already chosen by other
-    beams in the same time step.
-    Args:
-        diversity_penalty (`float`):
-            This value is subtracted from a beam's score if it generates a token same as any beam from other group at a
-            particular time. A higher `diversity_penalty` will enforce greater diversity among the beams. Adjusting
-            this value can help strike a balance between diversity and natural likelihood.
-        num_beams (`int`):
-            Number of beams for beam search. 1 means no beam search.
-        num_beam_groups (`int`):
-            Number of groups to divide `num_beams` into in order to ensure diversity among different groups of beams.
-            [this paper](https://huggingface.co/papers/1610.02424) for more details.
-    Examples:
-    ```python
-    >>> from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-    >>> import torch
-    >>> # Initialize the model and tokenizer
-    >>> tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-base")
-    >>> model = AutoModelForSeq2SeqLM.from_pretrained("google-t5/t5-base")
-    >>> # A long text about the solar system
-    >>> text = (
-    ...     "The Solar System is a gravitationally bound system comprising the Sun and the objects that orbit it, "
-    ...     "either directly or indirectly. Of the objects that orbit the Sun directly, the largest are the eight "
-    ...     "planets, with the remainder being smaller objects, such as the five dwarf planets and small Solar System "
-    ...     "bodies. The Solar System formed 4.6 billion years ago from the gravitational collapse of a giant "
-    ...     "interstellar molecular cloud."
-    ... )
-    >>> inputs = tokenizer("summarize: " + text, return_tensors="pt")
-    >>> # Generate diverse summary
-    >>> outputs_diverse = model.generate(
-    ...     **inputs,
-    ...     num_beam_groups=2,
-    ...     diversity_penalty=10.0,
-    ...     max_length=100,
-    ...     num_beams=4,
-    ...     num_return_sequences=2,
-    ... )
-    >>> summaries_diverse = tokenizer.batch_decode(outputs_diverse, skip_special_tokens=True)
-    >>> # Generate non-diverse summary
-    >>> outputs_non_diverse = model.generate(
-    ...     **inputs,
-    ...     max_length=100,
-    ...     num_beams=4,
-    ...     num_return_sequences=2,
-    ... )
-    >>> summary_non_diverse = tokenizer.batch_decode(outputs_non_diverse, skip_special_tokens=True)
-    >>> # With `diversity_penalty`, the resulting beams are much more diverse
-    >>> print(summary_non_diverse)
-    ['the solar system formed 4.6 billion years ago from the collapse of a giant interstellar molecular cloud. of the objects that orbit the Sun directly, the largest are the eight planets.',
-    'the Solar System formed 4.6 billion years ago from the collapse of a giant interstellar molecular cloud. of the objects that orbit the Sun directly, the largest are the eight planets.']
-    >>> print(summaries_diverse)
-    ['the solar system formed 4.6 billion years ago from the collapse of a giant interstellar molecular cloud. of the objects that orbit the Sun directly, the largest are the eight planets.',
-    'the solar system formed 4.6 billion years ago from the collapse of a giant interstellar molecular cloud. of the objects that orbit the Sun directly, the largest are the eight planets. the rest of the objects are smaller objects, such as the five dwarf planets and small solar system bodies.']
-    ```
-    """
-
-    def __init__(self, diversity_penalty: float, num_beams: int, num_beam_groups: int):
-        logger.warning_once(
-            "`HammingDiversityLogitsProcessor` is deprecated and will be removed in v4.62.0, as constrained beam search has been moved to the Hub: https://hf.co/transformers-community/constrained-beam-search."
-        )
-        if not isinstance(diversity_penalty, float) or (not diversity_penalty > 0.0):
-            raise ValueError("`diversity_penalty` should be a float strictly larger than 0.")
-        self._diversity_penalty = diversity_penalty
-        if not isinstance(num_beams, int) or num_beams < 2:
-            raise ValueError("`num_beams` should be an integer strictly larger than 1.")
-        self._num_beams = num_beams
-        if not isinstance(num_beam_groups, int) or num_beam_groups < 2:
-            raise ValueError("`num_beam_groups` should be an integer strictly larger than 1.")
-        if num_beam_groups > num_beams:
-            raise ValueError("`beam_groups` has to be smaller or equal to `num_beams`.")
-        self._num_sub_beams = num_beams // num_beam_groups
-
-    def __call__(
-        self,
-        input_ids: torch.LongTensor,
-        scores: torch.FloatTensor,
-        current_tokens: torch.LongTensor,
-        beam_group_idx: int,
-    ) -> torch.FloatTensor:
-        r"""
-        Args:
-            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
-                Indices of input sequence tokens in the vocabulary. [What are input IDs?](../glossary#input-ids)
-            scores (`torch.FloatTensor` of shape `(batch_size, config.vocab_size)`):
-                Prediction scores of a language modeling head. These can be logits for each vocabulary when not using
-                beam search or log softmax for each vocabulary token when using beam search
-            current_tokens (`torch.LongTensor` of shape `(batch_size)`):
-                Indices of input sequence tokens in the vocabulary, corresponding to the tokens selected by the other
-                beam groups in the current generation step.
-            beam_group_idx (`int`):
-                The index of the beam group currently being processed.
-        Return:
-            `torch.FloatTensor` of shape `(batch_size, config.vocab_size)`:
-                The processed prediction scores.
-        """
-        # hamming diversity: penalise using same token in current group which was used in previous groups at
-        # the same time step
-        batch_size = current_tokens.shape[0] // self._num_beams
-        group_start_idx = beam_group_idx * self._num_sub_beams
-        group_end_idx = min(group_start_idx + self._num_sub_beams, self._num_beams)
-        group_size = group_end_idx - group_start_idx
-        vocab_size = scores.shape[-1]
-
-        if group_start_idx == 0:
-            return scores
-
-        scores_processed = scores.clone()
-        for batch_idx in range(batch_size):
-            # predicted tokens of last time step of previous groups
-            previous_group_tokens = current_tokens[
-                batch_idx * self._num_beams : batch_idx * self._num_beams + group_start_idx
-            ]
-            token_frequency = torch.bincount(previous_group_tokens, minlength=vocab_size).to(scores.device)
-            scores_processed[batch_idx * group_size : (batch_idx + 1) * group_size] -= (
-                self._diversity_penalty * token_frequency
-            )
-
-        return scores_processed
-
-
 class ForcedBOSTokenLogitsProcessor(LogitsProcessor):
     r"""
     [`LogitsProcessor`] that enforces the specified token as the first generated token. Used with encoder-decoder

transformers/generation/streamers.py

@@ -1,4 +1,3 @@
-# coding=utf-8
 # Copyright 2023 The HuggingFace Inc. team.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

transformers/generation/utils.py

@@ -1,4 +1,3 @@
-# coding=utf-8
 # Copyright 2020 The Google AI Language Team Authors, Facebook AI Research authors and The HuggingFace Inc. team.
 # Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
 #
@@ -21,7 +20,7 @@ import warnings
 from collections.abc import Callable
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Any, Optional
 
 import torch
 import torch.distributed as dist
@@ -332,9 +331,9 @@ class GenerateBeamEncoderDecoderOutput(ModelOutput):
 
 
 # Typing shortcuts
-GenerateNonBeamOutput = Union[GenerateDecoderOnlyOutput, GenerateEncoderDecoderOutput]
-GenerateBeamOutput = Union[GenerateBeamDecoderOnlyOutput, GenerateBeamEncoderDecoderOutput]
-GenerateOutput = Union[GenerateNonBeamOutput, GenerateBeamOutput]
+GenerateNonBeamOutput = GenerateDecoderOnlyOutput | GenerateEncoderDecoderOutput
+GenerateBeamOutput = GenerateBeamDecoderOnlyOutput | GenerateBeamEncoderDecoderOutput
+GenerateOutput = GenerateNonBeamOutput | GenerateBeamOutput
 
 
 class GenerationMixin(ContinuousMixin):
@@ -1792,26 +1791,15 @@ class GenerationMixin(ContinuousMixin):
         generation_config = copy.deepcopy(generation_config)
 
         # First set values from the loaded `self.generation_config`, then set default values (BC)
-        # Do not update any values that aren't `None`, i.e. if set by users explicitly and passed
-        # to `generate()`. Thus the `defaults_only=True` is used
+        #
+        # Only update values that are `None`, i.e. these values were not explicitly set by users to `generate()`,
+        # or values that are not present in the current config, i.e. custom entries that were set via `**kwargs`.
+        # Thus we use the specific kwargs `defaults_only=True` (`None` values only) and `allow_custom_entries=True`
+        # (custom entries are carried over).
         global_defaults = self.generation_config._get_default_generation_params()
-        generation_config.update(**self.generation_config.to_dict(), defaults_only=True)
+        generation_config.update(**self.generation_config.to_dict(), defaults_only=True, allow_custom_entries=True)
         generation_config.update(**global_defaults, defaults_only=True)
 
-        # Due to some values being boolean and not `None`, we need additional logic to overwrite
-        # them explicitly (`defaults_only=False`) on the condition that it's only a previous
-        # default value
-        default_generation_config = GenerationConfig()
-        generation_config.update(
-            **{
-                k: v
-                for k, v in self.generation_config.to_dict().items()
-                if isinstance(v, bool)
-                and hasattr(default_generation_config, k)
-                and getattr(generation_config, k, None) == getattr(default_generation_config, k)
-            }
-        )
-
         # Finally, if there are any kwargs, update config with it -> highest priority at the end
         model_kwargs = generation_config.update(**kwargs)
 
@@ -1914,6 +1902,7 @@ class GenerationMixin(ContinuousMixin):
         # NOTE: remove xlnet/reformer when the models are deprecated, non-standard model architecture/cache name
         return not cls._is_stateful and all(
             special_model_name not in cls.__name__.lower()
+            or "minimaxm2" in cls.__name__.lower()  # name clash between minimax and minimax m2
             for special_model_name in [
                 "reformer",
                 "minimax",
@@ -2429,7 +2418,7 @@ class GenerationMixin(ContinuousMixin):
                 raise NotImplementedError(
                     f"assistant_model is not supported for continuous batching. Got {assistant_model = }"
                 )
-            if streamer is not None:  # TODO: actualy this could be supported
+            if streamer is not None:  # TODO: actually this could be supported
                 raise NotImplementedError(f"streaming is not supported for continuous batching. Got {streamer = }")
             if negative_prompt_ids is not None:
                 raise NotImplementedError(
@@ -2485,9 +2474,11 @@ class GenerationMixin(ContinuousMixin):
         generation_config, model_kwargs = self._prepare_generation_config(generation_config, **kwargs)
 
         generation_mode = generation_config.get_generation_mode(assistant_model)
+        deprecated_mode_repo = self._get_deprecated_gen_repo(generation_mode, trust_remote_code, custom_generate)
+
         if isinstance(custom_generate, Callable):
             decoding_method = custom_generate
-        else:
+        elif deprecated_mode_repo is None:
             # type() required to access the unbound class-level method
             decoding_method = getattr(type(self), GENERATION_MODES_MAPPING[generation_mode])
 
@@ -2498,7 +2489,7 @@ class GenerationMixin(ContinuousMixin):
         # NOTE: This must come after initializing generation_config, since we need it to determine if this is a deprecated mode.
         # It must also be before any preparation steps, since Hub repos expect to be loaded before preparation steps.
         # TODO joao, manuel: remove this in v4.62.0
-        if deprecated_mode_repo := self._get_deprecated_gen_repo(generation_mode, trust_remote_code, custom_generate):
+        if deprecated_mode_repo is not None:
             return GenerationMixin.generate(
                 self,
                 inputs=inputs,
@@ -3875,7 +3866,6 @@ class GenerationMixin(ContinuousMixin):
                 model_kwargs["cache_position"] = torch.arange(
                     past_length, current_length, dtype=torch.long, device=input_chunk.device
                 )
-                model_kwargs["position_ids"] = model_kwargs["cache_position"].unsqueeze(0)
                 model_inputs = self.prepare_inputs_for_generation(input_chunk, **model_kwargs)
 
                 outputs = model_forward(**model_inputs, return_dict=True)