Merge pull request #10 from eubinecto/issue-9

config.yaml

@@ -1,12 +1,13 @@
 # for training an idiomifier
 idiomifier:
-  ver: m-1-2
-  desc: just overfitting the model, but on the entire PIE dataset.
+  ver: m-1-3
+  desc: Just overfitting on PIE dataset, but now with <idiom> & </idiom> special tokens.
   bart: facebook/bart-base
-  lr: 0.0001
-  literal2idiomatic_ver: d-1-2
-  idioms_ver: d-1-2
-  max_epochs: 2
+  lr: 0.00005
+  literal2idiomatic_ver: d-1-3
+  idioms_ver: d-1-3
+  tokenizer_ver: t-1-1
+  max_epochs: 8
   batch_size: 40
   shuffle: true
   seed: 104

explore/explore_bart_tokenizer_decode_idiom_special_tokens.py

@@ -0,0 +1,14 @@
+from idiomify.fetchers import fetch_tokenizer
+
+
+def main():
+    tokenizer = fetch_tokenizer("t-1-1")
+    sent = "There will always be a <idiom> silver lining </idiom> even when things look pitch black"
+    ids = tokenizer(sent)['input_ids']
+    print(ids)
+    decoded = tokenizer.decode(ids)
+    print(decoded)
+
+
+if __name__ == '__main__':
+    main()

explore/explore_fetch_tokenizer.py

@@ -12,6 +12,9 @@ def main():
     print(tokenizer.unk_token)
     print(tokenizer.additional_special_tokens)  # this should have been added
 
+    # the size of the vocab
+    print(len(tokenizer))
+
 
 """
 <s>
@@ -22,6 +25,7 @@ def main():
 <pad>
 <unk>
 ['<idiom>', '</idiom>']
+50267
 """
 
 if __name__ == '__main__':

idiomify/fetchers.py

@@ -27,7 +27,7 @@ def fetch_idioms(ver: str, run: Run = None) -> pd.DataFrame:
         artifact = run.use_artifact(f"idioms:{ver}", type="dataset")
     else:
         artifact = wandb.Api().artifact(f"eubinecto/idiomify/idioms:{ver}", type="dataset")
-    artifact_dir = artifact.download(root=idioms_dir(ver))
+    artifact_dir = artifact.download(root=str(idioms_dir(ver)))
     tsv_path = path.join(artifact_dir, "all.tsv")
     return pd.read_csv(tsv_path, sep="\t")
 
