[#1] SRCBuilder, TGTBuilder implemented and tested

explore/explore_src_builder.py

@@ -0,0 +1,18 @@
+from transformers import BartTokenizer
+from idiomify.builders import SRCBuilder
+
+BATCH = [
+    ("I could die at any moment", "I could kick the bucket at any moment"),
+    ("Speak plainly", "Don't beat around the bush")
+]
+
+
+def main():
+    tokenizer = BartTokenizer.from_pretrained("facebook/bart-large")
+    builder = SRCBuilder(tokenizer)
+    src = builder(BATCH)
+    print(src)
+
+
+if __name__ == '__main__':
+    main()

explore/explore_tgt_builder.py

@@ -0,0 +1,19 @@
+from transformers import BartTokenizer
+from idiomify.builders import TGTBuilder
+
+BATCH = [
+    ("I could die at any moment", "I could kick the bucket at any moment"),
+    ("Speak plainly", "Don't beat around the bush")
+]
+
+
+def main():
+    tokenizer = BartTokenizer.from_pretrained("facebook/bart-large")
+    builder = TGTBuilder(tokenizer)
+    tgt_r, tgt = builder(BATCH)
+    print(tgt_r)
+    print(tgt)
+
+
+if __name__ == '__main__':
+    main()

idiomify/builders.py

@@ -45,40 +45,39 @@ class Idiom2SubwordsBuilder(TensorBuilder):
         return input_ids
 
 
-class Idiom2ContextBuilder(TensorBuilder):
-
-    def __call__(self, idiom2context: List[Tuple[str, str]]):
-        """
-            Given a list of tuples of idiom and context,
-            it returns a tensor of shape (batch_size, 3, max_seq_len)
-            :param idiom2context: List[Tuple[str, str]], a list of tuples of idiom and context
-            :type idiom2context: List[Tuple[str, str]]
-            :return: The input_ids, token_type_ids, and attention_mask for each context.
-        """
-        contexts = [context for _, context in idiom2context]
-        encodings = self.tokenizer(text=contexts,
+class SRCBuilder(TensorBuilder):
+    """
+    to be used for both training and inference
+    """
+    def __call__(self, literal2idiomatic: List[Tuple[str, str]]) -> torch.Tensor:
+        encodings = self.tokenizer(text=[literal for literal, _ in literal2idiomatic],
                                    return_tensors="pt",
-                                   add_special_tokens=True,
-                                   truncation=True,
                                    padding=True,
-                                   verbose=True)
-        return torch.stack([encodings['input_ids'],
-                            encodings['token_type_ids'],
-                            encodings['attention_mask']], dim=1)
+                                   truncation=True,
+                                   add_special_tokens=True)
+        src = torch.stack([encodings['input_ids'],
+                           encodings['attention_mask']], dim=1)   # (N, 2, L)
+        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([
+            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)
+        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
 
 
-class TargetsBuilder(TensorBuilder):
 
-    def __call__(self, idiom2sent: List[Tuple[str, str]], idioms: List[str]) -> torch.Tensor:
-        """
-            Given a list of idioms and a list of sentences, return a list of indices of the idioms in the sentences
-            :param idiom2sent: A list of tuples, where each tuple is an idiom and its corresponding sentence
-            :type idiom2sent: List[Tuple[str, str]]
-            :param idioms: A list of idioms
-            :type idioms: List[str]
-            :return: A tensor of indices of the idioms in the list of idioms.
-        """
-        return torch.LongTensor([
-            idioms.index(idiom)
-            for idiom, _ in idiom2sent
-        ])

idiomify/datamodules.py

@@ -3,7 +3,7 @@ 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, TargetsBuilder
+from idiomify.builders import Idiom2DefBuilder, Idiom2ContextBuilder, LabelsBuilder
 from transformers import BertTokenizer
 
 
@@ -67,7 +67,7 @@ class Idiom2DefDataModule(LightningDataModule):
         # --- set up the builders --- #
         # build the datasets
         X = Idiom2DefBuilder(self.tokenizer)(self.idiom2def, self.config['k'])
-        y = TargetsBuilder(self.tokenizer)(self.idiom2def, self.idioms)
+        y = LabelsBuilder(self.tokenizer)(self.idiom2def, self.idioms)
         self.dataset = IdiomifyDataset(X, y)
 
     def train_dataloader(self) -> DataLoader:
@@ -107,7 +107,7 @@ class Idiom2ContextsDataModule(LightningDataModule):
     def setup(self, stage: Optional[str] = None):
         # build the datasets
         X = Idiom2ContextBuilder(self.tokenizer)(self.idiom2context)
-        y = TargetsBuilder(self.tokenizer)(self.idiom2context, self.idioms)
+        y = LabelsBuilder(self.tokenizer)(self.idiom2context, self.idioms)
         self.dataset = IdiomifyDataset(X, y)
 
     def train_dataloader(self):