Update app.py

app.py

@@ -1,13 +1,10 @@
 import streamlit as st
-import requests
-import json
-import os
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Hugging Face API details
-API_URL = "https://api-inference.huggingface.co/models/meta-llama/Llama-Guard-3-1B"
-API_KEY = os.getenv("HF_API_KEY")  # Fetch API key from environment variable
-
-HEADERS = {"Authorization": f"Bearer {API_KEY}"}
+# Load model and tokenizer
+MODEL_NAME = "meta-llama/Llama-Guard-3-1B"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
 # Streamlit UI
 st.title("AI Safe Content Checker Tool")
@@ -16,30 +13,15 @@ st.write("Enter text below, and the model will check if it's safe.")
 # User input
 user_input = st.text_area("Enter your text here:")
 
-def query(payload):
-    try:
-        response = requests.post(API_URL, headers=HEADERS, json=payload)
-
-        # Check response status
-        if response.status_code != 200:
-            st.error(f"API Error {response.status_code}: {response.text}")
-            return None
-
-        return response.json()  # Properly parse JSON
-
-    except requests.exceptions.RequestException as e:
-        st.error(f"Request failed: {e}")
-        return None
+def check_content(text):
+    inputs = tokenizer(text, return_tensors="pt")
+    outputs = model.generate(**inputs, max_new_tokens=50)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 if st.button("Check Content"):
     if user_input:
-        payload = {"inputs": user_input}  # Llama Guard expects this format
-        response = query(payload)
-
-        if response:
-            st.subheader("Moderation Result:")
-            st.json(response)  # Display result as formatted JSON
-        else:
-            st.error("No valid response from the API.")
+        result = check_content(user_input)
+        st.subheader("Moderation Result:")
+        st.write(result)
     else:
         st.warning("Please enter some text.")