postgresml · montanalow · Jun 3, 2023 · Jun 3, 2023
diff --git a/pgml-extension/requirements.txt b/pgml-extension/requirements.txt
@@ -19,3 +19,4 @@ tqdm==4.65.0
 transformers==4.29.2
 xgboost==1.7.5
 langchain==0.0.180
+einops==0.6.1
diff --git a/pgml-extension/src/bindings/transformers.py b/pgml-extension/src/bindings/transformers.py
@@ -97,6 +97,10 @@ def transform(task, args, inputs):
    ensure_device(task)
    convert_dtype(task)

+    model = task.get("model", None)
+    if model and "tokenizer" not in task:
+        task["tokenizer"] = AutoTokenizer.from_pretrained(model)
+
    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]

diff --git a/pgml-extension/src/bindings/transformers.rs b/pgml-extension/src/bindings/transformers.rs
@@ -35,17 +35,24 @@ pub fn transform(
    let results = Python::with_gil(|py| -> String {
        let transform: Py<PyAny> = PY_MODULE.getattr(py, "transform").unwrap().into();

-        transform
+        let result = transform
            .call1(
                py,
                PyTuple::new(
                    py,
                    &[task.into_py(py), args.into_py(py), inputs.into_py(py)],
                ),
-            )
-            .unwrap()
-            .extract(py)
-            .unwrap()
+            );
+
+        let result = match result {
+            Err(e) => {
+                let traceback = e.traceback(py).unwrap().format().unwrap();
+                error!("{traceback} {e}")
+            }
+            Ok(o) => o.extract(py).unwrap(),
+        };
+
+        result
    });
    serde_json::from_str(&results).unwrap()
 }