fix messages strat loader

axolotl-ai-cloud · Oct 11, 2024 · 0e08627 · 0e08627
1 parent 14d4b16
commit 0e08627
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/axolotl/prompt_strategies/messages/__init__.py b/src/axolotl/prompt_strategies/messages/__init__.py
@@ -15,7 +15,9 @@ def load(tokenizer, cfg, ds_cfg, processor=None):
         if strategy.split(".")[-1].startswith("load_"):
             load_fn = strategy.split(".")[-1]
             strategy = ".".join(strategy.split(".")[:-1])
-        mod = importlib.import_module(f".{strategy}", "axolotl.prompt_strategies")
+        mod = importlib.import_module(
+            f".{strategy}", "axolotl.prompt_strategies.messages"
+        )
         func = getattr(mod, load_fn)
         load_kwargs = {}
         sig = inspect.signature(func)