added model_kwargs

SeldonIO · Sep 27, 2023 · e1a77d3 · e1a77d3
1 parent d41f24e
commit e1a77d3
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/runtimes/huggingface/mlserver_huggingface/common.py b/runtimes/huggingface/mlserver_huggingface/common.py
@@ -54,6 +54,7 @@ def load_pipeline_from_settings(
     hf_pipeline = pipeline(
         hf_settings.task_name,
         model=model,
+        model_kwargs=hf_settings.model_kwargs,
         tokenizer=tokenizer,
         device=hf_settings.device,
         batch_size=batch_size,

diff --git a/runtimes/huggingface/mlserver_huggingface/settings.py b/runtimes/huggingface/mlserver_huggingface/settings.py
@@ -61,7 +61,10 @@ class Config:
     """
     Name of the model that should be loaded in the pipeline.
     """
-
+    model_kwargs: Optional[dict] = None
+    """
+    model kwargs that should be loaded in the pipeline.
+    """
     pretrained_tokenizer: Optional[str] = None
     """
     Name of the tokenizer that should be loaded in the pipeline.