Repository: huggingface/transformers
Pull request: [CB] [Serve] Fix broken serve tests

This PR fixes two issues related to CB and serve:
- when the `max_length` argument as set to `None` when creating a request, CB had an error because it always expected max_length to be an int. This is now an expected behavior, and we have a hidden variable to avoid checking if the variable is None
- the CB streamer did not yield the last generated token when the request finished. This was not an issue before the last PR, because the result of each request was yielded twice (with the same amount of tokens). Now that this is fixed, it revealed the bug related to the streamer, which is now fixed.

Changed files:
filename: src/transformers/cli/serve.py
patch: @@ -839,6 +839,17 @@ def stream_chat_completion(request_id, decode_stream):
                 for result in self.running_continuous_batching_manager.request_id_iter(request_id):
                     n_tokens_generated += 1
 
+                    # Always yield the token content (even for the final FINISHED token)
+                    if result.generated_tokens:
+                        token_id = result.generated_tokens[-1]
+                        yield self.build_chat_completion_chunk(
+                            request_id=request_id,
+                            content=token_id,
+                            model=model_id_and_revision,
+                            decode_stream=decode_stream,
+                            tokenizer=tokenizer,
+                        )
+
                     if result.status == RequestStatus.FINISHED:
                         generated_all_tokens = n_tokens_generated >= generation_config.max_new_tokens
 
@@ -855,14 +866,6 @@ def stream_chat_completion(request_id, decode_stream):
                             model=model_id_and_revision,
                         )
                         break
-                    else:
-                        yield self.build_chat_completion_chunk(
-                            request_id=request_id,
-                            content=result.generated_tokens[-1],
-                            model=model_id_and_revision,
-                            decode_stream=decode_stream,
-                            tokenizer=tokenizer,
-                        )
 
             except Exception as e:
                 logger.error(str(e))

filename: src/transformers/generation/continuous_batching/requests.py
patch: @@ -129,7 +129,7 @@ class RequestState:
         position_offset (int): The current position in the sequence for position_ids.
         status (RequestStatus): The status of the request: can be one of PENDING, PREFILLING, PREFILLING_SPLIT,
                                 SPLIT_PENDING_REMAINDER, DECODING, FINISHED, FAILED
-        max_new_tokens (int): The maximum number of new tokens to generate.
+        max_new_tokens (int | None): The maximum number of new tokens to generate.
         eos_token_id (int): The ID of the end-of-sequence token.
         streaming (bool): Whether to stream tokens as they're generated
         created_time (float): The time the request was created.
@@ -149,7 +149,7 @@ class RequestState:
     allocated_blocks: int = 0  # Number of blocks allocated to the request
     position_offset: int = 0  # Current position in the sequence for position_ids
     _status: RequestStatus = RequestStatus.PENDING  # Status of the request, hidden behind a property
-    max_new_tokens: int = 20  # Maximum number of new tokens to generate
+    max_new_tokens: int | None = 20  # Maximum number of new tokens to generate. None means no limit. Default to 20.
     eos_token_id: int = -1  # ID of the end-of-sequence token
     streaming: bool = False  # Whether to stream tokens as they're generated
     created_time: float = field(default_factory=time.perf_counter)  # Time the request was created
@@ -158,6 +158,11 @@ class RequestState:
     _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
+    _new_tokens_limit: int = 2147483647  # An int to check the max number of new tokens w/out always comparing w/ None
+
+    def __post_init__(self):
+        # If no max length is set, we set an absurdly high value which will never be reached
+        self._new_tokens_limit = 2147483647 if self.max_new_tokens is None else self.max_new_tokens
 
     @property
     def status(self) -> RequestStatus:
@@ -218,14 +223,14 @@ def update_and_check_completion(self, token_id: int) -> bool:
 
         # Replace the temporary token if we're not finishing due to max length
         # (EOS tokens should still be added to the output)
-        if is_eos or (current_len < self.max_new_tokens):
+        if is_eos or (current_len < self._new_tokens_limit):
             self.generated_tokens[-1] = token_id
             current_len += 1
         else:
             logger.warning(f"Request {self.request_id} generated a useless token: {token_id}")
             self.generated_tokens.pop()
 
-        if is_eos or current_len >= self.max_new_tokens:
+        if is_eos or current_len >= self._new_tokens_limit:
             self.status = RequestStatus.FINISHED
             return True
         return False  # We still need to process more tokens
@@ -294,13 +299,14 @@ def create_equivalent_initial_request(self) -> "RequestState":
         # Remove the temporary token if it exists
         if self.generated_tokens and self.generated_tokens[-1] == TMP_TOKEN_ID:
             self.generated_tokens.pop()
+        max_new_tokens = None if self.max_new_tokens is None else (self.max_new_tokens - len(self.generated_tokens))
         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),
+            max_new_tokens=max_new_tokens,
             eos_token_id=self.eos_token_id,
             streaming=self.streaming,
         )