Update app.py

app.py

@@ -23,19 +23,22 @@ def load_llama_model(model_name):
 
     tokenizer = LlamaTokenizer.from_pretrained(model_name, token=HUGGINGFACE_TOKEN)
 
-    # Load the checkpoint manually
-    model_path = f"{model_name}/consolidated.00.pth"
-    state_dict = torch.load(model_path, map_location="cpu")  # Adjust for GPU if needed
+    # 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!")
-    
-    return tokenizer, state_dict 
+
+    # Initialize model and load state_dict
+    model = AutoModelForCausalLM.from_pretrained(model_name, state_dict=state_dict)
+
+    return tokenizer, model
 
 # Load the quantized Llama model
 tokenizer, model = load_llama_model(QUANTIZED_MODEL)
 
 # Load Llama Guard for content moderation
-guard_tokenizer, guard_model = load_llama_model(LLAMA_GUARD_NAME, is_guard=True)
+guard_tokenizer, guard_model = load_llama_model(LLAMA_GUARD_NAME)
 
 # Define Prompt Templates
 PROMPTS = {