Added reranking into the extension

SilasMarvin · SilasMarvin · commit da7df56208d7 · 2024-05-31T14:46:52.000-07:00
diff --git a/pgml-extension/Cargo.lock b/pgml-extension/Cargo.lock
diff --git a/pgml-extension/Cargo.toml b/pgml-extension/Cargo.toml
@@ -41,7 +41,7 @@ ndarray-stats = "0.5.1"
 parking_lot = "0.12"
 pgrx = "=0.11.3"
 pgrx-pg-sys = "=0.11.3"
-pyo3 = { version = "0.20.0", features = ["auto-initialize"], optional = true }
+pyo3 = { version = "0.20.0", features = ["anyhow", "auto-initialize"], optional = true }
 rand = "0.8"
 rmp-serde = { version = "1.1" }
 signal-hook = "0.3"
diff --git a/pgml-extension/src/api.rs b/pgml-extension/src/api.rs
@@ -603,7 +603,22 @@ pub fn embed_batch(
     kwargs: default!(JsonB, "'{}'"),
 ) -> SetOfIterator<'static, Vec<f32>> {
     match crate::bindings::transformers::embed(transformer, inputs, &kwargs.0) {
-        Ok(output) => SetOfIterator::new(output.into_iter()),
+        Ok(output) => SetOfIterator::new(output),
+        Err(e) => error!("{e}"),
+    }
+}
+
+#[cfg(all(feature = "python", not(feature = "use_as_lib")))]
+#[pg_extern(immutable, parallel_safe, name = "rank")]
+// pub fn rank(transformer: &str, query: &str, documents: Vec<&str>, kwargs: default!(JsonB, "'{}'")) -> Vec<pgrx::JsonB> {
+pub fn rank(
+    transformer: &str,
+    query: &str,
+    documents: Vec<&str>,
+    kwargs: default!(JsonB, "'{}'"),
+) -> SetOfIterator<'static, pgrx::JsonB> {
+    match crate::bindings::transformers::rank(transformer, query, documents, &kwargs.0) {
+        Ok(output) => SetOfIterator::new(output.into_iter().map(pgrx::JsonB)),
         Err(e) => error!("{e}"),
     }
 }
diff --git a/pgml-extension/src/bindings/transformers/mod.rs b/pgml-extension/src/bindings/transformers/mod.rs
@@ -6,7 +6,7 @@ use std::{collections::HashMap, path::Path};
 use anyhow::{anyhow, bail, Context, Result};
 use pgrx::*;
 use pyo3::prelude::*;
-use pyo3::types::PyTuple;
+use pyo3::types::{PyBool, PyDict, PyFloat, PyInt, PyList, PyString, PyTuple};
 use serde_json::Value;
 
 use crate::create_pymodule;
@@ -21,6 +21,57 @@ pub use transform::*;
 
 create_pymodule!("/src/bindings/transformers/transformers.py");
 
