Don't raise KeyError when popping in GGMLPipeline (#1464)

kczimm · web-flow · commit 6d061eda7c72 · 2024-05-14T12:01:08.000-05:00
diff --git a/pgml-extension/src/bindings/transformers/transformers.py b/pgml-extension/src/bindings/transformers/transformers.py
@@ -211,9 +211,9 @@ class GGMLPipeline(object):
     def __init__(self, model_name, **task):
         import ctransformers
 
-        task.pop("model")
-        task.pop("task")
-        task.pop("device")
+        task.pop("model", None)
+        task.pop("task", None)
+        task.pop("device", None)
         self.model = ctransformers.AutoModelForCausalLM.from_pretrained(
             model_name, **task
         )
