script now generates a small dataset

data_generator.py

@@ -153,14 +153,12 @@ def generate_image(directory: str, latex_path: str, filename: str, max_length=20
     assert (pr2.returncode == 0)
 
 
-def generate_dataset(
+def generate_data(
         filenames: iter(str),
-        directory: str = "/external2/dkkoshman/repos/ML2TransformerApp/data/",
-        latex_path: str = "/external2/dkkoshman/repos/ML2TransformerApp/resources/latex.json",
+        directory: str,
+        latex_path: str,
         overwrite: bool = False
 ) -> None:
-
-
     """
     Generates a latex dataset in given directory
     -------

data_preprocessing.py

@@ -67,26 +67,6 @@ class TexImageDataset(Dataset):
         else:
             self.image_transform = normalize
 
-    def subjoin_tex_tokenize_transform(self, texs, vocab_size=300):
-        """Returns a tokenizer trained on given tex strings"""
-
-        # os.environ['TOKENIZERS_PARALLELISM'] = 'false'
-        tokenizer = tokenizers.Tokenizer(tokenizers.models.BPE(unk_token="[UNK]"))
-        tokenizer_trainer = tokenizers.trainers.BpeTrainer(
-            vocab_size=vocab_size,
-            special_tokens=["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]"]
-        )
-        tokenizer.pre_tokenizer = tokenizers.pre_tokenizers.Whitespace()
-        tokenizer.train_from_iterator(texs, trainer=tokenizer_trainer)
-        tokenizer.post_processor = tokenizers.processors.TemplateProcessing(
-            single="$A [SEP]",
-            special_tokens=[("[SEP]", tokenizer.token_to_id("[SEP]"))]
-        )
-        tokenizer.enable_padding(pad_id=tokenizer.token_to_id("[PAD]"), pad_token="[PAD]")
-
-        self.tokenizer = tokenizer
-        return tokenizer
-
 
 class BatchCollator(object):
     """Image, tex batch collator"""
@@ -156,3 +136,23 @@ class ExtractEquationFromTexTransform(object):
         equation = equation.strip()
         equation = self.spaces.sub(' ', equation)
         return equation
+
+
+def generate_tex_tokenizer(texs, vocab_size=300):
+    """Returns a tokenizer trained on given tex strings"""
+
+    # os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+    tokenizer = tokenizers.Tokenizer(tokenizers.models.BPE(unk_token="[UNK]"))
+    tokenizer_trainer = tokenizers.trainers.BpeTrainer(
+        vocab_size=vocab_size,
+        special_tokens=["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]"]
+    )
+    tokenizer.pre_tokenizer = tokenizers.pre_tokenizers.Whitespace()
+    tokenizer.train_from_iterator(texs, trainer=tokenizer_trainer)
+    tokenizer.post_processor = tokenizers.processors.TemplateProcessing(
+        single="$A [SEP]",
+        special_tokens=[("[SEP]", tokenizer.token_to_id("[SEP]"))]
+    )
+    tokenizer.enable_padding(pad_id=tokenizer.token_to_id("[PAD]"), pad_token="[PAD]")
+
+    return tokenizer

train.py

@@ -1,23 +1,30 @@
+from data_generator import generate_data
 from data_preprocessing import TexImageDataset, RandomizeImageTransform, ExtractEquationFromTexTransform, \
-    generate_tex_tokenizer, BatchCollator
+    BatchCollator, generate_tex_tokenizer
 
 import torch
 from torch.utils.data import DataLoader
-import tqdm
+
+DATA_DIR = 'data'
+LATEX_PATH = 'resources/latex.json'
 
 if __name__ == '__main__':
+    generate_data(
+        filenames=map(str, range(1000)),
+        directory=DATA_DIR,
+        latex_path=LATEX_PATH,
+    )
+
     image_transform = RandomizeImageTransform()
     tex_transform = ExtractEquationFromTexTransform()
-    dataset = TexImageDataset('data', image_transform=image_transform, tex_transform=tex_transform)
+    dataset = TexImageDataset(DATA_DIR, image_transform=image_transform, tex_transform=tex_transform)
     dataset.subjoin_image_normalize_transform()
     train_dataset, test_dataset = torch.utils.data.random_split(
         dataset,
         [len(dataset) * 9 // 10, len(dataset) // 10]
     )
-    train_dataloader = DataLoader(train_dataset, batch_size=16, num_workers=16)
-    texs = list(tqdm.tqdm(batch['tex'] for batch in train_dataloader))
-    tokenizer = generate_tex_tokenizer(texs)
-    collate_fn = BatchCollator(tokenizer)
+    tex_tokenizer = generate_tex_tokenizer(dataset.texs)
+    collate_fn = BatchCollator(tex_tokenizer)
 
     train_dataloader = DataLoader(train_dataset, batch_size=16, shuffle=True, num_workers=16,
                                   collate_fn=collate_fn)