separate pipeline creation and usage for transform (#921)

kczimm · web-flow · commit e3892bc23372 · 2023-08-22T11:06:53.000-05:00
diff --git a/pgml-dashboard/Cargo.lock b/pgml-dashboard/Cargo.lock
diff --git a/pgml-extension/src/bindings/mod.rs b/pgml-extension/src/bindings/mod.rs
@@ -73,15 +73,18 @@ pub trait Bindings: Send + Sync + Debug {
         Self: Sized;
 }
 
-trait TracebackError<T> {
+pub trait TracebackError<T> {
     fn format_traceback(self, py: Python<'_>) -> Result<T>;
 }
 
 impl<T> TracebackError<T> for PyResult<T> {
     fn format_traceback(self, py: Python<'_>) -> Result<T> {
-        self.map_err(|e| {
-            let traceback = e.traceback(py).unwrap().format().unwrap();
-            anyhow!("{traceback} {e}")
+        self.map_err(|e| match e.traceback(py) {
+            Some(traceback) => match traceback.format() {
+                Ok(traceback) => anyhow!("{traceback} {e}"),
+                Err(format_e) => anyhow!("{e} {format_e}"),
+            },
+            None => anyhow!("{e}"),
         })
     }
 }
diff --git a/pgml-extension/src/bindings/transformers/transformers.py b/pgml-extension/src/bindings/transformers/transformers.py
@@ -184,28 +184,45 @@ def get_model_from(task):
         return model[ty][0]
 
 
+def create_pipeline(task):
+    if isinstance(task, str):
+        task = orjson.loads(task)
+    ensure_device(task)
+    convert_dtype(task)
+    model_name = task.get("model", None)
+    if model_name and "-ggml" in model_name:
+        pipe = GGMLPipeline(model_name, **task)
+    elif model_name and "-gptq" in model_name:
+        pipe = GPTQPipeline(model_name, **task)
+    else:
+        try:
+            pipe = StandardPipeline(model_name, **task)
+        except TypeError:
+            # some models fail when given "device" kwargs, remove and try again
+            task.pop("device")
+            pipe = StandardPipeline(model_name, **task)
+    return pipe
+
+
+def transform_using(pipeline, args, inputs):
+    args = orjson.loads(args)
+    inputs = orjson.loads(inputs)
+
+    if pipeline.task == "question-answering":
+        inputs = [orjson.loads(input) for input in inputs]
+    convert_eos_token(pipeline.tokenizer, args)
+
+    return orjson.dumps(pipeline(inputs, **args), default=orjson_default).decode()
+
+
 def transform(task, args, inputs):
     task = orjson.loads(task)
     args = orjson.loads(args)
     inputs = orjson.loads(inputs)
 
     key = ",".join([f"{key}:{val}" for (key, val) in sorted(task.items())])
     if key not in __cache_transform_pipeline_by_task:
-        ensure_device(task)
-        convert_dtype(task)
-        model_name = task.get("model", None)
-        if model_name and "-ggml" in model_name:
-            pipe = GGMLPipeline(model_name, **task)
-        elif model_name and "-gptq" in model_name:
-            pipe = GPTQPipeline(model_name, **task)
-        else:
-            try:
-                pipe = StandardPipeline(model_name, **task)
-            except TypeError:
-                # some models fail when given "device" kwargs, remove and try again
-                task.pop("device")
-                pipe = StandardPipeline(model_name, **task)
-
+        pipe = create_pipeline(task)
         __cache_transform_pipeline_by_task[key] = pipe
 
     pipe = __cache_transform_pipeline_by_task[key]

Original file line number	Diff line number	Diff line change
`@@ -73,15 +73,18 @@ pub trait Bindings: Send + Sync + Debug {`
`73`	`73`	`Self: Sized;`
`74`	`74`	`}`
`75`	`75`
`76`		`-trait TracebackError<T> {`
	`76`	`+pub trait TracebackError<T> {`
`77`	`77`	`fn format_traceback(self, py: Python<'_>) -> Result<T>;`
`78`	`78`	`}`
`79`	`79`
`80`	`80`	`impl<T> TracebackError<T> for PyResult<T> {`
`81`	`81`	`fn format_traceback(self, py: Python<'_>) -> Result<T> {`
`82`		`- self.map_err(\|e\| {`
`83`		`- let traceback = e.traceback(py).unwrap().format().unwrap();`
`84`		`- anyhow!("{traceback} {e}")`
	`82`	`+ self.map_err(\|e\| match e.traceback(py) {`
	`83`	`+ Some(traceback) => match traceback.format() {`
	`84`	`+ Ok(traceback) => anyhow!("{traceback} {e}"),`
	`85`	`+ Err(format_e) => anyhow!("{e} {format_e}"),`
	`86`	`+ },`
	`87`	`+ None => anyhow!("{e}"),`
`85`	`88`	`})`
`86`	`89`	`}`
`87`	`90`	`}`