Update app.py

app.py

@@ -34,27 +34,25 @@ def load_base_model():
     global base_model, tokenizer
 
     if base_model is None:
-        try:
             print(f"Loading base model: {LLAMA_CKPT_PATH}")
-            base_model = load_model_llamarecipes(
-                model_name=LLAMA_CKPT_PATH,
-                device_map="auto",
-                quantization=None,
-                use_fast_kernels=True,
-            )
-            tokenizer = AutoTokenizer.from_pretrained(LLAMA_CKPT_PATH)
-        except Exception as e:
-            # Try to get the local path using the download function
-            model_path = download_base_model("lucas-ventura/chapter-llama", local_dir=".")
-            model_path += "/" + LLAMA_CKPT_PATH
-            print(f"Model path: {model_path}")
-            base_model = load_model_llamarecipes(
-                model_name=model_path,
-                device_map="auto",
-                quantization=None,
-                use_fast_kernels=True,
-            )
-            tokenizer = AutoTokenizer.from_pretrained(model_path)
+        # base_model = load_model_llamarecipes(
+        #     model_name=LLAMA_CKPT_PATH,
+        #     device_map="auto",
+        #     quantization=None,
+        #     use_fast_kernels=True,
+        # )
+        # tokenizer = AutoTokenizer.from_pretrained(LLAMA_CKPT_PATH)
+        # Try to get the local path using the download function
+        model_path = download_base_model("lucas-ventura/chapter-llama", local_dir=".")
+        model_path += f"lucas-ventura/chapter-llama/{LLAMA_CKPT_PATH}"
+        print(f"Model path: {model_path}")
+        base_model = load_model_llamarecipes(
+            model_name=model_path,
+            device_map="auto",
+            quantization=None,
+            use_fast_kernels=True,
+        )
+        tokenizer = AutoTokenizer.from_pretrained(model_path)
 
         base_model.eval()
         tokenizer.pad_token = tokenizer.eos_token