From abc9160881c7455bb8592c7404b9849e6eab74b3 Mon Sep 17 00:00:00 2001
From: Frankie Primerano <frankieprimerano@gmail.com>
Date: Sun, 16 Apr 2023 18:12:38 +0000
Subject: [PATCH 1/2] Adds ability to cache models to make subsequent calls to
 `transform` faster.

---
 .../user_guides/transformers/pre_trained_models.md    | 10 ++++++----
 pgml-extension/Cargo.lock                             |  2 +-
 pgml-extension/src/api.rs                             |  6 ++++--
 pgml-extension/src/bindings/transformers.py           | 11 +++++++++--
 pgml-extension/src/bindings/transformers.rs           |  3 ++-
 5 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/pgml-docs/docs/user_guides/transformers/pre_trained_models.md b/pgml-docs/docs/user_guides/transformers/pre_trained_models.md
index fa9677b86..c21c86bd1 100644
--- a/pgml-docs/docs/user_guides/transformers/pre_trained_models.md
+++ b/pgml-docs/docs/user_guides/transformers/pre_trained_models.md
@@ -11,9 +11,10 @@ The Hugging Face [`Pipeline`](https://huggingface.co/docs/transformers/main_clas
 
 ```sql linenums="1" title="transformer.sql"
 pgml.transform(
-    task TEXT OR JSONB,      -- task name or full pipeline initializer arguments
-    call JSONB,              -- additional call arguments alongside the inputs
-    inputs TEXT[] OR BYTEA[] -- inputs for inference
+    task TEXT OR JSONB,       -- task name or full pipeline initializer arguments
+    call JSONB,               -- additional call arguments alongside the inputs
+    inputs TEXT[] OR BYTEA[], -- inputs for inference
+    cache_model BOOLEAN       -- if true, the model will be cached in memory. FALSE by default
 )
 ```
 
@@ -73,7 +74,8 @@ Sentiment analysis is one use of `text-classification`, but there are [many othe
         inputs => ARRAY[
             'I love how amazingly simple ML has become!', 
             'I hate doing mundane and thankless tasks. ☹️'
-        ]
+        ],
+        cache_model => TRUE
     ) AS positivity;
     ```
 
diff --git a/pgml-extension/Cargo.lock b/pgml-extension/Cargo.lock
index a731ecb6c..4abcdd0dc 100644
--- a/pgml-extension/Cargo.lock
+++ b/pgml-extension/Cargo.lock
@@ -1681,7 +1681,7 @@ dependencies = [
 
 [[package]]
 name = "pgml"
-version = "2.4.0"
+version = "2.5.0"
 dependencies = [
  "anyhow",
  "blas",
diff --git a/pgml-extension/src/api.rs b/pgml-extension/src/api.rs
index c332b7f98..dae7830ed 100644
--- a/pgml-extension/src/api.rs
+++ b/pgml-extension/src/api.rs
@@ -564,9 +564,10 @@ pub fn transform_json(
     task: JsonB,
     args: default!(JsonB, "'{}'"),
     inputs: default!(Vec<String>, "ARRAY[]::TEXT[]"),
+    cache_model: default!(bool, false)
 ) -> JsonB {
     JsonB(crate::bindings::transformers::transform(
-        &task.0, &args.0, &inputs,
+        &task.0, &args.0, &inputs, cache_model
     ))
 }
 
@@ -576,12 +577,13 @@ pub fn transform_string(
     task: String,
     args: default!(JsonB, "'{}'"),
     inputs: default!(Vec<String>, "ARRAY[]::TEXT[]"),
+    cache_model: default!(bool, false)
 ) -> JsonB {
     let mut task_map = HashMap::new();
     task_map.insert("task", task);
     let task_json = json!(task_map);
     JsonB(crate::bindings::transformers::transform(
-        &task_json, &args.0, &inputs,
+        &task_json, &args.0, &inputs, cache_model
     ))
 }
 
diff --git a/pgml-extension/src/bindings/transformers.py b/pgml-extension/src/bindings/transformers.py
index da109b9f2..f417aded9 100644
--- a/pgml-extension/src/bindings/transformers.py
+++ b/pgml-extension/src/bindings/transformers.py
@@ -39,6 +39,7 @@
 
 __cache_transformer_by_model_id = {}
 __cache_sentence_transformer_by_name = {}
+__cache_transform_pipeline_model_by_name = {}
 
 class NumpyJSONEncoder(json.JSONEncoder):
     def default(self, obj):
@@ -46,12 +47,18 @@ def default(self, obj):
             return float(obj)
         return super().default(obj)
 
-def transform(task, args, inputs):
+def transform(task, args, inputs, cache_model):
     task = json.loads(task)
     args = json.loads(args)
     inputs = json.loads(inputs)
 
-    pipe = transformers.pipeline(**task)
+    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 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
index 4791fab12..676835a82 100644
--- a/pgml-extension/src/bindings/transformers.rs
+++ b/pgml-extension/src/bindings/transformers.rs
@@ -25,6 +25,7 @@ pub fn transform(
     task: &serde_json::Value,
     args: &serde_json::Value,
     inputs: &Vec<String>,
+    cache_model: bool 
 ) -> serde_json::Value {
     let task = serde_json::to_string(task).unwrap();
     let args = serde_json::to_string(args).unwrap();
@@ -38,7 +39,7 @@ pub fn transform(
                 py,
                 PyTuple::new(
                     py,
-                    &[task.into_py(py), args.into_py(py), inputs.into_py(py)],
+                    &[task.into_py(py), args.into_py(py), inputs.into_py(py), cache_model.into_py(py)],
                 ),
             )
             .unwrap()

From eda180c39149d539343f32bf174147370a3d1da4 Mon Sep 17 00:00:00 2001
From: Frankie Primerano <frankieprimerano@gmail.com>
Date: Mon, 17 Apr 2023 04:29:13 +0000
Subject: [PATCH 2/2] Caching based on task parameters instead of just the
 model name

---
 .../transformers/pre_trained_models.md         |  4 ++--
 pgml-extension/Cargo.lock                      |  2 +-
 pgml-extension/src/api.rs                      |  8 ++++----
 pgml-extension/src/bindings/transformers.py    | 18 +++++++++---------
 pgml-extension/src/bindings/transformers.rs    |  4 ++--
 5 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/pgml-docs/docs/user_guides/transformers/pre_trained_models.md b/pgml-docs/docs/user_guides/transformers/pre_trained_models.md
index c21c86bd1..7eec2791c 100644
--- a/pgml-docs/docs/user_guides/transformers/pre_trained_models.md
+++ b/pgml-docs/docs/user_guides/transformers/pre_trained_models.md
@@ -14,7 +14,7 @@ pgml.transform(
     task TEXT OR JSONB,       -- task name or full pipeline initializer arguments
     call JSONB,               -- additional call arguments alongside the inputs
     inputs TEXT[] OR BYTEA[], -- inputs for inference
-    cache_model BOOLEAN       -- if true, the model will be cached in memory. FALSE by default
+    cache BOOLEAN             -- if TRUE, the model will be cached in memory. FALSE by default.
 )
 ```
 
