[#9] `fetch_pipeline` has been added. Fixed the bug where <pad> tokens will appear in the final output.

idiomify/fetchers.py

@@ -8,6 +8,7 @@ from idiomify.paths import CONFIG_YAML, idioms_dir, literal2idiomatic, idiomifie
 from idiomify.urls import PIE_URL
 from transformers import AutoModelForSeq2SeqLM, AutoConfig, BartTokenizer
 from idiomify.models import Idiomifier
+from idiomify.pipeline import Pipeline
 
 
 # --- from the web --- #
@@ -75,6 +76,20 @@ def fetch_tokenizer(ver: str, run: Run = None) -> BartTokenizer:
     return tokenizer
 
 
+def fetch_pipeline() -> Pipeline:
+    """
+    fetch a pipeline of the version stated in config.yaml
+    """
+    config = fetch_config()['idiomifier']
+    model = fetch_idiomifier(config['ver'])
+    tokenizer = fetch_tokenizer(config['tokenizer_ver'])
+    idioms = fetch_idioms(config['idioms_ver'])
+    model.eval()  # this is crucial to obtain consistent results
+    pipeline = Pipeline(model, tokenizer, idioms)
+    return pipeline
+
+
+# --- from local --- #
 def fetch_config() -> dict:
     with open(str(CONFIG_YAML), 'r', encoding="utf-8") as fh:
         return yaml.safe_load(fh)

idiomify/pipeline.py

@@ -1,4 +1,5 @@
 import re
+import pandas as pd
 from typing import List
 from transformers import BartTokenizer
 from idiomify.builders import SourcesBuilder
@@ -7,9 +8,10 @@ from idiomify.models import Idiomifier
 
 class Pipeline:
 
-    def __init__(self, model: Idiomifier, tokenizer: BartTokenizer):
+    def __init__(self, model: Idiomifier, tokenizer: BartTokenizer, idioms: pd.DataFrame):
         self.model = model
         self.builder = SourcesBuilder(tokenizer)
+        self.idioms = idioms
 
     def __call__(self, sents: List[str], max_length=100) -> List[str]:
         srcs = self.builder(literal2idiomatic=[(sent, "") for sent in sents])
@@ -19,9 +21,10 @@ class Pipeline:
             decoder_start_token_id=self.model.hparams['bos_token_id'],
             max_length=max_length,
         )  # -> (N, L_t)
+        # we don't skip special tokens because we have to keep <idiom> & </idiom> for highlighting idioms.
         tgts = self.builder.tokenizer.batch_decode(pred_ids, skip_special_tokens=False)
         tgts = [
-            re.sub(r"<s>|</s>", "", tgt)
+            re.sub(r"<s>|</s>|<pad>", "", tgt)
             for tgt in tgts
         ]
         return tgts

main_deploy.py

@@ -3,30 +3,24 @@ we deploy the pipeline via streamlit.
 """
 import re
 import streamlit as st
-from idiomify.fetchers import fetch_config, fetch_idiomifier, fetch_idioms, fetch_tokenizer
+from idiomify.fetchers import fetch_pipeline
 from idiomify.pipeline import Pipeline
 
 
 @st.cache(allow_output_mutation=True)
-def fetch_resources() -> tuple:
-    config = fetch_config()['idiomifier']
-    model = fetch_idiomifier(config['ver'])
-    tokenizer = fetch_tokenizer(config['tokenizer_ver'])
-    idioms = fetch_idioms(config['idioms_ver'])
-    return config, model, tokenizer, idioms
+def cache_pipeline() -> Pipeline:
+    return fetch_pipeline()
 
 
 def main():
     # fetch a pre-trained model
-    config, model, tokenizer, idioms = fetch_resources()
-    model.eval()
-    pipeline = Pipeline(model, tokenizer)
+    pipeline = cache_pipeline()
     st.title("Idiomify Demo")
     text = st.text_area("Type sentences here",
                         value="Just remember that there will always be a hope even when things look hopeless")
     with st.sidebar:
         st.subheader("Supported idioms")
-        idioms = [row["Idiom"] for _, row in idioms.iterrows()]
+        idioms = [row["Idiom"] for _, row in pipeline.idioms.iterrows()]
         st.write(" / ".join(idioms))
 
     if st.button(label="Idiomify"):

main_infer.py

@@ -2,9 +2,7 @@
 This is for just a simple sanity check on the inference.
 """
 import argparse
-from idiomify.pipeline import Pipeline
-from idiomify.fetchers import fetch_config, fetch_idiomifier, fetch_tokenizer
-from transformers import BartTokenizer
+from idiomify.fetchers import fetch_pipeline
 
 
 def main():
@@ -12,15 +10,9 @@ def main():
     parser.add_argument("--sent", type=str,
                         default="Just remember that there will always be a hope even when things look hopeless")
     args = parser.parse_args()
-    config = fetch_config()['idiomifier']
-    config.update(vars(args))
-    model = fetch_idiomifier(config['ver'])
-    tokenizer = fetch_tokenizer(config['tokenizer_ver'])
-    model.eval()  # this is crucial
-    pipeline = Pipeline(model, tokenizer)
-    src = config['sent']
-    tgts = pipeline(sents=[src])
-    print(src, "\n->", tgts[0])
+    pipeline = fetch_pipeline()
+    tgts = pipeline(sents=[args.sent])
+    print(args.sent, "\n->", tgts[0])
 
 
 if __name__ == '__main__':