Minor enhancement

Signed-off-by: Pengyun Lin <[email protected]>

Author: LinPoly

tensorrt_llm/commands/serve.py

@@ -371,7 +371,8 @@ def convert(self, value: Any, param: Optional["click.Parameter"],
 @click.option("--chat_template",
               type=str,
               default=None,
-              help="[Experimental] Specify the chat template.")
+              help="[Experimental] Specify a custom chat template. "
+              "Can be a file path or one-liner template string")
 def serve(
         model: str, tokenizer: Optional[str], host: str, port: int,
         log_level: str, backend: str, max_beam_width: int, max_batch_size: int,

tensorrt_llm/serve/chat_utils.py

@@ -259,7 +259,8 @@ def make_tool_call_id(id_type: str = "random", func_name=None, idx=None):
 
 # Adapted from
 # https://github.com/vllm-project/vllm/blob/44b5ce956d3cf28841615a58c1c0873af87bcfe2/vllm/entrypoints/chat_utils.py
-def _load_chat_template(
+@lru_cache
+def load_chat_template(
     chat_template: Path | str | None,
     *,
     is_literal: bool = False,
@@ -290,15 +291,4 @@ def _load_chat_template(
 
         # If opening a file fails, set chat template to be args to
         # ensure we decode so our escape are interpreted correctly
-        return _load_chat_template(chat_template, is_literal=True)
-
-
-_cached_load_chat_template = lru_cache(_load_chat_template)
-
-
-def load_chat_template(
-    chat_template: Path | str | None,
-    *,
-    is_literal: bool = False,
-) -> str | None:
-    return _cached_load_chat_template(chat_template, is_literal=is_literal)
+        return load_chat_template(chat_template, is_literal=True)

tests/unittest/llmapi/apps/test_chat_utils.py

@@ -1,5 +1,3 @@
-import os
-import tempfile
 from unittest.mock import MagicMock
 
 import pytest
@@ -187,17 +185,12 @@ def test_tool_message_without_tool_call_id(self, mock_mm_data_tracker):
 
 
 @pytest.fixture
-def chat_template_path():
+def chat_template_path(tmp_path):
     """Return the path to the chat template."""
-    temp_dir = tempfile.gettempdir()
-    temp_file_path = os.path.join(temp_dir, "chat_template.jinja")
-    try:
-        with open(temp_file_path, "w") as f:
-            f.write(TEMPLATE_CHATML)
-        yield temp_file_path
-    finally:
-        if os.path.exists(temp_file_path):
-            os.remove(temp_file_path)
+    temp_file_path = tmp_path / "chat_template.jinja"
+    with open(temp_file_path, "w") as f:
+        f.write(TEMPLATE_CHATML)
+    return temp_file_path
 
 
 class TestLoadChatTemplate: