guide.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# Evaluating and improving transformers pre-trained on ASTs for Code Completion: Code Guide\n",
    "This guide can be regarded as supplementary material to my thesis: `Evaluating and Improving Generative Pre-Trained Transformers for Code Completion`. The thesis introduces eight research questions, each of investigates another aspect of the AST based GPT-2 model.\n",
    "\n",
    "## Prerequisites\n",
    "- Python 3.8.5\n",
    "- PyTorch 1.8.1\n",
    "\n",
    "## Preparation\n",
    "- Clone the project from github.com/derochs/code-prediction-transformer\n",
    "- Download the py150k dataset and merge `python50k_eval` and `python100k_train` JSON files into a single file inside of the cloned repository, for example in `/data/python150k.json`. The resulting dataset should contain `150,000` lines, each line resembling a single AST."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RQ1: How does the baseline model perform when reproduced?\n",
    "This research question intends to establish a baseline model required for comparison with models from other research questions. As baseline model the `trav_trans` model from github.com/facebookresearch/code-prediction-transformer is used. The following scripts for this research question originate from this repository. In order to prepare the data, train and evaluate the model, the following instructions have to be followed:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generating new trees"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python generate_new_trees.py -i data/python150k.json -o tmp/new_trees.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As `INPUT` you can use the `data/python150k.json` dataset, `OUTPUT` is the JSON file which contains the processed ASTs, by default this path is set to `/tmp/new_trees.json`. You may have to specify the output path yourself because the default path is an absolute path, so probably `tmp/new_trees.json` is the better way to go, without the leading forward-slash.\n",
    "\n",
    "**Explanation**: This script will iterate through each line of the INPUT dataset. Remember, every line resembles an AST. For each Abstrac Syntax Tree, the script will iterate through all tree nodes. Currently a node can hold a `TYPE`, a `VALUE` or both. The `trav_trans` model however expects only **either** a `TYPE` node **or** a `VALUE` node. Threfore, if the script finds a node containing both, a type and a value, it will remove the value from the node and re-add it as child node. A simple visualization can be found [in the original project\\`s README](https://github.com/facebookresearch/code-prediction-transformer#modifying-the-ast)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generating data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python models\\trav_trans\\generate_data.py -a tmp/new_trees.json -o tmp/dps.txt -c 1000"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As `INPUT` the pre-processed ASTs will be used, in our example it's `tmp/new_trees.json`. `OUTPUT` is the output path of the sequenced ASTs, by default `/tmp/dps.txt`, again, probably you don't want to use the default one but rather `tmp/dps.txt`. The third argument `CONTEXTS` is the context size of your model: The default model is set to have a context of size `1,000`. This means that in one iteration the model can only process `1,000` elements. \n",
    "\n",
    "**Explanation**: The ASTs in the dataset are very likely to exceed the models context size by containing more thatn `1,000` nodes. Therefore, the ASTs are sliced into pieces that contain at most `1,000` nodes. One thing to note here is that the `trav_trans` authors split the ASTs with a sliding window of `CONTEXTS / 2`, in our case `1,000 / 2 = 500`. For example: For an AST with 1,700 elements, the script will output three slices: `[0:1000]`, `[500:1500]`, `[700:1700]`. The sliding window moves `500` elements. If there are no more elements in order to make a `500` element large slide, the last slice will go from `[highest node - CONTEXT:highest node]`, in this case `[1700 - 1000:1700]` or `[700:1700]`.\n",
    "\n",
    "The slices contain an additional piece of information: The first slice actually looks like this: `[ast[0:1000],0]`. The last element `0` specifies the last node that has already been visited. In this case it's `0` as it's the first slices. The second slice looks like this: `[ast[500:1500],1000]`. Now the last element is `1000` as the slice may go from `500` to `1500` but the previous slice has already visited element `1000`. The third and last slide would look like this: `[ast[700, 1700],1500]`, as node `1500` has already been visited by the second slice. This additional information is required during the model training process, as the model will only take into account nodes which have not yet been visited in order to avoid taking into account the same nodes twice.\n",
    "\n",
    "Finally, the data is still in a tree structure. In order to turn these trees into a sequence, which is required by the transformer model, the trees are traversed in pre-order sequence. As output, we therefore have a TXT file instead of a JSON. Each line of the TXT file contains at most `1,000` AST nodes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generating vocab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python generate_vocab.py -n 100000 -i tmp/new_trees.json -o tmp/vocab.pkl -t \"ast\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Vocab deals as a mapping between the actual node value and a machine learning interpretable representation, in this case integer representation. This allows human readable node values to be mapped into a datatype which the model can use and vice versa. The model output can therefore be mapped back to human readable values.\n",
    "\n",
    "By default the vocab size is limited to `100,000` tokens. The vocab can be seen as a dictionary: Given an `ID`, the dictionary returns the according string. \n",
    "\n",
    "**Explanation**: This script will iterate through all ASTs node-by-node and check its value. It will keep track on how often a specific value occurs throughout the entire dataset. Finally, the top `100,000` elements will be preserved, [pickled](https://docs.python.org/3/library/pickle.html) and stored for later use."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generating AST IDs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python models\\trav_trans\\generate_ast_ids.py -a tmp/new_trees.json -o tmp/ids.txt -c 1000 \"all\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The trained model is able to generate AST nodes. All generated nodes could be evaluated against the expected node in order to establish a score. However, this would end up in a single score for the entire model. The `trav_trans` authors decided to introduce two score-groups: One for node values and one for node types. This allows node value predictions to be evaluated against attributes, nums, names, params and strings and node type predictions to be evaluated against calls, assigns, returns, lists, dicts and raises. By having this vast variety of different evaluation categories, the model can be compared to other implementations much more thoroughly: One could determine how well a model performs in string predictions and compare the string-prediction capabilities with another model.\n",
    "\n",
    "**Explanation**: This script will iterate through all ASTs and their respective nodes and checks for certain keywords. If the script finds a node of type `type`, it will check which type the node resembles: For example, if the type is `Call`, the node ID will be stored in a `call_ids` list. If the script finds a type `ListComp, ListLoad or ListStore`, the node ID will be stored in the `list_ids` list. \n",
    "\n",
    "Same procedure for nodes of type `value`: In this case the values parent node type is checked: If the parent is of type `Str`, then the value node ID is stored in a `string_ids` list. If the parent nodes type is `NameLoad or NameStore`, the value node ID is stored in `name_ids`.\n",
    "\n",
    "Finally there are several lists containing the IDs of certain node categories. During evaluation these lists will be used in order to compute scores for each category. For example: The model makes a prediction for node ID `n`. It will then check if `n` is in any list. It may find, that ID `n` is in list `string_ids` and therefore the expected node should be of type `String`. The computed score will then be added to the exclusive `String` score."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training\n",
    "Usually the dataset gets split into training/evaluation datasets. This means that the `data/python150k.json` file could initially be split into `60:40` test-split. Afterwards all the previous scripts should be applied to each dataset respectively. The following training script will then of course only use the training dataset. The evaluation script will then use the evaluation dataset.\n",
    "\n",
    "You may want to execute this script on a high performance cluster or on a machine with sufficient GPU power."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python train.py --batch_size 4 --num_epoch 3 --dps tmp/dps.txt --ids tmp/ids.txt --suffix demo --save_on_epoch False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The training script allows the specification of batch size (instead of feeding the ASTs one by one to the model, multiple can be fed to the model at once, limited by GPU memory), number of epochs (if the model was trained on the entire dataset this means that it was trained on one epoch. Multiple epochs mean that the model will train on the entire dataset for multiple times), a suffix to adjust the model file names and `save_on_epoch` which, if enabled, stores a snapshot of the model after each epoch instead of storing only one model after finishing the entire training.\n",
    "\n",
    "**Explanation**: This script contains the training routine for the model which can be found in `model.py`. Training is essentially just a simple for loop which calls the same procedures over and over again: First, feed the model one batch from the dataset, the model returns a [Loss](https://en.wikipedia.org/wiki/Loss_function) and adjusts its weights accordingly.\n",
    "\n",
    "The training script will output a model `{suffix}-model-final.pt` which can be used for evaluation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python evaluate.py --model output/{suffix}-model-final.pt --dps tmp/eval_dps.txt --ids tmp/eval_ids.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluation is very similar to training with the core difference that model weights are not adjusted during the process. Batches of evaluation data are fed to the model which returns predictions as a result. Finally, the predictions are compared to the expected results with a scoring function, in this case \"Mean Reciprocal Rank\" or MRR. The scores are then displayed as console output."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Results\n",
    "\n",
    "The results for RQ1 can be found in a separate notebook [rq1.ipynb](https://github.com/derochs/code-prediction-transformer/blob/master/rq1.ipynb). It contains training times, scores and visualizes the results in charts."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RQ2: Can the basline model be improved by adjusting architectural settings?\n",
    "This research question investigates certain architectural aspects of the model itself. There's no new script for this research question. Instead, the `train.py` script was adjusted for each scenario: In [train.py line 22](https://github.com/derochs/code-prediction-transformer/blob/master/train.py#L22) the model is istantiated, including the two parameters relevant for this research question, namely *number of decoder blocks* and *embedding size*. \n",
    "\n",
    "This means that all steps of RQ1 were performed identically, with the difference being that for each test case, the according model was instantiated with different parameters. The parameters for this research questions were: Number of layers: `1`, `3 (default)`, `6`, `9`, and embedding sizes `120`, `240`, `300 (default)`, and `540`. After the training process, this resulted in eight different models which are then compared to each other."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Results\n",
    "\n",
    "The results for RQ2 can be found in a separate notebook [rq2.ipynb](https://github.com/derochs/code-prediction-transformer/blob/master/rq2.ipynb). It contains training times, scores for the different architectural settings and visualizes the results in charts."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RQ3: Should additional metrics be tracked for model evaluation?\n",
    "This research question does not introduce any custom scripts. It mainly focuses on the evaluation groups and whether or not all relevant metrics are captured by the categories introduced by the `trav_trans` authors."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Results\n",
    "\n",
    "The results for RQ3 can be found in a separate notebook [rq3.ipynb](https://github.com/derochs/code-prediction-transformer/blob/master/rq3.ipynb). It shows which metrics were added to the evaluation routine and how well these new metrics hold up against other code structures."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RQ4: Can the out-of-vocabulary issue be reduced by using an alternative tokenizer?\n",
    "\n",
    "This research question tries to replace the original tokenizer (a [WordLevel tokenizer](https://huggingface.co/docs/tokenizers/pipeline#model) which replaces a word with an integer) with a WordPiece tokenizer, which splits words into common subwords and maps those to integers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating the WordPiece tokenizer\n",
    "At first, the new tokenizer has to be \"trained\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python rq4_tokenizer.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This script will create a tokenizer which is stored as JSON file in order to be reusable. \n",
    "\n",
    "**Explanation**: The dataset input file for this script was hardcoded, so make sure that the dataset can be found at `data/python150k.json`. The tokenizer requires a _delimiter_ which specifies how two separate two words from each other. The delimiter used is the `,` character. However, an AST node value can theoretically contain a large body of String, therefore eventually contain such a `,` character. We want each node to be treated as a word, so there are no `,` characters allowed. The script will iterate through all AST nodes and remove any instance of `,` characters. The cleaned dataset is then stored as `data/python150k_rq4.json`.\n",
    "\n",
    "In the next step, all nodes are then extracted into a new file called `output/all_raw.json` which contains all cleaned node values, each node being separated with the delimiter `,` from each other. \n",
    "\n",
    "This file is then fed to a tokenizer pipeline, which will train a WordPiece tokenizer on the provided dataset and export a `output/tokenizer.json`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generating new trees\n",
    "Similar to RQ1, RQ4 requires AST nodes to either contain a type or a value, but never both. This is done by the next script:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python rq4_preprocess.py --file_path data/python150k.json --tokenizer output/tokenizer.json --suffix demo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This model will generated new trees just like the `generate_new_trees.py` script from RQ1, remove all `,` characters as required by the tokenizer, split and traverse the ASTs. The outputs are `output/[suffix]_dps.txt` and `output/[suffix]_ids.txt`, depending on the suffix that was chosen for the script call.\n",
    "\n",
    "**Explanation**: The new WordPiece tokenizer is used to convert the node values into integers. As a WordPiece tokenizer converts words into subwords, additional logic had to be performed when splitting the trees."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python rq4_train.py --batch_size 4 --num_epoch 3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This script will train a model. There are two hardcoded file paths that should exists, `output/tokenizer.json` should contain the trained tokenizer and `output/[suffix]_rq4_dps.txt` should contain the dataset that was generated in the previous step by the script `rq4_preprocess.py`.\n",
    "\n",
    "Just like in RQ1, the trained model is then stored for future evaluation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python rq4_evaluate.py --model output/modelname.pt --dps output/rq4_test_dps.txt --ids output/rq4_test_ids.txt --tokenizer output/tokenizer.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This script will evaluate the model with the new tokenizer just like the evaluation in RQ1."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Results\n",
    "\n",
    "The results for RQ4 can be found in a separate notebook [rq4.ipynb](https://github.com/derochs/code-prediction-transformer/blob/master/rq4.ipynb). It contains the evaluation results of a model trained on a subword tokenizer compared to the \"normal\" model trained on a wordlevel tokenizer. The most interesting aspect is the comparison of String predictions, as the subword tokenizer is able to make improved predictions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RQ5: Which impact does the overlap size have on model performance?\n",
    "\n",
    "This research question investigates the impact of the overlap size during the tree splitting process. By default, the sliding window has a stride of `50%`. The research question deals with testing out different overlap sizes and compares the results.\n",
    "\n",
    "The parameter `half_len` which resembles the sliding window in [utils.py line 95](https://github.com/derochs/code-prediction-transformer/blob/master/utils.py#L95) is set to `75%`, `50% (default)` and `30%` in different test cases. This results in three different datasets, on which a model is trained and evaluated on respectively, following the same process as RQ1.\n",
    "\n",
    "The scripts are all trailing a `rq6` despite being the scripts for RQ5 due to changes in the thesis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python rq6_preprocess.py --file_path data/python150k.json --suffix demo --overlap 2 --context_size 1000 --generate_vocab True --n_vocab 100000"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This script executes other scripts: At first it will run `rq6_generate_new_trees.py` which will prepare and split the trees. As overlap, it will use `1 / OVERLAP`, `OVERLAP` being the parameter passed to the `rq6_preprocess.py` script. This means for example, that a `2` means `1 / 2 = 0.5` as an overlap.\n",
    "\n",
    "Afterwards, a vocab is generated if specified in as argument and AST IDs are collected similar to RQ1. Training and Evaluating is performed with the same scripts from RQ1."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Results\n",
    "\n",
    "The results for RQ5 can be found in a separate notebook [rq6.ipynb](https://github.com/derochs/code-prediction-transformer/blob/master/rq6.ipynb) (research numbers were changed throughout the research, hence the file name differing from the actual RQ number). It contains the training times and scores for model trained on datasets split with different overlap sizes as well as a visual comparison between them."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RQ6 Does positional embedding improve the performance of the model?\n",
    "This research question introduces a change in the model architecture. Therefore, a new model `rq8_model.py` has been created which reintroduces a positional embedding layer. The dataset preparation steps are equal to the ones from RQ1. Only the testing and evaluation scripts differ. In this case there are two new scripts:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python rq8_train.py --batch_size 4 --num_epoch 3 --dps output/train_dps.txt --ids output/train_ids.txt --suffix demo --save_on_epoch True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This script is actually identical to the training script from RQ1, the only change being that the new modified model is trained."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python rq8_evaluate.py --model output/model.pt --dps output/eval_dps.txt --ids output/eval_ids.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This script is identical to the RQ1 evaluation script with the difference being that the new modified model is used for evaluation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Results\n",
    "\n",
    "The results for RQ6 can be found in a separate notebook [rq8.ipynb](https://github.com/derochs/code-prediction-transformer/blob/master/rq8.ipynb) (research numbers were changed throughout the research, hence the file name differing from the actual RQ number). It contains a comparison between the original `trav_trans` model and how a model with reintroduced positional embedding layer holds up against it. "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}