Update app.py

app.py

@@ -22,8 +22,8 @@ subprocess.run("~/ollama pull gemma2", shell=True)
 
 import copy
 import gradio as gr
-import ollama
-
+from ollama import Client
+client = Client(host='http://localhost:11434', timeout=60)
 
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 MODEL_ID = os.environ.get("MODEL_ID", "google/gemma-2-9b-it")
@@ -65,7 +65,7 @@ def stream_chat(message: str, history: list, temperature: float, context_window:
 
     print(f"Conversation is -\n{conversation}")
 
-    response = ollama.chat(
+    response = client.chat(
         model="gemma2",
         messages=conversation,
         stream=True