add examples

app.py

@@ -3,6 +3,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import requests
 from PIL import Image
 from io import BytesIO
+from datasets import load_dataset
 
 # 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.
@@ -28,6 +29,12 @@ def load_image_from_url(url):
     img = Image.open(BytesIO(response.content))
     return img
 
+# Load dataset
+@st.cache_data
+def load_example_dataset():
+    ds = load_dataset("walledai/XSTest")
+    return ds['train']['prompt'][:10]  # Get first 10 examples
+
 # Evaluation function
 def evaluate_text(user_input):
     if user_input:
@@ -57,8 +64,19 @@ st.title("Text Safety Evaluator")
 if 'model_and_tokenizer' not in st.session_state:
     st.session_state.model_and_tokenizer = load_model()
 
+# Load example dataset
+example_prompts = load_example_dataset()
+
+# Display example prompts
+st.subheader("Example Inputs:")
+for i, prompt in enumerate(example_prompts):
+    if st.button(f"Example {i+1}", key=f"example_{i}"):
+        st.session_state.user_input = prompt
+
 # User input
-user_input = st.text_area("Enter the text you want to evaluate:", height=100)
+user_input = st.text_area("Enter the text you want to evaluate:", 
+                          height=100, 
+                          value=st.session_state.get('user_input', ''))
 
 # Create an empty container for the result
 result_container = st.empty()
@@ -72,17 +90,17 @@ if st.button("Evaluate"):
         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
+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
+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