Update app.py

app.py

@@ -86,21 +86,25 @@ lottie_animation = load_lottie_url(LOTTIE_URL)
 # ---------------------------------------------------------------------
 @st.cache_resource
 def load_llama_pipeline(model_id: str, device: str, token: str):
-    tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        use_auth_token=token,
-        torch_dtype=torch.float16 if device == "auto" else torch.float32,
-        device_map=device,
-        low_cpu_mem_usage=True
-    )
-    text_gen_pipeline = pipeline(
-        "text-generation",
-        model=model,
-        tokenizer=tokenizer,
-        device_map=device
-    )
-    return text_gen_pipeline
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            use_auth_token=token,
+            torch_dtype=torch.float16 if device == "auto" else torch.float32,
+            device_map=device,
+            low_cpu_mem_usage=True
+        )
+        text_gen_pipeline = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            device_map=device
+        )
+        return text_gen_pipeline
+    except Exception as e:
+        st.error(f"Error loading Llama model: {e}")
+        raise
 
 # ---------------------------------------------------------------------
 # 5) GENERATE RADIO SCRIPT