/code2seq

PyTorch's implementation of the code2seq model.

Primary LanguagePythonMIT LicenseMIT

code2seq

JetBrains Research Github action: build Code style: black

PyTorch's implementation of code2seq model.

Installation

You can easily install model through the PIP:

pip install code2seq

Usage

Minimal code example to run the model:

from os.path import join

import hydra
from code2seq.dataset import PathContextDataModule
from code2seq.model import Code2Seq
from code2seq.utils.vocabulary import Vocabulary
from omegaconf import DictConfig
from pytorch_lightning import Trainer


@hydra.main(config_path="configs")
def train(config: DictConfig):
    vocabulary_path = join(config.data_folder, config.dataset.name, config.vocabulary_name)
    vocabulary = Vocabulary.load_vocabulary(vocabulary_path)
    model = Code2Seq(config, vocabulary)
    data_module = PathContextDataModule(config, vocabulary)

    trainer = Trainer(max_epochs=config.hyper_parameters.n_epochs)
    trainer.fit(model, datamodule=data_module)


if __name__ == "__main__":
    train()

Navigate to code2seq/configs to see examples of configs. If you had any questions then feel free to open the issue.