From 818b96cae6da354f0abaac767e31a74adbb39238 Mon Sep 17 00:00:00 2001
From: 1b5d <8110504+1b5d@users.noreply.github.com>
Date: Thu, 13 Apr 2023 21:57:15 +0200
Subject: [PATCH] update llama convert and migrate logic

---
 app/llms/llama/llama.py | 25 ++++++++++++++++---------
 requirements.txt        |  2 +-
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/app/llms/llama/llama.py b/app/llms/llama/llama.py
index b66e148..c00fd06 100644
--- a/app/llms/llama/llama.py
+++ b/app/llms/llama/llama.py
@@ -24,15 +24,10 @@ class LlamaLLM(BaseLLM):
     Llama LLM implementation
     """
 
-    def _setup(self):
-        model_path = os.path.join(
-            settings.models_dir,
-            f"ggml-{settings.model_family}-{settings.model_name}-q4.bin",
-        )
-
+    def _download(self, model_path):
         if os.path.exists(model_path):
             logger.info("found an existing model %s", model_path)
-            return model_path
+            return
 
         logger.info("downloading model to %s", model_path)
 
@@ -49,6 +44,14 @@ def _setup(self):
             model_path,
         )
 
+    def _setup(self):
+        model_path = os.path.join(
+            settings.models_dir,
+            f"ggml-{settings.model_family}-{settings.model_name}-q4.bin",
+        )
+
+        self._download(model_path=model_path)
+
         if settings.setup_params["convert"]:
             tokenizer_model_path = os.path.join(settings.models_dir, "tokenizer.model")
             logger.info("downloading tokenizer model %s", tokenizer_model_path)
@@ -67,14 +70,18 @@ def _setup(self):
                 )
             )
             logger.info("converting model %s", model_path)
-            convert_one_file(model_path, tokenizer)
+            try:
+                convert_one_file(model_path, tokenizer)
+            except Exception as exp:  # pylint: disable=broad-exception-caught
+                logger.warning("Could not convert the model %s", str(exp))
 
         if settings.setup_params["migrate"]:
             logger.info("migrating model %s", model_path)
             migrate(model_path)
             # clean up backed model since we won't need it
             logger.info("cleaning up ..")
-            os.remove(model_path + ".orig")
+            if os.path.exists(model_path + ".orig"):
+                os.remove(model_path + ".orig")
 
         logger.info("setup done successfully for %s", model_path)
         return model_path
diff --git a/requirements.txt b/requirements.txt
index 13dc5db..0cc6e5e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -7,7 +7,7 @@ filelock==3.10.7
 h11==0.14.0
 huggingface-hub==0.13.3
 idna==3.4
-llama-cpp-python==0.1.26
+llama-cpp-python==0.1.33
 numpy==1.24.2
 packaging==23.0
 pydantic==1.10.7