Update app.py

app.py

@@ -21,14 +21,18 @@ st.write("Enter text below, and the model will check if it's safe.")
 user_input = st.text_area("Enter your text here:")
 
 def check_content(text):
-    inputs = tokenizer(text, return_tensors="pt")
+    prompt = f"<|user|>\n{text}\n<|assistant|>\n"  # Ensure correct formatting for input
+    inputs = tokenizer(prompt, return_tensors="pt")
+    
     outputs = model.generate(**inputs, max_new_tokens=50)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    return response  # Check the response format to extract category
 
 if st.button("Check Content"):
     if user_input:
         result = check_content(user_input)
         st.subheader("Moderation Result:")
-        st.write(result)
+        st.write(result)  # Print raw result first to analyze format
     else:
         st.warning("Please enter some text.")