postgresml · SilasMarvin · Jun 5, 2024 · May 31, 2024 · Jun 3, 2024 · Jun 3, 2024
diff --git a/pgml-extension/Cargo.lock b/pgml-extension/Cargo.lock
diff --git a/pgml-extension/Cargo.toml b/pgml-extension/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "pgml"
-version = "2.8.5"
+version = "2.9.0"
 edition = "2021"
 
 [lib]
@@ -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/sql/pgml--2.8.5--2.9.0.sql b/pgml-extension/sql/pgml--2.8.5--2.9.0.sql
@@ -0,0 +1,15 @@
+-- src/api.rs:613
+-- pgml::api::rank
+CREATE  FUNCTION pgml."rank"(
+	"transformer" TEXT, /* &str */
+	"query" TEXT, /* &str */
+	"documents" TEXT[], /* alloc::vec::Vec<&str> */
+	"kwargs" jsonb DEFAULT '{}' /* pgrx::datum::json::JsonB */
+) RETURNS TABLE (
+	"corpus_id" bigint,  /* i64 */
+	"score" double precision,  /* f64 */
+	"text" TEXT  /* core::option::Option<alloc::string::String> */
+)
+IMMUTABLE STRICT PARALLEL SAFE 
+LANGUAGE c /* Rust */
+AS 'MODULE_PATHNAME', 'rank_wrapper';
diff --git a/pgml-extension/src/api.rs b/pgml-extension/src/api.rs
@@ -603,7 +603,21 @@ 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, "'{}'"),
+) -> TableIterator<'static, (name!(corpus_id, i64), name!(score, f64), name!(text, Option<String>))> {
+    match crate::bindings::transformers::rank(transformer, query, documents, &kwargs.0) {
+        Ok(output) => TableIterator::new(output.into_iter().map(|x| (x.corpus_id, x.score, x.text))),
         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,8 @@ 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::Deserialize;
 use serde_json::Value;
 
 use crate::create_pymodule;
@@ -21,6 +22,59 @@ 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));
+            }
+            Err(anyhow::anyhow!(
+                "Unsupported type for JSON conversion: {:?}",
+                ob.get_type()
+            ))?
+        }
+    }
+}
+
 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 +109,46 @@ pub fn embed(transformer: &str, inputs: Vec<&str>, kwargs: &serde_json::Value) -
     })
 }
 
+#[derive(Deserialize)]
+pub struct RankResult {
+    pub corpus_id: i64,
+    pub score: f64,
+    pub text: Option<String>,
+}
+
+pub fn rank(
+    transformer: &str,
+    query: &str,
+    documents: Vec<&str>,
+    kwargs: &serde_json::Value,
+) -> Result<Vec<RankResult>> {
+    let kwargs = serde_json::to_string(kwargs)?;
+    Python::with_gil(|py| -> Result<Vec<RankResult>> {
+        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)?;
+        out.into_iter()
+            .map(|x| {
+                let x: RankResult = serde_json::from_value(x.0)?;
+                Ok(x)
+            })
+            .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)