Text-to-code basic using Tensorflow

Introduction

Given the goal of improving software development productivity with machine learning methods, software intelligence research has attracted increasing attention in both academia and industries over the last decade. Software code intelligence techniques can help developers reduce tedious repetitive workloads, enhance the programming quality, and improve the overall software development productivity. This would reduce time spent writing software as well as reduce computational and operational costs.

Text-to-code generation is a task where we can generate code based on the natural language description. It can further be used to build an AI-powered coding assistant. Developers simply type the natural language description or the function signature to specify their intents, and the AI coding assistant can generate or complete the target function for them. This helps to accelerate implementation and also reduce their reliance on external resources.

CodeT5 by Salesforce is the first code-aware, encoder-decoder-based pre-trained programming language model, which enables a wide range of code intelligence applications including code understanding and generation tasks. CodeT5 achieves state-of-the-art performance on 14 sub-tasks in the CodeXGLUE code intelligence benchmark.

CodeT5 builds on an encoder-decoder framework with the same architecture as T5 by Google. The T5 model, pre-trained on C4, achieves state-of-the-art results on many NLP benchmarks while being flexible enough to be fine-tuned to a variety of important downstream tasks. T5 architecture employs denoising sequence-to-sequence (Seq2Seq) pre-training and has been shown to benefit out-of-the-box for both understanding and generation tasks in natural language.

In this notebook we will finetune CodeT5 on MBPP - Mostly Basic Python Problems by Google Research to generate code based on problem description. MBPP is a benchmark that consists of around 1,000 crowd-sourced Python programming problems, designed to be solvable by entry level programmers, covering programming fundamentals, standard library functionality, and so on. Each problem consists of a task description, code solution and 3 automated test cases.

The notebook demonstrates how to finetune CodeT5 on MBPP Dataset w/ TensorFlow.

T5

Explained

• T5 — Text-to-Text Transfer Transformer Model was proposed in the paper, Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer. This paper is essentially a survey of modern transfer learning techniques used in language understanding and hence proposes a unified framework that attempts to combine all language problems into a text-to-text format.

• The text-to-text framework suggests using the same model, same loss function, and the same hyperparameters on all the NLP tasks. In this approach, the inputs are modeled in such a way that the model shall recognize a task, and the output is simply the “text” version of the expected outcome.

• To avail the same model for all the downstream tasks, a task-specific text prefix is added to the original input that is fed to the model. This text prefix is also considered as a hyperparameter.

As an example,to ask the model to translate the sentence “That is good.” from English to German, the model would be fed the sequence “translate English to German: That is good.” and would be trained to output “Das ist gut.” — T5 Paper

Similarly, for classification tasks, the model predicts a single word corresponding to the target label.

For example, on the MNLI benchmark the goal is to predict whether a premise implies (“entailment”), contradicts (“contradiction”), or neither (“neutral”) a hypothesis. With our preprocessing, the input sequence becomes “mnli premise: I hate pigeons. hypothesis: My feelings towards pigeons are filled with animosity.” with the corresponding target word “entailment”. — T5 Paper

• T5 is pretrained using the denoising objective on C4— Colossal Clean Crawled Corpus - a 750GB dataset which is not just reasonably larger than the most pre-training datasets but also contains a relatively very clean text.

• The proposed model is essentially a Encoder-Decoder Transformer with some architectural changes (like applying Layer Normalization before a sub block and then adding the initial input to the sub-block output; also known as pre-norm). Moreover, the model configuration is similar to BERT base. T5 uses relative scalar embeddings. Encoder input padding can be done on the left and on the right.

• T5 comes in different sizes: t5-small, t5-base, t5-large, t5-3b, t5-11b.

Training

• T5 is an encoder-decoder model and converts all NLP problems into a text-to-text format. It is trained using teacher forcing. This means that for training, we always need an input sequence and a corresponding target sequence. The input sequence is fed to the model using input_ids. The target sequence is shifted to the right, i.e., prepended by a start-sequence token and fed to the decoder using the decoder_input_ids. In teacher-forcing style, the target sequence is then appended by the EOS token and corresponds to the labels. The PAD token is hereby used as the start-sequence token. T5 can be trained / fine-tuned both in a supervised and unsupervised fashion.

• One can use T5ForConditionalGeneration (or the Tensorflow/Flax variant) from HuggingFace transformers library, which includes the language modeling head on top of the decoder.

• T5 models need a slightly higher learning rate than the default one set in the Trainer when using the AdamW optimizer. Typically, 1e-4 and 3e-4 work well for most problems (classification, summarization, translation, question answering, question generation). Note that T5 was pre-trained using the AdaFactor optimizer.

• Task prefixes matter when (1) doing multi-task training (2) your task is similar or related to one of the supervised tasks used in T5’s pre-training mixture (see Appendix D of the paper for the task prefixes used).

