import streamlit as st
from datasets import load_dataset
import os 

HF_TOKEN = os.environ.get("HF_TOKEN", None)

st.set_page_config(page_title="SelfCheck", layout="wide")
st.title("SelfCheck scores")


@st.cache_data
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')
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
sample = ds[index]
st.markdown(f"**Passage Score:** {sample['passage_score']}, seed data: {sample['seed_data']}, format: {sample['format']}.")
st.markdown(sample['original_text'])