Update app.py

app.py

@@ -23,43 +23,44 @@ QLORA_ADAPTER = "meta-llama/Llama-3.2-1B-Instruct-QLORA_INT4_EO8"  # Ensure this
 LLAMA_GUARD_NAME = "meta-llama/Llama-Guard-3-1B-INT4"  # Ensure this is correct
 
 # Function to load Llama model
-def load_llama_model(model_path, is_guard=False):
-    print(f"Loading model: {model_path}")
-
-    try:
-        # Get token from secrets
-        token = os.getenv("HUGGINGFACE_TOKEN")
-        if not token:
-            print("Warning: HUGGINGFACE_TOKEN not found in environment variables")
-        else:
-            print("HUGGINGFACE_TOKEN found in environment")
-        
-        # Use the parameter name 'use_auth_token' instead of 'token'
-        tokenizer = AutoTokenizer.from_pretrained(
-            BASE_MODEL,
-            use_auth_token=token
-        )
-        
-        model = AutoModelForCausalLM.from_pretrained(
-            model_path,
-            use_auth_token=token,
-            torch_dtype=torch.float16,
-            low_cpu_mem_usage=True
-        )
-        
-        # Load QLoRA adapter if applicable
-        if not is_guard and "QLORA" in model_path:
-            print("Loading QLoRA adapter...")
-            model = PeftModel.from_pretrained(
-                model, 
-                model_path,
-                use_auth_token=token
-            )
-            print("Merging LoRA weights...")
-            model = model.merge_and_unload()
-        
-        model.eval()
-        return tokenizer, model
+def load_llama_model():
+    print(f"🔄 Loading Base Model: {BASE_MODEL}")
+    
+    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_auth_token=HUGGINGFACE_TOKEN)
+    model = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL, 
+        use_auth_token=HUGGINGFACE_TOKEN,
+        torch_dtype=torch.float16,
+        low_cpu_mem_usage=True
+    )
+
+    print(f"✅ Base Model Loaded Successfully")
+
+    # Load QLoRA adapter if available
+    print(f"🔄 Loading QLoRA Adapter: {QLORA_ADAPTER}")
+    model = PeftModel.from_pretrained(model, QLORA_ADAPTER, use_auth_token=HUGGINGFACE_TOKEN)
+    print("🔄 Merging LoRA Weights...")
+    model = model.merge_and_unload()
+    print("✅ QLoRA Adapter Loaded Successfully")
+
+    model.eval()
+    return tokenizer, model
+
+# Function to load Llama Guard Model for content moderation
+def load_llama_guard():
+    print(f"🔄 Loading Llama Guard Model: {LLAMA_GUARD_NAME}")
+    
+    tokenizer = AutoTokenizer.from_pretrained(LLAMA_GUARD_NAME, use_auth_token=HUGGINGFACE_TOKEN)
+    model = AutoModelForCausalLM.from_pretrained(
+        LLAMA_GUARD_NAME, 
+        use_auth_token=HUGGINGFACE_TOKEN,
+        torch_dtype=torch.float16,
+        low_cpu_mem_usage=True
+    )
+
+    model.eval()
+    print("✅ Llama Guard Model Loaded Successfully")
+    return tokenizer, model
 
     except Exception as e:
         print(f"❌ Error loading model {model_path}: {e}")