Upload app.py

app.py

@@ -25,16 +25,20 @@ def download_spacy_model(model="en"):
     return True
 
 
-@st.cache(suppress_st_warning=True, allow_output_mutation=True)
+@st.cache(allow_output_mutation=True)
 def get_model(model_name):
     return HappyTextToText("T5", model_name)
 
 
-@st.cache(suppress_st_warning=True, allow_output_mutation=True)
+@st.cache(allow_output_mutation=True)
 def get_annotator(lang: str):
     return errant.load(lang)
 
 
+def predict(model, args, text: str):
+    return model.generate_text(text, args=args).text
+
+
 def main():
     st.title("🤗 Writing Assistant")
     st.markdown(
@@ -59,7 +63,7 @@ def main():
         start = time.time()
         with st.spinner("Checking for errors 🔍"):
             prefixed_input_text = "Grammar: " + input_text
-            result = model.generate_text(prefixed_input_text, args=args).text
+            result = predict(model, args, prefixed_input_text)
 
             try:
                 show_highlights(annotator, input_text, result)