From acd3eb1140f0280b7630cd40dc1524c6dcb68b18 Mon Sep 17 00:00:00 2001
From: kingbri <bdashore3@proton.me>
Date: Sat, 7 Sep 2024 22:15:42 -0400
Subject: [PATCH] Model: Add model folder template support

Like tabby_config.yml in the model's folder, a custom template can
also be provided via tabby_template.yml in addition to the existing
templates folder. The config.yml always takes priority.

Signed-off-by: kingbri <bdashore3@proton.me>
---
 backends/exllamav2/model.py | 19 +++++++++++++++----
 common/templating.py        | 14 ++++++++++----
 2 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/backends/exllamav2/model.py b/backends/exllamav2/model.py
index 044bce0..6e0a8cc 100644
--- a/backends/exllamav2/model.py
+++ b/backends/exllamav2/model.py
@@ -401,19 +401,30 @@ class ExllamaV2Container:
         find_template_functions = [
             lambda: PromptTemplate.from_model_json(
                 pathlib.Path(self.config.model_dir) / "tokenizer_config.json",
-                "chat_template",
+                key="chat_template",
             ),
             lambda: PromptTemplate.from_file(find_template_from_model(model_directory)),
         ]
 
+        # Find the template in the model directory if it exists
+        model_dir_template_path = (
+            pathlib.Path(self.config.model_dir) / "tabby_template.jinja"
+        )
+        if model_dir_template_path.exists():
+            find_template_functions[:0] = [
+                lambda: PromptTemplate.from_file(model_dir_template_path)
+            ]
+
         # Add lookup from prompt template name if provided
         if prompt_template_name:
             find_template_functions[:0] = [
-                lambda: PromptTemplate.from_file(prompt_template_name),
+                lambda: PromptTemplate.from_file(
+                    pathlib.Path("templates") / prompt_template_name
+                ),
                 lambda: PromptTemplate.from_model_json(
                     pathlib.Path(self.config.model_dir) / "tokenizer_config.json",
-                    "chat_template",
-                    prompt_template_name,
+                    key="chat_template",
+                    name=prompt_template_name,
                 ),
             ]
 
diff --git a/common/templating.py b/common/templating.py
index 47299ff..30abb38 100644
--- a/common/templating.py
+++ b/common/templating.py
@@ -106,20 +106,26 @@ class PromptTemplate:
         self.template = self.compile(raw_template)
 
     @classmethod
-    def from_file(self, prompt_template_name: str):
+    def from_file(self, template_path: pathlib.Path):
         """Get a template from a jinja file."""
 
-        template_path = pathlib.Path(f"templates/{prompt_template_name}.jinja")
+        # Add the jinja extension if it isn't provided
+        if template_path.suffix.endswith(".jinja"):
+            template_name = template_path.name.split(".jinja")[0]
+        else:
+            template_name = template_path.name
+            template_path = template_path.with_suffix(".jinja")
+
         if template_path.exists():
             with open(template_path, "r", encoding="utf8") as raw_template_stream:
                 return PromptTemplate(
-                    name=prompt_template_name,
+                    name=template_name,
                     raw_template=raw_template_stream.read(),
                 )
         else:
             # Let the user know if the template file isn't found
             raise TemplateLoadError(
-                f'Chat template "{prompt_template_name}" not found in files.'
+                f'Chat template "{template_name}" not found in files.'
             )
 
     @classmethod