[#1] Alpha implemented

config.yaml

@@ -1,46 +1,8 @@
 alpha:
-  eng2eng:
-    bert: bert-base-uncased
-    desc:
-    seed: 410
-    idioms_ver: c
-    idiom2def_ver: c
-    k: 11
-    lr: 0.00001
-    max_epochs: 10
-    batch_size: 64
-    shuffle: true
-  kor2eng:
-    bert: bert-base-multilingual-uncased
-    desc:
-    seed: 410
-    idioms_ver: c
-    idiom2def_ver: d
-    k: 11
-    lr: 0.00001
-    max_epochs: 20
-    batch_size: 64
-    num_workers: 4
-    shuffle: true
-gamma:
-  eng2eng:
-    bert: bert-base-uncased
-    seed: 410
-    idioms_ver: c
-    idiom2def_ver: c
-    k: 11
-    lr: 0.00001
-    max_epochs: 50
-    batch_size: 64
-    shuffle: true
-  kor2eng:
-    bert: bert-base-multilingual-uncased
-    seed: 410
-    idioms_ver: c
-    idiom2def_ver: d
-    k: 11
-    lr: 0.00001
-    max_epochs: 50
-    batch_size: 64
-    num_workers: 4
+  overfit:
+    bart: facebook/bart-base
+    lr: 0.0001
+    literal2idiomatic_ver: pie_v0
+    max_epochs: 100
+    batch_size: 100
     shuffle: true

explore/explore_bart_logits_shape.py

@@ -0,0 +1,39 @@
+from transformers import BartTokenizer, BartForConditionalGeneration
+
+from data import IdiomifyDataModule
+
+
+CONFIG = {
+    "literal2idiomatic_ver": "pie_v0",
+    "batch_size": 20,
+    "num_workers": 4,
+    "shuffle": True
+}
+
+
+def main():
+    tokenizer = BartTokenizer.from_pretrained('facebook/bart-large')
+    bart = BartForConditionalGeneration.from_pretrained('facebook/bart-large')
+    datamodule = IdiomifyDataModule(CONFIG, tokenizer)
+    datamodule.prepare_data()
+    datamodule.setup()
+    for batch in datamodule.train_dataloader():
+        srcs, tgts_r, tgts = batch
+        input_ids, attention_mask = srcs[:, 0], srcs[:, 1]  # noqa
+        decoder_input_ids, decoder_attention_mask = tgts_r[:, 0], tgts_r[:, 1]
+        outputs = bart(input_ids=input_ids,
+                       attention_mask=attention_mask,
+                       decoder_input_ids=decoder_input_ids,
+                       decoder_attention_mask=decoder_attention_mask)
+        logits = outputs[0]
+        print(logits.shape)
+        """
+        torch.Size([20, 47, 50265])
+        (N, L, |V|)
+        """
+
+        break
+
+
+if __name__ == '__main__':
+    main()

explore/explore_idiomifydatamodule.py

@@ -1,5 +1,5 @@
 from transformers import BartTokenizer
-from idiomify.datamodules import IdiomifyDataModule
+from idiomify.data import IdiomifyDataModule
 
 
 CONFIG = {

idiomify/builders.py

@@ -81,8 +81,7 @@ class TargetsBuilder(TensorBuilder):
             idiomatic + self.tokenizer.eos_token  # no bos, but ends with eos
             for _, idiomatic in literal2idiomatic
         ], return_tensors="pt", add_special_tokens=False, padding=True, truncation=True)
-        tgts = torch.stack([encodings['input_ids'],
-                            encodings['attention_mask']], dim=1)  # (N, 2, L)
-        return tgts
+        tgts = encodings['input_ids']
+        return tgts  # (N, L)
 
 

idiomify/{datamodules.py → data.py}

@@ -2,6 +2,8 @@ import torch
 from typing import Tuple, Optional, List
 from torch.utils.data import Dataset, DataLoader
 from pytorch_lightning import LightningDataModule
+from wandb.sdk.wandb_run import Run
+
 from idiomify.fetchers import fetch_literal2idiomatic
 from idiomify.builders import SourcesBuilder, TargetsBuilder, TargetsRightShiftedBuilder
 from transformers import BartTokenizer
@@ -12,9 +14,9 @@ class IdiomifyDataset(Dataset):
                  srcs: torch.Tensor,
                  tgts_r: torch.Tensor,
                  tgts: torch.Tensor):
