Update extension for cloud streaming use (#1163)

SilasMarvin · web-flow · commit c2d69ad0af56 · 2023-11-16T11:54:22.000-08:00
diff --git a/pgml-extension/src/bindings/transformers/transformers.py b/pgml-extension/src/bindings/transformers/transformers.py
@@ -312,14 +312,16 @@ def create_pipeline(task):
     return pipe
 
 
-def transform_using(pipeline, args, inputs):
+def transform_using(pipeline, args, inputs, stream=False):
     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)
 
+    if stream:
+        return pipeline.stream(inputs, **args)
     return orjson.dumps(pipeline(inputs, **args), default=orjson_default).decode()
 
 
