[#7] training & fetching m-1-3 is ready

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
+  literal2idiomatic_ver: d-1-3
+  idioms_ver: d-1-3
+  tokenizer_ver: t-1-1
+  max_epochs: 3
   batch_size: 40
   shuffle: true
   seed: 104

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

@@ -60,6 +60,7 @@ def fetch_idiomifier(ver: str, run: Run = None) -> Idiomifier:
     artifact_dir = artifact.download(root=idiomifier_dir(ver))
     ckpt_path = path.join(artifact_dir, "model.ckpt")
     bart = AutoModelForSeq2SeqLM.from_config(AutoConfig.from_pretrained(config['bart']))
+    bart.resize_embeddings(config['vocab_size'])
     model = Idiomifier.load_from_checkpoint(ckpt_path, bart=bart)
     return model
 

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']])