Update app.py

app.py

@@ -11,7 +11,7 @@ from src.data.single_video import SingleVideo
 from src.data.utils_asr import PromptASR
 from src.models.llama_inference import inference
 from src.test.vidchapters import get_chapters
-from tools.download.models import download_model
+from tools.download.models import download_base_model, download_model
 
 # Set up proxies
 # from urllib.request import getproxies
@@ -26,7 +26,7 @@ tokenizer = None
 current_peft_model = None
 inference_model = None
 
-LLAMA_CKPT_PATH = "meta-llama/Llama-3.1-8B-Instruct"
+LLAMA_CKPT_PATH = "meta-llama/Meta-Llama-3.1-8B-Instruct"
 
 
 def load_base_model():
@@ -34,16 +34,28 @@ def load_base_model():
     global base_model, tokenizer
 
     if base_model is None:
-        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,
-        )
-        base_model.eval()
+        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(LLAMA_CKPT_PATH, local_dir=".")
+            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)
 
-        tokenizer = AutoTokenizer.from_pretrained(LLAMA_CKPT_PATH)
+        base_model.eval()
         tokenizer.pad_token = tokenizer.eos_token
 
         print("Base model loaded successfully")