Wootang01's picture
Update app.py
2199b57
raw
history blame
578 Bytes
import streamlit as st
from transformers import pipeline
st.title("Text Classifier - BART-Large-MNLI")
st.write("Input text and labels. Submit. The machine will classify the text according to the labels.")
classifier = pipeline(task="zero-shot-classification", device=0, model="facebook/bart-large-mnli")
input_text = st.text_input(label='Input Text')
labels = st.text_input(label='Labels')
if st.button("Submit"):
labs = labels.title().split(',')
text = input_text.title()
res = classifier(text, labs)
st.success(res['labels'][0])
st.info(res['scores'][0])