Building a Code Generative AI Model

Introduction

Automation is a core part of software engineering. Code Generative AI now makes it possible to have an AI write code for you. In this article, we’ll build a Code Generative AI model from scratch and answer some common questions along the way.

Can Generative AI write code?

Start with the obvious question: can Generative AI genuinely write code? Yes. Generative AI models, particularly those built on neural networks, have shown strong results generating human-like text, including code. These models train on large datasets covering various programming languages, which lets them produce code snippets that are both syntactically accurate and semantically meaningful.

What is Generative AI computer code?

Generative AI computer code is code that an artificial intelligence model, such as a neural network, produces from a natural language prompt. These models pick up the structure and the small conventions of code from large training sets, so what they produce looks close to what a human programmer would write. The output runs from simple functions to complex algorithms, depending on the prompt and on what the model was trained on.

How do I create a Generative AI for code?

Now the practical steps for building your Code Generative AI model. We’ll go through the process step by step so you can build your own AI code-writing assistant.

Step 1: Environment setup

To start, you’ll need a Python environment with the libraries and dependencies it needs. The code below is a basic setup for your AI model, with the imports and the configuration settings:

1
import logging
2
import torch
3
import peft
4
import transformers
5
from transformers import AutoTokenizer, AutoModelForCausalLM
6
from huggingface_hub.hf_api import HfFolder
7

8
# Configuration class for the Dumpster Copilot Generative model
9
class Configuration:
10
    ACCESS_TOKEN = 'ENTER YOUR HUGGINGFACE ACCESS TOKEN HERE'
11
    LOAD_IN_8BIT = False
12
    BASE_MODEL = 'meta-llama/Llama-2-7b-chat-hf'
13
    LORA_WEIGHTS = 'qblocks/llama2-7b-tiny-codes-code-generation'
14
    PROMPT = 'Write a Python function to divide 2 numbers and check for division by zero.'
15

16
# Exception classes for errors loading the model and generating text
17
class ModelLoadingError(Exception):
18
    pass
19

20
class DumpsterCopilotGenerativeError(Exception):
21
    pass
22

23
# Model loader class
24
class ModelLoader:
25
    @staticmethod
26
    def load_model() -> tuple:
27
        try:
28
            tokenizer = AutoTokenizer.from_pretrained(Configuration.LORA_WEIGHTS)
29
            model = AutoModelForCausalLM.from_pretrained(
30
                Configuration.BASE_MODEL,
31
                device_map='auto',
32
                torch_dtype=torch.float16,
33
                load_in_8bit=Configuration.LOAD_IN_8BIT
34
            )
35
            model = peft.PeftModel.from_pretrained(model, Configuration.LORA_WEIGHTS)
36
            return tokenizer, model
37
        except Exception as e:
38
            raise ModelLoadingError(f'Error loading tokenizer and model: {str(e)}')

In this code snippet, we’ve imported the libraries we need, transformers and torch, and added a Configuration class to hold the settings. The ModelLoader class is the one that loads the AI model.

Step 2: Loading your AI model

Now that your environment is set up, it’s time to load your Code Generative AI model. In the code snippet, we’ve defined a ModelLoader class with a load_model method that handles the loading. This method returns a tokenizer and a model instance. Remember to replace 'ENTER YOUR HUGGINGFACE ACCESS TOKEN HERE' with your actual Hugging Face access token.

Step 3: Generating code with your AI

With your model loaded, you can generate code snippets with your AI assistant. The DumpsterCopilotGenerative class below runs the generation from a prompt you pass in:

1
class DumpsterCopilotGenerative:
2
    def __init__(self, tokenizer, model):
3
        self.tokenizer = tokenizer
4
        self.model = model
5

6
    def dumpster_copilot_generative(self, prompt: str) -> str:
7
        try:
8
            logging.info(f'Generating text for prompt: {prompt}')
9

10
            generator = transformers.pipeline(
11
                'text-generation',
12
                model=self.model,
13
                tokenizer=self.tokenizer
14
            )
15

