🐛 trainer

Author: vpj

.labml.yaml

@@ -1,2 +0,0 @@
-check_repo_dirty: False
-web_api:  https://api.lab-ml.com/api/v1/track?labml_token=team-samples

python_autocomplete/create_dataset.py

@@ -13,20 +13,20 @@
 PRINTABLE = set(string.printable)
 
 
-class _PythonFile(NamedTuple):
+class PythonFile(NamedTuple):
     relative_path: str
     project: str
     path: Path
 
 
-class _GetPythonFiles:
+class GetPythonFiles:
     """
     Get list of python files and their paths inside `data/source` folder
     """
 
     def __init__(self):
         self.source_path = Path(lab.get_data_path() / 'source')
-        self.files: List[_PythonFile] = []
+        self.files: List[PythonFile] = []
         self.get_python_files(self.source_path)
 
         logger.inspect([f.path for f in self.files])
@@ -36,28 +36,21 @@ def add_file(self, path: Path):
         Add a file to the list of tiles
         """
         project = path.relative_to(self.source_path).parents
-        project = project[len(project) - 2]
-        relative_path = path.relative_to(self.source_path / project)
+        relative_path = path.relative_to(self.source_path / project[len(project) - 3])
 
-        self.files.append(_PythonFile(relative_path=str(relative_path),
-                                      project=str(project),
-                                      path=path))
+        self.files.append(PythonFile(relative_path=str(relative_path),
+                                     project=str(project[len(project) - 2]),
+                                     path=path))
 
     def get_python_files(self, path: Path):
         """
         Recursively collect files
         """
         for p in path.iterdir():
-            if p.is_symlink():
-                p.unlink()
-                continue
             if p.is_dir():
                 self.get_python_files(p)
             else:
-                if p.suffix == '.py':
-                    self.add_file(p)
-                else:
-                    p.unlink()
+                self.add_file(p)
 
 
 def _read_file(path: Path) -> str:
@@ -72,15 +65,15 @@ def _read_file(path: Path) -> str:
     return content
 
 
-def _load_code(path: PurePath, source_files: List[_PythonFile]):
+def _load_code(path: PurePath, source_files: List[PythonFile]):
     with open(str(path), 'w') as f:
         for i, source in monit.enum(f"Write {path.name}", source_files):
             f.write(f"# PROJECT: {source.project} FILE: {str(source.relative_path)}\n")
             f.write(_read_file(source.path) + "\n")
 
 
 def main():
-    source_files = _GetPythonFiles().files
+    source_files = GetPythonFiles().files
 
     np.random.shuffle(source_files)
 

python_autocomplete/evaluate.py

@@ -94,7 +94,7 @@ def eval(self):
                 continue
             else:
                 if next_char == self.text[i + 1]:
-                    logs.append((self.text[i + 1], Style.underline))
+                    logs.append((self.text[i + 1], [Text.success, Style.underline]))
                 else:
                     logs.append((self.text[i + 1], Text.subtle))
 
@@ -107,17 +107,17 @@ def eval(self):
 
 def main():
     conf = Configs()
-    experiment.create(name="source_code_eval",
-                      comment='lstm model')
+    experiment.evaluate()
 
     # Replace this with your training experiment UUID
-    conf_dict = experiment.load_configs('6f10a292e77211ea89d69979079dc3d6')
-    experiment.configs(conf, conf_dict, 'run')
+    conf_dict = experiment.load_configs('8d16abcc3f6211ebb0be67ed81588441')
+    experiment.configs(conf, conf_dict)
     experiment.add_pytorch_models(get_modules(conf))
-    experiment.load('6f10a292e77211ea89d69979079dc3d6')
+    experiment.load('8d16abcc3f6211ebb0be67ed81588441')
 
     experiment.start()
-    evaluator = Evaluator(conf.model, conf.text, conf.text.valid, False)
+    from python_autocomplete.models.transformer import TransformerModel
+    evaluator = Evaluator(conf.model, conf.text, conf.text.valid, not isinstance(conf.model, TransformerModel))
     evaluator.eval()
 
 

python_autocomplete/extract_downloads.py

@@ -1,24 +1,29 @@
 import zipfile
 from pathlib import Path
 
+from labml.internal.util import rm_tree
+
 from labml import lab, monit
 
 
-def main():
+def extract_zips(overwrite: bool = False):
     download = Path(lab.get_data_path() / 'download')
     source = Path(lab.get_data_path() / 'source')
 
+    if not source.exists():
+        source.mkdir(parents=True)
+
     for repo in download.iterdir():
         with monit.section(f"Extract {repo.stem}"):
             repo_source = source / repo.stem
             if repo_source.exists():
-                continue
+                if overwrite:
+                    rm_tree(repo_source)
+                else:
+                    continue
             with zipfile.ZipFile(repo, 'r') as repo_zip:
                 repo_zip.extractall(repo_source)
 
-    if not source.exists():
-        source.mkdir(parents=True)
-
 
 if __name__ == '__main__':
-    main()
+    extract_zips()

python_autocomplete/train.py

@@ -3,16 +3,17 @@
 
 import torch
 import torch.nn as nn
-from labml import lab, experiment, monit, logger
+
+from labml import lab, experiment, monit, logger, tracker
 from labml.configs import option
 from labml.logger import Text
 from labml.utils.pytorch import get_modules
-
 from labml_helpers.datasets.text import TextDataset, SequentialDataLoader
 from labml_helpers.device import DeviceConfigs
+from labml_helpers.metrics.accuracy import Accuracy
 from labml_helpers.module import Module
-from labml_helpers.optimizer import OptimizerConfigs
-from labml_helpers.train_valid import TrainValidConfigs
+from labml_helpers.train_valid import TrainValidConfigs, hook_model_outputs, BatchIndex
+from labml_nn.optimizers.configs import OptimizerConfigs
 from labml_nn.transformers import TransformerConfigs
 
 
@@ -26,7 +27,9 @@ def __init__(self, path: PurePath, tokenizer: Callable):
 
 
 class Configs(TrainValidConfigs):
-    device = DeviceConfigs()
+    optimizer: torch.optim.Adam
+    device: torch.device = DeviceConfigs()
+
     model: Module
     text: TextDataset
     batch_size: int = 16
@@ -44,32 +47,50 @@ class Configs(TrainValidConfigs):
 
     transformer: TransformerConfigs
 
-    def run(self):
-        for _ in self.training_loop:
-            prompt = 'def train('
-            log = [(prompt, Text.subtle)]
-            for i in monit.iterate('Sample', 25):
-                data = self.text.text_to_i(prompt).unsqueeze(-1)
-                data = data.to(self.device)
-                output, *_ = self.model(data)
-                output = output.argmax(dim=-1).squeeze()
-                prompt += '' + self.text.itos[output[-1]]
-                log += [('' + self.text.itos[output[-1]], Text.value)]
+    accuracy_func = Accuracy()
+    loss_func: 'CrossEntropyLoss'
+
+    def init(self):
+        tracker.set_queue("loss.*", 20, True)
+        tracker.set_scalar("accuracy.*", True)
+        hook_model_outputs(self.mode, self.model, 'model')
+        self.state_modules = [self.accuracy_func]
+
+    def step(self, batch: any, batch_idx: BatchIndex):
+        data, target = batch[0].to(self.device), batch[1].to(self.device)
+
+        if self.mode.is_train:
+            tracker.add_global_step(len(data))
+
+        with self.mode.update(is_log_activations=batch_idx.is_last):
+            output, *_ = self.model(data)
 
-            logger.log(log)
+        loss = self.loss_func(output, target)
+        self.accuracy_func(output, target)
+        tracker.add("loss.", loss)
 
-            self.run_step()
+        if self.mode.is_train:
+            loss.backward()
 
+            self.optimizer.step()
+            if batch_idx.is_last:
+                tracker.add('model', self.model)
+            self.optimizer.zero_grad()
 
-class SimpleAccuracyFunc(Module):
-    def __call__(self, output: torch.Tensor, target: torch.Tensor) -> int:
-        pred = output.argmax(dim=-1)
-        return pred.eq(target).sum().item() / target.shape[1]
+        tracker.save()
 
+    def sample(self):
+        prompt = 'def train('
+        log = [(prompt, Text.subtle)]
+        for i in monit.iterate('Sample', 25):
+            data = self.text.text_to_i(prompt).unsqueeze(-1)
+            data = data.to(self.device)
+            output, *_ = self.model(data)
+            output = output.argmax(dim=-1).squeeze()
+            prompt += '' + self.text.itos[output[-1]]
+            log += [('' + self.text.itos[output[-1]], Text.value)]
 
-@option(Configs.accuracy_func)
-def simple_accuracy():
-    return SimpleAccuracyFunc()
+        logger.log(log)
 
 
 @option(Configs.transformer)
@@ -126,7 +147,7 @@ def lstm_model(c: Configs):
 
 @option(Configs.model)
 def rhn_model(c: Configs):
-    from python_autocomplete.models import RhnModel
+    from python_autocomplete.models.highway import RhnModel
     m = RhnModel(n_tokens=c.n_tokens,
                  embedding_size=c.d_model,
                  hidden_size=c.rnn_size,
@@ -137,7 +158,7 @@ def rhn_model(c: Configs):
 
 @option(Configs.model)
 def transformer_model(c: Configs):
-    from python_autocomplete.models import TransformerModel
+    from python_autocomplete.models.transformer import TransformerModel
     m = TransformerModel(n_tokens=c.n_tokens,
                          d_model=c.d_model,
                          encoder=c.transformer.encoder,
@@ -189,7 +210,7 @@ def main():
         'optimizer.optimizer': 'Adam',
         'optimizer.learning_rate': 2.5e-4,
         'device.cuda_device': 1
-    }, 'run')
+    })
     experiment.add_pytorch_models(get_modules(conf))
     # experiment.load('d5ba7f56d88911eaa6629b54a83956dc')
     with experiment.start():

readme.md

@@ -12,9 +12,9 @@ This repo trains deep learning models on source code.
 4. Run `python_autocomplete/extract_downloads.py` to extract the downloaded zip files to `data/source`.
  You can directly copy any python code to `data/source` to train on them.
 5. Run `python_autocomplete/remove_non_source_files.py` to all files except `.py` files.
-6. Run `create_dataset.py` to collect all python files.
+6. Run `python_autocomplete/create_dataset.py` to collect all python files.
  The collected code will be written to `data/train.py` and, `data/eval.py`.
-7. Run `train.py` to train the model.
+7. Run `python_autocomplete/train.py` to train the model.
  *Try changing hyper-parameters like model dimensions and number of layers*.
 8. Run `evaluate.py` to evaluate the model.
 9. Enjoy!