import streamlit as st
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline

x = st.slider('Select a value')
st.write(x, 'squared is', x * x)

model_ids = {
    'Bart MNLI': 'facebook/bart-large-mnli',
    'Bart MNLI + Yahoo Answers': 'joeddav/bart-large-mnli-yahoo-answers',
    'XLM Roberta XNLI (cross-lingual)': 'joeddav/xlm-roberta-large-xnli'
}

MODEL_DESC = {
    'Bart MNLI': """Bart with a classification head trained on MNLI.\n\nSequences are posed as NLI premises and topic labels are turned into premises, i.e. `business` -> `This text is about business.`""",
    'Bart MNLI + Yahoo Answers': """Bart with a classification head trained on MNLI and then further fine-tuned on Yahoo Answers topic classification.\n\nSequences are posed as NLI premises and topic labels are turned into premises, i.e. `business` -> `This text is about business.`""",
    'XLM Roberta XNLI (cross-lingual)': """XLM Roberta, a cross-lingual model, with a classification head trained on XNLI. Supported languages include: _English, French, Spanish, German, Greek, Bulgarian, Russian, Turkish, Arabic, Vietnamese, Thai, Chinese, Hindi, Swahili, and Urdu_.
Note that this model seems to be less reliable than the English-only models when classifying longer sequences.
Examples were automatically translated and may contain grammatical mistakes.
Sequences are posed as NLI premises and topic labels are turned into premises, i.e. `business` -> `This text is about business.`""",
}

device = 0 if torch.cuda.is_available() else -1

@st.cache_resource
def load_models():
    return {id: AutoModelForSequenceClassification.from_pretrained(id) for id in model_ids.values()}

models = load_models()

@st.cache_resource
def load_tokenizer(tok_id):
    return AutoTokenizer.from_pretrained(tok_id)

def get_most_likely(nli_model_id, sequence, labels, hypothesis_template, multi_class):
    classifier = pipeline(
        'zero-shot-classification',
        model=models[nli_model_id],
        tokenizer=load_tokenizer(nli_model_id),
        device=device
    )
    outputs = classifier(
        sequence,
        candidate_labels=labels,
        hypothesis_template=hypothesis_template,
        multi_label=multi_class
    )
    return outputs['labels'], outputs['scores']

def main():

    hypothesis_template = "This text is about {}."

    model_desc = st.sidebar.selectbox('Model', list(MODEL_DESC.keys()), 0)
    st.sidebar.markdown('#### Model Description')
    st.sidebar.markdown(MODEL_DESC[model_desc])

    model_id = model_ids[model_desc]

if __name__ == '__main__':
    main()