• We must make sure that padding token id’s of the labels are not taken into account by the loss function. In PyTorch and Tensorflow, this can be done by replacing them with -100, which is the ignore_index of the CrossEntropyLoss. We also pass attention_mask as additional input to the model, which makes sure that padding tokens of the inputs are ignored.

Inference

• At inference time, it is recommended to use generate(). This method takes care of encoding the input and feeding the encoded hidden states via cross-attention layers to the decoder and auto-regressively generates the decoder output. Check out this blog post to know all the details about generating text with Transformers. There’s also this blog post which explains how generation works in general in encoder-decoder models.

References

• Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer
• Google AI Blog - Exploring Transfer Learning with T5: the Text-To-Text Transfer Transformer
• HuggingFace - T5
• Understanding Transformer-Based Self-Supervised Architectures
• T5 - A Detailed Explanation

CodeT5

Explained

• CodeT5 by Salesforce was proposed in CodeT5: Identifier-aware Unified Pre-trained Encoder-Decoder Models for Code Understanding and Generation and is an open-source model that can understand and readily generate code. It is an identifier-aware unified pre-trained coder-encoder tool that enables a wide range of code intelligence applications. CodeT5 possesses an uninformed model for natural language processing tasks, which reframes text-to-text with input, and output data always being strings of texts.

• CodeT5 builds on the similar architecture of T5 but incorporates code-specific knowledge to endow the model with better code understanding. It takes code and its accompanying comments as a sequence to build and generate upon. It aims to derive generic representations for programming language (PL) and natural language (NL) via pre-training on unlabeled source code.

• CodeT5 achieves state-of-the-art performance on multiple code-related downstream tasks including understanding tasks such as code defect detection and clone detection, and generation tasks across various directions including PL-NL, NL-PL, and PL-PL.

Pretraining

• CodeT5 was pretrained on the CodeSearchNet data that consists of both unimodal (PL-only) and bimodal (PL-NL) data on six PLs - Ruby, JavaScript, Go, Python, PHP, C, and C#. In addition to that, they further collect extra data of C/C# from open-source Github repositories. They further finetune CodeT5 on most tasks in the CodeXGLUE benchmark, including two understanding tasks: code defect detection and clone detection, and generation tasks such as code summarization, generation, translation, and refinement.

• Some of the pre-training tasks of CodeT5 include:

  • Masked Span Prediction (MSP) randomly masks spans with arbitrary lengths and requires the decoder to recover the original input. It captures the syntactic information of the NL-PL input and learns robust cross-lingual representations as we pre-train on multiple PLs with a shared model.

  • Identifier Tagging (IT) applied only to the encoder which distinguishes whether each code token is an identifier (e.g., variables or function names) or not. It works like the syntax highlighting feature in some developer-aided tools.

  • Masked Identifier Prediction (MIP), in contrast to MSP, only masks identifiers and employs the same mask placeholder for all occurrences of one unique identifier. It works like deobfuscation in software engineering and is a more challenging task that requires the model to comprehend the code semantics based on the obfuscated code.

  • Bimodal Dual Generation (dual-gen) jointly optimizes the conversion from code to its comments and vice versa. It encourages a better alignment between the NL and PL counterparts.

Code-specific Tokenizer

• In CodeT5 they train a Byte-level BPE tokenizer and set the vocabulary size to 32,000 as T5. They add additional special tokens ([PAD], [CLS], [SEP], [MASK0], ..., [MASK99]). This tokenzier is trained on all pre-training data with non-printable characters and low-frequent tokens (occurring <3 times) filtered.

• When compared to T5’s default tokenizer they find that the trained tokenizer largely reduces the length of tokenized code sequence by 30% - 45% on downstream tasks. This accelerates the training and especially benefits generation tasks due to the shorter sequence to predict.

• They spot a severe problem for applying the T5’s default tokenizer on source code, where it would encode some common code tokens such as brackets [‘{’, ‘}’] into unknown tokens.

References

• CodeT5: Identifier-aware Unified Pre-trained Encoder-Decoder Models for Code Understanding and Generation
• CodeT5 Blog by Salesforce
• CodeT5 GitHub Repository
• HuggingFace CodeT5-Base Model
• Salesforce’s CodeT5 system can understand and generate code
• Salesforce CodeT5 vs Github Copilot: A Comparative Guide to Auto-code Generators
• Fine_tune_CodeT5_for_generating_docstrings_from_Ruby_code
• CodeXGLUE: A Machine Learning Benchmark Dataset for Code Understanding and Generation
• CodeSearchNet
• OpenAI Codex

Imports

Let's start by importing required libraries to the environment:

• TensorFlow an end-to-end open source platform for machine learning.
• transformers provides APIs to easily download and train state-of-the-art pretrained models
• datasets a library for easily accessing and sharing datasets.