Skip to main content
tensorgirl
Projects
PythonCodeGenerator
Reports
Building a Coding Assistant using LangChain and CodeLlama with QLoRA

Building a Coding Assistant using LangChain and CodeLlama with QLoRA

Usha Rengaraju
Created on January 21|Last edited on January 21
Code generation has become a pivotal tool for boosting productivity and efficiency. One of the latest advancements in this domain is the use of Large Language Models (LLMs) for code generation. These models, such as OpenAI's GPT-3, have demonstrated remarkable capabilities in understanding and generating human-like text, making them a revolutionary force in the world of programming. In this tutorial, we'll explore the fascinating realm of code generation using LLMs and delve into the potential they hold for streamlining development processes.


source

Understanding Large Language Models:

Large Language Models are sophisticated artificial intelligence models that have been trained on massive amounts of textual data. They leverage deep learning techniques to understand and generate human-like text based on the patterns and structures present in the training data. These models excel at natural language understanding and can generate coherent and contextually relevant text across a wide range of topics.


source

The Evolution of Code Generation:

Code generation is not a new concept in software development. Programmers have long used tools and frameworks to automate the generation of repetitive or boilerplate code, saving time and reducing the likelihood of errors. However, traditional code generation tools often lack the flexibility and adaptability required for more complex tasks.
LLMs, with their ability to comprehend context and mimic human language, take code generation to a whole new level. Developers can now provide high-level instructions or prompts to these models, and they generate code snippets that align with the specified requirements. This approach not only accelerates the development process but also allows for more natural and expressive interactions with the model.


Use Cases for Code Generation with LLMs:



Auto-Completion and Suggestions:

LLMs can assist developers by providing auto-completions and intelligent suggestions as they write code. This not only speeds up the coding process but also helps prevent common syntax errors.


Boilerplate Code Generation:

Developers often find themselves writing repetitive boilerplate code. LLMs can be trained to understand the context and generate boilerplate code snippets based on high-level instructions, freeing up developers to focus on more complex aspects of their projects.
/

Natural Language Interfaces:

LLMs enable the creation of natural language interfaces for code. Developers can describe the functionality they need in plain English, and the model translates these descriptions into executable code.


Code Summarization:

LLMs can be used to generate concise and human-readable summaries of code. This can be particularly useful for documentation purposes, making codebases more accessible and maintainable.


Challenges and Considerations:



Fine-Tuning and Specialization:

While LLMs are powerful, fine-tuning them for specific programming languages or domains can enhance their performance. Specialized training helps the models understand the intricacies of different programming paradigms.


Code Quality and Security:

Code generated by LLMs may require careful review, as the models may not always produce optimal or secure solutions. Human oversight is crucial to ensure the generated code meets quality standards and security requirements.


Ethical Considerations:

As with any AI technology, ethical considerations arise. Ensuring that code generated by LLMs adheres to ethical standards and doesn't violate privacy or security norms is essential.

Understanding Langchain

The age of language models (LLMs) is upon us, and their potential to revolutionize how we interact with technology is undeniable. But how do we harness this power? That's where LangChain comes in, a framework making LLM integration accessible and intuitive for developers of all skill levels.
Think of LangChain as a Lego set for LLMs. It provides pre-built blocks (Chains) that you can easily snap together to create powerful applications. Each Chain performs a specific task, like summarizing text, translating languages, or writing code. You can combine these Chains in endless ways to build complex workflows, all driven by the magic of LLMs.


What makes LangChain so special?

/bContext-aware: LangChain isn't just throwing prompts at an LLM. It allows you to inject rich context, like user data or external files, to ensure your responses are relevant and accurate.
Reasoning on steroids: Don't just generate text; reason with it! LangChain lets you use LLMs to make decisions, solve problems, and even perform actions in your application.
Built for production: LangChain isn't just a toy. It's designed for real-world applications, with features like monitoring, testing, and debugging to ensure your creations are robust and reliable.
So, what can you build with LangChain?

