Fixed template imports

Signed-off-by: Mustafa Eyceoz <[email protected]>
instructlab · Jun 21, 2024 · 0db67f6 · 0db67f6
1 parent 07b6662
commit 0db67f6
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 10 deletions.
diff --git a/src/instructlab/training/chat_templates/ibm_generic_tmpl.py b/src/instructlab/training/chat_templates/ibm_generic_tmpl.py
@@ -1,5 +1,5 @@
-# Third Party
-from tokenizer_utils import SpecialTokens
+# First Party
+from instructlab.training.tokenizer_utils import SpecialTokens
 
 SPECIAL_TOKENS = SpecialTokens(
     system="<|system|>",

diff --git a/src/instructlab/training/chat_templates/mistral_tmpl.py b/src/instructlab/training/chat_templates/mistral_tmpl.py
@@ -1,5 +1,5 @@
-# Third Party
-from tokenizer_utils import SpecialTokens
+# First Party
+from instructlab.training.tokenizer_utils import SpecialTokens
 
 SPECIAL_TOKENS = SpecialTokens(
     bos="<s>",

diff --git a/src/instructlab/training/utils.py b/src/instructlab/training/utils.py
@@ -29,12 +29,15 @@ def retrieve_chat_template(chat_tmpl_path):
     import importlib.util
     import sys
 
-    spec = importlib.util.spec_from_file_location("spcl_chat_tmpl", chat_tmpl_path)
-    module = importlib.util.module_from_spec(spec)
-    sys.modules["spcl_chat_tmpl"] = module
-    spec.loader.exec_module(module)
-    SPECIAL_TOKENS = module.SPECIAL_TOKENS
-    CHAT_TEMPLATE = module.CHAT_TEMPLATE
+    try:
+        spec = importlib.util.spec_from_file_location("spcl_chat_tmpl", chat_tmpl_path)
+        module = importlib.util.module_from_spec(spec)
+        sys.modules["spcl_chat_tmpl"] = module
+        spec.loader.exec_module(module)
+        SPECIAL_TOKENS = module.SPECIAL_TOKENS
+        CHAT_TEMPLATE = module.CHAT_TEMPLATE
+    except:
+        sys.exit(f"Invalid chat template path: {chat_tmpl_path}")
     return CHAT_TEMPLATE, SPECIAL_TOKENS