Repository: huggingface/transformers
Pull request: [chat] Split chat/serve (built on top of lysandre's PR)

# What does this PR do?

Adds the following to #38443

- Streaming tool calls
- Persistent KV cache across server calls
- BC `chat` generation parameterization
- docs

Changed files:
filename: docs/source/en/conversations.md
patch: @@ -27,6 +27,9 @@ This guide shows you how to quickly start chatting with Transformers from the co
 
 ## transformers CLI
 
+
+### Interactive chat session
+
 After you've [installed Transformers](./installation.md), chat with a model directly from the command line as shown below. It launches an interactive session with a model, with a few base commands listed at the start of the session.
 
 ```bash
@@ -51,6 +54,68 @@ transformers chat -h
 
 The chat is implemented on top of the [AutoClass](./model_doc/auto), using tooling from [text generation](./llm_tutorial) and [chat](./chat_templating).
 
+
+### Serving a model and using MCP tools
+
+> [!WARNING]
+> This section is experimental and subject to changes in future versions
+
+Powering the `chat` interface, we have a server that takes user messages and returns completions. The server has a chat completion API compatible with the OpenAI SDK, so you can also quickly experiment with `transformers` models on existing aplications. To launch a server separately, use the `transformers serve` CLI:
+
+```bash
+transformers serve Menlo/Jan-nano
+```
+
+Under the hood, the `chat` CLI launches and uses `transformers serve`. This server is also an MCP client, which can receive information available MCP servers (i.e. tools), massage their information into the model prompt, and prepare calls to these tools when the model commands to do so. Naturally, this requires a model that is trained to use tools.
+
+At the moment, MCP tool usage in `transformers` has the following constraints:
+- `chat` can't handle tools, but the [`tiny-agents`](https://huggingface.co/blog/python-tiny-agents) CLI can;
+- Only the `qwen` family of models is supported.
+
+The first step to use MCP tools is to let the model know which tools are available. As an example, let's consider a `tiny-agents` configuration file with a reference to an [image generation MCP server](https://evalstate-flux1-schnell.hf.space/).
+
+> [!TIP]
+> Many Hugging Face Spaces can be used as MCP servers. You can find all compatible Spaces [here](https://huggingface.co/spaces?filter=mcp-server).
+
+```json
+{
+    "model": "http://localhost:8000",
+    "provider": "local",
+    "servers": [
+        {
+            "type": "sse",
+            "config": {
+                "url": "https://evalstate-flux1-schnell.hf.space/gradio_api/mcp/sse"
+            }
+        }
+    ]
+}
+```
+
+You can then launch your `tiny-agents` chat interface with the following command.
+
+```bash
+tiny-agents run path/to/your/config.json
+```
+
+If you have a server (from `transformers serve`) running in the background, you're ready to use MCP tools from a local model! For instance, here's the example of a chat session:
+
+```bash
+Agent loaded with 1 tools:
+ • flux1_schnell_infer
+»  Generate an image of a cat on the moon
+<Tool req_0_tool_call>flux1_schnell_infer {"prompt": "a cat on the moon", "seed": 42, "randomize_seed": true, "width": 1024, "height": 1024, "num_inference_steps": 4}
+
+Tool req_0_tool_call
+[Binary Content: Image image/webp, 57732 bytes]
+The task is complete and the content accessible to the User
+Image URL: https://evalstate-flux1-schnell.hf.space/gradio_api/file=/tmp/gradio/3dbddc0e53b5a865ed56a4e3dbdd30f3f61cf3b8aabf1b456f43e5241bd968b8/image.webp
+380576952
+
+I have generated an image of a cat on the moon using the Flux 1 Schnell Image Generator. The image is 1024x1024 pixels and was created with 4 inference steps. Let me know if you would like to make any changes or need further assistance!
+```
+
+
 ## TextGenerationPipeline
 
 [`TextGenerationPipeline`] is a high-level text generation class with a "chat mode". Chat mode is enabled when a conversational model is detected and the chat prompt is [properly formatted](./llm_tutorial#wrong-prompt-format).

filename: src/transformers/commands/chat.py
patch: @@ -20,7 +20,6 @@
 import re
 import string
 import time
-import warnings
 from argparse import ArgumentParser, Namespace
 from dataclasses import dataclass, field
 from threading import Thread
@@ -109,19 +108,6 @@
 - **!exit**: closes the interface
 """
 
-# format: (optional CLI arg being deprecated, its current default, corresponding `generate` flag)
-_DEPRECATION_MAP = [
-    ("max_new_tokens", 256, "max_new_tokens"),
-    ("do_sample", True, "do_sample"),
-    ("num_beams", 1, "num_beams"),
-    ("temperature", 1.0, "temperature"),
-    ("top_k", 50, "top_k"),
-    ("top_p", 1.0, "top_p"),
-    ("repetition_penalty", 1.0, "repetition_penalty"),
-    ("eos_tokens", None, "eos_token_id"),
-    ("eos_token_ids", None, "eos_token_id"),
-]
-
 
 class RichInterface:
     def __init__(self, model_name: Optional[str] = None, user_name: Optional[str] = None):
@@ -252,25 +238,6 @@ class ChatArguments:
             ),
         },
     )
-    # Deprecated CLI args start here
-    max_new_tokens: int = field(default=256, metadata={"help": "Maximum number of tokens to generate."})
-    do_sample: bool = field(default=True, metadata={"help": "Whether to sample outputs during generation."})
-    num_beams: int = field(default=1, metadata={"help": "Number of beams for beam search."})
-    temperature: float = field(default=1.0, metadata={"help": "Temperature parameter for generation."})
-    top_k: int = field(default=50, metadata={"help": "Value of k for top-k sampling."})
-    top_p: float = field(default=1.0, metadata={"help": "Value of p for nucleus sampling."})
-    repetition_penalty: float = field(default=1.0, metadata={"help": "Repetition penalty."})
-    eos_tokens: Optional[str] = field(
-        default=None,
-        metadata={
-            "help": "EOS tokens (text format) to stop the generation. If multiple they should be comma separated."
-        },
-    )
-    eos_token_ids: Optional[str] = field(
-        default=None,
-        metadata={"help": "EOS token IDs to stop the generation. If multiple they should be comma separated."},
-    )
-    # Deprecated CLI args end here
 
     # Model loading
     model_revision: str = field(
@@ -346,15 +313,14 @@ def register_subcommand(parser: ArgumentParser):
                 "Flags to pass to `generate`, using a space as a separator between flags. Accepts booleans, numbers, "
                 "and lists of integers, more advanced parameterization should be set through --generation-config. "
                 "Example: `transformers chat <model_repo> max_new_tokens=100 do_sample=False eos_token_id=[1,2]`. "
-                "If you're a new user, check this basic flag guide: https://huggingface.co/docs/transformers/llm_tutorial#common-options"
+                "If you're a new user, check this basic flag guide: "
+                "https://huggingface.co/docs/transformers/llm_tutorial#common-options"
             ),
             nargs="*",
         )
         chat_parser.set_defaults(func=chat_command_factory)
 
     def __init__(self, args):
-        args = self._handle_deprecated_args(args)
-
         if args.model_name_or_path_or_address is not None:
             name = args.model_name_or_path_or_address
             if name.startswith("http") or name.startswith("https") or name.startswith("localhost"):
@@ -386,32 +352,6 @@ def __init__(self, args):
 
         self.args = args
 
-    def _handle_deprecated_args(self, args: ChatArguments) -> ChatArguments:
-        """
-        Handles deprecated arguments and their deprecation cycle. To be removed after we fully migrated to the new
-        args.
-        """
-        has_warnings = False
-
-        # Named generate option args
-        for deprecated_arg, default_value, new_arg in _DEPRECATION_MAP:
-            value = getattr(args, deprecated_arg)
-            if value != default_value:
-                has_warnings = True
-                warnings.warn(
-                    f"The --{deprecated_arg} argument is deprecated will be removed in v4.54.0. There are two "
-                    "alternative solutions to specify this generation option: \n"
-                    "1. Pass `--generation-config <path_to_file/Hub repo>` to specify a generation config.\n"
-                    "2. Pass `generate` flags through positional arguments, e.g. `transformers chat <model_repo> "
-                    f"{new_arg}={value}`",
-                    FutureWarning,
-                )
-
-        if has_warnings:
-            print("\n(Press enter to continue)")
-            input()
-        return args
-
     # -----------------------------------------------------------------------------------------------------------------
     # Chat session methods
     @staticmethod
