Add other datasets

app.py

@@ -4,9 +4,6 @@ import streamlit as st
 from datasets import load_dataset
 from annotated_text import annotated_text
 
-# Load data
-ds = load_dataset("hs-knowledge/hateval_enriched")
-
 
 # Show highlighted ner entities in a tweet
 def display_ner(example):
@@ -69,9 +66,9 @@ def display_ner(example):
     annotated_text(*chunks)
 
 
-def display_text(example):
+def display_text(example, text_column):
     # Use annotated_text to show entities
-    text = example["text"]
+    text = example[text_column]
 
     # Sort entities by start
     entities = sorted(example["entities"], key=lambda x: x["start"])
@@ -110,7 +107,20 @@ def display_text(example):
     annotated_text(*chunks)
 
 
-# Get first 1000 examples
+# selectbox to choose dataset
+
+selected_dataset = st.sidebar.selectbox(
+    "Select dataset", ["hateval_enriched", "sbf-enriched", "hatecheck-enriched"]
+)
+
+# Load data
+ds = load_dataset(f"hs-knowledge/{selected_dataset}")
+
+text_column = {
+    "hateval_enriched": "text",
+    "sbf-enriched": "post",
+    "hatecheck-enriched": "test_case",
+}
 
 elements = random.choices(range(len(ds["train"])), k=50)
 ds["train"] = ds["train"].select(elements)