add file uploader & download functions

app.py

@@ -10,8 +10,7 @@ import tensorflow as tf
 
 def main():
     st.title("Interactive demo: T5 Multitasking Demo")
-    st.write("**Demo for T5's different tasks including machine translation, \
-        text summarization, document similarity, and grammatical correctness of sentences.**")
+    st.sidebar.image("https://i.gzn.jp/img/2020/02/25/google-ai-t5/01.png")
     saved_model_path = load_model_cache()
 
     # Model is loaded in st.session_state to remain stateless across reloading
@@ -33,43 +32,88 @@ def load_model_cache():
     snapshot_download(repo_id="stevekola/T5", cache_dir=CACHE_DIR)
     saved_model_path = os.path.join(CACHE_DIR, os.listdir(CACHE_DIR)[0])
     return saved_model_path
-        
+
 
 def dashboard(model):
-    """"Function to display the inputs and results
+    """Function to display the inputs and results
     params:
         model   stateless model to run inference from
     """
-    st.sidebar.write("**Select the Task Type over here**")
     task_type = st.sidebar.radio("Task Type",
                                  [
-                                  "Translate English to French",
-                                  "Translate English to German",
-                                  "Translate English to Romanian",
-                                  "Grammatical Correctness of Sentence",
-                                  "Text Summarization",
-                                  "Document Similarity Score"
+                                     "Translate English to French",
+                                     "Translate English to German",
+                                     "Translate English to Romanian",
+                                     "Grammatical Correctness of Sentence",
+                                     "Text Summarization",
+                                     "Document Similarity Score"
                                  ])
-    if task_type.startswith("Document Similarity"): # document similarity requires two documents
-        sentence1 = st.text_area("The first document/sentence",
-                            "I reside in the commercial capital city of Nigeria, which is Lagos.")
-        sentence2 = st.text_area("The second document/sentence",
-                            "I live in Lagos.")
-        sentence = sentence1 + "---" + sentence2
-    elif task_type.startswith("Text Summarization"): # text summarization's default input should be longer
-        sentence = st.text_area("Input sentence",
-                                "I don't care about those doing the comparison, but comparing the Ghanaian Jollof Rice \
-                                to Nigerian Jollof Rice is an insult to Nigerians.")
+
+    default_sentence = "I am Steven and I live in Lagos, Nigeria."
+    text_summarization_sentence = "I don't care about those doing the comparison, but comparing \
+        the Ghanaian Jollof Rice to Nigerian Jollof Rice is an insult to Nigerians."
+    doc_similarity_sentence1 = "I reside in the commercial capital city of Nigeria, which is Lagos."
+    doc_similarity_sentence2 = "I live in Lagos."
+    help_msg = "You could either type in the sentences to run inferences on or use the upload button to \
+        upload text files containing those sentences. The input sentence box, by default, displays sample \
+        texts or the texts in the files that you've uploaded. Feel free to erase them and type in new sentences."
+
+    if task_type.startswith("Document Similarity"):  # document similarity requires two documents
+        uploaded_file = upload_files(help_msg, text="Upload 2 documents for similarity check", accept_multiple_files=True)
+        if uploaded_file:
+            sentence1 = st.text_area("Enter first document/sentence", uploaded_file[0], help=help_msg)
+            sentence2 = st.text_area("Enter second document/sentence", uploaded_file[1], help=help_msg)
+        else:
+            sentence1 = st.text_area("Enter first document/sentence", doc_similarity_sentence1)
+            sentence2 = st.text_area("Enter second document/sentence", doc_similarity_sentence2)
+        sentence = sentence1 + "---" + sentence2  # to be processed like other tasks' single sentences
     else:
-        sentence = st.text_area("Input sentence",
-                                "I am Steven and I live in Lagos, Nigeria.")
+        uploaded_file = upload_files(help_msg)
+        if uploaded_file:
+            sentence = st.text_area("Enter sentence", uploaded_file, help=help_msg)
+        elif task_type.startswith("Text Summarization"):  # text summarization's default input should be longer
+            sentence = st.text_area("Enter sentence", text_summarization_sentence, help=help_msg)
+        else:
+            sentence = st.text_area("Enter sentence", default_sentence, help=help_msg)
 
     st.write("**Output Text**")
-    with st.spinner("Please wait..."): # spinner while model is running inferences
+    with st.spinner("Waiting for prediction..."):  # spinner while model is running inferences
         output_text = predict(task_type, sentence, model)
         st.write(output_text)
-        # st.download_button("Download output text", output_text)  # download_button is yet to be production-ready
-    
+        try:  # to workaround the environment's Streamlit version
+            st.download_button("Download output text", output_text)
+        except AttributeError:
+            st.text("File download not enabled for this Streamlit version \U0001F612")
+
+
+def upload_files(help_msg, text="Upload a text file here", accept_multiple_files=False):
+    """Function to upload text files and return as string text
+    params:
+        text                    Display label for the upload button
+        accept_multiple_files   params for the file_uploader function to accept more than a file
+    returns:
+        a string or a list of strings (in case of multiple files being uploaded)
+    """
+
+    def upload():
+        uploaded_files = st.file_uploader(label="Upload text files only", 
+                                          type="txt", help=help_msg,
+                                          accept_multiple_files=accept_multiple_files)
+        if st.button("Process"):
+            if not uploaded_files:
+                st.write("**No file uploaded!**")
+                return None
+            st.write("**Upload successful!**")
+            if type(uploaded_files) == list:
+                return [f.read().decode("utf-8") for f in uploaded_files]
+            return uploaded_files.read().decode("utf-8")
+
+    try:  # to workaround the environment's Streamlit version
+        with st.expander(text):
+            return upload()
+    except AttributeError:
+        return upload()
+
 
 def predict(task_type, sentence, model):
     """Function to parse the user inputs, run the parsed text through the

requirements.txt

@@ -1,3 +1,3 @@
 t5
 huggingface_hub
-streamlit
+streamlit==1.0.0