data_preprocessing, base train script

latex_generator.py → data_generator.py

@@ -11,14 +11,15 @@ class DotDict(dict):
     __getattr__ = dict.get
     __setattr__ = dict.__setitem__
     __delattr__ = dict.__delitem__
-    
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         if len(args) > 0 and isinstance(args[0], dict):
             for key, value in self.items():
                 if isinstance(value, dict):
                     self.__setitem__(key, DotDict(value))
-    
+
+
 def _generate_equation(size_left, depth_left, latex, tokens):
     if size_left <= 0:
         return ""
@@ -27,17 +28,17 @@ def _generate_equation(size_left, depth_left, latex, tokens):
     pairs, scopes, special = latex.pairs, latex.scopes, latex.special
     weights = [3, depth_left > 0, depth_left > 0]
     group, = random.choices([tokens, pairs, scopes], weights=weights)
-    
+
     if group is tokens:
         equation += ' '.join([
             random.choice(tokens),
             _generate_equation(size_left - 1, depth_left, latex, tokens)
         ])
         return equation
-    
+
     post_scope_size = round(abs(random.gauss(0, size_left / 2)))
     size_left -= post_scope_size + 1
-    
+
     if group is pairs:
         pair = random.choice(pairs)
         equation += ' '.join([
@@ -47,18 +48,18 @@ def _generate_equation(size_left, depth_left, latex, tokens):
             _generate_equation(post_scope_size, depth_left, latex, tokens)
         ])
         return equation
-            
+
     elif group is scopes:
         scope_type, scope_group = random.choice(list(scopes.items()))
         scope_operator = random.choice(scope_group)
         equation += scope_operator
-        
+
         if scope_type == 'single':
             equation += ' '.join([
                 special.left_bracket,
                 _generate_equation(size_left, depth_left - 1, latex, tokens)
             ])
