Caching based on task parameters instead of just the model name

f-prime · f-prime · commit d3ef781bcebf · 2023-04-17T04:29:13.000Z
diff --git a/pgml-extension/Cargo.lock b/pgml-extension/Cargo.lock
diff --git a/pgml-extension/src/api.rs b/pgml-extension/src/api.rs
@@ -564,10 +564,10 @@ pub fn transform_json(
     task: JsonB,
     args: default!(JsonB, "'{}'"),
     inputs: default!(Vec<String>, "ARRAY[]::TEXT[]"),
-    cache_model: default!(bool, false)
+    cache: default!(bool, false)
 ) -> JsonB {
     JsonB(crate::bindings::transformers::transform(
-        &task.0, &args.0, &inputs, cache_model
+        &task.0, &args.0, &inputs, cache
     ))
 }
 
diff --git a/pgml-extension/src/bindings/transformers.py b/pgml-extension/src/bindings/transformers.py
@@ -39,26 +39,26 @@
 
 __cache_transformer_by_model_id = {}
 __cache_sentence_transformer_by_name = {}
-__cache_transform_pipeline_model_by_name = {}
+__cache_transform_pipeline_by_task = {}
 
 class NumpyJSONEncoder(json.JSONEncoder):
     def default(self, obj):
         if isinstance(obj, np.float32):
             return float(obj)
         return super().default(obj)
 
-def transform(task, args, inputs, cache_model):
+def transform(task, args, inputs, cache):
     task = json.loads(task)
     args = json.loads(args)
     inputs = json.loads(inputs)
 
-    model = task.get("model")
-    cached_model = __cache_transform_pipeline_model_by_name.get(model) if model is not None else None
-    
-    pipe = cached_model or transformers.pipeline(**task)
-
-    if cache_model and cached_model is None and model is not None:
-        __cache_transform_pipeline_model_by_name[model] = pipe
+    if cache:
+        key = ",".join([f"{key}:{val}" for (key, val) in sorted(task.items())])
+        if key not in __cache_transform_pipeline_by_task:
+            __cache_transform_pipeline_by_task[key] = transformers.pipeline(**task)
+        pipe = __cache_transform_pipeline_by_task[key]
+    else:
+        pipe = transformers.pipeline(**task)
 
     if pipe.task == "question-answering":
         inputs = [json.loads(input) for input in inputs]
diff --git a/pgml-extension/src/bindings/transformers.rs b/pgml-extension/src/bindings/transformers.rs
@@ -25,7 +25,7 @@ pub fn transform(
     task: &serde_json::Value,
     args: &serde_json::Value,
     inputs: &Vec<String>,
-    cache_model: bool 
+    cache: bool 
 ) -> serde_json::Value {
     let task = serde_json::to_string(task).unwrap();
     let args = serde_json::to_string(args).unwrap();
@@ -39,7 +39,7 @@ pub fn transform(
                 py,
                 PyTuple::new(
                     py,
-                    &[task.into_py(py), args.into_py(py), inputs.into_py(py), cache_model.into_py(py)],
+                    &[task.into_py(py), args.into_py(py), inputs.into_py(py), cache.into_py(py)],
                 ),
             )
             .unwrap()
