Update main.py

main.py

@@ -1,5 +1,7 @@
+
 from flask import Flask, request, jsonify, Response, stream_with_context
 from huggingface_hub import InferenceClient
+import time
 
 # Initialize Flask app
 app = Flask(__name__)
@@ -18,7 +20,7 @@ def format_prompt(message, history):
     return prompt
 
 def generate(prompt, history, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0):
-    print(f"\nUser: {prompt}\n")
+    print(f"\nUser: {prompt}\n", flush=True)
 
     temperature = float(temperature)
     if temperature < 1e-2:
@@ -47,12 +49,22 @@ def generate(prompt, history, temperature=0.9, max_new_tokens=256, top_p=0.95, r
         )
 
         output = ""
+        buffer = []
+        buffer_size = 5  # Adjust the buffer size as needed
+
         for token in response:
-            output += token.token.text
-            yield token.token.text  # Yield each token for streaming
+            buffer.append(token.token.text)
+            if len(buffer) >= buffer_size:
+                chunk = ''.join(buffer)
+                yield chunk
+                buffer.clear()
+                time.sleep(0.1)  # Introduce a delay to manage the flow of data
+
+        if buffer:
+            yield ''.join(buffer)
 
         # Print AI response
-        print(f"\nSema AI: {output}\n")
+        print(f"\nSema AI: {output}\n, flush=True")
     except Exception as e:
         print(f"Exception during generation: {str(e)}")
         yield "Error occurred"