update stream_chat

src/llmtuner/extras/template.py

@@ -1,4 +1,4 @@
-from typing import List, Optional
+from typing import List, Optional, Tuple
 from dataclasses import dataclass
 
 
@@ -139,25 +139,33 @@ class Template:
         else:
             raise ValueError("Template {} does not exist.".format(self.name))
 
-    def get_prompt(self, query: str, history: Optional[list] = None, prefix: Optional[str] = "") -> str:
+    def get_prompt(
+        self, query: str, history: Optional[List[Tuple[str, str]]] = None, prefix: Optional[str] = ""
+    ) -> str:
         r"""
         Returns a string containing prompt without response.
         """
         return "".join(self._format_example(query, history, prefix))
 
-    def get_dialog(self, query: str, resp: str, history: Optional[list] = None, prefix: Optional[str] = "") -> List[str]:
+    def get_dialog(
+        self, query: str, resp: str, history: Optional[List[Tuple[str, str]]] = None, prefix: Optional[str] = ""
+    ) -> List[str]:
         r"""
         Returns a list containing 2 * n elements where the 2k-th is a query and the (2k+1)-th is a response.
         """
         return self._format_example(query, history, prefix) + [resp]
 
-    def _register_template(self, prefix: str, prompt: str, sep: str, use_history: Optional[bool] = True) -> None:
+    def _register_template(
+        self, prefix: str, prompt: str, sep: str, use_history: Optional[bool] = True
+    ) -> None:
         self.prefix = prefix
         self.prompt = prompt
         self.sep = sep
         self.use_history = use_history
 
-    def _format_example(self, query: str, history: Optional[list] = None, prefix: Optional[str] = "") -> List[str]:
+    def _format_example(
+        self, query: str, history: Optional[List[Tuple[str, str]]] = None, prefix: Optional[str] = ""
+    ) -> List[str]:
         prefix = prefix if prefix else self.prefix # use prefix if provided
         prefix = prefix + self.sep if prefix else "" # add separator for non-empty prefix
         history = history if (history and self.use_history) else []