Add download button

app.py

@@ -1,3 +1,5 @@
+from io import BytesIO
+
 import pandas as pd
 import streamlit as st
 import tokenizers
@@ -20,7 +22,7 @@ st.set_page_config(
         tokenizers.AddedToken: lambda _: None,
     },
     allow_output_mutation=True,
-    show_spinner=False
+    show_spinner=False,
 )
 def load_classifier() -> Pipeline:
     classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
@@ -35,18 +37,19 @@ 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 column named `text` in its header):")
+data = st.file_uploader(
+    "Upload Excel file (it should contain a column named `text` in its header):"
+)
 labels = st.text_input("Enter comma-separated labels:")
 
 # classify first N snippets only for faster inference
-N = 100
 
 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(N)
+        table = table.loc[table["text"].apply(len) > 10].reset_index(drop=True)
 
         prog_bar = st.progress(0)
         preds = []
@@ -59,6 +62,13 @@ if st.button("Calculate labels"):
 
         st.table(table[["text", "label"]])
 
+        buf = BytesIO()
+        table[["text", "label"]].to_excel(buf)
+
+        st.download_button(
+            label="Download table", data=buf.getvalue(), file_name="output.xlsx"
+        )
+
     except:
         st.error(
             "Something went wrong. Make sure you upload an Excel file containing a column named `text` and a set of comma-separated labels is provided"

requirements.txt

@@ -4,4 +4,5 @@ protobuf
 sentencepiece
 xlrd
 -f https://download.pytorch.org/whl/torch_stable.html 
-torch
+torch
+appdirs