Update app.py

app.py

@@ -48,7 +48,7 @@ class GenerateRequest(BaseModel):
     input_text: str
     task_type: str
     temperature: float = 1.0
-    stream: bool = True  # Enforce stream for this functionality
+    stream: bool = True
     top_p: float = 1.0
     top_k: int = 50
     repetition_penalty: float = 1.0
@@ -149,21 +149,20 @@ async def generate(request: GenerateRequest):
             if not model_loader.download_model_from_huggingface(model_name):
                 raise HTTPException(status_code=500, detail=f"Failed to load model: {model_name}")
 
-        pipe = pipeline(task_type, model=model_name, token=HUGGINGFACE_HUB_TOKEN, device_map="auto")
+        text_pipeline = pipeline(task_type, model=model_name, token=HUGGINGFACE_HUB_TOKEN, device_map="auto")
         token_streamer = TokenIteratorStreamer()
 
-        def generate_on_thread(pipe, input_text, token_streamer, generation_params):
+        def generate_on_thread(pipeline, input_text, streamer, generation_params):
             try:
-                for output in pipe(input_text,
-                                  max_new_tokens=int(1e9),  # Effectively infinite
-                                  return_full_text=False,
-                                  streamer=token_streamer,
-                                  **generation_params):
-                    pass
+                pipeline(input_text,
+                         max_new_tokens=int(1e9),  # Effectively infinite
+                         return_full_text=False,
+                         streamer=streamer,
+                         **generation_params)
             finally:
-                token_streamer.end()
+                streamer.end()
 
-        thread = Thread(target=generate_on_thread, args=(pipe, input_text, token_streamer, generation_params))
+        thread = Thread(target=generate_on_thread, args=(text_pipeline, input_text, token_streamer, generation_params))
         thread.start()
 
         async def event_stream() -> AsyncIterator[str]:
@@ -177,7 +176,7 @@ async def generate(request: GenerateRequest):
                 await asyncio.sleep(request.chunk_delay)
             if tokens_buffer:
                 yield f"data: {json.dumps({'tokens': tokens_buffer})}\n\n"
-            yield "\n\n"  # Ensure final newline
+            yield "\n\n"
 
         return StreamingResponse(event_stream(), media_type="text/event-stream")