[#1] IdiomifyDataModule implemented (srcs, tgts_r, tgts)

explore/{explore_fetch_literal2idiom.py → explore_fetch_literal2idiomatic.py}

@@ -1,8 +1,8 @@
-from idiomify.fetchers import fetch_literal2idiom
+from idiomify.fetchers import fetch_literal2idiomatic
 
 
 def main():
-    for src, tgt in fetch_literal2idiom("pie_v0"):
+    for src, tgt in fetch_literal2idiomatic("pie_v0"):
         print(src, "->", tgt)
 
 

explore/explore_idiomifydatamodule.py

@@ -0,0 +1,26 @@
+from transformers import BartTokenizer
+from idiomify.datamodules import IdiomifyDataModule
+
+
+CONFIG = {
+    "literal2idiomatic_ver": "pie_v0",
+    "batch_size": 20,
+    "num_workers": 4,
+    "shuffle": True
+}
+
+
+def main():
+    tokenizer = BartTokenizer.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
+        print(srcs.shape)
+        print(tgts_r.shape)
+        print(tgts.shape)
+
+
+if __name__ == '__main__':
+    main()

explore/explore_src_builder.py

@@ -1,5 +1,5 @@
 from transformers import BartTokenizer
-from idiomify.builders import SRCBuilder
+from idiomify.builders import SourcesBuilder
 
 BATCH = [
     ("I could die at any moment", "I could kick the bucket at any moment"),
@@ -9,7 +9,7 @@ BATCH = [
 
 def main():
     tokenizer = BartTokenizer.from_pretrained("facebook/bart-large")
-    builder = SRCBuilder(tokenizer)
+    builder = SourcesBuilder(tokenizer)
     src = builder(BATCH)
     print(src)
 

explore/explore_tgt_builder.py

@@ -1,5 +1,5 @@
 from transformers import BartTokenizer
-from idiomify.builders import TGTBuilder
+from idiomify.builders import TargetsBuilder
 
 BATCH = [
     ("I could die at any moment", "I could kick the bucket at any moment"),
@@ -9,7 +9,7 @@ BATCH = [
 
 def main():
     tokenizer = BartTokenizer.from_pretrained("facebook/bart-large")
-    builder = TGTBuilder(tokenizer)
+    builder = TargetsBuilder(tokenizer)
     tgt_r, tgt = builder(BATCH)
     print(tgt_r)
     print(tgt)

idiomify/builders.py

@@ -4,12 +4,12 @@ builders must accept device as one of the parameters.
 """
 import torch
 from typing import List, Tuple
-from transformers import BertTokenizer
+from transformers import BertTokenizer, BartTokenizer
 
 
 class TensorBuilder:
 
-    def __init__(self, tokenizer: BertTokenizer):
+    def __init__(self, tokenizer: BartTokenizer):
         self.tokenizer = tokenizer
 
     def __call__(self, *args, **kwargs) -> torch.Tensor:
@@ -45,7 +45,7 @@ class Idiom2SubwordsBuilder(TensorBuilder):
         return input_ids
 
 
-class SRCBuilder(TensorBuilder):
+class SourcesBuilder(TensorBuilder):
     """
     to be used for both training and inference
     """
@@ -60,24 +60,29 @@ class SRCBuilder(TensorBuilder):
         return src  # (N, 2, L)
 
 
-class TGTBuilder(TensorBuilder):
-    """
-    This is to be used only for training. As for inference, we don't need this.
-    """
-    def __call__(self, literal2idiomatic: List[Tuple[str, str]]) -> Tuple[torch.Tensor, torch.Tensor]:
-        encodings_r = self.tokenizer([
+class TargetsRightShiftedBuilder(TensorBuilder):
+
+    def __call__(self, literal2idiomatic: List[Tuple[str, str]]) -> torch.Tensor:
+        encodings = self.tokenizer([
             self.tokenizer.bos_token + idiomatic  # starts with bos, but does not end with eos (right-shifted)
             for _, idiomatic in literal2idiomatic
         ], return_tensors="pt", add_special_tokens=False, padding=True, truncation=True)
+        tgts_r = torch.stack([encodings['input_ids'],
+                              encodings['attention_mask']], dim=1)  # (N, 2, L)
+        return tgts_r
+
+
+class TargetsBuilder(TensorBuilder):
+    """
+    This is to be used only for training. As for inference, we don't need this.
+    """
+    def __call__(self, literal2idiomatic: List[Tuple[str, str]]) -> torch.Tensor:
         encodings = self.tokenizer([
             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)
-        tgt_r = torch.stack([encodings_r['input_ids'],
-                             encodings_r['attention_mask']], dim=1)  # (N, 2, L)
-        tgt = torch.stack([encodings['input_ids'],
-                           encodings['attention_mask']], dim=1)  # (N, 2, L)
-        return tgt_r, tgt
-
+        tgts = torch.stack([encodings['input_ids'],
+                            encodings['attention_mask']], dim=1)  # (N, 2, L)
+        return tgts
 
 

idiomify/datamodules.py

@@ -2,35 +2,38 @@ import torch
 from typing import Tuple, Optional, List
 from torch.utils.data import Dataset, DataLoader
 from pytorch_lightning import LightningDataModule
-from idiomify.fetchers import fetch_idiom2def, fetch_epie
-from idiomify.builders import Idiom2DefBuilder, Idiom2ContextBuilder, LabelsBuilder
-from transformers import BertTokenizer
+from idiomify.fetchers import fetch_literal2idiomatic
+from idiomify.builders import SourcesBuilder, TargetsBuilder, TargetsRightShiftedBuilder
+from transformers import BartTokenizer
 
 
 class IdiomifyDataset(Dataset):
     def __init__(self,
-                 X: torch.Tensor,
-                 y: torch.Tensor):
-        self.X = X
-        self.y = y
+                 srcs: torch.Tensor,
+                 tgts_r: torch.Tensor,
+                 tgts: torch.Tensor):
+        self.srcs = srcs
+        self.tgts_r = tgts_r
+        self.tgts = tgts
 
     def __len__(self) -> int:
         """
         Returning the size of the dataset
         :return:
         """
-        return self.y.shape[0]
+        assert self.srcs.shape[0] == self.tgts_r.shape[0] == self.tgts.shape[0]
+        return self.srcs.shape[0]
 
-    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.LongTensor]:
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.LongTensor]:
         """
         Returns features & the label
         :param idx:
         :return:
         """
-        return self.X[idx], self.y[idx]
+        return self.srcs[idx], self.tgts_r[idx], self.tgts[idx]
 
 
-class Idiom2DefDataModule(LightningDataModule):
+class IdiomifyDataModule(LightningDataModule):
 
     # boilerplate - just ignore these
     def test_dataloader(self):
@@ -44,21 +47,19 @@ class Idiom2DefDataModule(LightningDataModule):
 
     def __init__(self,
                  config: dict,
-                 tokenizer: BertTokenizer,
-                 idioms: List[str]):
+                 tokenizer: BartTokenizer):
         super().__init__()
         self.config = config
         self.tokenizer = tokenizer
-        self.idioms = idioms
         # --- to be downloaded & built --- #
-        self.idiom2def: Optional[List[Tuple[str, str]]] = None
+        self.literal2idiomatic: Optional[List[Tuple[str, str]]] = None
         self.dataset: Optional[IdiomifyDataset] = None
 
     def prepare_data(self):
         """
         prepare: download all data needed for this from wandb to local.
         """
-        self.idiom2def = fetch_idiom2def(self.config['idiom2def_ver'])
+        self.literal2idiomatic = fetch_literal2idiomatic(self.config['literal2idiomatic_ver'])
 
     def setup(self, stage: Optional[str] = None):
         """
@@ -66,50 +67,11 @@ class Idiom2DefDataModule(LightningDataModule):
         """
         # --- set up the builders --- #
         # build the datasets
-        X = Idiom2DefBuilder(self.tokenizer)(self.idiom2def, self.config['k'])
-        y = LabelsBuilder(self.tokenizer)(self.idiom2def, self.idioms)
-        self.dataset = IdiomifyDataset(X, y)
+        srcs = SourcesBuilder(self.tokenizer)(self.literal2idiomatic)
+        tgts_r = TargetsRightShiftedBuilder(self.tokenizer)(self.literal2idiomatic)
+        tgts = TargetsBuilder(self.tokenizer)(self.literal2idiomatic)
+        self.dataset = IdiomifyDataset(srcs, tgts_r, tgts)
 
     def train_dataloader(self) -> DataLoader:
         return DataLoader(self.dataset, batch_size=self.config['batch_size'],
                           shuffle=self.config['shuffle'], num_workers=self.config['num_workers'])
-
-
-class Idiom2ContextsDataModule(LightningDataModule):
-
-    # boilerplate - just ignore these
-    def test_dataloader(self):
-        pass
-
-    def val_dataloader(self):
-        pass
-
-    def predict_dataloader(self):
-        pass
-
-    def __init__(self, config: dict, tokenizer: BertTokenizer, idioms: List[str]):
-        super().__init__()
-        self.config = config
-        self.tokenizer = tokenizer
-        self.idioms = idioms
-        self.idiom2context: Optional[List[Tuple[str, str]]] = None
-        self.dataset: Optional[IdiomifyDataset] = None
-
-    def prepare_data(self):
-        """
-        prepare: download all data needed for this from wandb to local.
-        """
-        self.idiom2context = [
-            (idiom, context)
-            for idiom, _, context in fetch_epie()
-        ]
-
-    def setup(self, stage: Optional[str] = None):
-        # build the datasets
-        X = Idiom2ContextBuilder(self.tokenizer)(self.idiom2context)
-        y = LabelsBuilder(self.tokenizer)(self.idiom2context, self.idioms)
-        self.dataset = IdiomifyDataset(X, y)
-
-    def train_dataloader(self):
-        return DataLoader(self.dataset, batch_size=self.config['batch_size'],
-                          shuffle=self.config['shuffle'], num_workers=self.config['num_workers'])

idiomify/fetchers.py

@@ -5,10 +5,7 @@ import wandb
 import requests
 from typing import Tuple, List
 from wandb.sdk.wandb_run import Run
-from transformers import AutoModelForMaskedLM, AutoConfig, BertTokenizer
-from idiomify.builders import Idiom2SubwordsBuilder
-from idiomify.models import Alpha, RD
-from idiomify.paths import CONFIG_YAML, idioms_dir, alpha_dir, literal2idiom
+from idiomify.paths import CONFIG_YAML, idioms_dir, literal2idiomatic
 from idiomify.urls import (
     EPIE_IMMUTABLE_IDIOMS_URL,
     EPIE_IMMUTABLE_IDIOMS_CONTEXTS_URL,
@@ -74,40 +71,20 @@ def fetch_idioms(ver: str, run: Run = None) -> List[str]:
         return [line.strip() for line in fh]
 
 
-def fetch_literal2idiom(ver: str, run: Run = None) -> List[Tuple[str, str]]:
+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)
     else:
-        artifact = wandb.Api().artifact(f"eubinecto/idiomify/literal2idiom:{ver}", type="dataset")
-    artifact_dir = artifact.download(root=literal2idiom(ver))
+        artifact = wandb.Api().artifact(f"eubinecto/idiomify/literal2idiomatic:{ver}", type="dataset")
+    artifact_dir = artifact.download(root=literal2idiomatic(ver))
     tsv_path = path.join(artifact_dir, "all.tsv")
     with open(tsv_path, 'r') as fh:
         reader = csv.reader(fh, delimiter="\t")
         return [(row[0], row[1]) for row in reader]
 
 
-def fetch_rd(model: str, ver: str) -> RD:
-    artifact = wandb.Api().artifact(f"eubinecto/idiomify-demo/{model}:{ver}", type="model")
-    config = artifact.metadata
-    artifact_path = alpha_dir(ver)
-    artifact.download(root=str(artifact_path))
-    mlm = AutoModelForMaskedLM.from_config(AutoConfig.from_pretrained(config['bert']))
-    ckpt_path = artifact_path / "rd.ckpt"
-    idioms = fetch_idioms(config['idioms_ver'])
-    tokenizer = BertTokenizer.from_pretrained(config['bert'])
-    idiom2subwords = Idiom2SubwordsBuilder(tokenizer)(idioms, config['k'])
-    # if model == Alpha.name():
-    #     rd = Alpha.load_from_checkpoint(str(ckpt_path), mlm=mlm, idiom2subwords=idiom2subwords)
-    # elif model == Gamma.name():
-    #     rd = Gamma.load_from_checkpoint(str(ckpt_path), mlm=mlm, idiom2subwords=idiom2subwords)
-    # else:
-    #     raise ValueError
-    rd = ...
-    return rd
-
-
 def fetch_config() -> dict:
     with open(str(CONFIG_YAML), 'r', encoding="utf-8") as fh:
         return yaml.safe_load(fh)

idiomify/paths.py

@@ -9,8 +9,8 @@ def idioms_dir(ver: str) -> Path:
     return ARTIFACTS_DIR / f"idioms_{ver}"
 
 
-def literal2idiom(ver: str) -> Path:
-    return ARTIFACTS_DIR / f"literal2idiom_{ver}"
+def literal2idiomatic(ver: str) -> Path:
+    return ARTIFACTS_DIR / f"literal2idiomatic_{ver}"
 
 
 def alpha_dir(ver: str) -> Path:

main_train.py

@@ -6,7 +6,7 @@ import pytorch_lightning as pl
 from pytorch_lightning.loggers import WandbLogger
 from termcolor import colored
 from transformers import BertForMaskedLM, BertTokenizer
-from idiomify.datamodules import Idiom2DefDataModule
+from idiomify.datamodules import IdiomifyDataModule
 from idiomify.fetchers import fetch_config, fetch_idioms
 from idiomify.models import Alpha, Gamma
 from idiomify.paths import ROOT_DIR
@@ -40,7 +40,7 @@ def main():
     else:
         raise ValueError
     # prepare datamodule
-    datamodule = Idiom2DefDataModule(config, tokenizer, idioms)
+    datamodule = IdiomifyDataModule(config, tokenizer, idioms)
 
     with wandb.init(entity="eubinecto", project="idiomify-demo", config=config) as run:
         logger = WandbLogger(log_model=False)

main_upload_literal2idiom.py → main_upload_literal2idiomatic.py

@@ -31,7 +31,7 @@ def main():
         raise NotImplementedError
 
     with wandb.init(entity="eubinecto", project="idiomify", config=config) as run:
-        artifact = wandb.Artifact(name="literal2idiom", type="dataset")
+        artifact = wandb.Artifact(name="literal2idiomatic", type="dataset")
         tsv_path = ROOT_DIR / "all.tsv"
         with open(tsv_path, 'w') as fh:
             writer = csv.writer(fh, delimiter="\t")