Fix issue with loading non-transformer LLM models in Extractor/RAG pi…

…peline, closes #734
neuml · Jun 18, 2024 · 7073084 · 7073084
1 parent 42583f0
commit 7073084
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/python/txtai/pipeline/text/extractor.py b/src/python/txtai/pipeline/text/extractor.py
@@ -6,7 +6,7 @@
 
 from ..base import Pipeline
 from ..data import Tokenizer
-from ..llm import LLM
+from ..llm import GenerationFactory, LLM
 
 from .questions import Questions
 from .similarity import Similarity
@@ -159,7 +159,8 @@ def load(self, path, quantize, gpu, model, task, **kwargs):
             return path
 
         # Attempt to resolve task if not provided
-        task = task if task else Models.task(path, **kwargs)
+        task = GenerationFactory.method(path, task)
+        task = Models.task(path, **kwargs) if task == "transformers" else task
 
         # Load Questions pipeline
         if task == "question-answering":