transformers 5.0.0rc2__py3-none-any.whl → 5.1.0__py3-none-any.whl

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,25 +55,25 @@ 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] | None:
+        """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:
         """Checks if there are requests ready to be processed."""
-        return len(self.active_requests) or len(self.waiting_requests)
+        return bool(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,21 @@ 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.
+        After a request is scheduled, it is part of the next batch unless there is an error.
+        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
+        # This does not check the request state, but DECODING request already have children set to 0.
         if state.num_children > 0 and token_budget >= len(request_tokens) - 1:
             token_budget = len(request_tokens) - 1
             self._requests_to_fork.append(state)
@@ -179,62 +191,68 @@ class Scheduler(ABC):
             state.remaining_prefill_tokens = request_tokens[token_budget:]
             state.tokens_to_process = request_tokens[:token_budget]
 
-
-# TODO: further common-ize the two classes
-@attach_tracer()
-class FIFOScheduler(Scheduler):
-    """This scheduler processes requests in the order they arrive, meaning decoding requests has priority over
-    prefilling requests. Additionally, it includes a safety margin mechanism to prevent cache exhaustion. By default,
-    when 80% of the cache is full, new requests will not be scheduled to prioritize decoding active requests."""
-
-    def __init__(self, cache: PagedAttentionCache, retain_cache_on_finish: bool = False, safety_margin: float = 0.2):
-        """Initializes the FIFO scheduler. The safety margin is the percentage of free blocks under which we stop
-        scheduling new prefill requests, so safety_margin = 0.1 means that when there is less than 10% of free blocks,
-        or equivalently when more than 90% of blocks are already allocated, we stop scheduling new prefill requests.
+    def _process_candidates(
+        self,
+        candidates: list[RequestState],
+        token_budget: int,
+        cache_budget: int,
+        request_ids_to_remove_from_waiting: set[str],
+        safety_margin: float = 0.0,
+    ) -> tuple[list[RequestState], bool]:
+        """Schedules candidate requests for the current batch.
+
+        This method contains the common logic shared by all schedulers: it checks token and cache budgets, allocates
+        cache blocks if needed, updates request states, and tracks which waiting requests should be removed from the
+        waiting queue.
         """
-        super().__init__(cache, retain_cache_on_finish)
-        self.safety_margin = safety_margin
-
-    @traced
-    def schedule_batch(self, token_budget: int) -> list[RequestState]:
-        priority_states: list[RequestState] = []
-        second_priority_states: list[RequestState] = []
         scheduled_requests = []
-
-        for state in self.active_requests.values():
-            if state.status == RequestStatus.DECODING:
-                priority_states.append(state)
-            if state.status in [RequestStatus.SPLIT_PENDING_REMAINDER, RequestStatus.PREFILLING_SPLIT]:
-                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])
-
-        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
+        safety_margins = safety_margin * self.cache.num_blocks
 
         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()
+            # If we are out the safety margin, we only accept decoding requests or the first prefill request
             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 reached a waiting request and the cache is full, all subsequent waiting requests will need
+                # allocation as well, so we can safely break out of the scheduling 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 +266,72 @@ 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
 
+        return scheduled_requests, one_allocation_failed
+
+    def _cleanup_waiting_queue(self, request_ids_to_remove_from_waiting: set[str]) -> None:
+        """Removes processed requests from the waiting queue order."""
         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]
         )
 
