Update app.py

app.py

@@ -258,6 +258,8 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
             streamer=streamer
         )
 
+        buffer = ""
+
         def generate_text():
             with torch.no_grad():
                 model.generate(input_ids['input_ids'], **generate_kwargs)
@@ -267,10 +269,11 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
         thread.start()
 
         # Collect generated text in real-time
-        buffer = ""
         for new_text in streamer:
             buffer += new_text
-            yield buffer
+
+        # Wait for the generation thread to finish
+        thread.join()
 
         print("  ")
         print("------")
@@ -279,9 +282,7 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
         print("  ")
         print("------")
 
-        results_text = buffer
-
-        return PlainTextResponse(results_text)
+        return PlainTextResponse(buffer)
     except Exception as e:
         return PlainTextResponse(f"Error: {str(e)}")