Update LLM pipeline to support GPU parameter with llama.cpp backend, c…

…loses #724
neuml · May 28, 2024 · b3a99f0 · b3a99f0
1 parent 468ed7a
commit b3a99f0
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/python/txtai/pipeline/llm/llama.py b/src/python/txtai/pipeline/llm/llama.py
@@ -45,6 +45,9 @@ def __init__(self, path, template=None, **kwargs):
         # Check if this is a local path, otherwise download from the HF Hub
         path = path if os.path.exists(path) else self.download(path)
 
+        # Default GPU layers if not already set
+        kwargs["n_gpu_layers"] = kwargs.get("n_gpu_layers", -1 if kwargs.get("gpu", True) else 0)
+
         # Create llama.cpp instance
         self.llm = Llama(path, verbose=kwargs.pop("verbose", False), **kwargs)