Make model_kwargs dictionary by default

dottxt-ai · Mar 14, 2024 · c8566e8 · c8566e8
1 parent 043117f
commit c8566e8
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/outlines/models/llamacpp.py b/outlines/models/llamacpp.py
@@ -22,10 +22,11 @@ def __init__(self, model: "Llama"):
         self.tokenizer = LlamaCppTokenizer(model=model)
 
 
-def llamacpp(model_path: str, device: Optional[str] = None, **model_kwargs) -> LlamaCpp:
+def llamacpp(
+    model_path: str, device: Optional[str] = None, model_kwargs: dict = {}
+) -> LlamaCpp:
     from llama_cpp import Llama
 
-    model_kwargs = model_kwargs.get("model_kwargs", {})
     if device == "cuda":
         model_kwargs["n_gpu_layers"].setdefault(-1)