add remove_pipeline to transformers.py

kczimm · kczimm · commit f3fc5fb7af98 · 2023-08-16T16:31:37.000-05:00
diff --git a/pgml-extension/src/bindings/transformers/transformers.py b/pgml-extension/src/bindings/transformers/transformers.py
@@ -204,6 +204,11 @@ def create_pipeline(task):
     return pipe
 
 
+def remove_pipeline(task):
+    key = key_from_task(task)
+    del __cache_transform_pipeline_by_task[key]
+
+
 def transform_using(pipeline, args, inputs):
     args = orjson.loads(args)
     inputs = orjson.loads(inputs)
@@ -215,12 +220,16 @@ def transform_using(pipeline, args, inputs):
     return orjson.dumps(pipeline(inputs, **args), default=orjson_default).decode()
 
 
+def key_from_task(task): 
+    return ",".join([f"{key}:{val}" for (key, val) in sorted(task.items())])
+
+
 def transform(task, args, inputs):
     task = orjson.loads(task)
     args = orjson.loads(args)
     inputs = orjson.loads(inputs)
 
-    key = ",".join([f"{key}:{val}" for (key, val) in sorted(task.items())])
+    key = key_from_task(task)
     if key not in __cache_transform_pipeline_by_task:
         pipe = create_pipeline(task)
         __cache_transform_pipeline_by_task[key] = pipe
