Fine_tuning_CodeLlama_demo.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hqCTP5D55YuD"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/googlecolab/colabtools/blob/master/notebooks/Fine-tuning-CodeLlama_demo.ipynb)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b3HJ07DT_sgw"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "os.environ['MAX_SPLIT_SIZE_GB'] = '32'\n",
        "import torch\n",
        "device = torch.device(\"cuda\")  # Make sure you're using the correct device\n",
        "mem_info = torch.cuda.memory_stats(device=device)\n",
        "memory_usage_bytes = mem_info.get(\"allocated_bytes.all.current\")\n",
        "torch.cuda.empty_cache()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HDPjMVKT3yre"
      },
      "source": [
        "#Downloading Dependencies"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bz0-DGpL3s0r"
      },
      "outputs": [],
      "source": [
        "!pip install bitsandbytes\n",
        "!pip install git+https://github.com/huggingface/transformers.git@refs/pull/25740/head accelerate\n",
        "!pip install -q -U git+https://github.com/huggingface/peft.git\n",
        "!pip install git+https://github.com/huggingface/accelerate\n",
        "!pip install trl\n",
        "!pip install einops wandb\n",
        "import transformers\n",
        "from transformers import AutoModelForCausalLM , AutoTokenizer , BitsAndBytesConfig , HfArgumentParser,TrainingArguments, pipeline, logging\n",
        "from peft import PeftModel,LoraConfig,PeftModelForCausalLM\n",
        "from trl import SFTTrainer\n",
        "from datasets import load_dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U63W7WrtlZPt"
      },
      "source": [
        "#Load dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vGNiWxH8lYpD"
      },
      "outputs": [],
      "source": [
        "dataset = load_dataset(\"nickrosh/Evol-Instruct-Code-80k-v1\" , split = \"train\")\n",
        "# Some other datasets\n",
        "# dataset = load_dataset(\"sahil2801/code_instructions_120k\" , split = \"train\")\n",
        "# dataset = load_dataset(\"iamtarun/python_code_instructions_18k_alpaca\" , split = \"train\")\n",
        "#dataset = load_dataset(\"HuggingFaceH4/CodeAlpaca_20K\",split = \"train\")\n",
        "# dataset = load_dataset(\"WizardLM/WizardLM_evol_instruct_70k\", split = \"train\")\n",
        "# dataset = load_dataset(\"mlabonne/CodeLlama-2-20k\" , split = \"train\")\n",
        "# dataset = load_dataset(\"VMware/open-instruct-v1-oasst-dolly-hhrlhf\", split = \"train\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ExF_LPi9k_r4"
      },
      "source": [
        "#Login to Huggingface\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CTfJeeA8k-BB"
      },
      "outputs": [],
      "source": [
        "!huggingface-cli login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1ypGo28yFCD3"
      },
      "source": [
        "#Tokenize the LLM"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wRgUebO8_q72"
      },
      "outputs": [],
      "source": [
        "model_name = \"codellama/CodeLlama-7b-Instruct-hf\"\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
        "# Fine-tuned model name\n",
        "new_model = \"Luffy/codellama-2-7b-Instruct-hf-Fine-tuned\"\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "47lXlkWE_q3P"
      },
      "outputs": [],
      "source": [
        "training_config = {\n",
        "    \"model\": {\n",
        "        \"pretrained_name\": model_name,\n",
        "        \"max_length\" : 2048\n",
        "    },\n",
        "    \"datasets\": {\n",
        "        # \"use_hf\": use_hf,\n",
        "        \"path\": dataset\n",
        "    },\n",
        "    \"verbose\": True # a boolean indicating whether to output detailed information during the process.\n",
        "}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cV0Mhz-X_qy8"
      },
      "outputs": [],
      "source": [
        "def tokenize_function(examples):\n",
        "    if \"question\" in examples and \"answer\" in examples:\n",
        "      text = examples[\"question\"][0] + examples[\"answer\"][0]\n",
        "    elif \"input\" in examples and \"output\" in examples:\n",
        "      text = examples[\"input\"][0] + examples[\"output\"][0]\n",
        "    elif \"instruction\" in examples and \"response\" in examples:\n",
        "      text = examples[\"instruction\"][0] + examples[\"response\"][0]\n",
        "    elif \"instruction\" in examples and \"completion\" in examples:\n",
        "      text = examples[\"instruction\"][0] + examples[\"completion\"][0]\n",
        "    elif \"instruction\" in examples and \"output\" in examples:\n",
        "      text = examples[\"instruction\"][0] + examples[\"output\"][0]\n",
        "    else:\n",
        "      text = examples[\"text\"][0]\n",
        "    tokenizer.pad_token = tokenizer.eos_token\n",
        "    tokenized_inputs = tokenizer(\n",
        "        text,\n",
        "        return_tensors=\"np\",\n",
        "        padding=True,\n",
        "    )\n",
        "\n",
        "    max_length = min(\n",
        "        tokenized_inputs[\"input_ids\"].shape[1],\n",
        "        2048\n",
        "    )\n",
        "    tokenizer.truncation_side = \"left\"\n",
        "    tokenized_inputs = tokenizer(\n",
        "        text,\n",
        "        return_tensors=\"np\",\n",
        "        truncation=True,\n",
        "        max_length=max_length\n",
        "    )\n",
        "\n",
        "    return tokenized_inputs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "muwqq3dT_qs2"
      },
      "outputs": [],
      "source": [
        "tokenized_dataset = dataset.map(\n",
        "    tokenize_function,\n",
        "    batched=True,\n",
        "    batch_size=1,\n",
        "    drop_last_batch=True\n",
        ")\n",
        "total = sum(len(sequence) for sequence in tokenized_dataset)\n",
        "print(total)\n",
        "\n",
        "print(tokenized_dataset)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uS2u1VSKFJ3A"
      },
      "source": [
        "#Spliting the Dataset into Train and Test Dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "o-pmcKIr_qlw"
      },
      "outputs": [],
      "source": [
        "total_samples = len(tokenized_dataset)\n",
        "train_ratio = 0.6  # 60% for training\n",
        "val_ratio = 0.2    # 20% for validation\n",
        "test_ratio = 0.2   # 20% for testing\n",
        "\n",
        "train_size = int(total_samples * train_ratio)\n",
        "val_size = int(total_samples * val_ratio)\n",
        "test_size = int(total_samples * test_ratio)\n",
        "\n",
        "# Define ranges for each split\n",
        "train_range = range(train_size)\n",
        "val_range = range(train_size, train_size + val_size)\n",
        "test_range = range(train_size + val_size, total_samples)\n",
        "\n",
        "# Create datasets based on the defined ranges\n",
        "train_dataset = tokenized_dataset.select(train_range)\n",
        "val_dataset = tokenized_dataset.select(val_range)\n",
        "test_dataset = tokenized_dataset.select(test_range)\n",
        "\n",
        "# Print the sizes of each dataset\n",
        "print(\"Train dataset size:\", len(train_dataset))\n",
        "print(\"Validation dataset size:\", len(val_dataset))\n",
        "print(\"Test dataset size:\", len(test_dataset))\n",
        "\n",
        "print(train_dataset)\n",
        "print(val_dataset)\n",
        "print(test_dataset)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7-8OQyBZ_qhO"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "bnb_config = BitsAndBytesConfig(\n",
        "    load_in_4bit=True,\n",
        "    bnb_4bit_use_double_quant=True,\n",
        "    bnb_4bit_quant_type=\"nf4\",\n",
        "    bnb_4bit_compute_dtype=torch.bfloat16\n",
        ")\n",
        "model = AutoModelForCausalLM.from_pretrained(\"codellama/CodeLlama-7b-Instruct-hf\", quantization_config=bnb_config, device_map={\"\":0})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2nUrehKDECtQ"
      },
      "outputs": [],
      "source": [
        "model.config.quantization_config.to_dict()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SfKUqWtx_qZo"
      },
      "outputs": [],
      "source": [
        "# QLoRA parameters\n",
        "################################################################################\n",
        "\n",
        "# LoRA attention dimension\n",
        "lora_r = 8\n",
        "\n",
        "# Alpha parameter for LoRA scaling\n",
        "lora_alpha = 16\n",
        "\n",
        "# Dropout probability for LoRA layers\n",
        "lora_dropout = 0.05\n",
        "\n",
        "################################################################################\n",
        "# bitsandbytes parameters\n",
        "################################################################################\n",
        "\n",
        "# Activate 4-bit precision base model loading\n",
        "use_4bit = True\n",
        "\n",
        "# Compute dtype for 4-bit base models\n",
        "bnb_4bit_compute_dtype = \"float16\"\n",
        "\n",
        "# Quantization type (fp4 or nf4)\n",
        "bnb_4bit_quant_type = \"nf4\"\n",
        "\n",
        "# Activate nested quantization for 4-bit base models (double quantization)\n",
        "use_nested_quant = False\n",
        "\n",
        "################################################################################\n",
        "# TrainingArguments parameters\n",
        "################################################################################\n",
        "\n",
        "# Output directory where the model predictions and checkpoints will be stored\n",
        "output_dir = \"./results\"\n",
        "\n",
        "# Number of training epochs\n",
        "num_train_epochs = 1\n",
        "\n",
        "# Enable fp16/bf16 training (set bf16 to True with an A100)\n",
        "fp16 = True\n",
        "bf16 = False\n",
        "\n",
        "# Batch size per GPU for training\n",
        "per_device_train_batch_size = 1\n",
        "\n",
        "# Batch size per GPU for evaluation\n",
        "per_device_eval_batch_size = 1\n",
        "\n",
        "# Number of update steps to accumulate the gradients for\n",
        "gradient_accumulation_steps = 4\n",
        "\n",
        "# Enable gradient checkpointing\n",
        "gradient_checkpointing = True\n",
        "\n",
        "# Maximum gradient normal (gradient clipping)\n",
        "max_grad_norm = 0.3\n",
        "\n",
        "# Initial learning rate (AdamW optimizer)\n",
        "learning_rate = 2e-5\n",
        "\n",
        "# Weight decay to apply to all layers except bias/LayerNorm weights\n",
        "weight_decay = 0.001\n",
        "\n",
        "# Optimizer to use\n",
        "optim = \"paged_adamw_8bit\"\n",
        "#  \"paged_adamw_32bit\"\n",
        "\n",
        "# Learning rate schedule (constant a bit better than cosine)\n",
        "lr_scheduler_type = \"cosine\"\n",
        "\n",
        "# Number of training steps (overrides num_train_epochs)\n",
        "max_steps = 100\n",
        "\n",
        "# Ratio of steps for a linear warmup (from 0 to learning rate)\n",
        "warmup_ratio = 0.03\n",
        "\n",
        "# Group sequences into batches with same length\n",
        "# Saves memory and speeds up training considerably\n",
        "group_by_length = True\n",
        "\n",
        "# Save checkpoint every X updates steps\n",
        "save_steps = 50\n",
        "\n",
        "# Log every X updates steps\n",
        "logging_steps = 10\n",
        "\n",
        "################################################################################\n",
        "# SFT parameters\n",
        "################################################################################\n",
        "\n",
        "# Maximum sequence length to use\n",
        "max_seq_length = 600\n",
        "\n",
        "# Pack multiple short examples in the same input sequence to increase efficiency\n",
        "packing = False\n",
        "\n",
        "# Load the entire model on the GPU 0\n",
        "device_map = {\"\": 0}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UXbV5npMDPrD"
      },
      "outputs": [],
      "source": [
        "def print_trainable_parameters(model):\n",
        "    \"\"\"\n",
        "    Prints the number of trainable parameters in the model.\n",
        "    \"\"\"\n",
        "    trainable_params = 0\n",
        "    all_param = 0\n",
        "    for _, param in model.named_parameters():\n",
        "        all_param += param.numel()\n",
        "        if param.requires_grad:\n",
        "            trainable_params += param.numel()\n",
        "    print(\n",
        "        f\"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}\"\n",
        "    )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "atqEBChjDPnM"
      },
      "outputs": [],
      "source": [
        "from peft import get_peft_model\n",
        "model.config.use_cache = False\n",
        "model.config.pretraining_tp = 1\n",
        "tokenizer.pad_token = tokenizer.eos_token\n",
        "tokenizer.padding_side = \"right\" # Fix weird overflow issue with fp16 training\n",
        "# Load LoRA configuration\n",
        "peft_config = LoraConfig(\n",
        "    lora_alpha=lora_alpha,\n",
        "    lora_dropout=lora_dropout,\n",
        "    r=lora_r,\n",
        "    bias=\"none\",\n",
        "    task_type=\"CAUSAL_LM\",\n",
        ")\n",
        "model = get_peft_model(model, peft_config)\n",
        "print_trainable_parameters(model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EWvuLyx5DPjE"
      },
      "outputs": [],
      "source": [
        "# Set training parameters\n",
        "training_arguments = TrainingArguments(\n",
        "    output_dir=output_dir,\n",
        "    num_train_epochs=num_train_epochs,\n",
        "    per_device_train_batch_size=per_device_train_batch_size,\n",
        "    gradient_accumulation_steps=gradient_accumulation_steps,\n",
        "    optim=optim,\n",
        "    save_steps=save_steps,\n",
        "    logging_steps=logging_steps,\n",
        "    learning_rate=learning_rate,\n",
        "    weight_decay=weight_decay,\n",
        "    fp16=fp16,\n",
        "    bf16=bf16,\n",
        "    max_grad_norm=max_grad_norm,\n",
        "    max_steps=max_steps,\n",
        "    warmup_ratio=warmup_ratio,\n",
        "    group_by_length=group_by_length,\n",
        "    lr_scheduler_type=lr_scheduler_type,\n",
        "    report_to=\"tensorboard\",\n",
        "    warmup_steps=30\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dQvrt9sVDPgD"
      },
      "outputs": [],
      "source": [
        "# Set supervised fine-tuning parameters\n",
        "trainer = SFTTrainer(\n",
        "    model=model,\n",
        "    train_dataset=train_dataset,\n",
        "    peft_config=peft_config,\n",
        "    dataset_text_field=\"instruction\",\n",
        "    max_seq_length=max_seq_length,\n",
        "    tokenizer=tokenizer,\n",
        "    args=training_arguments,\n",
        "    packing=packing,\n",
        "    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NLUeNV6IFqgF"
      },
      "source": [
        "#Train the LLM"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jEp6Nyw6DPcQ"
      },
      "outputs": [],
      "source": [
        "torch.cuda.empty_cache()\n",
        "# Train model\n",
        "trainer.train()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "75t1DunTQ0sJ"
      },
      "outputs": [],
      "source": [
        "# Save trained model\n",
        "trainer.model.save_pretrained(new_model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I52n996jTvZZ"
      },
      "outputs": [],
      "source": [
        "%load_ext tensorboard\n",
        "%tensorboard --logdir results/runs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IUNtsJ4oWF1H"
      },
      "outputs": [],
      "source": [
        "# Reload model in FP16 and merge it with LoRA weights\n",
        "base_model = AutoModelForCausalLM.from_pretrained(\n",
        "    model_name,\n",
        "    low_cpu_mem_usage=True,\n",
        "    return_dict=True,\n",
        "    torch_dtype=torch.float16,\n",
        "    device_map=device_map,\n",
        ")\n",
        "model = PeftModel.from_pretrained(base_model, new_model)\n",
        "model = model.merge_and_unload()\n",
        "\n",
        "# Reload tokenizer to save it\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n",
        "tokenizer.add_special_tokens({'pad_token': '[PAD]'})\n",
        "tokenizer.pad_token = tokenizer.eos_token\n",
        "tokenizer.padding_side = \"right\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GRxh7w7kDPaG"
      },
      "outputs": [],
      "source": [
        "model.eval()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SetmCK0JZ5zI"
      },
      "outputs": [],
      "source": [
        "model.push_to_hub(new_model, use_temp_dir=False)\n",
        "tokenizer.push_to_hub(new_model, use_temp_dir=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PzPwG5jVZYak"
      },
      "source": [
        "# Creating UI using Gradio\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DYG--w1pDmKs"
      },
      "outputs": [],
      "source": [
        "import gradio as gr\n",
        "from transformers import pipeline\n",
        "\n",
        "def generate_response(prompt):\n",
        "\n",
        "    system_message =\"\"\"Below is an instruction that describes a task.Write a response that appropriately completes the request.Please wrap your code answer using ``` \"\"\"\n",
        "\n",
        "    # system_message = \"\"\"You are a helpful, respectful and honest assistant.Your job is to generate python code to solve the following coding problem that obeys the constraints and you also have to give some test cases as an example and show the output.\n",
        "    # Explain the code after the code completion.Ask the user for any another queries.Please wrap your code answer using ```\"\"\"\n",
        "\n",
        "    prompt_template= f'''\n",
        "    [INST]\n",
        "    <<sys>>\n",
        "    {system_message}\n",
        "    <</sys>>\n",
        "    {prompt}\n",
        "    [/INST]\n",
        "    '''\n",
        "    # Generate a response using the pipeline\n",
        "    pipe = pipeline(\n",
        "        \"text-generation\",\n",
        "        model=model,\n",
        "        tokenizer=tokenizer,\n",
        "        max_length=1024,\n",
        "        temperature=0.3,\n",
        "        top_p=0.95,\n",
        "        repetition_penalty=1.15\n",
        "    )\n",
        "\n",
        "    generated_text = pipe(prompt_template)[0]['generated_text']\n",
        "    # Extract content between triple backticks\n",
        "    code_start = generated_text.find(\"```\")\n",
        "    code_end = generated_text.rfind(\"```\")\n",
        "    if code_start != -1 and code_end != -1:\n",
        "        generated_text = generated_text[code_start + 3:code_end].strip()\n",
        "\n",
        "    # Remove any remaining unwanted text\n",
        "    generated_text = generated_text.replace(\"<</sys>>\", \"\").replace(\"[/INST]\", \"\").strip()\n",
        "    return generated_text\n",
        "title = \"CodeLlama-13B for Code Generation \"\n",
        "examples = [\n",
        "    'Write a python code to find the Fibonacci series.',\n",
        "    'Write a python code for Merge Sort.',\n",
        "    'Write a python code for Binary search.',\n",
        "    'Write a python code for the Longest subsequence.'\n",
        "]\n",
        "\n",
        "gr.Interface(\n",
        "    fn=generate_response,\n",
        "    inputs=gr.inputs.Textbox(label=\"Enter your prompt here...\"),\n",
        "    outputs=gr.outputs.Textbox(),\n",
        "    title=title,\n",
        "    examples=examples\n",
        ").launch()\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}