add comments to app file

app.py

@@ -12,13 +12,18 @@ 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.**")
-    model = load_model_cache()
-    dashboard(model)
+    saved_model_path = load_model_cache()
+
+    # Model is loaded in st.session_state to remain stateless across reloading
+    if 'model' not in st.session_state:
+        st.session_state.model = tf.saved_model.load(saved_model_path, ["serve"])
+
+    dashboard(st.session_state.model)
 
 
 @st.cache
 def load_model_cache():
-    """Function to retrieve the model from HuggingFace Hub, load it and cache it using st.cache wrapper
+    """Function to retrieve the model from HuggingFace Hub and cache it using st.cache wrapper
     """
     CACHE_DIR = "hfhub_cache"  # where the library's fork would be stored once downloaded
     if not os.path.exists(CACHE_DIR):
@@ -27,11 +32,15 @@ def load_model_cache():
     # download the files from huggingface repo and load the model with tensorflow
     snapshot_download(repo_id="stevekola/T5", cache_dir=CACHE_DIR)
     saved_model_path = os.path.join(CACHE_DIR, os.listdir(CACHE_DIR)[0])
-    model = tf.saved_model.load(saved_model_path, ["serve"])
-    return model
-    
+    return saved_model_path
+        
 
 def dashboard(model):
+    """"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",
@@ -41,22 +50,26 @@ def dashboard(model):
                                   "Text Summarization",
                                   "Document Similarity Score"
                                  ])
-    if task_type.startswith("Document Similarity"):
-        sentence1 = st.text("The first document/sentence.",
-                            "I reside in the commercial capital city of Nigeria, which is Lagos")
-        sentence2 = st.text("The second document/sentence.",
-                            "I live in Lagos")
+    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"):
-        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")
+    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.")
     else:
-        sentence = st.text_area("Input sentence.",
-                                "I am Steven and I live in Lagos, Nigeria")
-    st.write("**Output Text**")
-    st.write(predict(task_type, sentence, model))
+        sentence = st.text_area("Input sentence",
+                                "I am Steven and I live in Lagos, Nigeria.")
 
+    st.write("**Output Text**")
+    with st.spinner("Please wait..."): # 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
+    
 
 def predict(task_type, sentence, model):
     """Function to parse the user inputs, run the parsed text through the
@@ -64,6 +77,7 @@ def predict(task_type, sentence, model):
     params:
         task_type   sentence representing the type of task to run on T5 model
         sentence    sentence to get inference on
+        model       model to get inferences from
     returns:
         text decoded into a human-readable format.
     """
@@ -80,13 +94,14 @@ def predict(task_type, sentence, model):
     if task_type.startswith("Document Similarity"):
         sentences = sentence.split('---')
         question = f"{task_dict[task_type]} sentence1: {sentences[0]} sentence2: {sentences[1]}"
-    return predict_fn([question])[0].decode('utf-8')
+    return predict_fn([question], model)[0].decode('utf-8')
 
 
 def predict_fn(x, model):
     """Function to get inferences from model on live data points.
     params:
-        x    input text to run get output on
+        x       input text to run get output on
+        model   model to run inferences from
     returns:
         a numpy array representing the output
     """