[#2] deploying the model with `main_deploy.py`.

config.yaml

@@ -4,6 +4,7 @@ idiomifier:
   bart: facebook/bart-base
   lr: 0.0001
   literal2idiomatic_ver: d-1-2
+  idioms_ver: d-1-2
   max_epochs: 2
   batch_size: 40
   shuffle: true

idiomify/pipeline.py

@@ -1,7 +1,7 @@
 from typing import List
 from transformers import BartTokenizer
-from builders import SourcesBuilder
-from models import Idiomifier
+from idiomify.builders import SourcesBuilder
+from idiomify.models import Idiomifier
 
 
 class Pipeline:

main_deploy.py

@@ -0,0 +1,41 @@
+"""
+we deploy the pipeline via streamlit.
+"""
+from typing import Tuple, List
+import streamlit as st
+from transformers import BartTokenizer
+from idiomify.fetchers import fetch_config, fetch_idiomifier, fetch_idioms
+from idiomify.pipeline import Pipeline
+from idiomify.models import Idiomifier
+
+
+@st.cache(allow_output_mutation=True)
+def fetch() -> Tuple[Idiomifier, BartTokenizer, List[str]]:
+    config = fetch_config()['idiomifier']
+    model = fetch_idiomifier(config['ver'])
+    idioms = fetch_idioms(config['idioms_ver'])
+    tokenizer = BartTokenizer.from_pretrained(config['bart'])
+    return model, tokenizer, idioms
+
+
+def main():
+    # fetch a pre-trained model
+    model, tokenizer, idioms = fetch()
+    pipeline = Pipeline(model, tokenizer)
+    st.title("Idiomify Demo")
+    text = st.text_area("Type sentences here",
+                        value="Just remember there will always be a hope even when things look black")
+    with st.sidebar:
+        st.subheader("Supported idioms")
+        st.write(" / ".join(idioms))
+
+    if st.button(label="Idiomify"):
+        with st.spinner("Please wait..."):
+            sents = [sent for sent in text.split(".") if sent]
+            sents = pipeline(sents, max_length=200)
+            # highlight the rule & honorifics that were applied
+            st.write(". ".join(sents))
+
+
+if __name__ == '__main__':
+    main()

requirements.txt

@@ -2,4 +2,6 @@ pytorch-lightning==1.5.10
 transformers==4.16.2
 wandb==0.12.10
 scikit-learn==1.0.2
-pandas==1.4.1
+pandas==1.4.1
+streamlit==1.7.0
+watchdog==2.1.6