✨ download checkpoint

Author: vpj

.gitignore

@@ -12,4 +12,5 @@ labml_helpers
 einops
 *.egg-info
 build
-dist
+dist
+*.tar.gz

notebooks/evaluate.ipynb

@@ -28,7 +28,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "%%capture\n",
     "!pip install labml labml_python_autocomplete"
    ]
   },
@@ -69,7 +68,14 @@
     "\n",
     "[![View Run](https://img.shields.io/badge/labml-experiment-brightgreen)](https://web.lab-ml.com/run?uuid=39b03a1e454011ebbaff2b26e3148b3d)\n",
     "\n",
-    "*If you want to try this on Colab you need to run this on the same space where you run the training, because models are saved locally.*"
+    "If you have a locally trained model load it directly with:\n",
+    "\n",
+    "```python\n",
+    "run_uuid = 'RUN_UUID'\n",
+    "checkpoint = None # Get latest checkpoint\n",
+    "```\n",
+    "\n",
+    "`load_bundle` will download an archive with a saved checkpoint (pretrained model)."
    ]
   },
   {
@@ -78,7 +84,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "TRAINING_RUN_UUID = '39b03a1e454011ebbaff2b26e3148b3d'"
+    "run_uuid, checkpoint = experiment.load_bundle(\n",
+    "    lab.get_path() / 'saved_checkpoint.tar.gz',\n",
+    "    url='https://github.com/lab-ml/python_autocomplete/releases/download/0.0.4/transformer_checkpoint.tar.gz')"
    ]
   },
   {
@@ -145,7 +153,7 @@
     }
    ],
    "source": [
-    "custom_conf = experiment.load_configs(TRAINING_RUN_UUID)\n",
+    "custom_conf = experiment.load_configs(run_uuid)\n",
     "custom_conf"
    ]
   },
@@ -234,7 +242,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "experiment.load(TRAINING_RUN_UUID)"
+    "experiment.load(run_uuid, checkpoint)"
    ]
   },
   {
@@ -817,4 +825,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

python_autocomplete/bundle.py

@@ -0,0 +1,5 @@
+from labml import experiment, lab
+
+if __name__ == '__main__':
+    experiment.save_bundle(lab.get_path() / 'bundle.tar.gz', '39b03a1e454011ebbaff2b26e3148b3d',
+                           data_files=['cache/itos.json', 'cache/n_tokens.json', 'cache/stoi.json'])

python_autocomplete/evaluate.py

@@ -189,21 +189,33 @@ def complete(predictor: Predictor, text: str, completion: int):
     logger.log(logs)
 
 
-def main():
+def get_predictor():
     conf = Configs()
     experiment.evaluate()
 
-    # Replace this with your training experiment UUID
-    run_uuid = '39b03a1e454011ebbaff2b26e3148b3d'
+    # This will download a pretrained model checkpoint and some cached files.
+    # It will download the archive as `saved_checkpoint.tar.gz` and extract it.
+    #
+    # If you have a locally trained model load it directly with
+    # run_uuid = 'RUN_UUID'
+    # And for latest checkpoint
+    # checkpoint = None
+    run_uuid, checkpoint = experiment.load_bundle(
+        lab.get_path() / 'saved_checkpoint.tar.gz',
+        url='https://github.com/lab-ml/python_autocomplete/releases/download/0.0.4/transformer_checkpoint.tar.gz')
 
     conf_dict = experiment.load_configs(run_uuid)
     experiment.configs(conf, conf_dict)
     experiment.add_pytorch_models(get_modules(conf))
-    experiment.load(run_uuid)
+    experiment.load(run_uuid, checkpoint)
 
     experiment.start()
-    predictor = Predictor(conf.model, cache('stoi', lambda: conf.text.stoi), cache('itos', lambda: conf.text.itos))
     conf.model.eval()
+    return Predictor(conf.model, cache('stoi', lambda: conf.text.stoi), cache('itos', lambda: conf.text.itos))
+
+
+def main():
+    predictor = get_predictor()
 
     with open(str(lab.get_data_path() / 'sample.py'), 'r') as f:
         sample = f.read()

python_autocomplete/serve.py

@@ -4,32 +4,11 @@
 
 from flask import Flask, request, jsonify
 
-from labml import experiment, monit
-from labml.utils.cache import cache
-from labml.utils.pytorch import get_modules
-from python_autocomplete.evaluate import Predictor
-from python_autocomplete.train import Configs
+from labml import monit
+from python_autocomplete.evaluate import get_predictor
 
 TOKEN_CHARS = set(string.ascii_letters + string.digits + ' ' + '\n' + '\r' + '_')
 
-
-def get_predictor():
-    conf = Configs()
-    experiment.evaluate()
-
-    # Replace this with your training experiment UUID
-    run_uuid = '39b03a1e454011ebbaff2b26e3148b3d'
-
-    conf_dict = experiment.load_configs(run_uuid)
-    experiment.configs(conf, conf_dict)
-    experiment.add_pytorch_models(get_modules(conf))
-    experiment.load(run_uuid)
-
-    experiment.start()
-    conf.model.eval()
-    return Predictor(conf.model, cache('stoi', lambda: conf.text.stoi), cache('itos', lambda: conf.text.itos))
-
-
 app = Flask('python_autocomplete')
 predictor = get_predictor()
 lock = threading.Lock()

setup.py

@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name='labml_python_autocomplete',
-    version='0.0.3',
+    version='0.0.4',
     author="Varuna Jayasiri",
     author_email="vpjayasiri@gmail.com",
     description="A simple model that learns to predict Python source code",
@@ -19,7 +19,7 @@
                                                'labml_nn', 'labml_nn.*',
                                                'labml', 'labml.*',
                                                'test', 'test.*')),
-    install_requires=['labml>=0.4.86',
+    install_requires=['labml>=0.4.98',
                       'labml_helpers>=0.4.70',
                       'labml_nn>=0.4.70'
                       'torch',