{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "19fe0df6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip3 install https://github.com/casper-hansen/AutoAWQ/releases/download/v0.1.8/autoawq-0.1.8+cu118-cp310-cp310-linux_x86_64.whl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "47528ce3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "20861f3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from awq import AutoAWQForCausalLM\n",
    "from transformers import AutoConfig, AwqConfig, AutoTokenizer, AutoModelForCausalLM\n",
    "import torch\n",
    "\n",
    "model_path = 'mesolitica/malaysian-mistral-7b-32k-instructions-v4'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9939ad4e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "21aaf06156a0404fab259a4b45a92d6a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype = torch.bfloat16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fdb86f50",
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -rf test2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "72e76288",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_pretrained('./test2', safe_serialization = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "aa245150",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "97bc7ca67ff84ab5b302c7720ec2d0d7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "model = AutoAWQForCausalLM.from_pretrained('./test2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d3949cf4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "15a4bf485ecc475ca011da02fe421e8b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/470M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d006dc3bcd834c3d8bf89f20e079d83e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "AWQ: 100%|██████████| 32/32 [09:49<00:00, 18.41s/it]\n"
     ]
    }
   ],
   "source": [
    "quant_path = 'malaysian-mistral-7b-32k-instructions-v4-awq'\n",
    "quant_config = { \"zero_point\": True, \"q_group_size\": 128, \"w_bit\": 4, \"version\": \"GEMM\" }\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)\n",
    "model.quantize(tokenizer, quant_config=quant_config, calib_data = 'mesolitica/malaysian-calibration')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ee290c1e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:`quant_config.json` is being deprecated in the future in favor of quantization_config in config.json.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "('malaysian-mistral-7b-32k-instructions-v4-awq/tokenizer_config.json',\n",
       " 'malaysian-mistral-7b-32k-instructions-v4-awq/special_tokens_map.json',\n",
       " 'malaysian-mistral-7b-32k-instructions-v4-awq/tokenizer.json')"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.save_quantized(quant_path, safetensors = False)\n",
    "tokenizer.save_pretrained(quant_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "737f2403",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/mesolitica/malaysian-mistral-7b-32k-instructions-v4-AWQ/commit/cba1704e3977bd29352015ee3b4c2a26efa17666', commit_message='Upload tokenizer', commit_description='', oid='cba1704e3977bd29352015ee3b4c2a26efa17666', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.push_to_hub('mesolitica/malaysian-mistral-7b-32k-instructions-v4-AWQ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ed92c8ee",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/mesolitica/malaysian-mistral-7b-32k-instructions-v4-AWQ/commit/e72805b82ace8600987f5bae5e336b336d9cd7d0', commit_message='Upload config', commit_description='', oid='e72805b82ace8600987f5bae5e336b336d9cd7d0', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "quantization_config = AwqConfig(\n",
    "    bits=quant_config['w_bit'],\n",
    "    group_size=quant_config['q_group_size'],\n",
    "    zero_point=quant_config['zero_point'],\n",
    "    backend='autoawq',\n",
    "    version=quant_config['version'].lower(),\n",
    ")\n",
    "\n",
    "config = AutoConfig.from_pretrained(model_path)\n",
    "config.quantization_config = quantization_config\n",
    "\n",
    "config.push_to_hub('mesolitica/malaysian-mistral-7b-32k-instructions-v4-AWQ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c74b2f45",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "config.json\t\tquant_config.json\t tokenizer_config.json\r\n",
      "generation_config.json\tspecial_tokens_map.json\r\n",
      "pytorch_model.bin\ttokenizer.json\r\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    }
   ],
   "source": [
    "!ls malaysian-mistral-7b-32k-instructions-v4-awq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "2e0fb591",
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import HfApi\n",
    "\n",
    "api = HfApi()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "dd06cfa2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b38f5cc50f58466fa996e890a43b14a3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pytorch_model.bin:   0%|          | 0.00/4.15G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/mesolitica/malaysian-mistral-7b-32k-instructions-v4-AWQ/commit/be94bbf12947f371acbc8ebb374f1340f3a308cd', commit_message='Upload pytorch_model.bin with huggingface_hub', commit_description='', oid='be94bbf12947f371acbc8ebb374f1340f3a308cd', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "api.upload_file(\n",
    "    path_or_fileobj='malaysian-mistral-7b-32k-instructions-v4-awq/pytorch_model.bin',\n",
    "    path_in_repo=\"pytorch_model.bin\",\n",
    "    repo_id='mesolitica/malaysian-mistral-7b-32k-instructions-v4-AWQ',\n",
    "    repo_type=\"model\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "1383ff2c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/mesolitica/malaysian-mistral-7b-32k-instructions-v4-AWQ/commit/7235673cad93ca30015d0b1d66f96f658af70f2c', commit_message='Upload quant_config.json with huggingface_hub', commit_description='', oid='7235673cad93ca30015d0b1d66f96f658af70f2c', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "api.upload_file(\n",
    "    path_or_fileobj='malaysian-mistral-7b-32k-instructions-v4-awq/quant_config.json',\n",
    "    path_in_repo=\"quant_config.json\",\n",
    "    repo_id='mesolitica/malaysian-mistral-7b-32k-instructions-v4-AWQ',\n",
    "    repo_type=\"model\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5852ec02",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You have loaded an AWQ model on CPU and have a CUDA device available, make sure to set your model on a GPU device in order to run your model.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2f96880b405a46f48665f6b5ff9ca39e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pytorch_model.bin:   0%|          | 0.00/4.15G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "quantized_model = AutoModelForCausalLM.from_pretrained('mesolitica/malaysian-mistral-7b-32k-instructions-v4-AWQ')\n",
    "_ = quantized_model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66895e20",
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = [\n",
    "    {'role': 'user', 'content': 'KWSP tu apa'}\n",
    "]\n",
    "prompt = tokenizer.apply_chat_template(messages, tokenize = False)\n",
    "inputs = tokenizer([prompt], return_tensors='pt', add_special_tokens=False).to('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b320f33",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "generate_kwargs = dict(\n",
    "    inputs,\n",
    "    max_new_tokens=100,\n",
    "    top_p=0.95,\n",
    "    top_k=50,\n",
    "    temperature=0.9,\n",
    "    do_sample=True,\n",
    "    num_beams=1,\n",
    ")\n",
    "r = quantized_model.generate(**generate_kwargs)\n",
    "tokenizer.decode(r[0])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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": 5
}