Building a Python Code Generator. Using Transformers to convert English... | by Divyam Malay Shah | Towards Data Science

Can NLP techniques be used to generate actual code? How far are we from a world where AI is used to write software?
In this blog, I make an attempt to build a python code generator that can convert simple English problem statements into their corresponding python code.
We tackle this problem by treating it as a Sequence to Sequence(Seq2Seq) learning problem. Here, our English sentence will be our Input or SRC sequence, and our Python Code will be our Output or TRG sequence.
Over the last couple of years, Transformers have emerged as the dominant architecture for solving Seq2Seq problems. Most present-day SoA models like BERT or GPT-3 make use of transformers internally. Transformers, as we know them today were first introduced by Google in their “Attention Is All You Need” paper. The “English-to-Python” model that we discuss in this blog also finds its motivation in this paper.
Before we begin solving our problem let’s first briefly review Transformers.
Transformers can be understood in terms of their three components:
An Encoder that encodes an input sequence into state representation vectors.An Attention mechanism that enables our Transformer model to focus on the right aspects of the sequential input stream. This is used repeatedly within both the encoder and the decoder to help them contextualize the input data.A Decoder that decodes the state representation vector to generate the target output sequence.
An Encoder that encodes an input sequence into state representation vectors.
An Attention mechanism that enables our Transformer model to focus on the right aspects of the sequential input stream. This is used repeatedly within both the encoder and the decoder to help them contextualize the input data.
A Decoder that decodes the state representation vector to generate the target output sequence.
I have explained each of these components in detail along with a code walkthrough in my previous blog.
Let's now see how we can feed data into our transformer.
We will be using a custom dataset curated by The School of AI(TSAI) to train our model. This dataset contains about 5000 data points where each data point comprises an English problem statement and its corresponding Python code. You can follow my project repo to understand how to parse the data.
Sample data Point:
English statement: “write a function that adds two numbers”
Python Code:
def add_two_numbers (num1 ,num2 ):    sum =num1 +num2     return sum
Here the “English statement” is our Input or SRC sequence and the “Python Code” is our Output or TRG sequence for training.
Our Input(SRC) and Output(TRG) sequence exist in the form of single strings that need to be further tokenized in order to be sent into the transformer model.
To tokenize the Input(SRC) sequence we make use of spacy. This is implemented by default in PyTorch's torchtext.data.Field. We will be using torchtext.data.Field to tokenize our data using spacy.
Input = data.Field(tokenize = 'spacy',            init_token='<sos>',             eos_token='<eos>',             lower=True)
To tokenize our Output(TRG) sequence we make use of our custom tokenizer built upon Python’s source code tokenizer. Python’s tokenizer returns several attributes for each token. We only extract the token type and the corresponding string attribute in form of a tuple(i.e., (token_type_int, token_string)) as the final token.
Tokenized Input(SRC):
SRC = [' ', 'write', 'a', 'python', 'function', 'to', 'add', 'two', 'user', 'provided', 'numbers', 'and', 'return', 'the', 'sum']
Tokenized Output(TRG):
TRG = [(57, 'utf-8'), (1, 'def'), (1, 'add_two_numbers'), (53, '('), (1, 'num1'), (53, ','), (1, 'num2'), (53, ')'), (53, ':'), (4, '\n'), (5, '    '), (1, 'sum'), (53, '='), (1, 'num1'), (53, '+'), (1, 'num2'), (4, '\n'), (1, 'return'), (1, 'sum'), (4, ''), (6, ''), (0, '')]
Since our dataset contains mere 5000 data points, we make use of data augmentations to increase the size of our dataset. While tokenizing the python code, we mask the names of certain variables randomly(with ‘var_1, ‘var_2’ etc) to ensure that the model that we train does not merely fixate on the way the variables are named and actually tries to understand the inherent logic and syntax of the python code.
For example, consider the following program.
def add_two_numbers (num1 ,num2 ):    sum =num1 +num2     return sum
We can replace some of the above variables to create new data points. The following are valid augmentations.
1.
def add_two_numbers (var_1 ,num2 ):    sum =var_1 +num2     return sum
2.
def add_two_numbers (num1 ,var_1 ):   sum =num1 +var_1    return sum 
3.
def add_two_numbers (var_1 ,var_2 ):    sum = var_1 + var_2     return sum
In the above example, we have therefore expanded a single data point into 3 more data points using our random variable replacement technique.
We implement our augmentations at the time of generating our TRG tokens.
While randomly picking variables to mask we avoid keyword literals(keyword.kwlist), control structures(as can be seen in below skip_list), and object properties. We add all such literals that need to be skipped into the skip_list.
We now apply our augmentations and tokenization using Pytorch’s torchtext.data.Field.
Output = data.Field(tokenize = augment_tokenize_python_code,                    init_token='<sos>',                     eos_token='<eos>',                     lower=False)
Our tokenized Output(TRG) after applying tokenization:
TRG = [(57, 'utf-8'), (1, 'def'), (1, 'add_two_numbers'), (53, '('), (1, 'num1'), (53, ','), (1, 'var_1'), (53, ')'), (53, ':'), (4, '\n'), (5, '    '), (1, 'sum'), (53, '='), (1, 'num1'), (53, '+'), (1, 'var_1'), (4, '\n'), (1, 'return'), (1, 'sum'), (4, ''), (6, ''), (0, '')]
To feed data into our model we first create batches using Pytorch’s torchtext.data.BucketIterator. This ensures that Inputs that have similar lengths stay together in a single batch for ease of training. Then we feed in our tokenized Input(SRC) batches into the encoder and use the tokenized Output(TRG) batches in our decoder. Our goal is to use the encoder’s tokenized English Input(SRC) to predict the tokenized Python Output(TRG) via the decoder. The tokenized predictions are then untokenized via the untokenize function of Python’s source code tokenizer.
We have used augmentations in our dataset to mask variable literals. This means that our model can predict a variety of values for a particular variable and all of them are correct as long as the predictions are consistent through the code. This would mean that our training labels are not very certain and hence it would make more sense to treat them to be correct with probability 1- smooth_eps and incorrect otherwise. This is what label smoothening does. By adding label smoothening to Cross-Entropy we ensure that the model does not become too confident in predicting some of our variables that can be replaced via augmentations.
Now with all our components set we can train our model using backpropagation. We split our dataset into training and validation data. Our model is trained until our validation loss does not improve any further.
It is important to note that label smoothening leads to much higher loss values as compared to models that do not make use of label smoothening. But this is as expected as we do not intend to be certain with our label predictions. This is particularly the case with variables as there can be multiple correct options as long as the predictions are consistent through the target code sequence.
Please refer to my notebook for the complete training code.
Input: “program to sort a list of dictionaries by key”Output:
var_1 ={'Nikhil':{'roll':24 ,'marks':17 },'Akshat':{'roll':54 ,'marks':12 },'Akash':{'roll':15 },'marks':15 }}sort_key ='marks'res ='marks'res =var_2 (test_dict .items (),key =lambda x :x [1 ][sort_key ])print ("The sorted dictionary by marks is : "+str (res ))
Input: “function to sum odd elements of list”Output:
def sum_odd_elements (l :list ):    return sum ([i for i in l if i %2 ==1 ])
Input: “program to reverse a string”Output:
var_1 ='Today is bad day'var_1 [::-1 ]
For more samples take a look at my project repo.
We have successfully trained a model that is capable of converting simple problem statements(in English) into corresponding python code. Looks like we aren’t too far away from an era where an AI writes software instead of humans. Software Developers beware!
github.com
github.com
[1] Aston Zhang and Zachary C. Lipton and Mu Li and Alexander J. Smola, Dive into Deep Learning(2020).
[2] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, Illia Polosukhin, Attention Is All You Need (2017), 31st Conference on Neural Information Processing Systems (NIPS 2017), Long Beach, CA, USA
[3] Rafael Müller, Simon Kornblith, Geoffrey Hinton, When Does Label Smoothing Help? (2019), 33rd Conference on Neural Information Processing Systems (NeurIPS 2019), Vancouver, Canada
[4] Ian Goodfellow and Yoshua Bengio and Aaron Courville, Deep Learning Book(2016), MIT Press.