Update app.py

app.py

@@ -4,6 +4,7 @@ import threading
 import time
 import subprocess
 import spaces
+import asynico
 
 OLLAMA = os.path.expanduser("~/ollama")
 process = None
@@ -100,7 +101,7 @@ def launch():
     print("Giving ollama serve a moment")
     time.sleep(10)
   
-def stream_chat(message: str, history: list, model: str, temperature: float, max_new_tokens: int, top_p: float, top_k: int, penalty: float):
+async def stream_chat(message: str, history: list, model: str, temperature: float, max_new_tokens: int, top_p: float, top_k: int, penalty: float):
     print(f"message: {message}")
     conversation = []
     for prompt, answer in history:
@@ -114,6 +115,7 @@ def stream_chat(message: str, history: list, model: str, temperature: float, max
                
         response = client.chat(
             model=model,
+            stream=True,
             messages=conversation,
             keep_alive="60s",
             options={
@@ -127,8 +129,11 @@ def stream_chat(message: str, history: list, model: str, temperature: float, max
         )
         
         print(response)
-        return response['message']['content']
-
+        
+        buffer = ""
+            for chunk in response:
+                buffer += chunk["message"]["content"]
+                yield buffer
 
 
 def main(message: str, history: list, model: str, temperature: float, max_new_tokens: int, top_p: float, top_k: int, penalty: float):
@@ -141,8 +146,8 @@ def main(message: str, history: list, model: str, temperature: float, max_new_to
         else:
             if not process:
                 launch()
-                
-            response = stream_chat(
+               
+            response = await stream_chat(
                 message,
                 history,
                 model,
@@ -152,8 +157,7 @@ def main(message: str, history: list, model: str, temperature: float, max_new_to
                 top_k,
                 penalty             
             )
-            
-            terminate()
+
             yield response