Upload app.py

app.py

@@ -1,6 +1,7 @@
 import time
 
 import errant
+import spacy
 import streamlit as st
 from happytransformer import HappyTextToText, TTSettings
 
@@ -14,6 +15,12 @@ checkpoints = [
 ]
 
 
+@st.cache
+def initialize_spacy(model="en_core_web_sm"):
+    spacy.cli.download(model)  # type: ignore
+    return True
+
+
 @st.cache(suppress_st_warning=True, allow_output_mutation=True)
 def get_model(model_name):
     return HappyTextToText("T5", model_name)
@@ -24,30 +31,10 @@ def get_annotator(lang: str):
     return errant.load(lang)
 
 
-st.title("🤗 Writing Assistant")
-st.markdown(
-    """This writing assistant will proofread any text for you! See my [GitHub repo](https://github.com/aseifert/hf-writing-assistant) for implementation details."""
-)
-
-checkpoint = st.selectbox("Choose model", checkpoints)
-happy_tt = get_model(checkpoint)
-annotator = get_annotator("en_core_web_sm")
-args = TTSettings(num_beams=5, min_length=1, max_length=1024)
-
-default_text = "A dog is bigger then mouse."
-default_text = "A dog is bigger then mouse."
-input_text = st.text_area(
-    label="Original text",
-    value=default_text,
-    placeholder="Enter your text here",
-)
-button = st.button("✍️ Check")
-
-
-def output(input_text):
+def output(model, args, annotator, input_text):
     with st.spinner("Checking for errors 🔍"):
         prefixed_input_text = "Grammar: " + input_text
-        result = happy_tt.generate_text(prefixed_input_text, args=args).text
+        result = model.generate_text(prefixed_input_text, args=args).text
 
         try:
             st.success(result)
@@ -58,7 +45,32 @@ def output(input_text):
             st.stop()
 
 
-start = time.time()
-output(input_text)
-st.write("---")
-st.text(f"Built by Team Writing Assistant ❤️ – prediction took {time.time() - start:.2f}s")
+def main():
+    st.title("🤗 Writing Assistant")
+    st.markdown(
+        """This writing assistant will proofread any text for you! See my [GitHub repo](https://github.com/aseifert/hf-writing-assistant) for implementation details."""
+    )
+
+    initialize_spacy()
+    annotator = get_annotator("en")
+    checkpoint = st.selectbox("Choose model", checkpoints)
+    model = get_model(checkpoint)
+    args = TTSettings(num_beams=5, min_length=1, max_length=1024)
+
+    default_text = "A dog is bigger then mouse."
+    default_text = "what be the reason for everyone leave the comapny"
+    input_text = st.text_area(
+        label="Original text",
+        value=default_text,
+        placeholder="Enter your text here",
+    )
+
+    if st.button("✍️ Check"):
+        start = time.time()
+        output(model, args, annotator, input_text)
+        st.write("---")
+        st.text(f"Built by Team Writing Assistant ❤️ – prediction took {time.time() - start:.2f}s")
+
+
+if __name__ == "__main__":
+    main()