StableCode Tutorial: How to get started with StableCode

What is StableCode from Stability AI?

StableCode the latest offering from Stability AI, is an innovative generative AI product designed to enhance the coding experience for developers at all levels. It serves as a powerful tool for both experienced programmers seeking efficiency and newcomers looking to strengthen their coding skills.

Base Model:

The foundation of StableCode is a comprehensive model that underwent initial training on a wide range of programming languages, sourced from the stack-dataset (v1.2) from BigCode. To refine its capabilities, the base model was further trained using popular languages such as Python, Go, Java, JavaScript, C, Markdown, and C++. This training involved a substantial dataset, comprising a staggering 560 billion tokens of code. This robust foundation equips StableCode with a deep understanding of various programming languages and structures.

Instruction Model:

This model was meticulously fine-tuned for specific use cases, focusing on solving intricate programming challenges. By exposing it to around 120,000 pairs of code instruction and corresponding responses in Alpaca format, the instruction model has been sharpened to provide intelligent solutions for complex coding tasks.

Long-Context Window Model:

StableCode introduces an advanced long-context window model that excels at generating single and multi-line autocomplete suggestions. Compared to previous open models with limited context windows, this new model is designed to handle significantly more code at once—approximately 2 to 4 times more. As a result, developers can effortlessly review or edit the equivalent of multiple average-sized Python files concurrently. This extended context window is particularly beneficial for those eager to expand their coding expertise and take on larger coding challenges.

What we're going to do in this tutorial?

In this tutorial, we will learn how to use StableCode to generate code completion. We will go through each model and see how it works. We will also learn how to use StableCode in Google Colab and Hugging Face Inference API to run StableCode even if you don't have really good GPU.

Implementation in Google Colab

Step 1 - Setting up the project

Let's start by creating new Notebook in Google Colab.

Go to Google Colab and create new Notebook and name it StableCode Tutorial.

Step 2 - Install required packages

First of all, we need to set Runtime type to Python 3 and Hardware accelerator to GPU.

Let's install or update Python packages related to natural language processing (NLP) and machine learning:

!pip install transformers

Click Run button or CMD/CTRL + Enter to run the single code cell.

Wait until the installation is complete. Then, you are ready to go to the next step.

Step 3 - StableCode - Base Model

Let's try StableCode - Base Model.

Add new code cell.

from transformers import AutoModelForCausalLM, AutoTokenizer

# Define all available models
BASE_MODEL = "stabilityai/stablecode-completion-alpha-3b-4k"
INSTRUCTION_MODEL = "stabilityai/stablecode-instruct-alpha-3b"
LONG_CONTEXT_WINDOW_MODEL = "stabilityai/stablecode-completion-alpha-3b"

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    trust_remote_code=True,
    torch_dtype="auto",
)
model.cuda()

Click Run or CMD/CTRL + Enter.

Wait until the model is loaded. Usually it takes more longer depends on your internet connection.

Next, let's define a function to run the model. This function will take prompt as input and return the result generated by StableCode.

Add new code cell.

def run_model(prompt):
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    del inputs["token_type_ids"]
    tokens = model.generate(
        **inputs,
        max_new_tokens=96,  # feel free to change this value
        temperature=0.2,  # feel free to change this value
        do_sample=True,
    )

    return tokenizer.decode(tokens[0], skip_special_tokens=True)

Click Run or CMD/CTRL + Enter.

Give your desired prompt for completion.

Add new code cell.

r = run_model("import torch\nimport torch.nn as nn")
print(r)

Click Run or CMD/CTRL + Enter.

You should see something similar to the following output.

Step 4 - StableCode - Instruction Model

Let's try StableCode - Instruction Model. Just change BASE_MODEL to INSTRUCTION_MODEL in the from_pretrained() function.

Add new code cell.

from transformers import AutoModelForCausalLM, AutoTokenizer

# Define all available models
BASE_MODEL = "stabilityai/stablecode-completion-alpha-3b-4k"
INSTRUCTION_MODEL = "stabilityai/stablecode-instruct-alpha-3b"
LONG_CONTEXT_WINDOW_MODEL = "stabilityai/stablecode-completion-alpha-3b"

tokenizer = AutoTokenizer.from_pretrained(INSTRUCTION_MODEL)
model = AutoModelForCausalLM.from_pretrained(
    INSTRUCTION_MODEL,
    trust_remote_code=True,
    torch_dtype="auto",
)
model.cuda()

Click Run or CMD/CTRL + Enter.

Again, wait until the model is loaded. Then, give your desired prompt for completion.

Add new code cell.

r = run_model("Write a python program to perform binary search in a given list")
print(r)

Click Run or CMD/CTRL + Enter.

You should see something similar to the following output.

Step 5 - StableCode - Long Context Window Model

Let's try StableCode - Long Context Window Model. To do that, just change INSTRUCTION_MODEL to LONG_CONTEXT_WINDOW_MODEL in the from_pretrained() function.

Add new code cell.

from transformers import AutoModelForCausalLM, AutoTokenizer

# Define all available models
BASE_MODEL = "stabilityai/stablecode-completion-alpha-3b-4k"
INSTRUCTION_MODEL = "stabilityai/stablecode-instruct-alpha-3b"
LONG_CONTEXT_WINDOW_MODEL = "stabilityai/stablecode-completion-alpha-3b"

tokenizer = AutoTokenizer.from_pretrained(LONG_CONTEXT_WINDOW_MODEL)
model = AutoModelForCausalLM.from_pretrained(
    LONG_CONTEXT_WINDOW_MODEL,
    trust_remote_code=True,
    torch_dtype="auto",
)
model.cuda()

Click Run or CMD/CTRL + Enter.

Again, wait until the model is loaded. Then, give your desired prompt for completion.

Add new code cell.

r = run_model("def quicksort(arr, reverse=False):")
print(r)

Click Run or CMD/CTRL + Enter.

You should see something similar to the following output.

Implementation with Hugging Face Inference API

Alternatively you can use Hugging Face Inference API to run StableCode. This is cool when you don't have really good GPU.

Step 1 - Create an account in Hugging Face

Go to Hugging Face and create new account or login if you already have one.

Step 2 - Create a new token in Hugging Face

Next, you need a token to use Hugging Face Inference API. Go to your profile and click Access tokens in the left sidebar, then New token button.

Give your token a name, from dropdown select read and hit Generate a token.

Perfect! You are good to go to the next step.

Step 3 - Running the StableCode with Hugging Face Inference API

Go to StableCode model page and click Deploy.

Then from dropdown select Inference API. It will open a dialog with code snippet, just hit the Copy.

Cool! Now, you can use StableCode even if you don't have really good GPU.

Conclusion

Thank you for following along with this tutorial.

If you have any questions, feel free to reach out to me on LinkedIn or Twitter. I'd love to hear from you!

made with 💜 by abdibrokhim for lablab.ai tutorials.