-        self.srcs = srcs
-        self.tgts_r = tgts_r
-        self.tgts = tgts
+        self.srcs = srcs  # (N, 2, L)
+        self.tgts_r = tgts_r  # (N, 2, L)
+        self.tgts = tgts  # (N, L)
 
     def __len__(self) -> int:
         """
@@ -47,10 +49,12 @@ class IdiomifyDataModule(LightningDataModule):
 
     def __init__(self,
                  config: dict,
-                 tokenizer: BartTokenizer):
+                 tokenizer: BartTokenizer,
+                 run: Run = None):
         super().__init__()
         self.config = config
         self.tokenizer = tokenizer
+        self.run = run
         # --- to be downloaded & built --- #
         self.literal2idiomatic: Optional[List[Tuple[str, str]]] = None
         self.dataset: Optional[IdiomifyDataset] = None
@@ -59,12 +63,9 @@ class IdiomifyDataModule(LightningDataModule):
         """
         prepare: download all data needed for this from wandb to local.
         """
-        self.literal2idiomatic = fetch_literal2idiomatic(self.config['literal2idiomatic_ver'])
+        self.literal2idiomatic = fetch_literal2idiomatic(self.config['literal2idiomatic_ver'], self.run)
 
     def setup(self, stage: Optional[str] = None):
-        """
-        setup the builders.
-        """
         # --- set up the builders --- #
         # build the datasets
         srcs = SourcesBuilder(self.tokenizer)(self.literal2idiomatic)

idiomify/fetchers.py

@@ -62,7 +62,7 @@ def fetch_idioms(ver: str, run: Run = None) -> List[str]:
     # if run object is given, we track the lineage of the data.
     # if not, we get the dataset via wandb Api.
     if run:
-        artifact = run.use_artifact("idioms", type="dataset", aliases=ver)
+        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))
@@ -75,7 +75,7 @@ def fetch_literal2idiomatic(ver: str, run: Run = None) -> List[Tuple[str, str]]:
     # if run object is given, we track the lineage of the data.
     # if not, we get the dataset via wandb Api.
     if run:
-        artifact = run.use_artifact("literal2idiom", type="dataset", aliases=ver)
+        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))

idiomify/metrics.py

@@ -0,0 +1,4 @@
+"""
+you may want to include bleu score.
+and more metrics for paraphrasing.
+"""

idiomify/models.py

@@ -1,56 +1,56 @@
 """
 The reverse dictionary models below are based off of: https://github.com/yhcc/BertForRD/blob/master/mono/model/bert.py
 """
-from typing import Tuple, List, Optional
+from typing import Tuple
 import torch
 from torch.nn import functional as F
 import pytorch_lightning as pl
-from transformers import BertForMaskedLM
+from transformers import BartForConditionalGeneration
 
 
-class Idiomifier(pl.LightningModule):
+class Alpha(pl.LightningModule):  # noqa
     """
-    @eubinecto
-    The superclass of all the reverse-dictionaries. This class houses any methods that are required by
-    whatever reverse-dictionaries we define.
+    the baseline.
     """
-    # passing them to avoid warnings ---  #
-    def train_dataloader(self):
-        pass
+    def __init__(self, bart: BartForConditionalGeneration, lr: float, bos_token_id: int, pad_token_id: int):  # noqa
+        super().__init__()
+        self.bart = bart
+        self.save_hyperparameters(ignore=["bart"])
 
-    def test_dataloader(self):
-        pass
-
-    def val_dataloader(self):
-        pass
-
-    def predict_dataloader(self):
-        pass
-
-    def __init__(self, mlm: BertForMaskedLM, idiom2subwords: torch.Tensor, k: int, lr: float):  # noqa
-        """
-        :param mlm: a bert model for masked language modeling
-        :param idiom2subwords: (|W|, K)
-        :return: (N, K, |V|); (num samples, k, the size of the vocabulary of subwords)
+    def forward(self, srcs: torch.Tensor, tgts_r: torch.Tensor) -> torch.Tensor:
         """
-        pass
-
-    def forward(self, X: torch.Tensor) -> torch.Tensor:
-        """
-        given a batch, forward returns a batch of hidden vectors
-        :param X: (N, 3, L). input_ids, token_type_ids, and what was the last one...?
-        :return: (N, L, H)
+        as for using bart for CG, refer to:
+        https://huggingface.co/docs/transformers/model_doc/bart#transformers.BartForQuestionAnswering.forward
+        param srcs: (N, 2, L_s)
+        param tgts_r: (N, 2, L_t)
+        return: (N, L, |V|)
         """
-        pass
-
-    def step(self):
-        pass
-
-    def predict(self):
-        pass
-
-    def training_step(self):
-        pass
+        input_ids, attention_mask = srcs[:, 0], srcs[:, 1]
+        decoder_input_ids, decoder_attention_mask = tgts_r[:, 0], tgts_r[:, 1]
+        outputs = self.bart(input_ids=input_ids,
+                            attention_mask=attention_mask,
+                            decoder_input_ids=decoder_input_ids,
+                            decoder_attention_mask=decoder_attention_mask)
+        logits = outputs[0]  # (N, L, |V|)
+        return logits
+
+    def training_step(self, batch: Tuple[torch.Tensor, torch.Tensor, torch.Tensor]) -> dict:
+        srcs, tgts_r, tgts = batch  # (N, 2, L_s), (N, 2, L_t), (N, 2, L_t)
+        logits = self.forward(srcs, tgts_r)  # -> (N, L, |V|)
+        logits = logits.transpose(1, 2)  # (N, L, |V|) -> (N, |V|, L)
+        loss = F.cross_entropy(logits, tgts, ignore_index=self.hparams['pad_token_id'])\
+                .sum()  # (N, L, |V|), (N, L) -> (N,) -> (1,)
+        return {
+            "loss": loss
+        }
+
+    def predict(self, srcs: torch.Tensor) -> torch.Tensor:
+        pred_ids = self.bart.generate(
+            inputs=srcs[:, 0],  # (N, 2, L) -> (N, L)
+            attention_mask=srcs[:, 1],  # (N, 2, L) -> (N, L)
+            decoder_start_token_id=self.hparams['bos_token_id'],
+        )
+        return pred_ids  # (N, L)
 
     def configure_optimizers(self) -> torch.optim.Optimizer:
         """
