import streamlit as st
import pandas as pd
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)

classifier = pipeline("zero-shot-classification", model="typeform/distilbert-base-uncased-mnli")

st.title("Zero-shot classification from tabular data")
st.text("Upload a table and perform zero-shot classification on a set of custom labels")

data = st.file_uploader("Upload Excel file:")
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")