diff --git a/setup.py b/setup.py index 13c174e..17e7c9b 100755 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setuptools.setup( name="text2text", - version="1.6.8", + version="1.6.9", author="artitw", author_email="artitw@gmail.com", description="Text2Text: Crosslingual NLP/G toolkit", diff --git a/text2text/assistant.py b/text2text/assistant.py index b8d68a1..549c599 100644 --- a/text2text/assistant.py +++ b/text2text/assistant.py @@ -69,28 +69,22 @@ def load_model(self): self.ollama_serve_proc = subprocess.Popen(["ollama", "serve"]) time.sleep(1) - result = subprocess.check_output( - ["ollama", "-v"], - stderr=subprocess.STDOUT - ).decode("utf-8") - if not result.startswith("ollama version"): - raise Exception(result) + if not ollama_version(): + raise Exception("Cannot serve ollama") 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]) + if result["status"] == "success": + ollama_run_proc = subprocess.Popen(["ollama", "run", self.model_name]) + else: + raise Exception(f"Did not pull {self.model_name}. Try restarting.") 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.") + warnings.warn("No model loaded. Retrying...") + self.load_model() + return self.chat_completion(messages=messages, stream=stream, **kwargs) except Exception as e: warnings.warn(str(e)) warnings.warn("Retrying...")