Update type hints for decode_request and apply_chat_template methods

Lightning-AI · Feb 23, 2025 · 8171cca · 8171cca
1 parent 194dd02
commit 8171cca
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/litgpt/deploy/serve.py b/litgpt/deploy/serve.py
@@ -170,7 +170,7 @@ def setup(self, device: str):
             self.chat_template = chat_template
 
 
-    def decode_request(self, request: ChatCompletionRequest) -> Any:
+    def decode_request(self, request: 'ChatCompletionRequest') -> Any:
         prompt = self.apply_chat_template(request.messages)
         return prompt
 
@@ -185,7 +185,7 @@ def predict(self, inputs) -> Any:
             stream=True
         )
 
-    def apply_chat_template(self, messages:List[ChatMessage])-> str:
+    def apply_chat_template(self, messages:List['ChatMessage'])-> str:
         if not _JINJA2_AVAILABLE:
             raise ImportError(
                 "apply_chat_template requires jinja2 to be installed. Please install it using `pip install jinja2`."