Skip to content

Commit

Permalink
Make model_kwargs dictionary by default
Browse files Browse the repository at this point in the history
  • Loading branch information
rlouf committed Mar 14, 2024
1 parent 043117f commit c8566e8
Showing 1 changed file with 3 additions and 2 deletions.
5 changes: 3 additions & 2 deletions outlines/models/llamacpp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down

0 comments on commit c8566e8

Please sign in to comment.