code2vec

To support a new language, you must implement a new extractor (similar to JavaExtractor) and register it in preprocess.sh.

An extractor must output a single text file where each row is an example. Each example must be a space-delimited list of fields:

1. The first word is the target label, using | as an internal delimiter.
2. Subsequent words are contexts. Each context must consist of three comma-separated components: token,path,token.

Example Context Format: foo|Bar System,FIELD_ACCESS,out System.out,FIELD_ACCESS,println THE_METHOD,returns,void THE_METHOD,prints,"hello_world"

In this format, the 1st and 3rd components of a context are drawn from the 'tokens' vocabulary, while the 2nd component is drawn from the 'paths' vocabulary.

To begin using code2vec, clone the repository and enter the directory:

git clone https://github.com/tech-srl/code2vec
cd code2vec

To run code2vec on Ubuntu, ensure you have the following installed:

• Python 3 (>=3.6)
• TensorFlow 2.0.0
• CUDA 10.0 (if using a GPU)
• cuDNN >= 7.5 (if using a GPU)
• Java JDK (required for creating new datasets or manually examining a trained model to parse code examples)

Use these commands to verify your versions:

# Check Python version
python3 --version

# Check TensorFlow version
python3 -c 'import tensorflow as tf; print(tf.__version__)'

# Check CUDA version
nvcc --version

# Check cuDNN version
cat /usr/include/cudnn.h | grep CUDNN_MAJOR -A 2

The following preprocessed datasets are available for training:

• Java-small: ~700K examples (9 training projects, 1 validation, 1 testing).

  wget https://s3.amazonaws.com/code2vec/data/java-small_data.tar.gz

• Java-med: ~4M examples (800 training, 100 validation, 100 testing).

  wget https://s3.amazonaws.com/code2vec/data/java-med_data.tar.gz

• Java-large: ~16M examples (9000 training, 200 validation, 300 testing).

  wget https://s3.amazonaws.com/code2vec/data/java-large_data.tar.gz

Pre-trained Models:

• Java-large (Trainable): 3.5 GB

  wget https://code2vec.s3.amazonaws.com/model/java-large-model.tar.gz

• Java-large (Released/Inference only): 1.4 GB

  wget https://code2vec.s3.amazonaws.com/model/java-large-released-model.tar.gz

You can either download the preprocessed ~14M example dataset or create your own from Java sources.

Download preprocessed dataset

wget https://s3.amazonaws.com/code2vec/data/java14m_data.tar.gz
tar -xvzf java14m_data.tar.gz

This creates a data/java14m/ directory containing training, test, and validation sets.

Preprocess your own Java dataset

1. Edit preprocess.sh to point to your specific training, validation, and test directories.
2. Run the script:

source preprocess.sh

# Download preprocessed dataset
wget https://s3.amazonaws.com/code2vec/data/java14m_data.tar.gz
tar -xvzf java14m_data.tar.gz

# Or preprocess your own
source preprocess.sh

If you want to save a trained model for inference only (which prevents further training but reduces disk space by approximately 3x), use the --release flag. This creates a copy of the model with a .release suffix.

python3 code2vec.py --load models/java14_model/saved_model_iter8 --release

Hyperparameters are managed by editing the config.py file. Key parameters include:

The code2vec.py script supports two model implementations:

1. Pure TensorFlow: The default implementation (defined in tensorflow_model.py).
2. Keras: Uses TensorFlow's Keras API (defined in keras_model.py).

To switch implementations, use the --framework flag.

Important: When loading a trained model from a file, you must use the same framework that was used during its training.

python3 code2vec.py --framework tensorflow
# OR
python3 code2vec.py --framework keras

To interactively examine a model's predictions and attention scores:

1. Run the prediction command:

python3 code2vec.py --load models/java14_model/saved_model_iter8.release --predict

2. Follow the on-screen instructions to edit Input.java with a Java method or code snippet to see how the model reacts.

To evaluate a specific iteration of a trained model on a test set, use the code2vec.py script. If using a released model, ensure you append .release to the model path.

python3 code2vec.py --load models/java14_model/saved_model_iter8.release --test data/java14m/java14m.test.c2v

Results (each test example name and its prediction) are written to log.txt during evaluation.

Exported word2vec files can be loaded and inspected using the gensim Python package.

from gensim.models import KeyedVectors as word2vec

vectors_text_path = 'models/java14_model/targets.txt'
model = word2vec.load_word2vec_format(vectors_text_path, binary=False)

# Find most similar words
print(model.most_similar(positive=['equals', 'to|lower']))
print(model.most_similar(positive=['download', 'send'], negative=['receive']))