bpe training

Author: vpj

python_autocomplete/bpe.py

@@ -3,28 +3,133 @@
 from typing import List, Tuple
 
 from labml import lab, monit
+from labml.utils.cache import cache_set
 
 ID_CHARS = set(string.ascii_letters + string.digits + '_')
 
 
 class BPE:
+    def __init__(self, bpe_en_de: 'BPEEnDe', tokenizer):
+        self.bpe = bpe_en_de
+        self.tokenizer = tokenizer
+
+    @property
+    def n_tokens(self):
+        return len(self.bpe.bpe)
+
+    @property
+    def itos(self):
+        return self.bpe.bpe_itos
+
+    @property
+    def stoi(self):
+        return self.bpe.bpe_stoi
+
+    def encode(self, data: str):
+        words = self.tokenizer.tokenize(data)
+
+        res = []
+        for w in monit.iterate('Encode words', words):
+            res += self.bpe.encode(w)
+
+        return res
+
+    def __call__(self, data: str):
+        encoded = self.encode(data)
+        return [self.itos[c] for c in encoded]
+
+
+class _BPEEncoder:
+    def __init__(self, pairs):
+        self.pairs = pairs
+        self.codes = []
+        self.next_idx = []
+        self.prev_idx = []
+        self.heap = []
+
+    def encode(self, codes: List[int]):
+        self.codes = codes
+        self.next_idx = BPELearner.default_next_pointers(len(codes))
+        self.prev_idx = BPELearner.default_prev_pointers(len(codes))
+        self.heap = []
+
+        for i in range(len(self.codes) - 1):
+            self.add_pair((self.codes[i], self.codes[i + 1]), i)
+
+        while self.heap:
+            _, idx, pair = heappop(self.heap)
+
+        return [c for c in self.codes if c != -1]
+
+    def merge(self, p2, pair):
+        p3 = self.next_idx[p2]
+
+        if p3 == -1 or pair[0] != self.codes[p2] or pair[1] != self.codes[p3]:
+            return
+
+        self.codes[p2] = self.pairs[pair]
+        self.codes[p3] = -1
+        p1 = self.prev_idx[p2]
+        p4 = self.next_idx[p3]
+
+        if p1 != -1:
+            self.add_pair((self.codes[p1], self.codes[p2]), p1)
+        self.next_idx[p2] = p4
+        if p4 != -1:
+            self.prev_idx[p4] = p2
+            self.add_pair((self.codes[p2], self.codes[p4]), p2)
+
+    def add_pair(self, pair, idx):
+        if pair not in self.pairs:
+            return
+
+        heappush(self.heap, (self.pairs[pair], idx, pair))
+
+
+class BPEEnDe:
     def __init__(self):
         self.char_itos = []
         self.char_stoi = {}
-        self.bpe_itos = []
         self.bpe = []
-        self.common = {}
+        self.popular_words = {}
+
+        self.bpe_itos = []
+        self.bpe_stoi = {}
+        self.pairs = {}
+        self.encoder = None
 
+    def load(self, char_itos, char_stoi, bpe):
+        self.char_itos = char_itos
+        self.char_stoi = char_stoi
+        self.bpe = bpe
+
+        self.calc()
+
+    def set_popular_words(self, popular_words):
+        self.popular_words = popular_words
+
+    def calc(self):
         self.bpe_itos = self.calc_bpe_itos()
+        self.bpe_stoi = {s: i for i, s in enumerate(self.bpe_itos)}
+        self.pairs = {(p[0], p[1]): c for c, p in enumerate(self.bpe) if isinstance(p, tuple)}
+
+        self.encoder = _BPEEncoder(self.pairs)
 
     def to_char_stoi(self, w: str):
         return [self.char_stoi[c] for c in w]
 
     def calc_bpe_itos(self):
         itos = list(self.char_itos)
-        itos += [itos[p1] + itos[p2] for p1, p2 in self.bpe[len(self.char_itos):]]
+        for p1, p2 in self.bpe[len(self.char_itos):]:
+            itos.append(itos[p1] + itos[p2])
         return itos
 
+    def encode(self, word: str):
+        if word in self.popular_words:
+            return self.popular_words[word]
+
+        return self.encoder.encode([self.char_stoi[c] for c in word])
+
 
 class Tokenizer:
     def collect_words(self, data: str):
@@ -284,7 +389,7 @@ def main():
     path = lab.get_data_path() / 'train.py'
 
     with open(str(path), 'r') as f:
