bug fixes

Author: vpj

python_autocomplete/bpe.py

@@ -58,6 +58,7 @@ def encode(self, codes: List[int]):
 
         while self.heap:
             _, idx, pair = heappop(self.heap)
+            self.merge(idx, pair)
 
         return [c for c in self.codes if c != -1]
 
@@ -111,7 +112,7 @@ def set_popular_words(self, 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.pairs = {(p[0], p[1]): c for c, p in enumerate(self.bpe) if not isinstance(p, int)}
 
         self.encoder = _BPEEncoder(self.pairs)
 

python_autocomplete/evaluate.py

@@ -6,7 +6,7 @@
 import torch.nn
 from torch import nn
 
-from labml import experiment, logger, lab
+from labml import experiment, logger, lab, monit
 from labml.logger import Text, Style
 from labml.utils.pytorch import get_modules
 from labml_helpers.module import Module
@@ -55,7 +55,7 @@ def get_probabilities(self, prompt: str, state: Any) -> Tuple[np.ndarray, Any]:
 
         return prediction.detach().cpu().numpy(), state
 
-    def get_next_char(self, prompt: str, state: Any) -> Tuple[str, Any]:
+    def get_next_token(self, prompt: str, state: Any) -> Tuple[str, Any]:
         prediction, state = self.get_predictions(prompt, state)
         best = prediction.argmax(-1).squeeze().item()
         return self.itos[best], state
@@ -68,25 +68,26 @@ def get_start_state(self, prompt: str):
         if not self.is_token_by_token:
             return prompt, None
 
-        _, state = self.get_next_char(prompt[:-1], None)
+        _, state = self.get_next_token(prompt[:-1], None)
         return prompt[-1], state
 
-    def get_token(self, prompt: str, token_chars: Optional[Set[str]], state: Any) -> Tuple[str, Any]:
+    def get_next_word(self, prompt: str, token_chars: Optional[Set[str]], state: Any) -> Tuple[str, Any]:
         result = ''
         if token_chars is None:
             token_chars = set(string.ascii_letters + string.digits + ' ' + '\n' + '\r')
         while True:
-            next_char, state = self.get_next_char(prompt, state)
-            if len(result) > 2 and next_char not in token_chars or (next_char.strip() == '' and result.strip() != ''):
+            next_token, state = self.get_next_token(prompt, state)
+            if len(result) > 2 and next_token not in token_chars or (next_token.strip() == '' and result.strip() != ''):
                 if not result:
-                    result += next_char
+                    result += next_token
                 return result, state
-            result += next_char
+            result += next_token
             if len(result) > 20:
                 return result, state
-            prompt += next_char
             if self.is_token_by_token:
-                prompt = prompt[-1:]
+                prompt = next_token
+            else:
+                prompt += next_token
 
 
 def evaluate(predictor: Predictor, text: str):
@@ -99,7 +100,7 @@ def evaluate(predictor: Predictor, text: str):
     key_strokes = 0
 
     while i + 1 < len(text):
-        next_token, state = predictor.get_token(text[:i + 1], None, None)
+        next_token, state = predictor.get_next_word(text[:i + 1], None, None)
         if next_token == text[i + 1: i + 1 + len(next_token)]:
             correct += len(next_token)
             right = True
@@ -187,7 +188,7 @@ def complete(predictor: Predictor, text: str, completion: int):
         if len(text) > i + 1:
             c = text[i + 1]
         else:
-            c, _ = predictor.get_next_char(text[:i + 1], None)
+            c, _ = predictor.get_next_token(text[:i + 1], None)
 
         if c == '\n':
             logger.log(logs)
@@ -219,7 +220,8 @@ def get_predictor():
     # And for latest checkpoint
     # checkpoint = None
 
-    run_uuid = 'c45857026a2811eba16c27c69839e51f'
+    run_uuid = '41dc02106d1611eb9ab213fdf628e807' # bpe
+    # run_uuid = 'c45857026a2811eba16c27c69839e51f'  # xl
     checkpoint = None
     # run_uuid, checkpoint = experiment.load_bundle(
     #     lab.get_path() / 'saved_checkpoint.tar.gz',
@@ -242,7 +244,8 @@ def main():
 
     with open(str(lab.get_data_path() / 'sample.py'), 'r') as f:
         sample = f.read()
-    evaluate(predictor, sample)
+    with monit.section('Evaluate'):
+        evaluate(predictor, sample)
 
 
 if __name__ == '__main__':

python_autocomplete/serve.py

@@ -28,7 +28,7 @@ def autocomplete():
     with monit.section('Predict') as s:
         acquired = lock.acquire(blocking=False)
         if acquired:
-            res, state = predictor.get_token(prompt, TOKEN_CHARS, None)
+            res, state = predictor.get_next_word(prompt, TOKEN_CHARS, None)
             lock.release()
             s.message = f'{json.dumps(prompt[-5:])} -> {json.dumps(res)}'
             return jsonify({'success': True, 'prediction': res})

python_autocomplete/train.py

@@ -38,8 +38,8 @@ class BPESourceCodeDataset(TextDataset):
 
     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]
+            train = self.load(path / 'train.py')  # [:100_000]
+            valid = self.load(path / 'valid.py')  # [:100_000]
 
         super().__init__(path, bpe, train, valid, '',
                          n_tokens=bpe.n_tokens,