source
The possibilities are truly endless. Here are just a few examples:

  • Chatbots that understand you: Build chatbots that can carry on meaningful conversations, answer your questions, and even complete tasks for you.
  • Document summarization with a twist: Go beyond basic summaries. Use LLMs to analyze documents, extract key insights, and even generate creative interpretations.
  • Code generation made easy: Let LLMs help you write code, debug existing code, or even generate new algorithms.
  • Personalized content creation: Tailor content to individual users based on their preferences and past interactions.


Code LLAMA model

In this tutorial we will use the CodeLlama model and finetune it to use for our problem.
This model, developed by Meta AI, is designed to make the coding process more efficient, accurate, and even a little more fun. CodeLlama is a large language model (LLM) specifically trained on a massive dataset of code. It can generate code, translate languages, and answer your coding questions in an informative way. It's considered a state-of-the-art model for code generation and programming tasks.

source

Key Features and Capabilities:


Code Generation:
  • Generates code from natural language descriptions.
  • Fills in missing code blocks given a prompt and context.
  • Generates code in multiple programming languages.

Code Translation:
Translates code between different programming languages, preserving functionality.

Code Understanding:
  • Understands code semantics and can answer complex questions about code.
  • Explains code functionality and suggests improvements.

Zero-Shot Instruction Following:
Follows instructions to perform programming tasks without needing additional training examples.


Potential Applications:


  • Streamlining Development: Automating repetitive coding tasks and generating code from high-level descriptions, saving time and effort.
  • Code Education: Providing personalized code explanations and guidance, enhancing learning experiences.
  • Code Maintenance: Refactoring code, catching bugs, and suggesting improvements, boosting code quality.
  • Creative Coding: Generating new and innovative code ideas, expanding possibilities.


Using W&B

Create a W&B account and install W&B using 
pip install wandb
Then login using
wandb login


Understanding BitsAndBytes Library

Bitsandbytes is a powerful library that's making waves in the realm of AI, specifically in the training and deployment of large language models (LLMs). It offers a unique set of tools to optimize model performance and resource usage, making LLMs more accessible and efficient.
Here's a quick look at its key features:
1. Quantization:
Reduces the precision of model weights from 16-bit to 4-bit or even 8-bit, significantly decreasing memory footprint and computation time.
This means you can train and run LLMs using less hardware, opening up new possibilities for smaller devices and more affordable setups.

2. Optimized Optimizers:
Provides 8-bit optimizers that are specifically designed for quantized models.
These optimizers perform operations using 8-bit integers instead of 16-bit floats, leading to faster training and inference.

3. Matrix Multiplication:
Offers a custom 8-bit matrix multiplication function (LLM.int8()) that's highly optimized for speed and efficiency.
This function is a crucial building block for many AI computations, and its optimization can have a significant impact on overall performance.

4. Integration with Popular Frameworks:
Works seamlessly with PyTorch and Transformers, allowing easy integration into existing AI workflows.
This makes it accessible to a wide range of developers and researchers.

Benefits:
  • Reduced Memory Usage: Enables training and running larger models on devices with less memory.
  • Faster Training and Inference: Speeds up the training process and makes real-time applications more responsive.
  • Lower Energy Consumption: Can lead to lower energy consumption and longer battery life on mobile devices.
  • Increased Accessibility: Makes LLMs more accessible to researchers and developers with limited hardware resources.

Finetuning the CodeLLAMA model

For the purpose of finetuning we use the QLoRA technique. QLoRA stands for Quantization and Low-Rank Adapters, and it's a novel approach to fine-tuning large language models (LLMs). Here's a quick breakdown of its key points:
1. Problem: Fine-tuning large LLMs with their massive parameter count often requires immense computational resources, making it inaccessible or costly.
2. Solution: QLoRA tackles this by combining two techniques:
  • Quantization: Reducing the precision of model weights from 16-bit to 4-bit, significantly decreasing memory consumption.
  • Low-Rank Adapters: Introducing small, trainable matrices alongside the quantized weights to capture task-specific information, allowing for effective fine-tuning without retraining the entire model.
