Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import StreamingResponse
-from pydantic import BaseModel, field_validator
+from pydantic importBaseModel, field_validator
 from transformers import pipeline, AutoConfig, AutoTokenizer
 from transformers.utils import logging
 from google.cloud import storage
@@ -55,7 +55,7 @@ class GenerateRequest(BaseModel):
     num_return_sequences: int = 1
     do_sample: bool = False
     chunk_delay: float = 0.0
-    max_new_tokens: int = 512  # Initial max tokens, can be large
+    max_new_tokens: int = 512
 
     @field_validator("model_name")
     def model_name_cannot_be_empty(cls, v):
@@ -114,25 +114,6 @@ class GCSModelLoader:
 
 model_loader = GCSModelLoader(bucket)
 
-class TokenIteratorStreamer:
-    def __init__(self):
-        self.queue = asyncio.Queue()
-
-    def put(self, value):
-        self.queue.put_nowait(value)
-
-    def end(self):
-        self.queue.put_nowait(None)
-
-    def __aiter__(self):
-        return self
-
-    async def __anext__(self):
-        value = await self.queue.get()
-        if value is None:
-            raise StopAsyncIteration
-        return value
-
 @app.post("/generate")
 async def generate(request: GenerateRequest):
     model_name = request.model_name
@@ -151,34 +132,19 @@ async def generate(request: GenerateRequest):
 
         tokenizer = AutoTokenizer.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN)
 
-        async def generate_responses() -> AsyncIterator[Dict[str, str]]:
-            current_input_text = input_text
-
-            text_pipeline = pipeline(task_type, model=model_name, tokenizer=tokenizer, token=HUGGINGFACE_HUB_TOKEN)
-            token_streamer = TokenIteratorStreamer()
-            generated_text = ""
-
-            def generate_on_thread(pipeline, input_text, streamer, generation_params, max_new_tokens):
-                try:
-                    for output in pipeline(input_text,
-                                     max_new_tokens=max_new_tokens,
-                                     return_full_text=False,
-                                     streamer=streamer,
-                                     **generation_params):
-                        streamer.put(output)  # Put the output dictionary into the queue
-                finally:
-                    streamer.end()
-
-            thread = Thread(target=generate_on_thread, args=(text_pipeline, current_input_text, token_streamer, generation_params, initial_max_new_tokens))
-            thread.start()
+        async def generate_responses() -> AsyncIterator[Dict[str, List[Dict[str, str]]]]:
+            text_pipeline = pipeline(task_type, model=model_name, tokenizer=tokenizer, token=HUGGINGFACE_HUB_TOKEN, **generation_params, max_new_tokens=initial_max_new_tokens)
 
-            async for output_dict in token_streamer:
-                if isinstance(output_dict, dict) and "generated_text" in output_dict:
-                    token = output_dict["generated_text"]
-                    generated_text += token
-                    yield {"token": token, "generated_text": generated_text}
+            def generate_on_thread(pipeline, input_text, output_queue):
+                result = pipeline(input_text)
+                output_queue.put_nowait(result)
 
-            thread.join() # Ensure the thread finishes before exiting the generator
+            output_queue = asyncio.Queue()
+            thread = Thread(target=generate_on_thread, args=(text_pipeline, input_text, output_queue))
+            thread.start()
+            result = await output_queue.get()
+            thread.join()
+            yield {"response": result}
 
         async def text_stream():
             async for data in generate_responses():