LG3

app.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"
     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)
         

app_walledguard.py

@@ -0,0 +1,88 @@
+import streamlit as st
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import requests
+from PIL import Image
+from io import BytesIO
+
+# Define the template
+TEMPLATE = """[INST] Evaluate if the given text is harmful i.e. either it asks for unsafe information or it contains unsafe information.
+<START TEXT>
+{prompt}
+<END TEXT>
+
+Answer: [/INST]
+"""
+
+# Load the model and tokenizer
+@st.cache_resource
+def load_model():
+    model_name = "walledai/walledguard-c"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    return tokenizer, model
+
+# Function to load image from URL
+@st.cache_data()
+def load_image_from_url(url):
+    response = requests.get(url)
+    img = Image.open(BytesIO(response.content))
+    return img
+
+# Evaluation function
+def evaluate_text(user_input):
+    if user_input:
+        # Get model and tokenizer from session state
+        tokenizer, model = st.session_state.model_and_tokenizer
+        
+        # Prepare input
+        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)
+        
+        # Decode output
+        prompt_len = input_ids.shape[-1]
+        output_decoded = tokenizer.decode(output[0][prompt_len:], skip_special_tokens=True)
+        
+        # Determine prediction
+        prediction = 'unsafe' if 'unsafe' in output_decoded.lower() else 'safe'
+        
+        return prediction
+    return None
+
+# Streamlit app
+st.title("Text Safety Evaluator")
+
+# Load model and tokenizer once and store in session state
+if 'model_and_tokenizer' not in st.session_state:
+    st.session_state.model_and_tokenizer = load_model()
+
+# User input
+user_input = st.text_area("Enter the text you want to evaluate:", height=100)
+
+# Create an empty container for the result
+result_container = st.empty()
+
+if st.button("Evaluate"):
+    prediction = evaluate_text(user_input)
+    if prediction:
+        result_container.subheader("Evaluation Result:")
+        result_container.write(f"The text is evaluated as: **{prediction.upper()}**")
+    else:
+        result_container.warning("Please enter some text to evaluate.")
+
+# Add logo at the bottom center (only once)
+#if 'logo_displayed' not in st.session_state:
+col1, col2, col3 = st.columns([1,2,1])
+with col2:
+    logo_url = "https://github.com/walledai/walledeval/assets/32847115/d8b1d14f-7071-448b-8997-2eeba4c2c8f6"
+    logo = load_image_from_url(logo_url)
+    st.image(logo, use_column_width=True, width=500)  # Adjust the width as needed
+#st.session_state.logo_displayed = True
+
+# Add information about Walled Guard Advanced (only once)
+#if 'info_displayed' not in st.session_state:
+col1, col2, col3 = st.columns([1,2,1])
+with col2:
+    st.info("For a more performant version, check out Walled Guard Advanced. Connect with us at admin@walled.ai for more information.")
+#st.session_state.info_displayed = True