import streamlit as st
from datasets import load_dataset

dataset = load_dataset("reshinthadith/pairwise-code-review-instruct-critique-revision-python")
total_length = len(dataset["train"])

st.title('Pairwise Comparison')
with st.form("dataset_form"):
    len_index = st.slider('Dataset Index', 0, total_length,0)
    if st.form_submit_button("Load"):
        st.markdown("## Question")
        st.markdown(dataset["train"]["prompt"][len_index],unsafe_allow_html=True)
        st.markdown("## Chosen")
        chosen_text = dataset["train"]["chosen"][len_index]
        if "Critique" in chosen_text:
            st.text(chosen_text.split("Critique")[1])
        else:
            st.text(chosen_text)
        rejected_text = dataset["train"]["rejected"][len_index]
        st.markdown("## Rejected")
        if "Critique" in rejected_text:
            st.text(rejected_text.split("Critique")[1])
        else:
            st.text(rejected_text)