+// Need a wrapper so we can implement traits for it
+struct Json(Value);
+
+impl From<Json> for Value {
+    fn from(value: Json) -> Self {
+        value.0
+    }
+}
+
+impl FromPyObject<'_> for Json {
+    fn extract(ob: &PyAny) -> PyResult<Self> {
+        if ob.is_instance_of::<PyDict>() {
+            let dict: &PyDict = ob.downcast()?;
+            let mut json = serde_json::Map::new();
+            for (key, value) in dict.iter() {
+                let value = Json::extract(value)?;
+                json.insert(String::extract(key)?, value.0);
+            }
+            Ok(Self(serde_json::Value::Object(json)))
+        } else if ob.is_instance_of::<PyBool>() {
+            let value = bool::extract(ob)?;
+            Ok(Self(serde_json::Value::Bool(value)))
+        } else if ob.is_instance_of::<PyInt>() {
+            let value = i64::extract(ob)?;
+            Ok(Self(serde_json::Value::Number(value.into())))
+        } else if ob.is_instance_of::<PyFloat>() {
+            let value = f64::extract(ob)?;
+            let value =
+                serde_json::value::Number::from_f64(value).context("Could not convert f64 to serde_json::Number")?;
+            Ok(Self(serde_json::Value::Number(value)))
+        } else if ob.is_instance_of::<PyString>() {
+            let value = String::extract(ob)?;
+            Ok(Self(serde_json::Value::String(value)))
+        } else if ob.is_instance_of::<PyList>() {
+            let value = ob.downcast::<PyList>()?;
+            let mut json_values = Vec::new();
+            for v in value {
+                let v = v.extract::<Json>()?;
+                json_values.push(v.0);
+            }
+            Ok(Self(serde_json::Value::Array(json_values)))
+        } else {
+            if ob.is_none() {
+                return Ok(Self(serde_json::Value::Null));
+            }
+            eprintln!("\n\nTHE OBJ: {:?}\n\n", ob.get_type());
+            Err(anyhow::anyhow!("Unsupported type for JSON conversion"))?
+        }
+    }
+}
+
 pub fn get_model_from(task: &Value) -> Result<String> {
     Python::with_gil(|py| -> Result<String> {
         let get_model_from = get_module!(PY_MODULE)
@@ -55,6 +106,29 @@ pub fn embed(transformer: &str, inputs: Vec<&str>, kwargs: &serde_json::Value) -
     })
 }
 
+pub fn rank(transformer: &str, query: &str, documents: Vec<&str>, kwargs: &serde_json::Value) -> Result<Vec<Value>> {
+    let kwargs = serde_json::to_string(kwargs)?;
+    Python::with_gil(|py| -> Result<Vec<Value>> {
+        let embed: Py<PyAny> = get_module!(PY_MODULE).getattr(py, "rank").format_traceback(py)?;
+        let output = embed
+            .call1(
+                py,
+                PyTuple::new(
+                    py,
+                    &[
+                        transformer.to_string().into_py(py),
+                        query.into_py(py),
+                        documents.into_py(py),
+                        kwargs.into_py(py),
+                    ],
+                ),
+            )
+            .format_traceback(py)?;
+        let out: Vec<Json> = output.extract(py).format_traceback(py)?;
+        Ok(out.into_iter().map(|x| x.into()).collect())
+    })
+}
+
 pub fn finetune_text_classification(
     task: &Task,
     dataset: TextClassificationDataset,
diff --git a/pgml-extension/src/bindings/transformers/transformers.py b/pgml-extension/src/bindings/transformers/transformers.py
@@ -12,7 +12,7 @@
 import orjson
 from rouge import Rouge
 from sacrebleu.metrics import BLEU
-from sentence_transformers import SentenceTransformer
+from sentence_transformers import SentenceTransformer, CrossEncoder
 from sklearn.metrics import (
     mean_squared_error,
     r2_score,
@@ -500,6 +500,33 @@ def transform(task, args, inputs, stream=False):
     return orjson.dumps(pipe(inputs, **args), default=orjson_default).decode()
 
 
+def create_cross_encoder(transformer):
+    return CrossEncoder(transformer)
+
+
+def rank_using(model, query, documents, kwargs):
+    if isinstance(kwargs, str):
+        kwargs = orjson.loads(kwargs)
+
+    # The score is a numpy float32 before we convert it
+    return [
+        {"score": x.pop("score").item(), **x}
+        for x in model.rank(query, documents, **kwargs)
+    ]
+
+
+def rank(transformer, query, documents, kwargs):
+    kwargs = orjson.loads(kwargs)
+
+    if transformer not in __cache_sentence_transformer_by_name:
+        __cache_sentence_transformer_by_name[transformer] = create_cross_encoder(
+            transformer
+        )
+    model = __cache_sentence_transformer_by_name[transformer]
+
+    return rank_using(model, query, documents, kwargs)
+
+
 def create_embedding(transformer):
     return SentenceTransformer(transformer)
 
