postgresml · montanalow · Jun 19, 2023 · Jun 19, 2023
diff --git a/pgml-extension/src/api.rs b/pgml-extension/src/api.rs
@@ -580,7 +580,6 @@ pub fn embed_batch(
     crate::bindings::transformers::embed(transformer, inputs, &kwargs.0)
 }
 
-
 /// Clears the GPU cache.
 ///
 /// # Arguments
@@ -596,10 +595,9 @@ pub fn embed_batch(
 /// SELECT pgml.clear_gpu_cache(memory_usage => 0.5);
 /// ```
 #[pg_extern(immutable, parallel_safe, name = "clear_gpu_cache")]
-pub fn clear_gpu_cache(
-    memory_usage: default!(Option<f32>, "NULL")
-) -> bool {
-    let memory_usage: Option<f32> = memory_usage.map(|memory_usage| memory_usage.try_into().unwrap());
+pub fn clear_gpu_cache(memory_usage: default!(Option<f32>, "NULL")) -> bool {
+    let memory_usage: Option<f32> =
+        memory_usage.map(|memory_usage| memory_usage.try_into().unwrap());
     crate::bindings::transformers::clear_gpu_cache(memory_usage)
 }
 

diff --git a/pgml-extension/src/bindings/transformers.rs b/pgml-extension/src/bindings/transformers.rs
@@ -62,21 +62,27 @@ pub fn embed(transformer: &str, inputs: Vec<&str>, kwargs: &serde_json::Value) -
     let kwargs = serde_json::to_string(kwargs).unwrap();
     Python::with_gil(|py| -> Vec<Vec<f32>> {
         let embed: Py<PyAny> = PY_MODULE.getattr(py, "embed").unwrap().into();
-        embed
-            .call1(
+        let result = embed.call1(
+            py,
+            PyTuple::new(
                 py,
-                PyTuple::new(
-                    py,
-                    &[
-                        transformer.to_string().into_py(py),
-                        inputs.into_py(py),
-                        kwargs.into_py(py),
-                    ],
-                ),
-            )
-            .unwrap()
-            .extract(py)
-            .unwrap()
+                &[
+                    transformer.to_string().into_py(py),
+                    inputs.into_py(py),
+                    kwargs.into_py(py),
+                ],
+            ),
+        );
+
+        let result = match result {
+            Err(e) => {
+                let traceback = e.traceback(py).unwrap().format().unwrap();
+                error!("{traceback} {e}")
+            }
+            Ok(o) => o.extract(py).unwrap(),
+        };
+
+        result
     })
 }
 
@@ -312,25 +318,13 @@ pub fn load_dataset(
     num_rows
 }
 
-pub fn clear_gpu_cache(
-    memory_usage: Option<f32>
-) -> bool {
-
+pub fn clear_gpu_cache(memory_usage: Option<f32>) -> bool {
     Python::with_gil(|py| -> bool {
         let clear_gpu_cache: Py<PyAny> = PY_MODULE.getattr(py, "clear_gpu_cache").unwrap().into();
         clear_gpu_cache
-            .call1(
-                py,
-                PyTuple::new(
-                    py,
-                    &[
-                        memory_usage.into_py(py),
-                    ],
-                ),
-            )
+            .call1(py, PyTuple::new(py, &[memory_usage.into_py(py)]))
             .unwrap()
             .extract(py)
             .unwrap()
     })
 }
-