{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "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": 5,
   "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/mallam-1.1b-20k-instructions-rag'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9939ad4e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a7408b5508d3488da91b64410fded9f6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config.json:   0%|          | 0.00/655 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d03be70a280d4dee8f0467dc05eff53e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/2.25G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c8aba7f9cf0c4e76a6e7b9763e927d8a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "generation_config.json:   0%|          | 0.00/111 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype = torch.bfloat16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "72e76288",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_pretrained('./test', safe_serialization = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "aa245150",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = AutoAWQForCausalLM.from_pretrained('./test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d3949cf4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f650c1b900e748bd972fc144c30cf136",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/1.69k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "44912719d8d54135a1a4f291f7e9e261",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer.json:   0%|          | 0.00/1.34M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1cb98c3aaf5b4594acc24940a33fde68",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "special_tokens_map.json:   0%|          | 0.00/638 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2f50e02ec430412fba8b65ee0a79aae9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/351M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1c66c052f248428cba3e5a9c45165cc8",
       "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%|██████████| 22/22 [02:28<00:00,  6.75s/it]\n"
     ]
    }
   ],
   "source": [
    "quant_path = 'mallam-1.1b-20k-instructions-rag-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": 10,
   "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": [
       "('mallam-1.1b-20k-instructions-rag-awq/tokenizer_config.json',\n",
       " 'mallam-1.1b-20k-instructions-rag-awq/special_tokens_map.json',\n",
       " 'mallam-1.1b-20k-instructions-rag-awq/tokenizer.json')"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.save_quantized(quant_path, safetensors = False)\n",
    "tokenizer.save_pretrained(quant_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "737f2403",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/mesolitica/mallam-1.1b-20k-instructions-rag-AWQ/commit/ef848a832fb7b21f0bf4b808a83af300ed07beca', commit_message='Upload tokenizer', commit_description='', oid='ef848a832fb7b21f0bf4b808a83af300ed07beca', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.push_to_hub('mesolitica/mallam-1.1b-20k-instructions-rag-AWQ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "ed92c8ee",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/mesolitica/mallam-1.1b-20k-instructions-rag-AWQ/commit/40c9e6295642ca0e5cdc4b8ff8888a73c1ce25d3', commit_message='Upload config', commit_description='', oid='40c9e6295642ca0e5cdc4b8ff8888a73c1ce25d3', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 13,
     "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/mallam-1.1b-20k-instructions-rag-AWQ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "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 mallam-1.1b-20k-instructions-rag-awq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2e0fb591",
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import HfApi\n",
    "\n",
    "api = HfApi()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "dd06cfa2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5685dd9dac3b41f6a0d082201cfd610e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pytorch_model.bin:   0%|          | 0.00/778M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/mesolitica/mallam-1.1b-20k-instructions-rag-AWQ/commit/4caae4a226e3dd569ac467c5099cd159ac5fa3e0', commit_message='Upload pytorch_model.bin with huggingface_hub', commit_description='', oid='4caae4a226e3dd569ac467c5099cd159ac5fa3e0', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "api.upload_file(\n",
    "    path_or_fileobj='mallam-1.1b-20k-instructions-rag-awq/pytorch_model.bin',\n",
    "    path_in_repo=\"pytorch_model.bin\",\n",
    "    repo_id='mesolitica/mallam-1.1b-20k-instructions-rag-AWQ',\n",
    "    repo_type=\"model\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "1383ff2c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/mesolitica/mallam-1.1b-20k-instructions-rag-AWQ/commit/81a4cf5cdf2ecd38d74354f558af0d2896a24b7d', commit_message='Upload quant_config.json with huggingface_hub', commit_description='', oid='81a4cf5cdf2ecd38d74354f558af0d2896a24b7d', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "api.upload_file(\n",
    "    path_or_fileobj='mallam-1.1b-20k-instructions-rag-awq/quant_config.json',\n",
    "    path_in_repo=\"quant_config.json\",\n",
    "    repo_id='mesolitica/mallam-1.1b-20k-instructions-rag-AWQ',\n",
    "    repo_type=\"model\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "5852ec02",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "51b30c41eab64d23ae86cf451c6bf196",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config.json:   0%|          | 0.00/909 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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": "26fe778711d8474699ee792ce27e3092",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pytorch_model.bin:   0%|          | 0.00/778M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "quantized_model = AutoModelForCausalLM.from_pretrained('mesolitica/mallam-1.1b-20k-instructions-rag-AWQ')\n",
    "_ = quantized_model.cuda()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "6adc9240",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'<s>[INST] kwsp tu apa [/INST]'"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "messages = [\n",
    "    {'role': 'user', 'content': 'kwsp tu apa'}\n",
    "]\n",
    "tokenizer.apply_chat_template(messages, tokenize = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "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": 27,
   "id": "4b320f33",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:2 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 4.65 s, sys: 3.73 ms, total: 4.65 s\n",
      "Wall time: 4.65 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'<s> [INST] KWSP tu apa? [/INST]Kumpulan Wang Simpanan Pekerja (KWSP) atau EPF adalah tabung simpanan wang pekerja di Malaysia. KWSP ialah institusi keselamatan sosial yang disediakan oleh kerajaan Malaysia untuk membantu pekerja Malaysia menguruskan dana persaraan mereka.\\n\\nKWSP bertujuan untuk menyediakan sumber kewangan untuk membantu pekerja Malaysia melabur dan menyimpan untuk masa depan mereka. Ia mempunyai objektif utama untuk mencapai matlamat ini, iaitu memastikan ahli KWSP mempunyai simpanan yang mencukupi untuk memenuhi keperluan persaraan mereka.\\n\\nKWSP menyediakan faedah persaraan yang berbeza, seperti KWSP Prihatin dan KWSP i-Sinar. Simpanan KWSP yang dikhaskan boleh digunakan untuk tujuan yang berbeza, seperti membiayai pendidikan anak-anak, membeli rumah, atau memenuhi keperluan perubatan.\\n\\nSecara amnya, ahli KWSP layak untuk pelbagai jenis manfaat persaraan yang mereka boleh perolehi, bergantung pada umur, tahap kecacatan, dan sektor pekerjaan mereka. Ini termasuk faedah KWSP, bantuan bayaran pencen, serta kemudahan simpanan tambahan.\\n\\nSekiranya anda ingin mengetahui lebih lanjut mengenai KWSP atau menyertai KWSP, anda boleh menghubungi cawangan KWSP terdekat atau layari laman web KWSP untuk mendapatkan maklumat lanjut.</s>'"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "generate_kwargs = dict(\n",
    "    inputs,\n",
    "    max_new_tokens=1024,\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])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9a93555",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}