@@ -39,7 +39,7 @@ def fetch_literal2idiomatic(ver: str, run: Run = None) -> Tuple[pd.DataFrame, pd
         artifact = run.use_artifact(f"literal2idiomatic:{ver}", type="dataset")
     else:
         artifact = wandb.Api().artifact(f"eubinecto/idiomify/literal2idiomatic:{ver}", type="dataset")
-    artifact_dir = artifact.download(root=literal2idiomatic(ver))
+    artifact_dir = artifact.download(root=str(literal2idiomatic(ver)))
     train_path = path.join(artifact_dir, "train.tsv")
     test_path = path.join(artifact_dir, "test.tsv")
     train_df = pd.read_csv(train_path, sep="\t")
@@ -57,9 +57,10 @@ def fetch_idiomifier(ver: str, run: Run = None) -> Idiomifier:
     else:
         artifact = wandb.Api().artifact(f"eubinecto/idiomify/idiomifier:{ver}", type="model")
     config = artifact.metadata
-    artifact_dir = artifact.download(root=idiomifier_dir(ver))
+    artifact_dir = artifact.download(root=str(idiomifier_dir(ver)))
     ckpt_path = path.join(artifact_dir, "model.ckpt")
     bart = AutoModelForSeq2SeqLM.from_config(AutoConfig.from_pretrained(config['bart']))
+    bart.resize_token_embeddings(config['vocab_size'])
     model = Idiomifier.load_from_checkpoint(ckpt_path, bart=bart)
     return model
 
@@ -69,7 +70,7 @@ def fetch_tokenizer(ver: str, run: Run = None) -> BartTokenizer:
         artifact = run.use_artifact(f"tokenizer:{ver}", type="other")
     else:
         artifact = wandb.Api().artifact(f"eubinecto/idiomify/tokenizer:{ver}", type="other")
-    artifact_dir = artifact.download(root=tokenizer_dir(ver))
+    artifact_dir = artifact.download(root=str(tokenizer_dir(ver)))
     tokenizer = BartTokenizer.from_pretrained(artifact_dir)
     return tokenizer
 

idiomify/models.py

@@ -71,4 +71,3 @@ class Idiomifier(pl.LightningModule):  # noqa
         """
         # The authors used Adam, so we might as well use it as well.
         return torch.optim.AdamW(self.parameters(), lr=self.hparams['lr'])
-

idiomify/paths.py

@@ -19,4 +19,3 @@ def idiomifier_dir(ver: str) -> Path:
 
 def tokenizer_dir(ver: str) -> Path:
     return ARTIFACTS_DIR / f"tokenizer_{ver}"
-

idiomify/pipeline.py

@@ -1,3 +1,4 @@
+import re
 from typing import List
 from transformers import BartTokenizer
 from idiomify.builders import SourcesBuilder
@@ -18,5 +19,9 @@ class Pipeline:
             decoder_start_token_id=self.model.hparams['bos_token_id'],
             max_length=max_length,
         )  # -> (N, L_t)
-        tgts = self.builder.tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
+        tgts = self.builder.tokenizer.batch_decode(pred_ids, skip_special_tokens=False)
+        tgts = [
+            re.sub(r"<s>|</s>", "", tgt)
+            for tgt in tgts
+        ]
         return tgts

idiomify/preprocess.py

@@ -59,4 +59,3 @@ def stratified_split(df: pd.DataFrame, ratio: float, seed: int) -> Tuple[pd.Data
                                           test_size=other_size, random_state=seed,
                                           shuffle=True)
     return ratio_df, other_df
-

main_deploy.py

@@ -1,20 +1,18 @@
 """
 we deploy the pipeline via streamlit.
 """
-from typing import Tuple, List
+import re
 import streamlit as st
-from transformers import BartTokenizer
-from idiomify.fetchers import fetch_config, fetch_idiomifier, fetch_idioms
+from idiomify.fetchers import fetch_config, fetch_idiomifier, fetch_idioms, fetch_tokenizer
 from idiomify.pipeline import Pipeline
-from idiomify.models import Idiomifier
 
 
 @st.cache(allow_output_mutation=True)
-def fetch_resources() -> Tuple[dict, Idiomifier, BartTokenizer, List[str]]:
+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'])
-    tokenizer = BartTokenizer.from_pretrained(config['bart'])
     return config, model, tokenizer, idioms
 
 
@@ -24,20 +22,21 @@ def main():
     model.eval()
     pipeline = Pipeline(model, tokenizer)
     st.title("Idiomify Demo")
-    st.markdown(f"Author: `Eu-Bin KIM`")
-    st.markdown(f"Version: `{config['ver']}`")
     text = st.text_area("Type sentences here",
-                        value="Just remember there will always be a hope even when things look black")
+                        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()]
         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)
+            preds = pipeline(sents, max_length=200)
             # highlight the rule & honorifics that were applied
-            st.write(". ".join(sents))
+            preds = [re.sub(r"<idiom>|</idiom>", "`", pred)
+                     for pred in preds]
+            st.markdown(". ".join(preds))
 
 
 if __name__ == '__main__':

main_eval.py

@@ -6,7 +6,7 @@ import pytorch_lightning as pl
 from pytorch_lightning.loggers import WandbLogger
 from transformers import BartTokenizer
 from idiomify.datamodules import IdiomifyDataModule
-from idiomify.fetchers import fetch_config, fetch_idiomifier
+from idiomify.fetchers import fetch_config, fetch_idiomifier, fetch_tokenizer
 from idiomify.paths import ROOT_DIR
 
 
@@ -17,10 +17,10 @@ def main():
     args = parser.parse_args()
     config = fetch_config()['idiomifier']
     config.update(vars(args))
-    tokenizer = BartTokenizer.from_pretrained(config['bart'])
     # prepare the datamodule
     with wandb.init(entity="eubinecto", project="idiomify", config=config) as run:
         model = fetch_idiomifier(config['ver'], run)  # fetch a pre-trained model
+        tokenizer = fetch_tokenizer(config['tokenizer_ver'], run)
         datamodule = IdiomifyDataModule(config, tokenizer, run)
         logger = WandbLogger(log_model=False)
         trainer = pl.Trainer(fast_dev_run=config['fast_dev_run'],

main_infer.py

@@ -3,25 +3,24 @@ 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
+from idiomify.fetchers import fetch_config, fetch_idiomifier, fetch_tokenizer
 from transformers import BartTokenizer
 
 
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--sent", type=str,
-                        default="If there's any good to loosing my job,"
-                                " it's that I'll now be able to go to school full-time and finish my degree earlier.")
+                        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
-    tokenizer = BartTokenizer.from_pretrained(config['bart'])
     pipeline = Pipeline(model, tokenizer)
     src = config['sent']
-    tgt = pipeline(sents=[config['sent']])
-    print(src, "\n->", tgt)
+    tgts = pipeline(sents=[src])
+    print(src, "\n->", tgts[0])
 
 
 if __name__ == '__main__':

main_train.py

@@ -5,9 +5,9 @@ import argparse
 import pytorch_lightning as pl
 from termcolor import colored
 from pytorch_lightning.loggers import WandbLogger
-from transformers import BartTokenizer, BartForConditionalGeneration
+from transformers import BartForConditionalGeneration
 from idiomify.datamodules import IdiomifyDataModule
-from idiomify.fetchers import fetch_config
+from idiomify.fetchers import fetch_config, fetch_tokenizer
 from idiomify.models import Idiomifier
 from idiomify.paths import ROOT_DIR
 
@@ -23,12 +23,13 @@ def main():
     config.update(vars(args))
     if not config['upload']:
         print(colored("WARNING: YOU CHOSE NOT TO UPLOAD. NOTHING BUT LOGS WILL BE SAVED TO WANDB", color="red"))
-    # prepare the model
+    # prepare a pre-trained BART
     bart = BartForConditionalGeneration.from_pretrained(config['bart'])
-    tokenizer = BartTokenizer.from_pretrained(config['bart'])
-    model = Idiomifier(bart, config['lr'], tokenizer.bos_token_id, tokenizer.pad_token_id)
     # prepare the datamodule
     with wandb.init(entity="eubinecto", project="idiomify", config=config) as run:
+        tokenizer = fetch_tokenizer(config['tokenizer_ver'], run)
+        bart.resize_token_embeddings(len(tokenizer))  # because new tokens are added, this process is necessary
+        model = Idiomifier(bart, config['lr'], tokenizer.bos_token_id, tokenizer.pad_token_id)
         datamodule = IdiomifyDataModule(config, tokenizer, run)
         logger = WandbLogger(log_model=False)
         trainer = pl.Trainer(max_epochs=config['max_epochs'],
@@ -44,6 +45,7 @@ def main():
         if not config['fast_dev_run'] and trainer.current_epoch == config['max_epochs'] - 1:
             ckpt_path = ROOT_DIR / "model.ckpt"
             trainer.save_checkpoint(str(ckpt_path))
+            config['vocab_size'] = len(tokenizer)  # this will be needed to fetch a pretrained idiomifier later
             artifact = wandb.Artifact(name="idiomifier", type="model", metadata=config)
             artifact.add_file(str(ckpt_path))
             run.log_artifact(artifact, aliases=["latest", config['ver']])

requirements.txt

@@ -1,7 +1,7 @@
 pytorch-lightning==1.5.10
-transformers==4.16.2
-wandb==0.12.10
+transformers==4.17.0
+wandb==0.12.11
 scikit-learn==1.0.2
-pandas==1.3.5
+pandas==1.4.1
 streamlit==1.7.0
 watchdog==2.1.6