-        data = f.read()[:100_000]
+        data = f.read()
 
     tokenizer = SourceCodeTokenizer()
     tokenizer.collect_words(data)
@@ -295,6 +400,15 @@ def main():
     print(bpe.bpe_itos()[len(bpe.char_itos):])
     print(len(data), bpe.get_length())
 
+    cache_set('bpe', {
+        'char_itos': bpe.char_itos,
+        'char_stoi': bpe.char_stoi,
+        'bpe': bpe.bpe
+    })
+
+    bpe_en_de = BPEEnDe()
+    bpe_en_de.load(bpe.char_itos, bpe.char_stoi, bpe.bpe)
+
 
 if __name__ == '__main__':
     main()

python_autocomplete/train.py

@@ -16,6 +16,7 @@
 from labml_helpers.train_valid import TrainValidConfigs, hook_model_outputs, BatchIndex
 from labml_nn.optimizers.configs import OptimizerConfigs
 from labml_nn.transformers import TransformerConfigs
+from python_autocomplete.bpe import BPE, SourceCodeTokenizer
 
 
 class SourceCodeDataset(TextDataset):
@@ -32,6 +33,23 @@ def __init__(self, path: PurePath, tokenizer: Callable):
                          stoi=cache_get('stoi'))
 
 
+class BPESourceCodeDataset(TextDataset):
+    tokenizer: BPE
+
+    def __init__(self, path: PurePath, bpe: BPE):
+        with monit.section("Load data"):
+            train = self.load(path / 'train.py')  # [:1000_000]
+            valid = self.load(path / 'valid.py')  # [:1000_000]
+
+        super().__init__(path, bpe, train, valid, '',
+                         n_tokens=bpe.n_tokens,
+                         itos=bpe.itos,
+                         stoi=bpe.stoi)
+
+    def text_to_i(self, text: str) -> torch.Tensor:
+        return torch.tensor(self.tokenizer.encode(text))
+
+
 class Configs(TrainValidConfigs):
     optimizer: torch.optim.Adam
     device: torch.device = DeviceConfigs()
@@ -268,6 +286,22 @@ def source_code(c: Configs):
     return SourceCodeDataset(lab.get_data_path(), c.tokenizer)
 
 
+@option(Configs.text)
+def source_code_bpe(c: Configs):
+    from labml.utils.cache import cache_get
+    from python_autocomplete.bpe import BPEEnDe
+    bpe_cache = cache_get('bpe')
+
+    if bpe_cache:
+        bpe_en_de = BPEEnDe()
+        bpe_en_de.load(**bpe_cache)
+    else:
+        raise RuntimeError('BPE not cached')
+
+    tokenizer = BPE(bpe_en_de, SourceCodeTokenizer())
+    return BPESourceCodeDataset(lab.get_data_path(), tokenizer)
+
+
 @option(Configs.train_loader)
 def sequential_train_loader(c: Configs):
     return SequentialDataLoader(text=c.text.train,
@@ -316,26 +350,27 @@ def main():
     conf = Configs()
     # Assign one of transformer_mode, lstm_model, or rhn_model
     experiment.create(name="source_code",
-                      comment='transformer xl model')
+                      comment='bpe')
     experiment.configs(conf, {
-        # 'model': 'transformer_model',
-        'model': 'transformer_xl_model',
+        # 'text': 'source_code',
+        'text': 'source_code_bpe',
+        'model': 'transformer_model',
+        # 'model': 'transformer_xl_model',
         'n_layers': 6,
         'batch_size': 12,
         'epochs': 32,
         'optimizer.optimizer': 'Noam',
         'optimizer.learning_rate': 1.0,
         'device.cuda_device': 0,
         'seq_len': 512,
-        'is_token_by_token': True,
-        # 'train_loader': 'shuffled_train_loader',
-        # 'valid_loader': 'shuffled_valid_loader',
-        'train_loader': 'sequential_train_loader',
-        'valid_loader': 'sequential_valid_loader',
+        'is_token_by_token': False,
+        'state_updater': 'simple',
+        'train_loader': 'shuffled_train_loader',
+        'valid_loader': 'shuffled_valid_loader',
+        # 'train_loader': 'sequential_train_loader',
+        # 'valid_loader': 'sequential_valid_loader',
     })
     experiment.add_pytorch_models(model=conf.model)
-    # experiment.load('70df7f86450911eb887b25e3927208f3')
-    experiment.load('c45857026a2811eba16c27c69839e51f')
     with experiment.start():
         conf.run()