16
            generation_config = transformers.GenerationConfig(
17
                temperature=0.4,
18
                top_p=0.99,
19
                top_k=40,
20
                num_beams=2,
21
                max_new_tokens=400,
22
                repetition_penalty=1.3
23
            )
24

25
            t = generator(prompt, generation_config=generation_config)
26

27
            generated_text = t[0]['generated_text']
28

29
            logging.info(f'Generated text: {generated_text}')
30
            return generated_text
31
        except Exception as e:
32
            raise DumpsterCopilotGenerativeError(f'Error generating text: {str(e)}')

The class has one method, dumpster_copilot_generative, which takes a prompt as input and returns the generated code. What you get back depends on the prompt, so make the prompt explicit and specific.

Step 4: Running your Code Generative AI

Now that all the pieces are in place, you can run your Code Generative AI model and generate code. Here’s how:

1
if __name__ == '__main__':
2
    try:
3
        if Configuration.ACCESS_TOKEN:
4
            HfFolder.save_token(Configuration.ACCESS_TOKEN)
5

6
        logging.info('Initiating the text generation process.')
7

8
        tokenizer, model = ModelLoader.load_model()
9
        generator = DumpsterCopilotGenerative(tokenizer, model)
10
        generated_text = generator.dumpster_copilot_generative(Configuration.PROMPT)
11

12
        logging.info('Generated text:')
13
        logging.info(generated_text)
14

15
        logging.info('Successful completion of the text generation process.')
16
    except (ModelLoadingError, DumpsterCopilotGenerativeError) as e:
17
        logging.error(f'An error occurred: {str(e)}')

This block initializes your AI model, generates code from the prompt (in this case, “Write a Python function to divide 2 numbers and check for division by zero.”), and logs the resulting code.

Frequently asked questions

Now that you have the basics of building a Code Generative AI model, here are some questions that come up often.

Q1: How does Generative AI understand programming languages?

Generative AI models learn programming languages by training on large amounts of code written in them. They pick up the syntax, the semantics, and the recurring patterns from many datasets, so the code they produce follows the conventions of the language you asked for.

Q2: Can Generative AI replace human programmers?

Generative AI can automate parts of coding, such as writing boilerplate or completing code as you type. It does not replace human programmers. You still need people to design complex algorithms, to debug, and to make the decisions that matter in software development.

Q3: Are there ethical concerns related to AI-generated code?

Yes, there are. Two of them: bias in the training data, and people using generated code for harmful purposes. Use AI-generated code with care, and check that what it hands you meets your ethical standards.

Q4: What are some practical applications of Code Generative AI?

Code Generative AI is useful for code autocompletion, code refactoring, generating documentation, and helping with code reviews. It can make developers faster and it can help with code quality.

Q5: How can I fine-tune my Code Generative AI model?

Fine-tuning a Code Generative AI model means training it further on a specific dataset or domain so it gets better at that narrower job. You fine-tune an existing model with transfer learning techniques and domain-specific data.

Conclusion

In software engineering, AI, particularly Code Generative AI, could change how developers write code. By following the steps in this article, you can build your own Code Generative AI model to assist with your coding work. That said, AI is a tool, and human expertise and ethical considerations should still guide software development.

Building your own Code Generative AI shows you exactly where it helps in your work and where it doesn’t.

References

1. Hugging Face Transformers Documentation, https://huggingface.co/docs/transformers/index
2. PEFT - Parameter-Efficient Fine-Tuning of Billion-Scale Models on Low-Resource Hardware, https://github.com/huggingface/peft
3. Llama 2 - Meta AI, https://ai.meta.com/llama/
4. “Attention Is All You Need” (Transformer paper) - Vaswani, A., et al. (2017.), https://arxiv.org/abs/1706.03762
5. “Evaluating Large Language Models Trained on Code” - Chen, M., et al. (OpenAI Codex.), https://arxiv.org/abs/2107.03374
6. Hugging Face Model Hub, https://huggingface.co/models (Specifically search for code generation models like meta-llama/Llama-2-7b-chat-hf and qblocks/llama2-7b-tiny-codes-code-generation)
7. PyTorch Official Website, https://pytorch.org/