zeroshotcat / app.py
Davide Fiocco
Improve messaging about model loading
7c7d7c6
raw
history blame
1.37 kB
import pandas as pd
import streamlit as st
from transformers import pipeline
st.set_page_config(
page_title="Zero-shot classification from tabular data",
page_icon=None,
layout="wide",
initial_sidebar_state="auto",
menu_items=None,
)
with st.spinner("Setting stuff up..."):
classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
st.title("Zero-shot classification from tabular data")
st.text(
"Upload an Excel table and perform zero-shot classification on a set of custom labels"
)
data = st.file_uploader("Upload Excel file (it should contain a `text` column):")
labels = st.text_input("Enter comma-separated labels:")
if st.button("Calculate labels"):
try:
labels_list = labels.split(",")
table = pd.read_excel(data)
table = table.loc[table["text"].apply(len) > 10].reset_index(drop=True).head(50)
prog_bar = st.progress(0)
preds = []
for i in range(len(table)):
preds.append(classifier(table.loc[i, "text"], labels)["labels"][0])
prog_bar.progress((i + 1) / len(table))
table["label"] = preds
st.table(table[["text", "label"]])
except:
st.error(
"File load didn't work. Make sure you upload a file containing a `text` column and a set of comma-separated labels is provided"
)