Walledguard bring back

app.py

@@ -16,7 +16,7 @@ Answer: [/INST]
 # Load the model and tokenizer
 @st.cache_resource
 def load_model():
-    model_name = "meta-llama/Llama-Guard-3-8B-INT8"
+    model_name = "walledai/walledguard-c"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(model_name)
     return tokenizer, model
@@ -35,8 +35,8 @@ def evaluate_text(user_input):
         tokenizer, model = st.session_state.model_and_tokenizer
         
         # Prepare input
-        input_ids = tokenizer.apply_chat_template(chat, return_tensors="pt").to(device)        
-
+        input_ids = tokenizer.encode(TEMPLATE.format(prompt=user_input), return_tensors="pt")
+        
         # Generate output
         output = model.generate(input_ids=input_ids, max_new_tokens=20, pad_token_id=0)
         

app_walledguard.py → app_LG3_try.py

@@ -16,7 +16,7 @@ Answer: [/INST]
 # Load the model and tokenizer
 @st.cache_resource
 def load_model():
-    model_name = "walledai/walledguard-c"
+    model_name = "meta-llama/Llama-Guard-3-8B-INT8"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(model_name)
     return tokenizer, model
@@ -35,8 +35,8 @@ def evaluate_text(user_input):
         tokenizer, model = st.session_state.model_and_tokenizer
         
         # Prepare input
-        input_ids = tokenizer.encode(TEMPLATE.format(prompt=user_input), return_tensors="pt")
-        
+        input_ids = tokenizer.apply_chat_template(chat, return_tensors="pt").to(device)        
+
         # Generate output
         output = model.generate(input_ids=input_ids, max_new_tokens=20, pad_token_id=0)