import streamlit as st
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from ferret import Benchmark


@st.cache()
def get_model(model_name):
    return AutoModelForSequenceClassification.from_pretrained(model_name)


def get_tokenizer(tokenizer_name):
    return AutoTokenizer.from_pretrained(tokenizer_name, use_fast=True)


def body():

    st.title("Evaluate using *ferret* !")

    st.markdown(
        """
    
        ### 👋 Hi!

        Insert down below your text, choose a model and fire up ferret. We will use
        *ferret* to:
        1. produce explanations with all supported methods
        2. evaluate explanations on state-of-the-art **faithfulness metrics**. 
        """
    )

    col1, col2 = st.columns([1, 1])
    with col1:
        model_name = st.text_input("HF Model", "g8a9/bert-base-cased_ami18")
    with col2:
        tokenizer_name = st.text_input("HF Tokenizer", "bert-base-cased")

    text = st.text_input("Text")

    compute = st.button("Compute")

    if compute and model_name and tokenizer_name:
        model = get_model(model_name)
        tokenizer = get_tokenizer(tokenizer_name)

        bench = Benchmark(model, tokenizer)
        explanations = bench.explain(text)
        st.dataframe(bench.show_table(explanations))