Spaces:
Runtime error
Runtime error
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 | |
def load_models(): | |
return {id: AutoModelForSequenceClassification.from_pretrained(id) for id in model_ids.values()} | |
models = load_models() | |
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() |