import streamlit as st  
from parse import retrieve

  
def main():  
    st.title("Streamlit App")  
  
    model_name = st.selectbox(  
        "Select Model",  
        options=["T5", "Vicuna", "UL2", "ChatGPT"],  
        index=0,  
    )  
  
    dataset_name = st.selectbox(  
        "Select Dataset",  
        options=[  
            "SST-2", "CoLA", "QQP", "MRPC", "MNLI", "QNLI",  
            "RTE", "WNLI", "MMLU", "SQuAD V2", "IWSLT 2017", "UN Multi", "Math"  
        ],  
        index=0,  
    )

    attack_name = st.selectbox(
        "Select Attack",
        options=[
            "BertAttack", "CheckList", "DeepWordBug", "StressTest", "TextFooler", "TextBugger", "Semantic"
        ],
        index=0,
    )

    prompt_type = st.selectbox(  
        "Select Prompt Type",  
        options=["zeroshot-task", "zeroshot-role", "fewshot-task", "fewshot-role"],  
        index=0,  
    )  
  
    st.write(f"Model: {model_name}")  
    st.write(f"Dataset: {dataset_name}")  
    st.write(f"Prompt Type: {prompt_type}")  
  
    if st.button("Retrieve"):  
        results = retrieve(model_name, dataset_name, attack_name, prompt_type)
        
        for result in results:
            st.write("Original prompt: {}".format(result["origin prompt"]))
            st.write("Original acc: {}".format(result["origin acc"]))
            st.write("Attack prompt: {}".format(result["attack prompt"]))
            st.write("Attack acc: {}".format(result["attack acc"]))
  
if __name__ == "__main__":  
    main()