Update app.py

app.py

@@ -16,6 +16,7 @@ import uvicorn
 import asyncio
 from io import BytesIO
 from transformers import pipeline
+import json
 
 AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
 AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
@@ -193,15 +194,17 @@ async def stream_text(model, tokenizer, input_text,
             chunked_output = [final_output[i:i+10] for i in range(0, len(final_output), 10)]
             
             for chunk in chunked_output:
-                yield chunk
+              
+                yield json.dumps({"text": chunk, "is_end": False}) + "\n"
                 await asyncio.sleep(chunk_delay)
                 
+            yield json.dumps({"text": "", "is_end": True}) + "\n"
             break
                 
         else:
             chunked_output = [new_text[i:i+10] for i in range(0, len(new_text), 10)]
             for chunk in chunked_output:
-                yield chunk
+                yield json.dumps({"text": chunk, "is_end": False}) + "\n"
                 await asyncio.sleep(chunk_delay)
                 
         
@@ -210,8 +213,9 @@ async def stream_text(model, tokenizer, input_text,
             chunked_output = [output_text[i:i+10] for i in range(0, len(output_text), 10)]
             
             for chunk in chunked_output:
-                yield chunk
+                yield json.dumps({"text": chunk, "is_end": False}) + "\n"
                 await asyncio.sleep(chunk_delay)
+            yield json.dumps({"text": "", "is_end": True}) + "\n"
             break
 
         encoded_input = tokenizer(output_text,