Update app.py

app.py

@@ -92,15 +92,14 @@ def ollama_func(command):
     else:
         return "No supported command."
         
-      
+@spaces.GPU()      
 def launch():
     global OLLAMA_SERVICE_THREAD
     OLLAMA_SERVICE_THREAD = threading.Thread(target=ollama_service_thread)
     OLLAMA_SERVICE_THREAD.start()
     print("Giving ollama serve a moment")
     time.sleep(10)
-
-@spaces.GPU()  
+  
 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 = []
@@ -116,6 +115,8 @@ def stream_chat(message: str, history: list, model: str, temperature: float, max
         response = client.chat(
             model=model,
             messages=conversation,
+            keep_alive="60s",
+            stream=True,
             options={
                 'num_predict': max_new_tokens,
                 'temperature': temperature,
@@ -123,11 +124,13 @@ def stream_chat(message: str, history: list, model: str, temperature: float, max
                 'top_k': top_k,
                 'repeat_penalty': penalty,
                 'low_vram': True,
-                "keep_alive": "60s",
             },
         )
+        
+        terminate()
+        
         print(response)
-        return response['message']['content']
+        yield response['message']['content']
 
 
 
@@ -152,8 +155,11 @@ def main(message: str, history: list, model: str, temperature: float, max_new_to
                 top_k,
                 penalty             
             )
-            print(response)
-            yield response
+            
+            buffer = ""
+            for chunk in response:
+                buffer += chunk
+                yield buffer