-            
+
         elif scope_type == 'double_no_delimiters':
             equation += ' '.join([
                 special.left_bracket,
@@ -66,7 +67,7 @@ def _generate_equation(size_left, depth_left, latex, tokens):
                 special.right_bracket + special.left_bracket,
                 _generate_equation(size_left // 2, depth_left - 1, latex, tokens)
             ])
-            
+
         elif scope_type == 'double_with_delimiters':
             equation += ' '.join([
                 special.caret,
@@ -77,14 +78,15 @@ def _generate_equation(size_left, depth_left, latex, tokens):
                 special.left_bracket,
                 _generate_equation(size_left // 2, depth_left - 1, latex, tokens)
             ])
-            
+
         equation += ' '.join([
             special.right_bracket,
             _generate_equation(post_scope_size, depth_left, latex, tokens)
         ])
         return equation
-    
-def generate_equation(latex: dict, size, depth=3):
+
+
+def generate_equation(latex: DotDict, size, depth=3):
     """
     Generates a random latex equation
     -------
@@ -98,6 +100,7 @@ def generate_equation(latex: dict, size, depth=3):
     equation = _generate_equation(size, depth, latex, tokens)
     return equation
 
+
 def generate_image(directory: str, latex_path: str, filename: str, max_length=20):
     """
     Generates a random tex file and corresponding image
@@ -108,29 +111,29 @@ def generate_image(directory: str, latex_path: str, filename: str, max_length=20
     :filename: -- name for the generated files
     :max_length: -- max size of equation
     """
-    #TODO ARGPARSE, path parse
+    # TODO ARGPARSE, path parse
     filepath = directory + filename
-    
+
     with open(latex_path) as file:
         latex = json.load(file)
     latex = DotDict(latex)
-    
+
     template = string.Template(latex.template)
     font, font_options = random.choice(latex.fonts)
     font_option = random.choice([''] + font_options)
     fontsize = random.choice(latex.fontsizes)
-    equation = generate_equation(latex, 20)
+    equation = generate_equation(latex, max_length)
     tex = template.substitute(font=font, font_option=font_option, fontsize=fontsize, equation=equation)
-    
+
     files_before = set(os.listdir(directory))
     with open(f"{filepath}.tex", mode='w') as file:
         file.write(tex)
-        
+
     pr1 = subprocess.run(
         f"pdflatex -output-directory={directory} {filepath}.tex".split(),
         stderr=subprocess.PIPE,
     )
-    
+
     files_after = set(os.listdir(directory))
     if pr1.returncode != 0:
         files_to_delete = files_after - files_before
@@ -138,41 +141,43 @@ def generate_image(directory: str, latex_path: str, filename: str, max_length=20
             subprocess.run(['rm'] + [directory + file for file in files_to_delete])
         print(pr1.stderr.decode(), tex)
         return
-        
+
     pr2 = subprocess.run(
         f"gs -sDEVICE=png16m -dTextAlphaBits=4 -r200 -dSAFER -dBATCH -dNOPAUSE -o {filepath}.png {filepath}.pdf".split(),
         stderr=subprocess.PIPE,
     )
-    
-    files_to_delete = files_after - files_before - set([filename + '.png', filename + '.tex'])
+
+    files_to_delete = files_after - files_before - {filename + '.png', filename + '.tex'}
     if files_to_delete:
         subprocess.run(['rm'] + [directory + file for file in files_to_delete])
-    assert(pr2.returncode == 0)
-    
+    assert (pr2.returncode == 0)
+
+
 def generate_dataset(
-    filenames,
-    directory="/external2/dkkoshman/repos/ML2TransformerApp/data/",
-    latex_path="/external2/dkkoshman/repos/ML2TransformerApp/resources/latex.json",
-    overwrite: bool=False
-):
+        filenames: iter(str),
+        directory: str = "/external2/dkkoshman/repos/ML2TransformerApp/data/",
+        latex_path: str = "/external2/dkkoshman/repos/ML2TransformerApp/resources/latex.json",
+        overwrite: bool = False
+) -> None:
+
+
     """
-    Generates a latex dataset
+    Generates a latex dataset in given directory
     -------
     params:
     :filenames: - iterable of filenames to create, without extension
     :directory: - where to create
     :latex_path: - full path to latex json
-    :ovewrite: - whether to overwrite exsisting files
+    :overwrite: - whether to overwrite existing files
     """
-    
+
     filenames = set(filenames)
     if not overwrite:
         existing = set(file.split('.')[0] for file in os.listdir(directory) if file.endswith('.png'))
         filenames -= existing
-        
+
     while filenames:
         with Pool() as pool:
             pool.starmap(generate_image, ((directory, latex_path, name) for name in filenames))
         existing = set(file.split('.')[0] for file in os.listdir(directory) if file.endswith('.png'))
         filenames -= existing
-    

data_preprocessing.py

@@ -1,96 +1,139 @@
+import einops
 import os
-import re
 import tokenizers
 import torch
 import torchvision
 import torchvision.transforms as T
+from torch.utils.data import Dataset
 import tqdm
-import PIL
-from torch.utils.data import Dataset, DataLoader
-
+import re
 
-directory = "/external2/dkkoshman/repos/ML2TransformerApp/data/"
 
 class TexImageDataset(Dataset):
     """Image to tex dataset."""
 
-    def __init__(self, root_dir, image_preprocessing=None, tex_preprocessing=None):
+    def __init__(self, root_dir, image_transform=None, tex_transform=None):
         """
         Args:
             root_dir (string): Directory with all the images and tex files.
-            transform (callable, optional): Optional transform to be applied
-                on a sample.
-                
-            image_preprocessing: callable image preprocessing
-            
-            tex_preprocessing: callable tex preprocessing
+            image_transform: callable image preprocessing
+            tex_transform: callable tex preprocessing
         """
-        
+
         torch.multiprocessing.set_sharing_strategy('file_system')
         self.root_dir = root_dir
-        filenames = sorted(
-            set(os.path.splitext(filename)[0] for filename in os.listdir(root_dir) if filename.endswith('png'))
-        )
-        self.data = []
-        
-        for filename in tqdm.tqdm(filenames):
-            tex_path = self.root_dir + filename + '.tex'
-            image_path = self.root_dir + filename + '.png'
-            
-            with open(tex_path) as file:
-                tex = file.read()
-            if tex_preprocessing:
-                tex = tex_preprocessing(tex)
-                
-            image = torchvision.io.read_image(image_path)
-            if image_preprocessing:
-                image = image_preprocessing(image)
-                
-            self.data.append((image, tex))
+        self.filenames = sorted(set(
+            os.path.splitext(filename)[0] for filename in os.listdir(root_dir) if filename.endswith('png')
+        ))
+        self.image_transform = image_transform
+        self.tex_transform = tex_transform
 
     def __len__(self):
-        return len(self.data)
+        return len(self.filenames)
 
     def __getitem__(self, idx):
-        image, tex = self.data[idx]
+        filename = self.filenames[idx]
+        tex_path = self.root_dir + filename + '.tex'
+        image_path = self.root_dir + filename + '.png'
+
+        with open(tex_path) as file:
+            tex = file.read()
+        if self.tex_transform:
+            tex = self.tex_transform(tex)
+
+        image = torchvision.io.read_image(image_path)
+        if self.image_transform:
+            image = self.image_transform(image)
+
         return {"image": image, "tex": tex}
-    
-    
-class StandardizeImage(object):
-    """Pad and crop image to a given size, invert and normalize"""
+
+    def subjoin_normalize_layer(self):
+        """Appends a normalize layer with mean and std computed after iterating over dataset"""
+        mean = 0
+        std = 0
+        for item in tqdm.tqdm(self):
+            image = item['image']
+            mean += image.mean()
+            std += image.std()
+
+        mean /= len(self)
+        std /= len(self)
+        normalize = T.Normalize(mean, std)
+
+        if self.image_transform:
+            self.image_transform = T.Compose((self.image_transform, normalize))
+        else:
+            self.image_transform = normalize
+
+    @staticmethod
+    def collate_batch(batch):
+        images = [i['image'] for i in batch]
+        images = einops.rearrange(images, 'b c h w -> b c h w')
+
+        texs = [item['tex'] for item in batch]
+        texs = tokenizer.encode_batch(texs)
+        tex_ids = torch.Tensor([encoding.ids for encoding in texs])
+        attention_masks = torch.Tensor([encoding.attention_mask for encoding in texs])
+
+        return {'images': images, 'tex_ids': tex_ids, 'tex_attention_masks': attention_masks}
+
+
+class StandardizeImageTransform(object):
+    """Pad and crop image to a given size, grayscale and invert"""
 
     def __init__(self, width=1024, height=128):
-        self.transform = T.Compose((
+        self.standardize = T.Compose((
             T.Resize(height),
             T.Grayscale(),
             T.functional.invert,
-            T.CenterCrop((height, width))
+            T.CenterCrop((height, width)),
+            T.ConvertImageDtype(torch.float32)
         ))
 
     def __call__(self, image):
-        image = self.transform(image)
+        image = self.standardize(image)
         return image
-    
-    
-class RandomTransformImage(object):
+
+
+class RandomizeImageTransform(object):
     """Standardize image and randomly augment"""
 
-    def __init__(self, standardize, random_magnitude=5):
-        self.brighten = T.ColorJitter(brightness=(1/random_magnitude, 1 + 1/random_magnitude))
-        self.standardize = standardize
-        self.rand_aug = T.RandAugment(magnitude=random_magnitude)
+    def __init__(self, width=1024, height=128, random_magnitude=5):
+        self.transform = T.Compose((
+            T.ColorJitter(brightness=random_magnitude / 10),
+            T.Resize(height),
+            T.Grayscale(),
+            T.functional.invert,
+            T.CenterCrop((height, width)),
+            torch.Tensor.contiguous,
+            T.RandAugment(magnitude=random_magnitude),
+            T.ConvertImageDtype(torch.float32)
+        ))
 
     def __call__(self, image):
-        image = self.brighten(image)
-        image = self.standardize(image)
-        image = image.contiguous()
-        image = self.rand_aug(image)
+        image = self.transform(image)
         return image
 
 
-def generate_tex_tokenizer(texs):
-    """Returns a tokeniser trained on tex strings from dataset"""
-    
+class ExtractEquationFromTexTransform(object):
+    """Extracts ...\[ equation \]... from tex file"""
+
+    def __init__(self):
+        self.equation_pattern = re.compile(r'\\\[(?P<equation>.*)\\\]', flags=re.DOTALL)
+        self.spaces = re.compile(r' +')
+
+    def __call__(self, tex):
+        equation = self.equation_pattern.search(tex)
+        equation = equation.group('equation')
+        equation = equation.strip()
+        equation = self.spaces.sub(' ', equation)
+        return equation
+
+
+def generate_tex_tokenizer(texs: iter(str)):
+    """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=300,
@@ -103,5 +146,5 @@ def generate_tex_tokenizer(texs):
         special_tokens=[("[SEP]", tokenizer.token_to_id("[SEP]"))]
     )
     tokenizer.enable_padding(pad_id=tokenizer.token_to_id("[PAD]"), pad_token="[PAD]")
-    
+
     return tokenizer

resources/latex.json

@@ -1 +1,257 @@
-{"special": {"dollar": "$", "underscore": "_", "caret": "^", "left_bracket": "{", "right_bracket": "}", "ampersand": "&"}, "chars": "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!\"'()*+,-./:;<=>?@[]`|~", "greek": ["\\alpha", "\\beta", "\\gamma", "\\delta", "\\epsilon", "\\varepsilon", "\\zeta", "\\eta", "\\theta", "\\vartheta", "\\iota", "\\kappa", "\\lambda", "\\mu", "\\nu", "\\xi", "\\pi", "\\varpi", "\\rho", "\\varrho", "\\sigma", "\\varsigma", "\\tau", "\\upsilon", "\\phi", "\\varphi", "\\chi", "\\psi", "\\omega", "\\Gamma", "\\Delta", "\\Theta", "\\Lambda", "\\Xi", "\\Pi", "\\Sigma", "\\Upsilon", "\\Phi", "\\Psi", "\\Omega"], "functions": ["\\forall", "\\exists", "\\arccos", "\\arcsin", "\\arctan", "\\cos", "\\cosh", "\\cot", "\\coth", "\\csc", "\\deg", "\\det", "\\dim", "\\exp", "\\gcd", "\\hom", "\\inf", "\\ker", "\\lg", "\\lim", "\\liminf", "\\limsup", "\\ln", "\\log", "\\max", "\\min", "\\sec", "\\sin", "\\sinh", "\\sup", "\\tan", "\\tanh"], "operators": ["--", "---", "\\pm", "\\mp", "\\times", "\\div", "\\ast", "\\star", "\\bullet", "\\circ", "\\cdot", "\\leq", "\\ll", "\\subset", "\\geq", "\\gg", "\\equiv", "\\sim", "\\simeq", "\\approx", "\\neq", "\\propto", "\\not", "\\mid", "\\leftarrow", "\\Leftarrow", "\\longleftarrow", "\\Longleftarrow", "\\rightarrow", "\\Rightarrow", "\\longrightarrow", "\\Longrightarrow", "\\leftrightarrow", "\\Leftrightarrow", "\\longleftrightarrow", "\\uparrow", "\\downarrow", "\\Uparrow", "\\cdots", "\\ddots", "\\ldots", "\\vdots"], "pairs": [["\\left(", "\\right)"], ["\\left[", "\\right]"], ["\\left\\{", "\\right\\}"], ["\\langle", "\\rangle"]], "spaces": ["\\;", "\\:", "\\,", "\\!"], "fonts": [["sfmath", []], ["lmodern", []], ["eulervm", []], ["euler", []], ["beton", []], ["drm", []], ["boisik", []], ["gfsartemisia-euler", []], ["gfsartemisia", []], ["arev", []], ["anttor", ["math", "light,math", "condensed,math", "light,condensed,math"]]], "fontsizes": [6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20], "template": "\\documentclass[preview]{standalone}\n\\usepackage[$font_option]{$font}\n\\usepackage[T1]{fontenc}\n\\begin{document}\n{\\fontsize{$fontsize pt}{12 pt}\\selectfont \n\\[\n$equation\n\\]\n}\n\\end{document}", "scopes": {"single": ["^", "_", "\\sqrt", "\\underbrace", "\\underline", "\\boldmath", "\\hat", "\\widehat", "\\check", "\\tilde", "\\widetilde", "\\acute", "\\grave", "\\dot", "\\ddot", "\\breve", "\\bar", "\\vec"], "double_with_delimiters": ["\"\\sum", "\\prod", "\\int", "\\bigcup", "\\bigcap"], "double_no_delimiters": ["\\frac", "\\stackrel"]}}
+{
+  "special": {
+    "dollar": "$",
+    "underscore": "_",
+    "caret": "^",
+    "left_bracket": "{",
+    "right_bracket": "}",
+    "ampersand": "&"
+  },
+  "chars": "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!\"'()*+,-./:;<=>?@[]`|~",
+  "greek": [
+    "\\alpha",
+    "\\beta",
+    "\\gamma",
+    "\\delta",
+    "\\epsilon",
+    "\\varepsilon",
+    "\\zeta",
+    "\\eta",
+    "\\theta",
+    "\\vartheta",
+    "\\iota",
+    "\\kappa",
+    "\\lambda",
+    "\\mu",
+    "\\nu",
+    "\\xi",
+    "\\pi",
+    "\\varpi",
+    "\\rho",
+    "\\varrho",
+    "\\sigma",
+    "\\varsigma",
+    "\\tau",
+    "\\upsilon",
+    "\\phi",
+    "\\varphi",
+    "\\chi",
+    "\\psi",
+    "\\omega",
+    "\\Gamma",
+    "\\Delta",
+    "\\Theta",
+    "\\Lambda",
+    "\\Xi",
+    "\\Pi",
+    "\\Sigma",
+    "\\Upsilon",
+    "\\Phi",
+    "\\Psi",
+    "\\Omega"
+  ],
+  "functions": [
+    "\\forall",
+    "\\exists",
+    "\\arccos",
+    "\\arcsin",
+    "\\arctan",
+    "\\cos",
+    "\\cosh",
+    "\\cot",
+    "\\coth",
+    "\\csc",
+    "\\deg",
+    "\\det",
+    "\\dim",
+    "\\exp",
+    "\\gcd",
+    "\\hom",
+    "\\inf",
+    "\\ker",
+    "\\lg",
+    "\\lim",
+    "\\liminf",
+    "\\limsup",
+    "\\ln",
+    "\\log",
+    "\\max",
+    "\\min",
+    "\\sec",
+    "\\sin",
+    "\\sinh",
+    "\\sup",
+    "\\tan",
+    "\\tanh"
+  ],
+  "operators": [
+    "--",
+    "---",
+    "\\pm",
+    "\\mp",
+    "\\times",
+    "\\div",
+    "\\ast",
+    "\\star",
+    "\\bullet",
+    "\\circ",
+    "\\cdot",
+    "\\leq",
+    "\\ll",
+    "\\subset",
+    "\\geq",
+    "\\gg",
+    "\\equiv",
+    "\\sim",
+    "\\simeq",
+    "\\approx",
+    "\\neq",
+    "\\propto",
+    "\\not",
+    "\\mid",
+    "\\leftarrow",
+    "\\Leftarrow",
+    "\\longleftarrow",
+    "\\Longleftarrow",
+    "\\rightarrow",
+    "\\Rightarrow",
+    "\\longrightarrow",
+    "\\Longrightarrow",
+    "\\leftrightarrow",
+    "\\Leftrightarrow",
+    "\\longleftrightarrow",
+    "\\uparrow",
+    "\\downarrow",
+    "\\Uparrow",
+    "\\cdots",
+    "\\ddots",
+    "\\ldots",
+    "\\vdots"
+  ],
+  "pairs": [
+    [
+      "\\left(",
+      "\\right)"
+    ],
+    [
+      "\\left[",
+      "\\right]"
+    ],
+    [
+      "\\left\\{",
+      "\\right\\}"
+    ],
+    [
+      "\\langle",
+      "\\rangle"
+    ]
+  ],
+  "spaces": [
+    "\\;",
+    "\\:",
+    "\\,",
+    "\\!"
+  ],
+  "fonts": [
+    [
+      "sfmath",
+      []
+    ],
+    [
+      "lmodern",
+      []
+    ],
+    [
+      "eulervm",
+      []
+    ],
+    [
+      "euler",
+      []
+    ],
+    [
+      "beton",
+      []
+    ],
+    [
+      "drm",
+      []
+    ],
+    [
+      "boisik",
+      []
+    ],
+    [
+      "gfsartemisia-euler",
+      []
+    ],
+    [
+      "gfsartemisia",
+      []
+    ],
+    [
+      "arev",
+      []
+    ],
+    [
+      "anttor",
+      [
+        "math",
+        "light,math",
+        "condensed,math",
+        "light,condensed,math"
+      ]
+    ]
+  ],
+  "fontsizes": [
+    6,
+    7,
+    8,
+    9,
+    10,
+    11,
+    12,
+    13,
+    14,
+    15,
+    16,
+    17,
+    18,
+    19,
+    20
+  ],
+  "template": "\\documentclass[preview]{standalone}\n\\usepackage[$font_option]{$font}\n\\usepackage[T1]{fontenc}\n\\begin{document}\n{\\fontsize{$fontsize pt}{12 pt}\\selectfont \n\\[\n$equation\n\\]\n}\n\\end{document}",
+  "scopes": {
+    "single": [
+      "^",
+      "_",
+      "\\sqrt",
+      "\\underbrace",
+      "\\underline",
+      "\\boldmath",
+      "\\hat",
+      "\\widehat",
+      "\\check",
+      "\\tilde",
+      "\\widetilde",
+      "\\acute",
+      "\\grave",
+      "\\dot",
+      "\\ddot",
+      "\\breve",
+      "\\bar",
+      "\\vec"
+    ],
+    "double_with_delimiters": [
+      "\"\\sum",
+      "\\prod",
+      "\\int",
+      "\\bigcup",
+      "\\bigcap"
+    ],
+    "double_no_delimiters": [
+      "\\frac",
+      "\\stackrel"
+    ]
+  }
+}

train.py

@@ -0,0 +1,19 @@
+from data_preprocessing import TexImageDataset, RandomizeImageTransform, ExtractEquationFromTexTransform
+
+import torch
+from torch.utils.data import DataLoader
+
+if __name__ == '__main__':
+    image_transform = RandomizeImageTransform()
+    tex_transform = ExtractEquationFromTexTransform()
+    dataset = TexImageDataset('data', image_transform=image_transform, tex_transform=tex_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, shuffle=True, num_workers=16,
+                                  collate_fn=train_dataset.collate_fn)
+    batch = next(iter(train_dataloader))
+    print(batch['texs'])