Add model caching and bit code clean up

app.py

@@ -2,35 +2,45 @@ import streamlit as st
 from happytransformer import HappyTextToText, TTSettings
 
 checkpoint = "team-writing-assistant/t5-base-c4jfleg"
-happy_tt = HappyTextToText("T5", checkpoint)
+
+
+@st.cache(suppress_st_warning=True, allow_output_mutation=True)
+def get_happy_text(model_name):
+    st.write(f"Loading the HappyTextToText model {model_name}, please wait...")
+    return HappyTextToText("T5", model_name)
+
+
+happy_tt = get_happy_text(checkpoint)
 args = TTSettings(num_beams=5, min_length=1)
 
 st.title("Check & Improve English Grammar")
 st.markdown("This writing assistant detects 🔍 and corrects ✍️ grammatical mistakes for you!")
 
 st.subheader("Example text: ")
-col1, col2 = st.columns([1,1])
+col1, col2 = st.columns([1, 1])
 with col1:
-  example_1 = st.button("Speed of light is fastest then speed of sound")
+    example_1 = st.button("Speed of light is fastest then speed of sound")
 with col2:
-  example_2 = st.button("Who are the president?")
+    example_2 = st.button("Who are the president?")
 
 input_text = st.text_area('Enter your text here')
 button = st.button('Submit')
 
+
 def output(input_text):
-  with st.spinner('Detecting 🔍..'):
-    input_text = "grammar: " + input_text
-    result = happy_tt.generate_text(input_text, args=args)
-    st.markdown("## " + result.text)
+    with st.spinner('Detecting 🔍..'):
+        input_text = "grammar: " + input_text
+        result = happy_tt.generate_text(input_text, args=args)
+        st.markdown("## " + result.text)
+
 
 if example_1:
-  output("Speed of light is fastest then speed of sound")
+    output("Speed of light is fastest then speed of sound")
 if example_2:
-  output("Who are the president?")  
+    output("Who are the president?")
 if input_text:
-  output(input_text)
+    output(input_text)
   
 st.subheader("")
 
-st.text("Built by Team Writing Assistant ❤️")
+st.text("Built by Team Writing Assistant ❤️")