ENDGAME — Python code generation from natural language (2021)

A sequence-to-sequence model that turns a plain-English description of a task into working Python, trained on a hand-built dataset with custom-trained code embeddings.

I built this in 2021, before LLM code assistants existed — no Copilot, no ChatGPT, no API to call. If you wanted a model that wrote code from a description, you trained one yourself: cleaned the data by hand, built the vocabulary, trained the embeddings, and debugged a seq2seq network until it produced something that ran. This repository is that work, kept as-is, including the experiments that didn't pan out.

This is early, hands-on research, not a polished library. I'm keeping it public because the interesting part is the process: what worked, what didn't, and why.

---

The problem

Given a natural-language prompt like "write a python program to count uppercase and lowercase letters in a string", generate Python that actually runs. The model is a sequence-to-sequence network (PyTorch) trained on ~4,600 description-to-code pairs.

Building the dataset

The dataset was the hardest part, and most of the time went here.

• Started from a file of 4,600+ raw code examples and cleaned them by hand and with tooling. Many had syntax issues (e.g. Python 2 print statements) that had to be fixed before the code would even parse.
• Used regex to separate each code snippet from its description. Inline comments on their own lines made this messy and forced a lot of manual cleanup.
• Ran autopep8 over every snippet to normalize formatting and convert 4-space indentation to tabs. Neither autopep8 nor black reliably fixes multi-level indentation, so a fair amount was corrected by hand.
• Kept all spaces, tabs, and newlines as explicit tokens rather than stripping them. Python, unlike English, is whitespace-significant, so I left the structure in and let the model learn indentation as part of the language.
• Limited training to snippets under 251 characters to keep the sequences tractable.

The full cleaning pipeline is in crc check.ipynb.

Custom embeddings (the experiment that didn't work)

The plan was to give the model a head start with embeddings trained specifically on Python code, using Stanford's GloVe.

• Trained GloVe on the Python 3.8 standard library (cpython/Lib) at embedding dimension 256, after converting whitespace and trailing characters into tokens to match GloVe's whitespace-based splitting. (Corpora saved as lib_tokens.pkl and lib_tokens1.pkl.)
• The custom embeddings made predictions worse, not better.
• I suspected the standard-library corpus was too complex for the simple problems in the training set, so I retrained on simpler code from LeetCode-style repositories (one, two).
• That was worse too. So I dropped the pretrained embeddings entirely and let the model learn its own from the training data.

This is the part I'd point a reviewer to: the honest answer was that a reasonable-sounding idea didn't survive contact with the data, and the right call was to drop it rather than force it.

Loss function

I tried several objectives — NLLLoss, KLDivLoss, and others — and, somewhat anticlimactically, plain cross-entropy worked best.

Results

The model does genuinely well on simple functions and programs, and a few cases were more interesting than I expected.

It often finds a more concise solution than the reference. Asked to count uppercase and lowercase letters, the training example used explicit ASCII range checks; the model produced the cleaner idiomatic version:

# Prompt: count uppercase and lowercase letters in a string
sentence = 'The Quick Brown Fox'
lowercase = 0
uppercase = 0
for c in sentence:
    if c.isupper():
        uppercase += 1
    elif c.islower():
        lowercase += 1
print(f'Lowercase: {lowercase}, Uppercase: {uppercase}')

A few behaviors stood out:

• It inferred intent the training example left implicit. On a prompt whose reference solution computed a result but never printed it, the model added a print() — it had learned from other examples that producing output was expected.
• It started learning Python's structure. The model picked up that an indentation-level change should follow a colon, even in cases where the reference solution had an indentation error.

It is far from perfect — it struggles with longer functions and class definitions, and stray indentation errors throw it off. But for a from-scratch 2021 model trained on 4,600 examples, it learned more about the structure of Python than I expected going in.

Limitations and next steps

• Weak on large functions and class definitions.
• Indentation errors specifically in class-level problems need work.
• The most promising next direction is more data rather than pretrained embeddings — expanding the dataset with LeetCode-style repositories, many of which include the class definitions the current model handles poorly.

Stack

PyTorch · sequence-to-sequence · custom GloVe embeddings · Python tokenization · autopep8 / black

---

Built in 2021 as independent research into neural code generation, before LLM code assistants were available. Preserved here in its original form.