fix sampling when training

Author: vpj

python_autocomplete/dataset/dataset.py

@@ -1,4 +1,5 @@
 from pathlib import PurePath
+from typing import Optional
 
 import torch
 from torch.utils.data import Dataset, DataLoader
@@ -48,8 +49,8 @@ def load(path: PurePath):
     def get_train_valid(path: PurePath, is_load_data: bool):
         if is_load_data:
             with monit.section("Load data"):
-                train = TextDataset.load(path / 'train.py')[:1000_000]
-                valid = TextDataset.load(path / 'valid.py')[:1000_000]
+                train = TextDataset.load(path / 'train.py')
+                valid = TextDataset.load(path / 'valid.py')
         else:
             train = ''
             valid = ''
@@ -67,6 +68,7 @@ def __repr__(self):
 
 class SourceCodeDataConfigs(BaseConfigs):
     dataset: SourceCodeDataset
+    truncate_data: int = 0
     is_load_data: bool = True
     tokenizer: Tokenizer
     retrain_tokenizer: bool = True
@@ -84,6 +86,8 @@ def text_to_i(self, text: str, *, is_silent: bool = True) -> torch.Tensor:
 @option(SourceCodeDataConfigs.dataset, 'default')
 def _dataset(c: SourceCodeDataConfigs):
     train, valid = SourceCodeDataset.get_train_valid(lab.get_data_path(), c.is_load_data)
+    if c.truncate_data:
+        train, valid = train[:c.truncate_data], valid[:c.truncate_data]
     if not c.tokenizer.is_trained:
         c.tokenizer.train(train + valid)
     return SourceCodeDataset(c.tokenizer, train, valid)
@@ -120,7 +124,7 @@ def __init__(self, *,
         self.seq_len = seq_len
         self.data = data
         self.n_samples = (self.data.shape[0] - 1) // self.seq_len
-        if drop_last:
+        if not drop_last:
             self.n_batches = (self.n_samples + batch_size - 1) // batch_size
         else:
             self.n_batches = self.n_samples // batch_size
@@ -152,17 +156,21 @@ def transpose_batch(batch):
 def _train_loader(c: SourceCodeDataConfigs):
     return DataLoader(TokenDataset(data=c.text_to_i(c.dataset.train, is_silent=False),
                                    batch_size=c.batch_size,
-                                   seq_len=c.seq_len),
+                                   seq_len=c.seq_len,
+                                   drop_last=True),
                       batch_size=c.batch_size,
                       collate_fn=transpose_batch,
-                      shuffle=c.is_shuffle)
+                      shuffle=c.is_shuffle,
+                      drop_last=True)
 
 
 @option(SourceCodeDataConfigs.valid_loader)
 def _valid_loader(c: SourceCodeDataConfigs):
     return DataLoader(TokenDataset(data=c.text_to_i(c.dataset.valid, is_silent=False),
                                    batch_size=c.batch_size,
-                                   seq_len=c.seq_len),
+                                   seq_len=c.seq_len,
+                                   drop_last=True),
                       batch_size=c.batch_size,
                       collate_fn=transpose_batch,
-                      shuffle=c.is_shuffle)
+                      shuffle=c.is_shuffle,
+                      drop_last=True)

python_autocomplete/train.py

@@ -51,7 +51,7 @@ 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))
+            tracker.add_global_step(target.shape[0] * target.shape[1])
 
         with self.mode.update(is_log_activations=batch_idx.is_last):
             state = self.state.get()
@@ -87,7 +87,9 @@ def sample(self):
             output = output.argmax(dim=-1).squeeze(1)
             prompt += '' + self.text.tokenizer.itos[output[-1]]
             if self.is_token_by_token:
-                prompt = prompt[-1:]
+                prompt = self.text.tokenizer.itos[output[-1]]
+            else:
+                prompt += '' + self.text.tokenizer.itos[output[-1]]
             log += [('' + self.text.tokenizer.itos[output[-1]], Text.value)]
             state = self.state_updater(state, new_state)
 
@@ -260,22 +262,25 @@ def main():
     experiment.create(name="source_code",
                       comment='bpe')
     experiment.configs(conf, {
-        'model': 'transformer_model',
-        # 'model': 'transformer_xl_model',
+        # 'model': 'transformer_model',
+        'model': 'transformer_xl_model',
         'n_layers': 6,
         'epochs': 32,
-        'optimizer.optimizer': 'Noam',
-        'optimizer.learning_rate': 1.0,
+        'optimizer.optimizer': 'AdamW',
+        'optimizer.learning_rate': 1.25e-4,
         'device.cuda_device': 0,
-        'is_token_by_token': False,
-        'state_updater': 'simple',
+
+        'is_token_by_token': True,
+        'state_updater': 'transformer_memory',
+        'mem_len': 256,
 
         'text.is_shuffle': False,
-        'text.tokenizer': 'char',
+        'text.tokenizer': 'bpe',
         'text.batch_size': 12,
-        'text.seq_len': 512,
-
-        'inner_iterations': 10,
+        'text.seq_len': 256,
+        #
+        # 'inner_iterations': 10,
+        # 'text.truncate_data': 100_000,
     })
     experiment.add_pytorch_models(model=conf.model)
     with experiment.start():