Create embedding API (#578)

montanalow · web-flow · commit d5ed9de16c48 · 2023-03-29T09:28:26.000-07:00
diff --git a/pgml-extension/Dockerfile b/pgml-extension/Dockerfile
@@ -32,7 +32,7 @@ RUN apt update && apt-fast install -y \
 	python3-pip \
 	libpython3.10-dev \
 	python3.10-dev
-RUN pip3 install xgboost scikit-learn torch lightgbm transformers datasets
+RUN pip3 install xgboost scikit-learn torch lightgbm transformers datasets sentence_transformers
 RUN useradd postgresml -m -s /bin/bash -G sudo
 RUN echo 'postgresml ALL=(ALL) NOPASSWD: ALL' >> /etc/sudoers
 USER postgresml
diff --git a/pgml-extension/Dockerfile.local b/pgml-extension/Dockerfile.local
@@ -11,7 +11,7 @@ RUN cat /etc/apt/sources.list
 RUN apt-get update && apt-get install -y postgresql-pgml-14
 
 # Cache this, quicker
-RUN pip3 install xgboost scikit-learn diptest torch lightgbm transformers datasets sentencepiece sacremoses sacrebleu rouge
+RUN pip3 install xgboost scikit-learn diptest torch lightgbm transformers datasets sentencepiece sentence_transformers sacremoses sacrebleu rouge
 
 COPY --chown=postgres:postgres . /app
 WORKDIR /app
diff --git a/pgml-extension/src/api.rs b/pgml-extension/src/api.rs
@@ -553,6 +553,11 @@ fn load_dataset(
     TableIterator::new(vec![(name, rows)].into_iter())
 }
 
+#[pg_extern]
+pub fn embed(transformer: &str, text: &str, kwargs: default!(JsonB, "'{}'")) -> Vec<f32> {
+    crate::bindings::transformers::embed(transformer, &text, &kwargs.0)
+}
+
 #[cfg(feature = "python")]
 #[pg_extern(name = "transform")]
 pub fn transform_json(
diff --git a/pgml-extension/src/bindings/transformers.py b/pgml-extension/src/bindings/transformers.py
@@ -8,6 +8,7 @@
 import datasets
 from rouge import Rouge
 from sacrebleu.metrics import BLEU
+from sentence_transformers import SentenceTransformer
 from sklearn.metrics import (
     mean_squared_error,
     r2_score,
@@ -37,6 +38,7 @@
 )
 
 __cache_transformer_by_model_id = {}
+__cache_sentence_transformer_by_name = {}
 
 def transform(task, args, inputs):
     task = json.loads(task)
@@ -50,6 +52,13 @@ def transform(task, args, inputs):
 
     return json.dumps(pipe(inputs, **args))
 
+def embed(transformer, text, kwargs):
+    kwargs = json.loads(kwargs)
+    if transformer not in __cache_sentence_transformer_by_name:
+        __cache_sentence_transformer_by_name[transformer] = SentenceTransformer(transformer)
+    model = __cache_sentence_transformer_by_name[transformer]
+    return model.encode(text, **kwargs)
+
 def load_dataset(name, subset, limit: None, kwargs: "{}"):
     kwargs = json.loads(kwargs)
 
diff --git a/pgml-extension/src/bindings/transformers.rs b/pgml-extension/src/bindings/transformers.rs
@@ -48,6 +48,28 @@ pub fn transform(
     serde_json::from_str(&results).unwrap()
 }
 
+pub fn embed(transformer: &str, text: &str, kwargs: &serde_json::Value) -> Vec<f32> {
+    let kwargs = serde_json::to_string(kwargs).unwrap();
+    Python::with_gil(|py| -> Vec<f32> {
+        let embed: Py<PyAny> = PY_MODULE.getattr(py, "embed").unwrap().into();
+        embed
+            .call1(
+                py,
+                PyTuple::new(
+                    py,
+                    &[
+                        transformer.to_string().into_py(py),
+                        text.to_string().into_py(py),
+                        kwargs.into_py(py),
+                    ],
+                ),
+            )
+            .unwrap()
+            .extract(py)
+            .unwrap()
+    })
+}
+
 pub fn tune(
     task: &Task,
     dataset: TextDataset,
