{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LLM Compressor Workbench -- Getting Started\n",
    "\n",
    "This notebook will demonstrate how common [LLM Compressor](https://github.com/vllm-project/llm-compressor) flows can be run on the Alauda AI.\n",
    "\n",
    "We will show how a user can compress a Large Language Model, with a calibration dataset.\n",
    "\n",
    "The notebook will detect if a GPU is available. If one is not available, it will demonstrate an abbreviated run, so users without GPU access can still get a feel for `llm-compressor`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calibrated Compression with a Dataset\n",
    "\n",
    "Some more advanced compression algorithms require a small dataset of calibration samples that are meant to be a representative random subset of the language the model will see at inference.\n",
    "\n",
    "We will show how the previous section can be augmented with a calibration dataset and GPTQ, one of the first published LLM compression algorithms.\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "<b>Note:</b> This will take several minutes if no GPU is available\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "use_gpu = torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We will use a new recipe running GPTQ (https://arxiv.org/abs/2210.17323)\n",
    "# to reduce error caused by quantization. GPTQ requires a calibration dataset.\n",
    "from llmcompressor.modifiers.quantization import GPTQModifier\n",
    "\n",
    "# model to compress\n",
    "model_id = \"./TinyLlama-1.1B-Chat-v1.0\"\n",
    "recipe = GPTQModifier(targets=\"Linear\", scheme=\"W4A16\", ignore=[\"lm_head\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load up model using huggingface API\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_id, device_map=\"auto\", torch_dtype=\"auto\"\n",
    ")\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "# Create the calibration dataset, using Huggingface datasets API\n",
    "dataset_id = \"./ultrachat_200k\"\n",
    "\n",
    "# Select number of samples. 512 samples is a good place to start.\n",
    "# Increasing the number of samples can improve accuracy.\n",
    "num_calibration_samples = 512 if use_gpu else 4\n",
    "max_sequence_length = 2048 if use_gpu else 16\n",
    "\n",
    "# Load dataset\n",
    "ds = load_dataset(dataset_id, split=\"train_sft\")\n",
    "# Shuffle and grab only the number of samples we need\n",
    "ds = ds.shuffle(seed=42).select(range(num_calibration_samples))\n",
    "\n",
    "\n",
    "# Preprocess and tokenize into format the model uses\n",
    "def preprocess(example):\n",
    "    text = tokenizer.apply_chat_template(\n",
    "        example[\"messages\"],\n",
    "        tokenize=False,\n",
    "    )\n",
    "    return tokenizer(\n",
    "        text,\n",
    "        padding=False,\n",
    "        max_length=max_sequence_length,\n",
    "        truncation=True,\n",
    "        add_special_tokens=False,\n",
    "    )\n",
    "\n",
    "\n",
    "ds = ds.map(preprocess, remove_columns=ds.column_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run oneshot, with dataset\n",
    "from llmcompressor import oneshot\n",
    "\n",
    "model = oneshot(\n",
    "    model=model,\n",
    "    dataset=ds,\n",
    "    recipe=recipe,\n",
    "    max_seq_length=max_sequence_length,\n",
    "    num_calibration_samples=num_calibration_samples,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save model and tokenizer\n",
    "model_dir = \"./\" + model_id.split(\"/\")[-1] + \"-GPTQ-W4A16\"\n",
    "model.save_pretrained(model_dir)\n",
    "tokenizer.save_pretrained(model_dir);"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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": 2
}