+
+# TODO: further common-ize the two classes
+@attach_tracer()
+class FIFOScheduler(Scheduler):
+    """This scheduler processes requests in the order they arrive, meaning decoding requests has priority over
+    prefilling requests. Additionally, it includes a safety margin mechanism to prevent cache exhaustion. By default,
+    when 80% of the cache is full, new requests will not be scheduled to prioritize decoding active requests."""
+
+    def __init__(self, cache: PagedAttentionCache, retain_cache_on_finish: bool = False, safety_margin: float = 0.2):
+        """Initializes the FIFO scheduler. The safety margin is the percentage of free blocks under which we stop
+        scheduling new prefill requests, so safety_margin = 0.1 means that when there is less than 10% of free blocks,
+        or equivalently when more than 90% of blocks are already allocated, we stop scheduling new prefill requests.
+        """
+        super().__init__(cache, retain_cache_on_finish)
+        self.safety_margin = safety_margin
+
+    @traced
+    def schedule_batch(self, token_budget: int, cache_budget: int) -> list[RequestState] | None:
+        priority_states: list[RequestState] = []
+        second_priority_states: list[RequestState] = []
+
+        for state in self.active_requests.values():
+            if state.status == RequestStatus.DECODING:
+                priority_states.append(state)
+            if state.status in [RequestStatus.SPLIT_PENDING_REMAINDER, RequestStatus.PREFILLING_SPLIT]:
+                second_priority_states.append(state)
+
+        # Add waiting requests to second priority
+        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()
+        scheduled_requests, one_allocation_failed = self._process_candidates(
+            candidates,
+            token_budget,
+            cache_budget,
+            request_ids_to_remove_from_waiting,
+            safety_margin=self.safety_margin,
+        )
+
+        # We remove waiting requests before checking requests were scheduled, because there might have been prefill matches
+        self._cleanup_waiting_queue(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,10 +339,9 @@ 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 = []
 
         for state in self.active_requests.values():
             # XXX: when cache is full, state can stay on `PREFILLING_SPLIT` so we need to take those into account
@@ -280,47 +351,26 @@ 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()
+        scheduled_requests, one_allocation_failed = self._process_candidates(
+            candidates,
+            token_budget,
+            cache_budget,
+            request_ids_to_remove_from_waiting,
+            safety_margin=0.0,
+        )
 
-        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:
-                    break
-                continue
-
-            # Add the request to the scheduled requests
-            scheduled_requests.append(state)
-
-            # Update the token budget
-            token_budget -= request_len
-            # 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
-                tokens_after_forward = tokens_in_current_block + request_len
-                complete_blocks = tokens_after_forward // self.cache.block_size
-                self.cache.blocks_to_complete[state.request_id] = complete_blocks
-
-            # 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]
-                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:
-                break
+        # We remove waiting requests before checking requests were scheduled, because there might have been prefill matches
+        self._cleanup_waiting_queue(request_ids_to_remove_from_waiting)
 
-        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");
@@ -21,7 +20,6 @@ from typing import TYPE_CHECKING
 import numpy as np
 import torch
 
-from ..pytorch_utils import isin_mps_friendly
 from ..utils import add_start_docstrings
 from ..utils.logging import get_logger
 
@@ -94,6 +92,12 @@ class LogitsProcessorList(list):
 
         return scores
 
+    def set_continuous_batching_context(self, logits_indices: torch.Tensor, cu_seq_lens_q: torch.Tensor) -> None:
+        """Forwards the continuous batching metadata to all logit processors that need it."""
+        for processor in self:
+            if hasattr(processor, "set_continuous_batching_context"):
+                processor.set_continuous_batching_context(logits_indices, cu_seq_lens_q)
+
 
 class MinLengthLogitsProcessor(LogitsProcessor):
     r"""
@@ -149,7 +153,7 @@ class MinLengthLogitsProcessor(LogitsProcessor):
     @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
         vocab_tensor = torch.arange(scores.shape[-1], device=scores.device)
-        eos_token_mask = isin_mps_friendly(vocab_tensor, self.eos_token_id)
+        eos_token_mask = torch.isin(vocab_tensor, self.eos_token_id)
         scores_processed = scores.clone()
         if input_ids.shape[-1] < self.min_length:
             scores_processed = torch.where(eos_token_mask, -math.inf, scores)
@@ -221,7 +225,7 @@ class MinNewTokensLengthLogitsProcessor(LogitsProcessor):
         new_tokens_length = input_ids.shape[-1] - self.prompt_length_to_skip
         scores_processed = scores.clone()
         vocab_tensor = torch.arange(scores.shape[-1], device=scores.device)
-        eos_token_mask = isin_mps_friendly(vocab_tensor, self.eos_token_id)
+        eos_token_mask = torch.isin(vocab_tensor, self.eos_token_id)
         if new_tokens_length < self.min_new_tokens:
             scores_processed = torch.where(eos_token_mask, -math.inf, scores)
 
@@ -1543,133 +1547,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
@@ -1975,7 +1852,7 @@ class SuppressTokensAtBeginLogitsProcessor(LogitsProcessor):
     @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
         vocab_tensor = torch.arange(scores.shape[-1], device=scores.device)
-        suppress_token_mask = isin_mps_friendly(vocab_tensor, self.begin_suppress_tokens)
+        suppress_token_mask = torch.isin(vocab_tensor, self.begin_suppress_tokens)
         scores_processed = scores
         if input_ids.shape[-1] == self.begin_index:
             scores_processed = torch.where(suppress_token_mask, -float("inf"), scores)
@@ -2018,7 +1895,7 @@ class SuppressTokensLogitsProcessor(LogitsProcessor):
     @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
         vocab_tensor = torch.arange(scores.shape[-1], device=scores.device)
-        suppress_token_mask = isin_mps_friendly(vocab_tensor, self.suppress_tokens.to(scores.device))
+        suppress_token_mask = torch.isin(vocab_tensor, self.suppress_tokens.to(scores.device))
         scores = torch.where(suppress_token_mask, -float("inf"), scores)
         return scores
 

transformers/generation/stopping_criteria.py

@@ -8,7 +8,6 @@ import numpy as np
 import torch
 from torch.nn import functional as F
 
-from ..pytorch_utils import isin_mps_friendly
 from ..tokenization_utils_base import PreTrainedTokenizerBase
 from ..utils import add_start_docstrings, logging
 
@@ -468,7 +467,7 @@ class EosTokenCriteria(StoppingCriteria):
     @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING)
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.BoolTensor:
         self.eos_token_id = self.eos_token_id.to(input_ids.device)
-        is_done = isin_mps_friendly(input_ids[:, -1], self.eos_token_id)
+        is_done = torch.isin(input_ids[:, -1], self.eos_token_id)
         return is_done
 
 

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");