ashrestha's picture
Update app.py
06ba6f1
raw
history blame
955 Bytes
import streamlit as st
import json
from transformers import pipeline
classifier = pipeline("zero-shot-classification",
model="valhalla/distilbart-mnli-12-1")
with st.form('inputs'):
input_text = st.text_area("Input text")
input_label = st.text_input("Labels", placeholder="support, help, important")
input_multi = st.checkbox('Allow multiple true classes', value=False)
submit_button = st.form_submit_button(label='Submit')
if submit_button:
labels = list(l.strip() for l in input_label.split(','))
pred = classifier(input_text, labels, multi_class=input_multi)
if input_multi:
st.vega_lite_chart(json.dumps(pred),
{'mark': 'bar', 'encoding': {'y': 'labels', 'x': 'scores'}}
)
else:
out = f"Top predicted labels are {', '.join(p for p in pred['labels'][0:2])}"
st.success(out)
# st.markdown(pred)