File size: 1,764 Bytes
b88c922 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | import threading
import queue
import tiktoken
from langchain.chat_models import ChatOpenAI
from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from openai.error import InvalidRequestError
class ThreadedGenerator:
def __init__(self):
self.queue = queue.Queue()
def __iter__(self):
return self
def __next__(self):
item = self.queue.get()
if item is StopIteration: raise item
return item
def send(self, data):
self.queue.put(data)
def close(self):
self.queue.put(StopIteration)
class ChainStreamHandler(StreamingStdOutCallbackHandler):
def __init__(self, gen):
super().__init__()
self.gen = gen
def on_llm_new_token(self, token: str, **kwargs):
self.gen.send(token)
def run(g, prompts, retries = 0):
try:
chat = ChatOpenAI(
verbose=True,
streaming=True,
callback_manager=CallbackManager([ChainStreamHandler(g)]),
)
chat(prompts)
except InvalidRequestError as e:
if retries >= 10:
raise e
else:
retries += 1
del prompts[-1]
run(g, prompts, retries)
def llm_thread(g, prompts):
encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")
messages = []
tokens = 0
for prompt in prompts:
tokens += len(encoding.encode(prompt.content))
if tokens >= 4096:
break
messages.append(prompt)
try:
run(g, messages)
finally:
g.close()
def chat(prompts):
g = ThreadedGenerator()
threading.Thread(target=llm_thread, args=(g, prompts)).start()
return g |