From 8ac5a020e56a6f36fd012cdb0f3a916aa18207f6 Mon Sep 17 00:00:00 2001
From: SilasMarvin <19626586+SilasMarvin@users.noreply.github.com>
Date: Tue, 5 Dec 2023 14:02:23 -0800
Subject: [PATCH 1/2] Added correct timeout
---
.../src/bindings/transformers/transformers.py | 19 ++++++++++++-------
1 file changed, 12 insertions(+), 7 deletions(-)
diff --git a/pgml-extension/src/bindings/transformers/transformers.py b/pgml-extension/src/bindings/transformers/transformers.py
index c738fe1f5..d1dac1dd1 100644
--- a/pgml-extension/src/bindings/transformers/transformers.py
+++ b/pgml-extension/src/bindings/transformers/transformers.py
@@ -106,6 +106,7 @@ def __init__(self, tokenizer, skip_prompt=False, timeout=None, **decode_kwargs):
self.text_queue = queue.Queue()
self.token_cache = []
self.text_index_cache = []
+ print("\n the python timeout ", timeout, "\n", file=sys.stderr)
def put(self, values):
if self.skip_prompt and self.next_tokens_are_prompt:
@@ -131,7 +132,7 @@ def put(self, values):
self.text_index_cache[i] += len(printable_text)
output.append(printable_text)
if any(output):
- self.text_queue.put(output, self.timeout)
+ self.text_queue.put(output)
def end(self):
self.next_tokens_are_prompt = True
@@ -139,8 +140,8 @@ def end(self):
for i, tokens in enumerate(self.token_cache):
text = self.tokenizer.decode(tokens, **self.decode_kwargs)
output.append(text[self.text_index_cache[i] :])
- self.text_queue.put(output, self.timeout)
- self.text_queue.put(self.stop_signal, self.timeout)
+ self.text_queue.put(output)
+ self.text_queue.put(self.stop_signal)
def __iter__(self):
return self
@@ -264,12 +265,13 @@ def __init__(self, model_name, **kwargs):
if self.tokenizer.pad_token is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
- def stream(self, input, **kwargs):
+ def stream(self, input, timeout=None, **kwargs):
streamer = None
generation_kwargs = None
if self.task == "conversational":
streamer = TextIteratorStreamer(
self.tokenizer,
+ timeout=timeout,
skip_prompt=True,
)
if "chat_template" in kwargs:
@@ -286,7 +288,10 @@ def stream(self, input, **kwargs):
input = self.tokenizer(input, return_tensors="pt").to(self.model.device)
generation_kwargs = dict(input, streamer=streamer, **kwargs)
else:
- streamer = TextIteratorStreamer(self.tokenizer)
+ streamer = TextIteratorStreamer(
+ self.tokenizer,
+ timeout=timeout,
+ )
input = self.tokenizer(input, return_tensors="pt", padding=True).to(
self.model.device
)
@@ -355,7 +360,7 @@ def create_pipeline(task):
return pipe
-def transform_using(pipeline, args, inputs, stream=False):
+def transform_using(pipeline, args, inputs, stream=False, timeout=None):
args = orjson.loads(args)
inputs = orjson.loads(inputs)
@@ -364,7 +369,7 @@ def transform_using(pipeline, args, inputs, stream=False):
convert_eos_token(pipeline.tokenizer, args)
if stream:
- return pipeline.stream(inputs, **args)
+ return pipeline.stream(inputs, timeout=timeout, **args)
return orjson.dumps(pipeline(inputs, **args), default=orjson_default).decode()
From f55de0cce7011b5688bfd292d81dc923394aed74 Mon Sep 17 00:00:00 2001
From: SilasMarvin <19626586+SilasMarvin@users.noreply.github.com>
Date: Tue, 5 Dec 2023 14:03:04 -0800
Subject: [PATCH 2/2] Removed print
---
pgml-extension/src/bindings/transformers/transformers.py | 1 -
1 file changed, 1 deletion(-)
diff --git a/pgml-extension/src/bindings/transformers/transformers.py b/pgml-extension/src/bindings/transformers/transformers.py
index d1dac1dd1..7ff3bdc18 100644
--- a/pgml-extension/src/bindings/transformers/transformers.py
+++ b/pgml-extension/src/bindings/transformers/transformers.py
@@ -106,7 +106,6 @@ def __init__(self, tokenizer, skip_prompt=False, timeout=None, **decode_kwargs):
self.text_queue = queue.Queue()
self.token_cache = []
self.text_index_cache = []
- print("\n the python timeout ", timeout, "\n", file=sys.stderr)
def put(self, values):
if self.skip_prompt and self.next_tokens_are_prompt:
pFad - Phonifier reborn
Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.
Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.
Alternative Proxies:
Alternative Proxy
pFad Proxy
pFad v3 Proxy
pFad v4 Proxy