Update app.py

app.py

@@ -9,14 +9,17 @@ st.title("SelfCheck scores")
 
 
 @st.cache_data
-def load_data(min_score=0.4):
+def load_data(min_score=0.4, exclude_stories=True):
     ds = load_dataset("HuggingFaceTB/hallucinations_450_samples_scores", split="train", token=HF_TOKEN, num_proc=2)
     ds = ds.filter(lambda x: x["passage_score"] >= min_score)
+    if exclude_stories:
+        ds = ds.filter(lambda x: "story" not in x["format"])
     return ds
 
 
 min_value = st.slider('Select minimum selfcheck score', 0.0, 1.0, 0.1, key='min_score')
-ds = load_data(min_score=min_value)
+exclude_stories = st.checkbox("Exclude stories", False)
+ds = load_data(min_score=min_value, exclude_stories=exclude_stories)
 index = st.number_input(f'Found {len(ds)} samples, choose one', min_value=0, max_value=len(ds)-1, value=0, step=1)
 
 # Load data based on slider values and checkbox status