{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lab 2: Fine-Tuning LLM with LoRA (ID2223)\n",
    "\n",
    "This notebook is structured into separate sections:\n",
    "\n",
    "**PART 1: Installation and Configuration** - Always run this first\n",
    "\n",
    "**PART 2: Training Pipeline** - Run this section to fine-tune a new model from scratch\n",
    "\n",
    "**PART 3: Load Existing LoRA and Export** - Run this section to load previously trained LoRA weights and export to HuggingFace/GGUF\n",
    "\n",
    "**PART 4: Inference and Testing** - Test your model\n",
    "\n",
    "---\n",
    "\n",
    "Base model: `unsloth/Llama-3.2-3B-Instruct`\n",
    "\n",
    "Dataset: [FineTome-100k](https://huggingface.co/datasets/mlabonne/FineTome-100k)\n",
    "\n",
    "---\n",
    "\n",
    "Features:\n",
    "1. Uses Maxime Labonne's FineTome 100K dataset\n",
    "2. Convert ShareGPT to HuggingFace format via `standardize_sharegpt`\n",
    "3. Train on Completions / Assistant only via `train_on_responses_only`\n",
    "4. Checkpoint saving every 500 steps for resumable training\n",
    "5. Export to FP16 and GGUF formats for deployment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# PART 1: Installation and Configuration\n",
    "---\n",
    "\n",
    "Run this section first regardless of whether you are training or loading an existing model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.1 Install Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "%uv pip install unsloth\n",
    "# Also get the latest nightly Unsloth!\n",
    "%uv pip uninstall unsloth -y && pip install --upgrade --no-cache-dir --no-deps git+https://github.com/unslothai/unsloth.git@nightly git+https://github.com/unslothai/unsloth-zoo.git"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.2 Configuration\n",
    "\n",
    "Set your model parameters and paths here. These are used throughout the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "# Model configuration\n",
    "BASE_MODEL_NAME = \"unsloth/Llama-3.2-3B-Instruct\"\n",
    "max_seq_length = 2048\n",
    "dtype = \"float16\"  # Float16 for Tesla T4, V100. Use None for auto detection.\n",
    "load_in_4bit = True\n",
    "\n",
    "# LoRA configuration\n",
    "LORA_R = 16\n",
    "LORA_ALPHA = 16\n",
    "LORA_DROPOUT = 0\n",
    "LORA_TARGET_MODULES = [\n",
    "    \"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n",
    "    \"gate_proj\", \"up_proj\", \"down_proj\",\n",
    "]\n",
    "\n",
    "# Paths for saving/loading\n",
    "OUTPUT_DIR = \"/vol\"  # Training checkpoints directory\n",
    "LORA_ADAPTER_PATH = \"/vol/checkpoint-10688\"  # Path to saved LoRA adapter (for loading)\n",
    "MERGED_MODEL_DIR = \"/vol/merged-model\"  # Path for merged FP16 model\n",
    "GGUF_MODEL_DIR = \"/vol/gguf-model\"  # Path for GGUF model\n",
    "\n",
    "# HuggingFace configuration\n",
    "HF_REPO_ID = \"Jeppcode/ScalableLab2\"  # Your HuggingFace repo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.3 Supported Models\n",
    "\n",
    "We support Llama, Mistral, Phi-3, Gemma, Yi, DeepSeek, Qwen, TinyLlama, Vicuna, Open Hermes etc.\n",
    "We support 16bit LoRA or 4bit QLoRA. Both 2x faster.\n",
    "`max_seq_length` can be set to anything, since we do automatic RoPE Scaling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 4bit pre quantized models we support for 4x faster downloading + no OOMs.\n",
    "fourbit_models = [\n",
    "    \"unsloth/Meta-Llama-3.1-8B-bnb-4bit\",\n",
    "    \"unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit\",\n",
    "    \"unsloth/Meta-Llama-3.1-70B-bnb-4bit\",\n",
    "    \"unsloth/Meta-Llama-3.1-405B-bnb-4bit\",\n",
    "    \"unsloth/Mistral-Small-Instruct-2409\",\n",
    "    \"unsloth/mistral-7b-instruct-v0.3-bnb-4bit\",\n",
    "    \"unsloth/Phi-3.5-mini-instruct\",\n",
    "    \"unsloth/Phi-3-medium-4k-instruct\",\n",
    "    \"unsloth/gemma-2-9b-bnb-4bit\",\n",
    "    \"unsloth/gemma-2-27b-bnb-4bit\",\n",
    "    \"unsloth/Llama-3.2-1B-bnb-4bit\",\n",
    "    \"unsloth/Llama-3.2-1B-Instruct-bnb-4bit\",\n",
    "    \"unsloth/Llama-3.2-3B-bnb-4bit\",\n",
    "    \"unsloth/Llama-3.2-3B-Instruct-bnb-4bit\",\n",
    "    \"unsloth/Llama-3.3-70B-Instruct-bnb-4bit\"\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# PART 2: Training Pipeline\n",
    "---\n",
    "\n",
    "**Run this section to fine-tune a new model from scratch.**\n",
    "\n",
    "Skip this section if you already have trained LoRA weights and want to load/export them (go to PART 3)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.1 Load Base Model for Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from unsloth import FastLanguageModel\n",
    "\n",
    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
    "    model_name = BASE_MODEL_NAME,\n",
    "    max_seq_length = max_seq_length,\n",
    "    dtype = dtype,\n",
    "    load_in_4bit = load_in_4bit,\n",
    "    # token = \"hf_...\",  # Use if accessing gated models\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.2 Add LoRA Adapters\n",
    "\n",
    "We add LoRA adapters so we only need to update 1-10% of all parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = FastLanguageModel.get_peft_model(\n",
    "    model,\n",
    "    r = LORA_R,\n",
    "    target_modules = LORA_TARGET_MODULES,\n",
    "    lora_alpha = LORA_ALPHA,\n",
    "    lora_dropout = LORA_DROPOUT,\n",
    "    bias = \"none\",\n",
    "    use_gradient_checkpointing = \"unsloth\",\n",
    "    random_state = 3407,\n",
    "    use_rslora = False,\n",
    "    loftq_config = None,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.3 Data Preparation\n",
    "\n",
    "We use the Llama-3.1 format for conversation style finetunes. We use Maxime Labonne's FineTome-100k dataset in ShareGPT style and convert it to HuggingFace's normal multiturn format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from unsloth.chat_templates import get_chat_template\n",
    "\n",
    "tokenizer = get_chat_template(\n",
    "    tokenizer,\n",
    "    chat_template = \"llama-3.1\",\n",
    ")\n",
    "\n",
    "def formatting_prompts_func(examples):\n",
    "    convos = examples[\"conversations\"]\n",
    "    texts = [tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False) for convo in convos]\n",
    "    return { \"text\" : texts, }\n",
    "\n",
    "from datasets import load_dataset\n",
    "dataset = load_dataset(\"mlabonne/FineTome-100k\", split = \"train\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Standardize ShareGPT Format\n",
    "\n",
    "Convert ShareGPT style datasets into HuggingFace's generic format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from unsloth.chat_templates import standardize_sharegpt\n",
    "dataset = standardize_sharegpt(dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Split Dataset\n",
    "\n",
    "Split into train (85%), validation (5%), and test (10%) sets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First split: train+val (90%), test (10%)\n",
    "train_val_split = dataset.train_test_split(test_size=0.10, seed=42)\n",
    "test_dataset = train_val_split[\"test\"]\n",
    "train_val_dataset = train_val_split[\"train\"]\n",
    "\n",
    "# Second split: train (95%), val (5%) of the remaining 90%\n",
    "train_valid = train_val_dataset.train_test_split(test_size=0.05, seed=42)\n",
    "train_dataset = train_valid[\"train\"]\n",
    "valid_dataset = train_valid[\"test\"]\n",
    "\n",
    "# Apply formatting\n",
    "train_dataset = train_dataset.map(formatting_prompts_func, batched=True)\n",
    "valid_dataset = valid_dataset.map(formatting_prompts_func, batched=True)\n",
    "test_dataset  = test_dataset.map(formatting_prompts_func, batched=True)\n",
    "\n",
    "print(f\"Train: {len(train_dataset)}, Valid: {len(valid_dataset)}, Test: {len(test_dataset)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inspect Dataset (Optional)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View conversation structure\n",
    "print(\"Conversation structure:\")\n",
    "print(train_dataset[5][\"conversations\"])\n",
    "print(\"\\n\" + \"=\"*50 + \"\\n\")\n",
    "print(\"Formatted text:\")\n",
    "print(train_dataset[5][\"text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.4 Setup Trainer\n",
    "\n",
    "Using HuggingFace TRL's SFTTrainer with checkpointing enabled."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from trl import SFTTrainer\n",
    "from transformers import TrainingArguments, DataCollatorForSeq2Seq\n",
    "from unsloth import is_bfloat16_supported\n",
    "\n",
    "trainer = SFTTrainer(\n",
    "    model = model,\n",
    "    tokenizer = tokenizer,\n",
    "    train_dataset = train_dataset,\n",
    "    eval_dataset = valid_dataset,\n",
    "    dataset_text_field = \"text\",\n",
    "    max_seq_length = max_seq_length,\n",
    "    data_collator = DataCollatorForSeq2Seq(tokenizer = tokenizer),\n",
    "    dataset_num_proc = 2,\n",
    "    packing = False,\n",
    "    args = TrainingArguments(\n",
    "        per_device_train_batch_size = 2,\n",
    "        gradient_accumulation_steps = 4,\n",
    "        warmup_steps = 5,\n",
    "        num_train_epochs = 1,\n",
    "        # max_steps = 60,  # Uncomment for quick test runs\n",
    "        learning_rate = 2e-4,\n",
    "        fp16 = not is_bfloat16_supported(),\n",
    "        bf16 = is_bfloat16_supported(),\n",
    "        logging_steps = 50,\n",
    "        optim = \"adamw_8bit\",\n",
    "        weight_decay = 0.01,\n",
    "        lr_scheduler_type = \"linear\",\n",
    "        seed = 3407,\n",
    "        output_dir = OUTPUT_DIR,\n",
    "        report_to = \"none\",\n",
    "        \n",
    "        # Checkpointing - saves every 500 steps\n",
    "        save_strategy = \"steps\",\n",
    "        save_steps = 500,\n",
    "        save_total_limit = 5,\n",
    "        \n",
    "        # Evaluation\n",
    "        eval_steps = 500,\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train on Responses Only\n",
    "\n",
    "Only train on the assistant outputs, ignore the loss on user inputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from unsloth.chat_templates import train_on_responses_only\n",
    "trainer = train_on_responses_only(\n",
    "    trainer,\n",
    "    instruction_part = \"<|start_header_id|>user<|end_header_id|>\\n\\n\",\n",
    "    response_part = \"<|start_header_id|>assistant<|end_header_id|>\\n\\n\",\n",
    "    num_proc = 1,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Verify Masking (Optional)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check that system and instruction prompts are masked\n",
    "space = tokenizer(\" \", add_special_tokens = False).input_ids[0]\n",
    "print(\"Original:\")\n",
    "print(tokenizer.decode(trainer.train_dataset[5][\"input_ids\"]))\n",
    "print(\"\\n\" + \"=\"*50 + \"\\n\")\n",
    "print(\"Masked (spaces show masked tokens):\")\n",
    "print(tokenizer.decode([space if x == -100 else x for x in trainer.train_dataset[5][\"labels\"]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.5 Train the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show current memory stats\n",
    "gpu_stats = torch.cuda.get_device_properties(0)\n",
    "start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)\n",
    "max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)\n",
    "print(f\"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.\")\n",
    "print(f\"{start_gpu_memory} GB of memory reserved.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer_stats = trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show final memory and time stats\n",
    "used_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)\n",
    "used_memory_for_lora = round(used_memory - start_gpu_memory, 3)\n",
    "used_percentage = round(used_memory / max_memory * 100, 3)\n",
    "lora_percentage = round(used_memory_for_lora / max_memory * 100, 3)\n",
    "print(f\"{trainer_stats.metrics['train_runtime']} seconds used for training.\")\n",
    "print(f\"{round(trainer_stats.metrics['train_runtime']/60, 2)} minutes used for training.\")\n",
    "print(f\"Peak reserved memory = {used_memory} GB.\")\n",
    "print(f\"Peak reserved memory for training = {used_memory_for_lora} GB.\")\n",
    "print(f\"Peak reserved memory % of max memory = {used_percentage} %.\")\n",
    "print(f\"Peak reserved memory for training % of max memory = {lora_percentage} %.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.6 Save LoRA Adapters\n",
    "\n",
    "Save the trained LoRA adapters locally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_pretrained(f\"{OUTPUT_DIR}/lora-final\")\n",
    "tokenizer.save_pretrained(f\"{OUTPUT_DIR}/lora-final\")\n",
    "print(f\"LoRA adapters saved to {OUTPUT_DIR}/lora-final\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Upload LoRA Adapters to HuggingFace (Optional)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import login, HfApi\n",
    "login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "api = HfApi()\n",
    "api.upload_folder(\n",
    "    folder_path=f\"{OUTPUT_DIR}/lora-final\",\n",
    "    repo_id=HF_REPO_ID,\n",
    "    path_in_repo=\"lora_adapters\",\n",
    ")\n",
    "print(f\"Uploaded to {HF_REPO_ID}/lora_adapters\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# PART 3: Load Existing LoRA and Export\n",
    "---\n",
    "\n",
    "**Run this section if you already have trained LoRA weights and want to:**\n",
    "- Load the base model + LoRA adapters\n",
    "- Merge into a full FP16 model\n",
    "- Export to GGUF format for CPU inference\n",
    "- Upload to HuggingFace\n",
    "\n",
    "**Skip this section if you just trained a model in PART 2 and it is still in memory - go directly to PART 4 for inference.**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.1 Load Base Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from unsloth import FastLanguageModel\n",
    "\n",
    "base_model, tokenizer = FastLanguageModel.from_pretrained(\n",
    "    model_name = BASE_MODEL_NAME,\n",
    "    max_seq_length = max_seq_length,\n",
    "    dtype = dtype,\n",
    "    load_in_4bit = load_in_4bit,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.2 Recreate LoRA Structure and Load Weights\n",
    "\n",
    "We need to recreate the same LoRA structure that was used during training, then load the saved weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Recreate LoRA layers with the same configuration used during training\n",
    "model = FastLanguageModel.get_peft_model(\n",
    "    base_model,\n",
    "    r = LORA_R,\n",
    "    target_modules = LORA_TARGET_MODULES,\n",
    "    lora_alpha = LORA_ALPHA,\n",
    "    lora_dropout = LORA_DROPOUT,\n",
    "    bias = \"none\",\n",
    "    use_gradient_checkpointing = \"unsloth\",\n",
    "    random_state = 3407,\n",
    ")\n",
    "\n",
    "print(\"Empty LoRA structure recreated.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the saved LoRA adapter weights\n",
    "model.load_adapter(LORA_ADAPTER_PATH, adapter_name=\"default\")\n",
    "model.set_adapter(\"default\")\n",
    "\n",
    "print(f\"LoRA adapter loaded from: {LORA_ADAPTER_PATH}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.3 Merge LoRA into Full Model (FP16)\n",
    "\n",
    "Merge the LoRA adapters into the base model to create a standalone FP16 model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_pretrained_merged(\n",
    "    MERGED_MODEL_DIR,\n",
    "    tokenizer,\n",
    "    save_method=\"merged_16bit\",\n",
    ")\n",
    "\n",
    "print(f\"Merged FP16 model saved at: {MERGED_MODEL_DIR}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.4 Export to GGUF Format\n",
    "\n",
    "Export to GGUF format for CPU inference (e.g., with llama.cpp, Ollama, GPT4All).\n",
    "\n",
    "Quantization options:\n",
    "- `q8_0` - Fast conversion, high quality\n",
    "- `q4_k_m` - Recommended balance of size and quality\n",
    "- `q5_k_m` - Better quality than q4_k_m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_pretrained_gguf(\n",
    "    GGUF_MODEL_DIR,\n",
    "    tokenizer,\n",
    "    quantization_method=\"q4_k_m\",\n",
    ")\n",
    "\n",
    "print(f\"GGUF model saved at: {GGUF_MODEL_DIR}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.5 Upload to HuggingFace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import login, HfApi\n",
    "login()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Upload Merged FP16 Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "api = HfApi()\n",
    "api.upload_folder(\n",
    "    folder_path=MERGED_MODEL_DIR,\n",
    "    repo_id=HF_REPO_ID,\n",
    "    path_in_repo=\"merged-model-fp16\",\n",
    ")\n",
    "print(f\"Merged model uploaded to {HF_REPO_ID}/merged-model-fp16\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Upload GGUF Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "api = HfApi()\n",
    "api.upload_folder(\n",
    "    folder_path=GGUF_MODEL_DIR,\n",
    "    repo_id=HF_REPO_ID,\n",
    "    path_in_repo=\"gguf-model\",\n",
    ")\n",
    "print(f\"GGUF model uploaded to {HF_REPO_ID}/gguf-model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# PART 4: Inference and Testing\n",
    "---\n",
    "\n",
    "Test your fine-tuned model. This works with the model in memory from either PART 2 or PART 3."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4.1 Basic Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from unsloth.chat_templates import get_chat_template\n",
    "\n",
    "tokenizer = get_chat_template(\n",
    "    tokenizer,\n",
    "    chat_template = \"llama-3.1\",\n",
    ")\n",
    "FastLanguageModel.for_inference(model)\n",
    "\n",
    "messages = [\n",
    "    {\"role\": \"user\", \"content\": \"Continue the fibonacci sequence: 1, 1, 2, 3, 5, 8,\"},\n",
    "]\n",
    "inputs = tokenizer.apply_chat_template(\n",
    "    messages,\n",
    "    tokenize = True,\n",
    "    add_generation_prompt = True,\n",
    "    return_tensors = \"pt\",\n",
    ").to(\"cuda\")\n",
    "\n",
    "outputs = model.generate(input_ids = inputs, max_new_tokens = 64, use_cache = True,\n",
    "                         temperature = 1.5, min_p = 0.1)\n",
    "tokenizer.batch_decode(outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4.2 Streaming Inference\n",
    "\n",
    "Use TextStreamer to see generation token by token."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "FastLanguageModel.for_inference(model)\n",
    "\n",
    "messages = [\n",
    "    {\"role\": \"user\", \"content\": \"Explain what machine learning is in simple terms.\"},\n",
    "]\n",
    "inputs = tokenizer.apply_chat_template(\n",
    "    messages,\n",
    "    tokenize = True,\n",
    "    add_generation_prompt = True,\n",
    "    return_tensors = \"pt\",\n",
    ").to(\"cuda\")\n",
    "\n",
    "from transformers import TextStreamer\n",
    "text_streamer = TextStreamer(tokenizer, skip_prompt = True)\n",
    "_ = model.generate(input_ids = inputs, streamer = text_streamer, max_new_tokens = 128,\n",
    "                   use_cache = True, temperature = 1.5, min_p = 0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Additional Options\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Alternative: Load LoRA from Local Directory for Inference\n",
    "\n",
    "If you saved LoRA adapters and want to load them directly for inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if False:  # Set to True to run\n",
    "    from unsloth import FastLanguageModel\n",
    "    model, tokenizer = FastLanguageModel.from_pretrained(\n",
    "        model_name = \"lora_model\",  # Path to saved LoRA model\n",
    "        max_seq_length = max_seq_length,\n",
    "        dtype = dtype,\n",
    "        load_in_4bit = load_in_4bit,\n",
    "    )\n",
    "    FastLanguageModel.for_inference(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Alternative: Save/Upload with Different Methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Merge to 16bit\n",
    "if False: model.save_pretrained_merged(\"model\", tokenizer, save_method = \"merged_16bit\",)\n",
    "if False: model.push_to_hub_merged(\"hf/model\", tokenizer, save_method = \"merged_16bit\", token = \"\")\n",
    "\n",
    "# Merge to 4bit\n",
    "if False: model.save_pretrained_merged(\"model\", tokenizer, save_method = \"merged_4bit\",)\n",
    "if False: model.push_to_hub_merged(\"hf/model\", tokenizer, save_method = \"merged_4bit\", token = \"\")\n",
    "\n",
    "# Just LoRA adapters\n",
    "if False: model.save_pretrained_merged(\"model\", tokenizer, save_method = \"lora\",)\n",
    "if False: model.push_to_hub_merged(\"hf/model\", tokenizer, save_method = \"lora\", token = \"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Alternative: GGUF Export Options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save to 8bit Q8_0\n",
    "if False: model.save_pretrained_gguf(\"model\", tokenizer,)\n",
    "if False: model.push_to_hub_gguf(\"hf/model\", tokenizer, token = \"\")\n",
    "\n",
    "# Save to 16bit GGUF\n",
    "if False: model.save_pretrained_gguf(\"model\", tokenizer, quantization_method = \"f16\")\n",
    "if False: model.push_to_hub_gguf(\"hf/model\", tokenizer, quantization_method = \"f16\", token = \"\")\n",
    "\n",
    "# Save to q4_k_m GGUF\n",
    "if False: model.save_pretrained_gguf(\"model\", tokenizer, quantization_method = \"q4_k_m\")\n",
    "if False: model.push_to_hub_gguf(\"hf/model\", tokenizer, quantization_method = \"q4_k_m\", token = \"\")\n",
    "\n",
    "# Save to multiple GGUF options\n",
    "if False:\n",
    "    model.push_to_hub_gguf(\n",
    "        \"hf/model\",\n",
    "        tokenizer,\n",
    "        quantization_method = [\"q4_k_m\", \"q8_0\", \"q5_k_m\",],\n",
    "        token = \"\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Resources\n",
    "\n",
    "- [Unsloth GitHub](https://github.com/unslothai/unsloth)\n",
    "- [TRL SFT docs](https://huggingface.co/docs/trl/sft_trainer)\n",
    "- [FineTome-100k Dataset](https://huggingface.co/datasets/mlabonne/FineTome-100k)\n",
    "- [GGUF Quantization Options](https://github.com/unslothai/unsloth/wiki#gguf-quantization-options)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}