model_name = "codellama/CodeLlama-7b-Python-hf"
dataset_name = "lucasmccabe-lmi/CodeAlpaca-20k"

os.environ["WANDB_PROJECT"] = "PythonCodeGenerator"  # name your W&B project
os.environ["WANDB_LOG_MODEL"] = "checkpoint"  # log all model checkpoints

dataset = load_dataset(dataset_name, split="train")

# Load tokenizer and model with QLoRA configuration
compute_dtype = getattr(torch, bnb_4bit_compute_dtype)

bnb_config = BitsAndBytesConfig(
    load_in_4bit=use_4bit,
    bnb_4bit_quant_type=bnb_4bit_quant_type,
    bnb_4bit_compute_dtype=compute_dtype,
    bnb_4bit_use_double_quant=use_nested_quant,
)

# Check GPU compatibility with bfloat16
if compute_dtype == torch.float16 and use_4bit:
    major, _ = torch.cuda.get_device_capability()
    if major >= 8:
        print("=" * 80)
        print("Your GPU supports bfloat16: accelerate training with bf16=True")
        print("=" * 80)

# Load base model
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map=device_map
)
model.config.use_cache = False
model.config.pretraining_tp = 1

# Load LLaMA tokenizer
tokenizer = CodeLlamaTokenizerFast.from_pretrained("hf-internal-testing/llama-tokenizer")
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right" # Fix weird overflow issue with fp16 training

# Load LoRA configuration
peft_config = LoraConfig(
    lora_alpha=lora_alpha,
    lora_dropout=lora_dropout,
    r=lora_r,
    bias="none",
    task_type="CAUSAL_LM",
)

# Set training parameters
training_arguments = TrainingArguments(
    output_dir=output_dir,
    num_train_epochs=num_train_epochs,
    per_device_train_batch_size=per_device_train_batch_size,
    gradient_accumulation_steps=gradient_accumulation_steps,
    optim=optim,
    save_steps=save_steps,
    logging_steps=logging_steps,
    learning_rate=learning_rate,
    weight_decay=weight_decay,
    fp16=fp16,
    bf16=bf16,
    max_grad_norm=max_grad_norm,
    max_steps=max_steps,
    warmup_ratio=warmup_ratio,
    group_by_length=group_by_length,
    lr_scheduler_type=lr_scheduler_type,
    report_to="wandb"
)

def formatting_prompts_func(example):
    output_texts = []
    for i in range(len(example['instruction'])):
        text = f"### Question: {example['instruction'][i]}\n ### Answer: {example['output'][i]}"
        output_texts.append(text)
    return output_texts

response_template = "\n ### Answer:"
collator = DataCollatorForCompletionOnlyLM(response_template, tokenizer=tokenizer)


# Set supervised fine-tuning parameters
trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    peft_config=peft_config,
    formatting_func=formatting_prompts_func,
    data_collator=collator,
    args=training_arguments,
    packing=False,
)

# Train model
trainer.train()

# Save trained model
trainer.model.save_pretrained(new_model)



Building the Code Generator using Langchain

prompt = "Create a function that takes a specific input and produces a specific output using any mathematical operators. Write corresponding code in Python."
pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, max_length=200)

os.environ["LANGCHAIN_WANDB_TRACING"] = "True"
os.environ["WANDB_PROJECT"] = "PythonCodeGenerator"
hf = HuggingFacePipeline(pipeline=pipe)



template = """Create a function according to the following input. Write corresponding code in Python.
 {question}

Answer: Here is the code"""
prompt = PromptTemplate.from_template(template)

chain = prompt | hf

question = "write the code to solve a quadratic equation"

print(chain.invoke({"question": question}))

Run set
7


References

Created with ❤️ on Weights & Biases.

https://wandb.ai/tensorgirl/PythonCodeGenerator/reports/Building-a-Coding-Assistant-using-LangChain-and-CodeLlama-with-QLoRA---Vmlldzo2NTg0NTE0
Made with Weights & Biases. Sign up or log in to create reports like this one.