@@ -75,7 +75,7 @@ Sentiment analysis is one use of `text-classification`, but there are [many othe
             'I love how amazingly simple ML has become!', 
             'I hate doing mundane and thankless tasks. ☹️'
         ],
-        cache_model => TRUE
+        cache => TRUE
     ) AS positivity;
     ```
 
diff --git a/pgml-extension/Cargo.lock b/pgml-extension/Cargo.lock
index 4abcdd0dc..a731ecb6c 100644
--- a/pgml-extension/Cargo.lock
+++ b/pgml-extension/Cargo.lock
@@ -1681,7 +1681,7 @@ dependencies = [
 
 [[package]]
 name = "pgml"
-version = "2.5.0"
+version = "2.4.0"
 dependencies = [
  "anyhow",
  "blas",
diff --git a/pgml-extension/src/api.rs b/pgml-extension/src/api.rs
index dae7830ed..7f09fb8c8 100644
--- 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
     ))
 }
 
@@ -577,13 +577,13 @@ pub fn transform_string(
     task: String,
     args: default!(JsonB, "'{}'"),
     inputs: default!(Vec<String>, "ARRAY[]::TEXT[]"),
-    cache_model: default!(bool, false)
+    cache: default!(bool, false)
 ) -> JsonB {
     let mut task_map = HashMap::new();
     task_map.insert("task", task);
     let task_json = json!(task_map);
     JsonB(crate::bindings::transformers::transform(
-        &task_json, &args.0, &inputs, cache_model
+        &task_json, &args.0, &inputs, cache
     ))
 }
 
diff --git a/pgml-extension/src/bindings/transformers.py b/pgml-extension/src/bindings/transformers.py
index f417aded9..0ef690410 100644
--- a/pgml-extension/src/bindings/transformers.py
+++ b/pgml-extension/src/bindings/transformers.py
@@ -39,7 +39,7 @@
 
 __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):
@@ -47,18 +47,18 @@ def default(self, obj):
             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
index 676835a82..504202ba8 100644
--- 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()

<!DOCTYPE html PUBLIC '-//W3C//DTD XHTML 1.0 Transitional//EN' 'http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd'>
<html xmlns='http://www.w3.org/1999/xhtml'>
<head>
<title>pFad - Phonifier reborn</title>
<meta http-equiv='Content-Type' content='text/html; charset=utf-8' />
</head>
<body>
<h1>Pfad - The Proxy pFad of &#169; 2024 Garber Painting. All rights reserved.</h1>


<!-- Disclaimer -->
<p>Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.</p>
<br>
<p>Alternative Proxies:</p><p><a href="http://rainy.clevelandohioweatherforecast.com/php-proxy/index.php?q=https://patch-diff.githubusercontent.com/raw/postgresml/postgresml/pull/593.patch" target="_blank">Alternative Proxy</a></p><p><a href="http://rainy.clevelandohioweatherforecast.com/pFad/index.php?u=https://patch-diff.githubusercontent.com/raw/postgresml/postgresml/pull/593.patch" target="_blank">pFad Proxy</a></p><p><a href="http://rainy.clevelandohioweatherforecast.com/pFad/v3index.php?u=https://patch-diff.githubusercontent.com/raw/postgresml/postgresml/pull/593.patch" target="_blank">pFad v3 Proxy</a></p><p><a href="http://rainy.clevelandohioweatherforecast.com/pFad/v4index.php?u=https://patch-diff.githubusercontent.com/raw/postgresml/postgresml/pull/593.patch" target="_blank">pFad v4 Proxy</a></p></body>
</html>