expose the cabability of choosing DL framework of the HF pipelines mo… (

#1066)
SeldonIO · Apr 4, 2023 · 7d8a207 · 7d8a207
1 parent 84cbe8b
commit 7d8a207
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/runtimes/huggingface/mlserver_huggingface/common.py b/runtimes/huggingface/mlserver_huggingface/common.py
@@ -53,6 +53,7 @@ class Config:
     task_suffix: str = ""
     pretrained_model: Optional[str] = None
     pretrained_tokenizer: Optional[str] = None
+    framework: Optional[str] = None
     optimum_model: bool = False
     device: int = -1
 
@@ -140,6 +141,7 @@ def load_pipeline_from_settings(
         tokenizer=tokenizer,
         device=device,
         batch_size=batch_size,
+        framework=hf_settings.framework,
     )
 
     # If max_batch_size > 0 we need to ensure tokens are padded