akhtet's picture
Upload existing code from zero-shot-example
164e239
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()