@@ -59,21 +59,3 @@ class Idiomifier(pl.LightningModule):
         """
         # The authors used Adam, so we might as well use it as well.
         return torch.optim.AdamW(self.parameters(), lr=self.hparams['lr'])
-
-    @classmethod
-    def name(cls) -> str:
-        return cls.__name__.lower()
-
-
-class Alpha(Idiomifier):
-    """
-    @eubinecto
-    The first prototype.
-    S_wisdom = S_wisdom_literal
-    trained on: wisdom2def only.
-    """
-
-    def S_wisdom(self, H_all: torch.Tensor) -> torch.Tensor:
-        H_k = self.H_k(H_all)  # (N, L, H) -> (N, K, H)
-        S_wisdom = self.S_wisdom_literal(H_k)  # (N, K, H) -> (N, |W|)
-        return S_wisdom

main_train.py

@@ -3,20 +3,19 @@ import torch.cuda
 import wandb
 import argparse
 import pytorch_lightning as pl
-from pytorch_lightning.loggers import WandbLogger
 from termcolor import colored
-from transformers import BertForMaskedLM, BertTokenizer
-from idiomify.datamodules import IdiomifyDataModule
-from idiomify.fetchers import fetch_config, fetch_idioms
-from idiomify.models import Alpha, Gamma
+from pytorch_lightning.loggers import WandbLogger
+from transformers import BartTokenizer, BartForConditionalGeneration
+from idiomify.data import IdiomifyDataModule
+from idiomify.fetchers import fetch_config
+from idiomify.models import Alpha
 from idiomify.paths import ROOT_DIR
-from idiomify import tensors as T
 
 
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--model", type=str, default="alpha")
-    parser.add_argument("--ver", type=str, default="eng2eng")
+    parser.add_argument("--ver", type=str, default="overfit ")
     parser.add_argument("--num_workers", type=int, default=os.cpu_count())
     parser.add_argument("--log_every_n_steps", type=int, default=1)
     parser.add_argument("--fast_dev_run", action="store_true", default=False)
@@ -27,22 +26,17 @@ def main():
     if not config['upload']:
         print(colored("WARNING: YOU CHOSE NOT TO UPLOAD. NOTHING BUT LOGS WILL BE SAVED TO WANDB", color="red"))
 
-    # prepare arguments
-    mlm = BertForMaskedLM.from_pretrained(config['bert'])
-    tokenizer = BertTokenizer.from_pretrained(config['bert'])
-    idioms = fetch_idioms(config['idioms_ver'])
-    idiom2subwords = T.idiom2subwords(idioms, tokenizer, config['k'])
-    # choose the model to train
-    if config['model'] == Alpha.name():
-        rd = Alpha(mlm, idiom2subwords, config['k'], config['lr'])
-    elif config['model'] == Gamma.name():
-        rd = Gamma(mlm, idiom2subwords, config['k'], config['lr'])
+    # prepare the model
+    bart = BartForConditionalGeneration.from_pretrained(config['bart'])
+    tokenizer = BartTokenizer.from_pretrained(config['bart'])
+    if config['model'] == "alpha":
+        model = Alpha(bart, config['lr'], tokenizer.bos_token_id, tokenizer.pad_token_id)
     else:
-        raise ValueError
-    # prepare datamodule
-    datamodule = IdiomifyDataModule(config, tokenizer, idioms)
+        raise NotImplementedError
+    # prepare the datamodule
 
-    with wandb.init(entity="eubinecto", project="idiomify-demo", config=config) as run:
+    with wandb.init(entity="eubinecto", project="idiomify", config=config) as run:
+        datamodule = IdiomifyDataModule(config, tokenizer, run)
         logger = WandbLogger(log_model=False)
         trainer = pl.Trainer(max_epochs=config['max_epochs'],
                              fast_dev_run=config['fast_dev_run'],
@@ -52,10 +46,10 @@ def main():
                              enable_checkpointing=False,
                              logger=logger)
         # start training
-        trainer.fit(model=rd, datamodule=datamodule)
+        trainer.fit(model=model, datamodule=datamodule)
         # upload the model to wandb only if the training is properly done  #
         if not config['fast_dev_run'] and trainer.current_epoch == config['max_epochs'] - 1:
-            ckpt_path = ROOT_DIR / "rd.ckpt"
+            ckpt_path = ROOT_DIR / "model.ckpt"
             trainer.save_checkpoint(str(ckpt_path))
             artifact = wandb.Artifact(name=config['model'], type="model", metadata=config)
             artifact.add_file(str(ckpt_path))