{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4f8941fe"
   },
   "source": [
    "# Exercise: Fine-Tuning a Language Model with Unsloth\n",
    "\n",
    "This notebook demonstrates how to fine-tune a **Qwen2-0.5B-Instruct** language model using the **Unsloth** library.\n",
    "We will use LoRA (Low-Rank Adaptation) for efficient fine-tuning on a custom dataset.\n",
    "\n",
    "### Steps Covered:\n",
    "1. **Install Dependencies**: Ensure all required libraries are installed.\n",
    "2. **Load Model**: Download and prepare the base model from Unsloth.\n",
    "3. **Apply LoRA**: Modify the model with LoRA for parameter-efficient training.\n",
    "4. **Prepare Dataset**: Load and preprocess the dataset for fine-tuning.\n",
    "5. **Train the Model**: Fine-tune the model using **SFTTrainer**.\n",
    "6. **Save the Model**: Export the fine-tuned model for future use.\n",
    "\n",
    "---\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tXBHLQYs32m6"
   },
   "outputs": [],
   "source": [
    "# Check GPU availability\n",
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VfVgN29liGUW"
   },
   "outputs": [],
   "source": [
    "# Necessary to save into Drive\n",
    "from google.colab import drive\n",
    "drive.mount('/content/drive')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Tvv1F32oBf5p"
   },
   "outputs": [],
   "source": [
    "#Install required dependencies for Unsloth and related libraries\n",
    "!pip3 install torch==2.4.0 torchvision==0.19.0 unsloth[cu124-torch240]==2025.2.4 --extra-index-url git+https://github.com/unslothai/unsloth.git\n",
    "!pip3 install unsloth_zoo==2025.2.3 peft==0.14.0\n",
    "!pip3 install datasets==3.2.0 bitsandbytes==0.45.2\n",
    "!pip3 install trl==0.14.0 transformers==4.48.3\n",
    "!pip3 install pyopenssl==25.0.0 accelerate==1.3.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Q30kIwhEBf5s"
   },
   "outputs": [],
   "source": [
    "# Import necessary libraries for model training and dataset preparation\n",
    "from unsloth import FastLanguageModel\n",
    "\n",
    "# For dataset\n",
    "from datasets import load_dataset\n",
    "from unsloth.chat_templates import get_chat_template\n",
    "\n",
    "# For training\n",
    "from trl import SFTTrainer, SFTConfig\n",
    "from unsloth import is_bfloat16_supported"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "E5h3o9c0Bf5s"
   },
   "outputs": [],
   "source": [
    "# Define model configuration parameters\n",
    "max_seq_length = 4096  # Choose any! We auto support RoPE Scaling internally!\n",
    "dtype = None  # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+\n",
    "load_in_4bit = True # Use 4bit quantization to reduce memory usage. Can be False."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "id": "XBEmIg66Bf5t"
   },
   "outputs": [],
   "source": [
    "# Define the repository and model name for loading the pre-trained model\n",
    "repo_name = \"unsloth\"  # Repository containing the model\n",
    "model_name = \"Qwen2.5-1.5B\" # Model Name\n",
    "# Load the pre-trained language model and tokenizer\n",
    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
    "    model_name=f\"{repo_name}/{model_name}\",  # Construct the full model path\n",
    "    max_seq_length=max_seq_length,  # Set the maximum sequence length for tokenization\n",
    "    dtype=dtype,  # Define the data type (e.g., float16, bfloat16) or auto-detect\n",
    "    load_in_4bit=load_in_4bit  # Enable 4-bit quantization to reduce memory usage\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9xZOBVR0Bf5t"
   },
   "outputs": [],
   "source": [
    "# Apply LoRA (Low-Rank Adaptation) to the model for efficient fine-tuning.\n",
    "# LoRA reduces the number of trainable parameters, making fine-tuning\n",
    "# more memory-efficient while preserving model performance.\n",
    "model = FastLanguageModel.get_peft_model(\n",
    "    model,\n",
    "\n",
    "    # LoRA rank: Determines the number of learnable parameters per layer.\n",
    "    # Higher values increase expressiveness but also memory usage.\n",
    "    r=128,  # Common values: 8, 16, 32, 64, 128\n",
    "\n",
    "    # List of model layers to which LoRA will be applied.\n",
    "    # These layers are typically key components in transformer-based models.\n",
    "    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n",
    "                    \"gate_proj\", \"up_proj\", \"down_proj\",\n",
    "                    \"embed_tokens\", \"lm_head\"],\n",
    "\n",
    "    # Scaling factor for LoRA. It influences the learning rate adaptation.\n",
    "    lora_alpha=256,\n",
    "\n",
    "    # Dropout rate for LoRA layers. Setting it to 0 is optimized for stability.\n",
    "    lora_dropout=0.1,  # A non-zero value can improve generalization in some cases.\n",
    "\n",
    "    # Whether to train bias parameters. \"none\" is optimized for efficiency.\n",
    "    bias=\"none\",\n",
    "\n",
    "    # Enable gradient checkpointing to reduce memory usage for long context models.\n",
    "    # \"unsloth\" is a specialized version optimized for Unsloth models.\n",
    "    use_gradient_checkpointing=\"unsloth\",\n",
    "\n",
    "    # Whether to use Rank-Stabilized LoRA (rslora), which adapts LoRA dynamically.\n",
    "    use_rslora=False,\n",
    "\n",
    "    # Configuration for LoftQ (Low-rank Quantization), which reduces model size.\n",
    "    # Setting it to None disables LoftQ.\n",
    "    loftq_config=None,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "X2RD6stNBf5u"
   },
   "outputs": [],
   "source": [
    "# Prepare dataset by formatting prompts for training\n",
    "tokenizer = get_chat_template(tokenizer, chat_template=\"qwen-2.5\")\n",
    "\n",
    "def formatting_prompts_func(examples):\n",
    "    prompt = examples[\"input\"]\n",
    "    output = examples[\"output\"]\n",
    "\n",
    "    messages = []\n",
    "    for p, o in zip(prompt, output):\n",
    "        el = [\n",
    "              # {\"role\": \"system\", \"content\": \"Eres un asistente experto en Leonés\"},\n",
    "              {\"role\": \"user\", \"content\": p},\n",
    "              {\"role\": \"assistant\", \"content\": o},\n",
    "            ]\n",
    "        messages.append(el)\n",
    "\n",
    "    texts = [tokenizer.apply_chat_template(ele, tokenize=False, add_generation_prompt = False).strip() + tokenizer.eos_token for ele in messages]\n",
    "    return {\"text\": texts, }\n",
    "\n",
    "\n",
    "dataset = load_dataset(\"unileon-robotics/lliones-dict-tr\", split=\"train\")\n",
    "dataset_p = dataset.map(formatting_prompts_func, batched=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "T3eUm1GEeJ2i"
   },
   "outputs": [],
   "source": [
    "dataset_p[0][\"text\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oKHDa7-vBf5u"
   },
   "outputs": [],
   "source": [
    "# Configure the training process using SFTTrainer (Supervised Fine-Tuning Trainer)\n",
    "trainer = SFTTrainer(\n",
    "    model=model,  # The model to be fine-tuned\n",
    "    tokenizer=tokenizer,  # Tokenizer used for text processing\n",
    "\n",
    "    # Datasets for training and evaluation\n",
    "    train_dataset=dataset_p,  # Training dataset\n",
    "\n",
    "    # Training arguments\n",
    "    args=SFTConfig(\n",
    "        dataset_text_field=\"text\",  # Field name containing text data in the dataset\n",
    "        max_seq_length=max_seq_length,  # Maximum sequence length for input text\n",
    "\n",
    "        dataset_num_proc=4,  # Number of CPU processes for dataset preprocessing\n",
    "        packing=False,  # Whether to concatenate multiple examples into a single sequence\n",
    "\n",
    "        # Training batch size per GPU/TPU/CPU\n",
    "        per_device_train_batch_size=8,\n",
    "\n",
    "        # Number of steps to accumulate gradients before performing a backward pass\n",
    "        gradient_accumulation_steps=1,\n",
    "\n",
    "        # Number of training epochs\n",
    "        num_train_epochs=3,\n",
    "        # Number of steps\n",
    "        # max_steps=200,\n",
    "\n",
    "        # Number of warmup steps for the learning rate scheduler\n",
    "        warmup_ratio=0.1,\n",
    "\n",
    "        # Learning rate for the optimizer\n",
    "        learning_rate=2e-4,\n",
    "\n",
    "\n",
    "        # Use 16-bit floating-point precision (FP16) if BFloat16 is not supported\n",
    "        fp16=not is_bfloat16_supported(),\n",
    "        bf16=is_bfloat16_supported(),  # Use BFloat16 if the hardware supports it\n",
    "\n",
    "        # Frequency of logging training progress\n",
    "        logging_steps=500,\n",
    "\n",
    "        # Optimizer type; \"adamw_8bit\" reduces memory usage\n",
    "        optim=\"paged_adamw_8bit\",\n",
    "\n",
    "        # Weight decay for regularization (helps prevent overfitting)\n",
    "        weight_decay=0.01,\n",
    "\n",
    "        # Type of learning rate scheduler (linear decay in this case)\n",
    "        lr_scheduler_type=\"linear\",\n",
    "\n",
    "        # Directory to save training outputs (e.g., checkpoints, logs)\n",
    "        output_dir=f\"outputs/train/{model_name}\",\n",
    "\n",
    "        save_strategy=\"no\",\n",
    "\n",
    "        # Disable reporting to external loggers (e.g., WandB, TensorBoard)\n",
    "        report_to=\"none\",\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qWFTyqQTBf5v"
   },
   "outputs": [],
   "source": [
    "# Start the fine-tuning process\n",
    "trainer_stats = trainer.train(resume_from_checkpoint = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1jywPsc9s-Os"
   },
   "outputs": [],
   "source": [
    "!pip install mistral-common"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SPX0H4b2aj2b"
   },
   "outputs": [],
   "source": [
    "# Save the quantized fine-tuned model for later use\n",
    "model.save_pretrained_gguf(f\"outputs/gguf/{model_name}\", tokenizer, quantization_method = \"q5_k_m\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HrLxqNhftpIq"
   },
   "outputs": [],
   "source": [
    "FastLanguageModel.for_training(model) # Enable native 2x faster inference\n",
    "\n",
    "# Save the LoRA model for later use\n",
    "model.save_pretrained_merged(f\"outputs/loras/{model_name}\", tokenizer, save_method=\"lora\")"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}