Fix Assistant

artitw · Oct 7, 2024 · 0570d3e · 0570d3e
1 parent ce160df
commit 0570d3e
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/setup.py b/setup.py
@@ -5,7 +5,7 @@
 
 setuptools.setup(
   name="text2text",
-  version="1.6.7",
+  version="1.6.8",
   author="artitw",
   author_email="artitw@gmail.com",
   description="Text2Text: Crosslingual NLP/G toolkit",

diff --git a/text2text/assistant.py b/text2text/assistant.py
@@ -63,7 +63,7 @@ def load_model(self):
       result = os.system(
         "curl -fsSL https://ollama.com/install.sh | sh"
       )
-      if result!=0:
+      if result != 0:
         raise Exception("Cannot install ollama")
 
       self.ollama_serve_proc = subprocess.Popen(["ollama", "serve"])
@@ -79,13 +79,16 @@ def load_model(self):
     result = ollama.pull(self.model_name)
     if result["status"] != "success":
       raise Exception(f"Did not pull {self.model_name}.")
+
+    ollama_run_proc = subprocess.Popen(["ollama", "run", self.model_name])
 
   def chat_completion(self, messages=[{"role": "user", "content": "hello"}], stream=False, schema=None, **kwargs):
     try:
       result = ollama.ps()
       if not result or not result.get("models"):
         result = ollama.pull(self.model_name)
         if result["status"] == "success":
+          ollama_run_proc = subprocess.Popen(["ollama", "run", self.model_name])
           return self.chat_completion(messages=messages, stream=stream, **kwargs)
         raise Exception(f"Did not pull {self.model_name}. Try restarting.")
     except Exception as e: