demo completed, disabled check pointing

idiomify/datamodules.py

@@ -2,9 +2,9 @@ import torch
 from typing import Tuple, Optional, List
 from torch.utils.data import Dataset, DataLoader
 from pytorch_lightning import LightningDataModule
-from transformers import BertTokenizer
 from idiomify.fetchers import fetch_idiom2def
 from idiomify import tensors as T
+from transformers import BertTokenizer
 
 
 class IdiomifyDataset(Dataset):
@@ -66,8 +66,8 @@ class IdiomifyDataModule(LightningDataModule):
         """
         # --- set up the builders --- #
         # build the datasets
-        X = T.inputs([definition for _, definition in self.idiom2def], self.tokenizer, self.config['k'])
-        y = T.targets(self.idioms)
+        X = T.inputs(self.idiom2def, self.tokenizer, self.config['k'])
+        y = T.targets(self.idiom2def, self.idioms)
         self.dataset = IdiomifyDataset(X, y)
 
     def train_dataloader(self) -> DataLoader:

idiomify/fetchers.py

@@ -2,10 +2,10 @@ import csv
 import yaml
 import wandb
 from typing import Tuple, List
-from transformers import AutoModelForMaskedLM, AutoConfig, BertTokenizer
 from idiomify.models import Alpha, Gamma, RD
 from idiomify.paths import idiom2def_dir, CONFIG_YAML, idioms_dir, alpha_dir
 from idiomify import tensors as T
+from transformers import AutoModelForMaskedLM, AutoConfig, BertTokenizer
 
 
 # dataset
@@ -16,7 +16,6 @@ def fetch_idiom2def(ver: str) -> List[Tuple[str, str]]:
     tsv_path = artifact_path / "all.tsv"
     with open(tsv_path, 'r') as fh:
         reader = csv.reader(fh, delimiter="\t")
-        next(reader)
         return [
             (row[0], row[1])
             for row in reader
@@ -32,7 +31,7 @@ def fetch_idioms(ver: str) -> List[str]:
         reader = csv.reader(fh, delimiter="\t")
         next(reader)
         return [
-            (row[0])
+            row[0]
             for row in reader
         ]
 

idiomify/models.py

@@ -3,9 +3,9 @@ The reverse dictionary models below are based off of: https://github.com/yhcc/Be
 """
 from typing import Tuple, List, Optional
 import torch
-import pytorch_lightning as pl
-from transformers.models.bert.modeling_bert import BertForMaskedLM
 from torch.nn import functional as F
+import pytorch_lightning as pl
+from transformers import BertForMaskedLM
 
 
 class RD(pl.LightningModule):
@@ -135,7 +135,7 @@ class RD(pl.LightningModule):
 
     def training_epoch_end(self, outputs: List[dict]) -> None:
         # to see an average performance over the batches in this specific epoch
-        avg_loss = torch.stack([output['loss'] for output in outputs]).mean()
+        avg_loss = torch.stack([output['loss'].detach() for output in outputs]).mean()
         self.log("Train/Average Loss", avg_loss)
 
     def validation_step(self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int) -> dict:
@@ -146,7 +146,7 @@ class RD(pl.LightningModule):
 
     def validation_epoch_end(self, outputs: List[dict]) -> None:
         # to see an average performance over the batches in this specific epoch
-        avg_loss = torch.stack([output['loss'] for output in outputs]).mean()
+        avg_loss = torch.stack([output['loss'].detach() for output in outputs]).mean()
         self.log("Validation/Average Loss", avg_loss)
 
     def configure_optimizers(self) -> torch.optim.Optimizer:

idiomify/tensors.py

@@ -3,18 +3,18 @@ all the functions for building tensors are defined here.
 builders must accept device as one of the parameters.
 """
 import torch
-from typing import List
+from typing import List, Tuple
 from transformers import BertTokenizer
 
 
 def idiom2subwords(idioms: List[str], tokenizer: BertTokenizer, k: int) -> torch.Tensor:
     mask_id = tokenizer.mask_token_id
     pad_id = tokenizer.pad_token_id
-    # temporarily disable single-token status of the wisdoms
-    wisdoms = [idiom.split(" ") for idiom in idioms]
-    encodings = tokenizer(text=wisdoms,
+    # temporarily disable single-token status of the idioms
+    idioms = [idiom.split(" ") for idiom in idioms]
+    encodings = tokenizer(text=idioms,
                           add_special_tokens=False,
-                          # should set this to True, as we already have the wisdoms split.
+                          # should set this to True, as we already have the idioms split.
                           is_split_into_words=True,
                           padding='max_length',
                           max_length=k,  # set to k
@@ -24,10 +24,11 @@ def idiom2subwords(idioms: List[str], tokenizer: BertTokenizer, k: int) -> torch
     return input_ids
 
 
-def inputs(definitions: List[str], tokenizer: BertTokenizer, k: int) -> torch.Tensor:
-    lefts = [" ".join(["[MASK]"] * k)] * len(definitions)
+def inputs(idiom2def: List[Tuple[str, str]], tokenizer: BertTokenizer, k: int) -> torch.Tensor:
+    defs = [definition for _, definition in idiom2def]
+    lefts = [" ".join(["[MASK]"] * k)] * len(defs)
     encodings = tokenizer(text=lefts,
-                          text_pair=definitions,
+                          text_pair=defs,
                           return_tensors="pt",
                           add_special_tokens=True,
                           truncation=True,
@@ -47,9 +48,9 @@ def inputs(definitions: List[str], tokenizer: BertTokenizer, k: int) -> torch.Te
                         desc_mask], dim=1)
 
 
-def targets(idioms: List[str]) -> torch.Tensor:
+def targets(idiom2def: List[Tuple[str, str]], idioms: List[str]) -> torch.Tensor:
     return torch.LongTensor([
         idioms.index(idiom)
-        for idiom in idioms
+        for idiom, _ in idiom2def
     ])
 

main_infer.py

@@ -1,6 +1,6 @@
 import argparse
-from idiomify.fetchers import fetch_config, fetch_idioms, fetch_rd
 from idiomify import tensors as T
+from idiomify.fetchers import fetch_config, fetch_rd, fetch_idioms
 from transformers import BertTokenizer
 
 
@@ -11,14 +11,15 @@ def main():
         parser.add_argument("--ver", type=str,
                             default="eng2eng")
         parser.add_argument("--sent", type=str,
-                            default="avoid getting to the point")
+                            default="to avoid getting to the point")
         args = parser.parse_args()
         config = fetch_config()[args.model][args.ver]
         config.update(vars(args))
-        tokenizer = BertTokenizer.from_pretrained(config['bert'])
         idioms = fetch_idioms(config['idioms_ver'])
-        X = T.inputs([config['sent']], tokenizer, config['k'])
         rd = fetch_rd(config['model'], config['ver'])
+        rd.eval()
+        tokenizer = BertTokenizer.from_pretrained(config['bert'])
+        X = T.inputs([("", config['sent'])], tokenizer, config['k'])
         probs = rd.P_wisdom(X).squeeze().tolist()
         wisdom2prob = [
                 (wisdom, prob)
@@ -26,6 +27,7 @@ def main():
         ]
         # sort and append
         res = list(sorted(wisdom2prob, key=lambda x: x[1], reverse=True))
+        print(f"query: {config['sent']}")
         for idx, (idiom, prob) in enumerate(res):
             print(idx, idiom, prob)
 

main_train.py

@@ -49,6 +49,7 @@ def main():
                              log_every_n_steps=config['log_every_n_steps'],
                              gpus=torch.cuda.device_count(),
                              default_root_dir=str(ROOT_DIR),
+                             enable_checkpointing=False,
                              logger=logger)
         # start training
         trainer.fit(model=rd, datamodule=datamodule)