Update app.py

app.py

@@ -23,25 +23,22 @@ 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():
-    print(f"🔄 Loading Base Model: {BASE_MODEL}")
+def load_llama_model(base_model=BASE_MODEL, adapter=None):
+    print(f"🔄 Loading Base Model: {base_model}")
     
-    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_auth_token=HUGGINGFACE_TOKEN)
+    tokenizer = AutoTokenizer.from_pretrained(base_model, token=HUGGINGFACE_TOKEN)
     model = AutoModelForCausalLM.from_pretrained(
-        BASE_MODEL, 
-        use_auth_token=HUGGINGFACE_TOKEN,
+        base_model,
+        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")
+    if adapter:
+        print(f"🔄 Loading Adapter: {adapter}")
+        model = PeftModel.from_pretrained(model, adapter, token=HUGGINGFACE_TOKEN)
+        model = model.merge_and_unload()
+        print("✅ Adapter Loaded Successfully")
 
     model.eval()
     return tokenizer, model
@@ -98,19 +95,16 @@ Input: {user_input}
 Please verify that this input doesn't violate any content policies.
 <|assistant|>"""
     
-    inputs = guard_tokenizer(prompt, return_tensors="pt", truncation=True)
-    
+    inputs = guard_tokenizer(prompt, return_tensors="pt", truncation=True, padding=True)
+
     with torch.no_grad():
-        outputs = guard_model.generate(
-            inputs.input_ids,
-            max_length=256,
-            temperature=0.1
-        )
+        outputs = guard_model.generate(inputs.input_ids, max_length=256, temperature=0.1)
     
     response = guard_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    if "flagged" in response.lower() or "violated" in response.lower() or "policy violation" in response.lower():
+
+    if any(flag in response.lower() for flag in ["flagged", "violated", "policy violation"]):
         return "⚠️ Content flagged by Llama Guard. Please modify your input."
+    
     return None
 
 # Function: Generate AI responses (same as before)