@@ -519,7 +459,10 @@ def get_generation_parameterization(self, args: ChatArguments) -> tuple[Generati
             else:
                 generation_config = GenerationConfig.from_pretrained(args.generation_config)
         else:
-            generation_config = GenerationConfig()
+            # !!!!!!!!!
+            # This is a chat session, so we have a few non-standard defaults
+            # !!!!!!!!!
+            generation_config = GenerationConfig(do_sample=True, max_new_tokens=256)
 
         # Finally: parse and apply `generate_flags`
         parsed_generate_flags = self.parse_generate_flags(args.generate_flags)

filename: src/transformers/commands/serving.py
patch: @@ -11,13 +11,23 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import functools
 import json
+import re
 import time
 from argparse import ArgumentParser, Namespace
 from dataclasses import dataclass, field
 from threading import Thread
-from typing import Dict, Optional
+from typing import Any, Optional
+
+from huggingface_hub import (
+    ChatCompletionStreamOutputChoice,
+    ChatCompletionStreamOutputDelta,
+    ChatCompletionStreamOutputDeltaToolCall,
+    ChatCompletionStreamOutputFunction,
+    ModelInfo,
+    model_info,
+)
 
 from transformers.utils.import_utils import is_fastapi_available, is_pydantic_available, is_uvicorn_available
 
@@ -42,7 +52,7 @@
 if is_pydantic_available() and is_fastapi_available() and is_uvicorn_available():
     import uvicorn
     from fastapi import FastAPI
-    from fastapi.responses import StreamingResponse
+    from fastapi.responses import JSONResponse, StreamingResponse
     from pydantic import BaseModel
 
     class Message(BaseModel):
@@ -55,7 +65,7 @@ class ChatCompletionInput(BaseModel):
         stream: Optional[bool] = False
         model: Optional[str] = None
         request_id: Optional[str] = None
-        extra_body: Optional[Dict] = None
+        extra_body: Optional[dict] = None
         frequency_penalty: Optional[float] = None
         logit_bias: Optional[list[float]] = None
         max_tokens: Optional[int] = None
@@ -68,7 +78,7 @@ class ChatCompletionInput(BaseModel):
         # that aren't yet supported here.
 
         # logprobs: Optional[bool] = None
-        # tools: Any = None
+        tools: Any = None
         # n: Optional[int] = None
         # presence_penalty: Optional[float] = None
         # response_format: Optional[ChatCompletionInputGrammarType] = None
@@ -81,6 +91,17 @@ class ChatCompletionInput(BaseModel):
 logger = logging.get_logger(__name__)
 
 
+# Possible tokens that indicate the start/end of a tool call
+# TODO (joao, matt): streamline tool token detection logic
+_TOOL_CALL_TOKENS = {
+    "qwen": {
+        "start": "<tool_call>",
+        "end": "</tool_call>",
+    },
+}
+_MODELS_WITH_TOOL_SUPPORT = list(_TOOL_CALL_TOKENS.keys())
+
+
 def serve_command_factory(args: Namespace):
     """
     Factory function used to instantiate serving server from provided command line arguments.
@@ -90,7 +111,17 @@ def serve_command_factory(args: Namespace):
     return ServeCommand(args)
 
 
-def create_generation_config_from_req(req: "ChatCompletionInput"):
+def create_generation_config_from_req(req: "ChatCompletionInput") -> "GenerationConfig":
+    """
+    Creates a generation config from the parameters of the request. Note that we can pass a `GenerationConfig`
+    (serialized into a `dict`) in `extra_body`, for full `generate` parameterization.
+
+    Args:
+        req (`ChatCompletionInput`): The request which may optionally contain generation parameters.
+
+    Returns:
+        The prepared `GenerationConfig` object.
+    """
     if req.extra_body is not None and "generation_config" in req.extra_body:
         for key in req.extra_body["generation_config"].keys():
             if key in ChatCompletionInput.base_field_names.keys():
@@ -117,6 +148,20 @@ def create_generation_config_from_req(req: "ChatCompletionInput"):
     return generation_config
 
 
+class ToolState:
+    """Lightweight class to keep track of the tool call state."""
+
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        """Reset the tool call state (assumes we're outside a tool call)."""
+        self.inside_tool_call = False
+        self.has_tool_name_defined = False
+        self.arg_nesting_level = 0
+        self.buffer = ""
+
+
 @dataclass
 class ServeArguments:
     r"""
@@ -126,11 +171,6 @@ class ServeArguments:
     `transformers serve --help`
     """
 
-    # Model loading
-    model_revision: str = field(
-        default="main",
-        metadata={"help": "Specific model version to use (can be a branch name, tag name or commit id)."},
-    )
     device: str = field(default="cpu", metadata={"help": "Device to use for inference."})
     torch_dtype: Optional[str] = field(
         default="auto",
@@ -172,6 +212,10 @@ class ServeArguments:
 
 
 class ServeCommand(BaseTransformersCLICommand):
+    loaded_model: Optional[str] = None
+    model: PreTrainedModel
+    tokenizer: PreTrainedTokenizerFast
+
     @staticmethod
     def register_subcommand(parser: ArgumentParser):
         """
@@ -182,9 +226,6 @@ def register_subcommand(parser: ArgumentParser):
         """
         dataclass_types = (ServeArguments,)
         serve_parser = parser.add_parser("serve", dataclass_types=dataclass_types)
-
-        group = serve_parser.add_argument_group("Positional arguments")
-        group.add_argument("model_name_or_path", type=str, default=None, help="Name of the pre-trained model.")
         serve_parser.set_defaults(func=serve_command_factory)
 
     def __init__(self, args: ServeArguments):
@@ -194,29 +235,44 @@ def __init__(self, args: ServeArguments):
             )
 
         self.args = args
-        self.model, self.tokenizer = self.load_model_and_tokenizer(args)
         self.use_continuous_batching = self.args.attn_implementation == "sdpa_paged"
 
+        # State: preserves information about the last call and last KV cache, to determine whether we can reuse the KV
+        # cache and avoid re-running prefil
+        self.last_messages = None
+        self.last_kv_cache = None
+
         transformers_logger = logging.get_logger("transformers")
         transformers_logger.setLevel(logging.log_levels[self.args.log_level.lower()])
 
         cb_logger = logging.get_logger("transformers.generation.continuous_batching")
         cb_logger.setLevel(logging.log_levels[self.args.log_level.lower()])
 
-    def build_chunk(self, content: str, request_id: str, finish_reason: Optional[str] = None) -> str:
+    def build_chunk(
+        self,
+        content: str,
+        request_id: str,
+        role: Optional[str] = None,
+        finish_reason: Optional[str] = None,
+        tool_calls: Optional[list[ChatCompletionStreamOutputDeltaToolCall]] = None,
+    ) -> str:
         payload = {
             "object": "chat.completion.chunk",
             "id": request_id,
             "created": int(time.time()),
-            "model": self.args.model_name_or_path,
+            "model": self.loaded_model,
             "system_fingerprint": "",
             "choices": [
-                {
-                    "index": 0,
-                    "delta": {"role": "assistant", "content": content, "tool_calls": []},
-                    "logprobs": None,
-                    "finish_reason": finish_reason,
-                }
+                ChatCompletionStreamOutputChoice(
+                    delta=ChatCompletionStreamOutputDelta(
+                        role=role,
+                        content=content,
+                        tool_calls=tool_calls,
+                    ),
+                    index=0,
+                    logprobs=None,
+                    finish_reason=finish_reason,
+                ),
             ],
         }
         return f"data: {json.dumps(payload)}\n\n"
@@ -229,6 +285,37 @@ def run(self):
         else:
             self.generate(app)
 
+        @functools.lru_cache(maxsize=None)
+        def get_text_gen_models() -> list[ModelInfo]:
+            return [
+                model_info("Menlo/Jan-nano"),
+                model_info("Menlo/Jan-nano-128k"),
+                model_info("Qwen/Qwen2.5-0.5B-Instruct"),
+                model_info("Qwen/Qwen2.5-3B-Instruct"),
+                model_info("Qwen/Qwen2.5-7B-Instruct"),
+                model_info("Qwen/Qwen2.5-14B-Instruct"),
+                model_info("meta-llama/Llama-3.1-8B-Instruct"),
+                model_info("meta-llama/Llama-3.2-1B-Instruct"),
+                model_info("meta-llama/Llama-3.3-70B-Instruct"),
+            ]
+
+        @app.get("/v1/models")
+        def get_all_models():
+            return JSONResponse(
+                {
+                    "object": "list",
+                    "data": [
+                        {
+                            "id": model.id,
+                            "object": "model",
+                            "crated": model.created_at.timestamp(),
+                            "owned_by": model.author,
+                        }
+                        for model in get_text_gen_models()
+                    ],
+                }
+            )
+
         uvicorn.run(app, host=self.args.host, port=self.args.port, log_level=self.args.log_level)
 
     def continuous_batching(self, app):
@@ -253,6 +340,11 @@ def _serve(req: ChatCompletionInput):
             if not req.stream:
                 return {"error": "Only streaming mode is supported."}
 
+            update_model = req.model != self.loaded_model
+
+            if update_model:
+                self.model, self.tokenizer = self.load_model_and_tokenizer(req.model, self.args)
+
             chat = req.messages
 
             inputs = self.tokenizer.apply_chat_template(chat, return_tensors="pt", add_generation_prompt=True).to(
@@ -265,7 +357,6 @@ def stream_response(_inputs):
                 try:
                     max_new_tokens = req.max_tokens or generation_config.max_new_tokens or 256
                     request_id = manager.add_request(_inputs, request_id=req.request_id, max_new_tokens=max_new_tokens)
-                    print(request_id)
                     queue_is_flushed = False
 
                     for result in manager:
@@ -288,13 +379,68 @@ def stream_response(_inputs):
 
             return StreamingResponse(stream_response(inputs[0]), media_type="text/event-stream")
 
+    def is_continuation(self, req: ChatCompletionInput) -> bool:
+        """
+        Determines whether the current request is a continuation of the last request. In other words, if it is the
+        same chat session.
+
+        Args:
+            req (`ChatCompletionInput`): The request to check.
+
+        Returns:
+            `True` if the request is a continuation of the last request, `False` otherwise.
+        """
+        req_continues_last_messages = True
+
+        # No cached messages: this is a new request
+        if self.last_messages is None:
+            req_continues_last_messages = False
+        # The new request has fewer rounds of conversation: this is a new request
+        elif len(self.last_messages) > len(req.messages):
+            req_continues_last_messages = False
+        # Otherwise, check that the last messages are a subset of the new request
+        else:
+            for i in range(len(self.last_messages)):
+                if self.last_messages[i] != req.messages[i]:
+                    req_continues_last_messages = False
+                    break
+
+        self.last_messages = req.messages
+        return req_continues_last_messages
+
     def generate(self, app):
         @app.post("/v1/chat/completions")
         def _serve(req: ChatCompletionInput):
+            update_model = req.model != self.loaded_model
+
+            if update_model:
+                self.model, self.tokenizer = self.load_model_and_tokenizer(req.model, self.args)
+
             if not req.stream:
                 return {"error": "Only streaming mode is supported."}
 
-            text = self.tokenizer.apply_chat_template(req.messages, add_generation_prompt=True, tokenize=False)
+            # HACK for tiny-agents: it sends a request after the assistant message (???). Let's assume we can't have a
+            # request whose last message is from the assistant.
+            if req.messages[-1].role == "assistant":
+                return
+
+            # ====== TOOL PREPROCESSING LOGIC ======
+            tool_model_family = None
+            for supported_model_families in _MODELS_WITH_TOOL_SUPPORT:
+                if supported_model_families in self.model.config.architectures[0].lower():
+                    tool_model_family = supported_model_families
+                    break
+            # TODO: trigger 2 constrained generations after the tool call start token is emitted:
+            # 1. force generation to pick from the tool names
+            # 2. force generation to pick from that tool's arguments
+            # ====== END OF TOOL PREPROCESSING LOGIC ======
+
+            if tool_model_family is not None:
+                text = self.tokenizer.apply_chat_template(
+                    req.messages, add_generation_prompt=True, tokenize=False, tools=req.tools
+                )
+            else:
+                text = self.tokenizer.apply_chat_template(req.messages, add_generation_prompt=True, tokenize=False)
 
             inputs = self.tokenizer(text, return_tensors="pt").to(self.model.device)["input_ids"]
             request_id = req.request_id if req.request_id is not None else "req_0"
@@ -305,28 +451,107 @@ def _serve(req: ChatCompletionInput):
             max_new_tokens = req.max_tokens or generation_config.max_new_tokens or 256
             generation_config.max_new_tokens = max_new_tokens
 
-            print(generation_config)
+            last_kv_cache = None
+            if self.is_continuation(req) and not update_model:
+                last_kv_cache = self.last_kv_cache
 
             generation_kwargs = {
                 "inputs": inputs,
                 "attention_mask": torch.ones_like(inputs),
                 "streamer": generation_streamer,
                 "generation_config": generation_config,
+                "return_dict_in_generate": True,
+                "past_key_values": last_kv_cache,
             }
 
             def stream_response(streamer, _request_id):
-                thread = Thread(target=self.model.generate, kwargs=generation_kwargs)
+                # Thin wrapper to save the KV cache after generation
+                def generate_with_cache(**kwargs):
+                    generate_output = self.model.generate(**kwargs)
+                    self.last_kv_cache = generate_output.past_key_values
+
+                thread = Thread(target=generate_with_cache, kwargs=generation_kwargs)
 
                 try:
                     thread.start()
+                    tool_state = ToolState()
 
                     for result in streamer:
-                        yield self.build_chunk(result, _request_id)
-                    yield "data: [DONE]\n\n"
+                        # ====== TOOL CALL LOGIC ======
+                        if tool_model_family is not None:
+                            # Start of a tool call: reset state variables, set `inside_tool_call`
+                            if result.strip() == _TOOL_CALL_TOKENS[tool_model_family]["start"]:
+                                tool_state.inside_tool_call = True
+                                continue
+
+                            # End of tool call: reset `inside_tool_call`, emit a `finish_reason`
+                            if result.strip() == _TOOL_CALL_TOKENS[tool_model_family]["end"]:
+                                tool_state.reset()
+                                yield self.build_chunk("", _request_id, role=None, finish_reason="tool_calls")
+                                continue
+
+                            # Inside a tool call
+                            if tool_state.inside_tool_call:
+                                tool_state.buffer += result
+
+                                # First step: extract the tool name (may need several tokens, and we can't emit a delta
+                                # until we have the full name)
+                                if not tool_state.has_tool_name_defined:
+                                    tool_name = re.search(r"\"name\": \"(.*?)\"", tool_state.buffer)
+                                    if tool_name is None:
+                                        continue
+                                    else:
+                                        tool_name = tool_name.group(1)
+                                    tool_state.has_tool_name_defined = True
+                                    tool = ChatCompletionStreamOutputDeltaToolCall(
+                                        function=ChatCompletionStreamOutputFunction(
+                                            name=tool_name,
+                                            arguments=None,
+                                        ),
+                                        index=0,
+                                        type="function",
+                                        id=_request_id + "_tool_call",  # Only the first tool call delta has an id
+                                    )
+
+                                # Second step: extract tool arguments. The tool arguments can be seen as a json string
+                                # within the tool json string. We emit a delta for the arguments.
+                                else:
+                                    # Empty text: skip
+                                    if result == "":
+                                        continue
+                                    # Until we see the `"arguments": {` in the buffer, we skip
+                                    # TODO: other models will likely need more elaborate processing here
+                                    if '"arguments": {' not in tool_state.buffer:
+                                        continue
+
+                                    # Handle nesting. We want to exclude the last } from the emitted arguments (it's
+                                    # closing the outermost nesting level, outside the arguments block)
+                                    tool_state.arg_nesting_level += result.count("{")
+                                    tool_state.arg_nesting_level -= result.count("}")
+                                    if tool_state.arg_nesting_level < 0:
+                                        result = "".join(result.split("}")[:-2]) + "}"  # e.g. "4}}\n" -> "4}"
+
+                                    tool = ChatCompletionStreamOutputDeltaToolCall(
+                                        function=ChatCompletionStreamOutputFunction(
+                                            arguments=result,
+                                        ),
+                                        index=0,
+                                        type="function",
+                                        id=None,
+                                    )
+
+                                yield self.build_chunk(None, _request_id, role=None, tool_calls=[tool])
+                                continue
+                        # ====== END OF TOOL CALL LOGIC ======
+
+                        # All non-tool related tokens are emitted as assistant messages
+                        yield self.build_chunk(result, _request_id, role="assistant")
+                    yield self.build_chunk(None, _request_id, role=None, finish_reason="stop")
 
                     thread.join()
                 except Exception as e:
                     logger.error(str(e))
+                    raise
                     yield f'data: {{"error": "{str(e)}"}}'
 
                 finally:
@@ -354,33 +579,45 @@ def get_quantization_config(model_args: ServeArguments) -> Optional["BitsAndByte
 
         return quantization_config
 
-    def load_model_and_tokenizer(self, args: ServeArguments) -> tuple[PreTrainedModel, PreTrainedTokenizerFast]:
+    def load_model_and_tokenizer(
+        self, model_id_and_revision: str, args: ServeArguments
+    ) -> tuple[PreTrainedModel, PreTrainedTokenizerFast]:
+        logger.warning(f"Loading {model_id_and_revision}")
+
+        if "@" in model_id_and_revision:
+            model_id, revision = model_id_and_revision.split("@", 1)
+        else:
+            model_id, revision = model_id_and_revision, "main"
+
         tokenizer = AutoTokenizer.from_pretrained(
-            args.model_name_or_path,
-            revision=args.model_revision,
+            model_id,
+            revision=revision,
             trust_remote_code=args.trust_remote_code,
         )
 
         torch_dtype = args.torch_dtype if args.torch_dtype in ["auto", None] else getattr(torch, args.torch_dtype)
         quantization_config = self.get_quantization_config(args)
 
         model_kwargs = {
-            "revision": args.model_revision,
+            "revision": revision,
             "attn_implementation": args.attn_implementation,
             "torch_dtype": torch_dtype,
             "device_map": "auto",
             "quantization_config": quantization_config,
             "trust_remote_code": args.trust_remote_code,
         }
 
-        model = AutoModelForCausalLM.from_pretrained(args.model_name_or_path, **model_kwargs)
+        model = AutoModelForCausalLM.from_pretrained(model_id, **model_kwargs)
 
         if model.generation_config.max_new_tokens is not None and model.generation_config.max_new_tokens < 256:
             model.generation_config.max_new_tokens = 256
 
         if getattr(model, "hf_device_map", None) is None:
             model = model.to(args.device)
 
+        self.loaded_model = model_id_and_revision
+
+        print("Loaded model", model_id_and_revision)
         return model, tokenizer