Update app.py

app.py

@@ -175,6 +175,7 @@ async def generate(request: GenerateRequest):
                 repetition_penalty=repetition_penalty,
                 do_sample=do_sample,
                 num_return_sequences=num_return_sequences,
+                eos_token_id = tokenizer.eos_token_id
             )
             if stream:
                 return StreamingResponse(
@@ -245,27 +246,29 @@ async def stream_text(model, tokenizer, input_text,
             outputs.sequences[0][len(encoded_input["input_ids"][0]):],
             skip_special_tokens=True
         )
+
+        if len(new_text) == 0:
+            if not stop_criteria(outputs.sequences, None):
+                for text in output_text.split():
+                    yield json.dumps({"text": text, "is_end": False}) + "\n"
+                yield json.dumps({"text": "", "is_end": True}) + "\n"
+            break
         
         output_text += new_text
-
+        
         for text in new_text.split():
               yield json.dumps({"text": text, "is_end": False}) + "\n"
-
+              
         if stop_criteria(outputs.sequences, None):
             yield json.dumps({"text": "", "is_end": True}) + "\n"
             break
-      
-        if len(new_text) == 0:
-            for text in output_text.split():
-                 yield json.dumps({"text": text, "is_end": False}) + "\n"
-            yield json.dumps({"text": "", "is_end": True}) + "\n"
-            break
-
+        
         encoded_input = tokenizer(
             output_text, return_tensors="pt",
             truncation=True
         ).to(device)
         output_text = ""
+        
 
 async def generate_text(model, tokenizer, input_text,
                         generation_config, stop_sequences,
@@ -299,6 +302,7 @@ async def generate_text(model, tokenizer, input_text,
     
     return generated_text
 
+
 @app.post("/generate-image")
 async def generate_image(request: GenerateRequest):
     try: