Update app.py

app.py

@@ -22,16 +22,18 @@ def load_llama_model(model_name):
     print(f"🔄 Loading Model: {model_name}")
 
     tokenizer = LlamaTokenizer.from_pretrained(model_name, token=HUGGINGFACE_TOKEN)
+    
+    try:
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name, 
+            token=HUGGINGFACE_TOKEN,
+            trust_remote_code=True  # Allows loading non-standard model formats
+        )
+    except Exception as e:
+        print(f"❌ Error loading model: {e}")
+        raise ValueError(f"❌ Model {model_name} may not have valid weight files. Check the Hugging Face repository.")
 
-    # Manually load `.pth` state dictionary
-    model_url = f"https://huggingface.co/{model_name}/resolve/main/consolidated.00.pth"
-    state_dict = torch.hub.load_state_dict_from_url(model_url, map_location="cpu")
-
-    print("✅ Model state dictionary loaded successfully!")
-
-    # Initialize model and load state_dict
-    model = AutoModelForCausalLM.from_pretrained(model_name, state_dict=state_dict)
-
+    print("✅ Model loaded successfully!")
     return tokenizer, model
 
 # Load the quantized Llama model