[#7] tokenizer:t-1-1. The tokenizer with the idiom special tokens is now fetchable directly from wandb

config.yaml

@@ -1,3 +1,4 @@
+# for training an idiomifier
 idiomifier:
   ver: m-1-2
   desc: just overfitting the model, but on the entire PIE dataset.
@@ -21,3 +22,9 @@ literal2idiomatic:
   seed: 104
   boi_token: <idiom>
   eoi_token: </idiom>
+tokenizer:
+  ver: t-1-1
+  description: A pretrained BartTokenizer. The idiom special tokens are pre-added.
+  bart: facebook/bart-base
+  boi_token: <idiom>
+  eoi_token: </idiom>

explore/explore_fetch_tokenizer.py

@@ -0,0 +1,28 @@
+from idiomify.fetchers import fetch_tokenizer
+
+
+def main():
+    tokenizer = fetch_tokenizer("t-1-1")
+    print(tokenizer.bos_token)
+    print(tokenizer.cls_token)
+    print(tokenizer.eos_token)
+    print(tokenizer.sep_token)
+    print(tokenizer.mask_token)
+    print(tokenizer.pad_token)
+    print(tokenizer.unk_token)
+    print(tokenizer.additional_special_tokens)  # this should have been added
+
+
+"""
+<s>
+<s>
+</s>
+</s>
+<mask>
+<pad>
+<unk>
+['<idiom>', '</idiom>']
+"""
+
+if __name__ == '__main__':
+    main()

idiomify/fetchers.py

@@ -2,11 +2,11 @@ import yaml
 import wandb
 from os import path
 import pandas as pd
-from typing import Tuple, List
+from typing import Tuple
 from wandb.sdk.wandb_run import Run
-from idiomify.paths import CONFIG_YAML, idioms_dir, literal2idiomatic, idiomifier_dir
+from idiomify.paths import CONFIG_YAML, idioms_dir, literal2idiomatic, idiomifier_dir, tokenizer_dir
 from idiomify.urls import PIE_URL
-from transformers import AutoModelForSeq2SeqLM, AutoConfig
+from transformers import AutoModelForSeq2SeqLM, AutoConfig, BartTokenizer
 from idiomify.models import Idiomifier
 
 
@@ -64,6 +64,16 @@ def fetch_idiomifier(ver: str, run: Run = None) -> Idiomifier:
     return model
 
 
+def fetch_tokenizer(ver: str, run: Run = None) -> BartTokenizer:
+    if run:
+        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))
+    tokenizer = BartTokenizer.from_pretrained(artifact_dir)
+    return tokenizer
+
+
 def fetch_config() -> dict:
     with open(str(CONFIG_YAML), 'r', encoding="utf-8") as fh:
         return yaml.safe_load(fh)

idiomify/paths.py

@@ -15,3 +15,8 @@ def literal2idiomatic(ver: str) -> Path:
 
 def idiomifier_dir(ver: str) -> Path:
     return ARTIFACTS_DIR / f"idiomifier_{ver}"
+
+
+def tokenizer_dir(ver: str) -> Path:
+    return ARTIFACTS_DIR / f"tokenizer_{ver}"
+

main_upload_tokenizer.py

@@ -0,0 +1,29 @@
+import wandb
+import shutil
+from transformers import BartTokenizer
+from idiomify.fetchers import fetch_config
+from idiomify.paths import ROOT_DIR
+
+
+def main():
+    config = fetch_config()['tokenizer']
+    tokenizer = BartTokenizer.from_pretrained(config['bart'])
+    tokenizer.add_special_tokens({
+        "additional_special_tokens": ["<idiom>", "</idiom>"],  # beginning and end of an idiom
+    })
+
+    with wandb.init(entity="eubinecto", project="idiomify") as run:
+        # the paths to write datasets in
+        tok_dir = ROOT_DIR / "tokenizer"
+        tokenizer.save_pretrained(tok_dir)
+        artifact = wandb.Artifact(name="tokenizer", type="other", description=config['description'],
+                                  metadata=config)
+        artifact.add_dir(tok_dir)
+        # then, we just log them here.
+        run.log_artifact(artifact, aliases=["latest", config['ver']])
+        # don't forget to remove them
+        shutil.rmtree(tok_dir)
+
+
+if __name__ == '__main__':
+    main()