Update app.py

app.py

@@ -258,18 +258,44 @@ def simple_chat(message, temperature: float = 0.8, max_length: int = 4096, top_p
     )
     gen_kwargs = {**input_ids, **generate_kwargs}
 
-    with torch.no_grad():
-        thread = Thread(target=model.generate, kwargs=gen_kwargs)
-        thread.start()
-        buffer = ""
-        for new_text in streamer:
-            buffer += new_text
-            yield buffer
+       buffer = ""
+
+    def generate_text():
+        nonlocal buffer
+        with torch.no_grad():
+            thread = Thread(target=model.generate, kwargs=gen_kwargs)
+            thread.start()
+            for new_text in streamer:
+                buffer += new_text
+
+    # Start the generation in a separate thread
+    generate_text()
+    
+    # Wait for the generation to finish
+    thread.join()
+
+    print("---------")
+    print("Text: ")
+    print(buffer)
+    print("---------")
+
+    return PlainTextResponse(buffer)
+
+    # with torch.no_grad():
+    #     thread = Thread(target=model.generate, kwargs=gen_kwargs)
+    #     thread.start()
+    #     buffer = ""
+    #     for new_text in streamer:
+    #         buffer += new_text
+    #         yield buffer
         
-        print("---------")
-        print("Text: ")
-        print(buffer)
-        print("---------")
+    #     print("---------")
+    #     print("Text: ")
+    #     print(buffer)
+    #     print("---------")
+
+    # return PlainTextResponse(buffer)
+