{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3008bf1f-df10-484b-9662-96cda681b93b",
   "metadata": {},
   "source": [
    "```bash\n",
    "git clone https://github.com/kIshizaki-sci/AutoAWQ.git\n",
    "pip install -U transformers\n",
    "pip install -e ./AutoAWQ\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a388d190-a611-46b0-aa8c-dcf5c97b0c1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from awq import AutoAWQForCausalLM\n",
    "import torch\n",
    "import transformers\n",
    "from transformers import AutoProcessor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "de7f888d-fe64-47f1-a106-75a7b911354f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch version :  2.4.1+cu124\n",
      "transformers version :  4.51.3\n"
     ]
    }
   ],
   "source": [
    "print('torch version : ', torch.__version__)\n",
    "print('transformers version : ', transformers.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "61623ebd-d8a4-45b6-aff3-3c16a78ecfa5",
   "metadata": {},
   "outputs": [],
   "source": [
    "quant_path = 'kishizaki-sci/Llama-4-Scout-17B-16E-Instruct-AWQ'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c13b72e4-f6cd-4642-a110-040844127541",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4f6c7b795fb24940b752dded2f51dea0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config.json:   0%|          | 0.00/3.55k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/workspace/AutoAWQ/awq/models/llama4.py:312: UserWarning: Multimodal input has not been implemented in Llama4AWQForConditionalGeneration yet.\n",
      "  warnings.warn(\"Multimodal input has not been implemented in Llama4AWQForConditionalGeneration yet.\", UserWarning)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "02743a30540244bbb4c05585becc4cc8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Fetching 25 files:   0%|          | 0/25 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3200dd87c1df4ad2a0d1ce0a300f094c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "llama4_inference.ipynb:   0%|          | 0.00/10.5k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b132c3a11c024af78dd8fd4d30b579cc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "generation_config.json:   0%|          | 0.00/239 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1f505cb27bc04c92b48ca9e27202cb4b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       ".gitattributes:   0%|          | 0.00/1.57k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "43b1b86a0c304e56b1870d08598b374f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00002-of-00013.safetensors:   0%|          | 0.00/4.98G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7babbb9a6b154f509c88a56b4fee3d0e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "README.md:   0%|          | 0.00/138 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "405cd2691a844317b6570eb80c622d24",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00001-of-00013.safetensors:   0%|          | 0.00/5.00G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e979f0b6efbe40baa318dccb57bffe24",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00004-of-00013.safetensors:   0%|          | 0.00/5.00G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "efccd71f8aa2492bad03ebadbbef596d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00005-of-00013.safetensors:   0%|          | 0.00/5.00G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a3028a64430744beb181ace86753bee6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "chat_template.json:   0%|          | 0.00/5.18k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "448ff54d8a6642e685474d3bc49de09b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00003-of-00013.safetensors:   0%|          | 0.00/5.00G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "63cb05c948ef476c98df712735133fa7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00006-of-00013.safetensors:   0%|          | 0.00/4.98G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1afb22e38b7c46c0ab496bd67c949402",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00007-of-00013.safetensors:   0%|          | 0.00/5.00G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ded5ae0f710a4124b00cb0b061383962",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00008-of-00013.safetensors:   0%|          | 0.00/5.00G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0db9d375464c415ebb33fc0b268a9950",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00009-of-00013.safetensors:   0%|          | 0.00/5.00G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e49ff7276a7b493cb5ab7dc222e9d3ec",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00010-of-00013.safetensors:   0%|          | 0.00/5.00G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cbb0557242c546708bfb5ba4c7789499",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00011-of-00013.safetensors:   0%|          | 0.00/4.98G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2c603da8fcc34e0584963099695e7142",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00012-of-00013.safetensors:   0%|          | 0.00/5.00G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "645fee17e47d44e1b7fb51865a00fc73",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00013-of-00013.safetensors:   0%|          | 0.00/2.79G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a022c5cb1f174ac5bd9fad91cda9352a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors.index.json:   0%|          | 0.00/1.13M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f68405f7857d4db1b40b264f25175f43",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "preprocessor_config.json:   0%|          | 0.00/636 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "81bc3a84390f4c388f1012044d4ce6ec",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "processor_config.json:   0%|          | 0.00/128 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "842746981e6e478aae858d3daf844460",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "special_tokens_map.json:   0%|          | 0.00/448 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "067b4e51f41840cdb7305972f1aa0dc6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer.json:   0%|          | 0.00/27.9M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5ceb2f619f9f4513bc307c5846fae2e5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/237k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You have set `use_cache` to `False`, but cache_implementation is set to hybrid. cache_implementation will have no effect.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "db4fddb7315c464ea87c8ad507f9a651",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Replacing MoE Block...:   0%|          | 0/48 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Replacing layers...: 100%|██████████| 48/48 [00:07<00:00,  6.18it/s]\n",
      "/workspace/AutoAWQ/awq/models/base.py:539: UserWarning: Skipping fusing modules because AWQ extension is not installed.No module named 'awq_ext'\n",
      "  warnings.warn(\"Skipping fusing modules because AWQ extension is not installed.\" + msg)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1h 44min 26s, sys: 26min 26s, total: 2h 10min 53s\n",
      "Wall time: 28min 55s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "model = AutoAWQForCausalLM.from_quantized(quant_path, torch_dtype=torch.float16, use_cache=True, device_map='auto')\n",
    "processor = AutoProcessor.from_pretrained(quant_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "665ae24c-cb73-489d-bb03-35d760460070",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sat Apr 19 22:53:42 2025       \n",
      "+-----------------------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 565.57.01              Driver Version: 565.57.01      CUDA Version: 12.7     |\n",
      "|-----------------------------------------+------------------------+----------------------+\n",
      "| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |\n",
      "|                                         |                        |               MIG M. |\n",
      "|=========================================+========================+======================|\n",
      "|   0  NVIDIA A100-SXM4-80GB          On  |   00000000:87:00.0 Off |                    0 |\n",
      "| N/A   30C    P0             71W /  400W |   60415MiB /  81920MiB |      0%      Default |\n",
      "|                                         |                        |             Disabled |\n",
      "+-----------------------------------------+------------------------+----------------------+\n",
      "                                                                                         \n",
      "+-----------------------------------------------------------------------------------------+\n",
      "| Processes:                                                                              |\n",
      "|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |\n",
      "|        ID   ID                                                               Usage      |\n",
      "|=========================================================================================|\n",
      "+-----------------------------------------------------------------------------------------+\n"
     ]
    }
   ],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "39c8154a-08c4-485e-b136-955d1b4fbec9",
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = [\n",
    "    {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": [\n",
    "            {\"type\": \"text\", \"text\": \"What does means the torsion free in the general relativit?\"},\n",
    "        ]\n",
    "    },\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "44c4ae97-96cf-47f7-a8e4-16d2c34bdc1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = processor.apply_chat_template(\n",
    "    messages,\n",
    "    add_generation_prompt=True,\n",
    "    tokenize=True,\n",
    "    return_dict=True,\n",
    "    return_tensors=\"pt\",\n",
    ").to(model.model.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c6dddd7b-5382-45e8-8dd1-6ca719200f64",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 4min 28s, sys: 3min 39s, total: 8min 7s\n",
      "Wall time: 8min 8s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "outputs = model.generate(\n",
    "    **inputs,\n",
    "    max_new_tokens=2048,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ab1d4fe0-e315-40ee-b3e0-10db1e6b2023",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A question from the realm of differential geometry and general relativity!\n",
      "\n",
      "In general relativity, \"torsion-free\" refers to a property of a connection on a manifold, specifically in the context of Riemannian geometry.\n",
      "\n",
      "**Torsion** is a measure of how much a connection \"twists\" or \"turns\" a vector as it is parallel-transported around a closed loop. In other words, it's a measure of how much the connection deviates from being \"flat\" or \"Euclidean\".\n",
      "\n",
      "A **torsion-free connection**, also known as a **symmetric connection**, is a connection that has zero torsion. This means that when you parallel-transport a vector around a closed loop, it returns to its original orientation, without any twisting or turning.\n",
      "\n",
      "In mathematical terms, a torsion-free connection satisfies the following condition:\n",
      "\n",
      "$$\\Gamma^i_{jk} = \\Gamma^i_{kj}$$\n",
      "\n",
      "where $\\Gamma^i_{jk}$ are the Christoffel symbols of the second kind, which define the connection.\n",
      "\n",
      "In general relativity, the Levi-Civita connection is a fundamental example of a torsion-free connection. It's the unique connection that is:\n",
      "\n",
      "1. **Metric-compatible**: it preserves the metric tensor under parallel transport.\n",
      "2. **Torsion-free**: it has zero torsion.\n",
      "\n",
      "The Levi-Civita connection is used to define the covariant derivative, which is essential for describing the curvature of spacetime.\n",
      "\n",
      "The assumption of a torsion-free connection is crucial in general relativity, as it allows us to:\n",
      "\n",
      "1. **Define a unique covariant derivative**: which is necessary for formulating the Einstein field equations.\n",
      "2. **Ensure geodesic equation**: which describes the shortest path in curved spacetime, is well-defined.\n",
      "\n",
      "However, it's worth noting that there are some alternative theories, such as Einstein-Cartan theory, which consider torsion as a fundamental aspect of spacetime geometry.\n",
      "\n",
      "I hope this explanation helps! Do you have any follow-up questions?<|eot|>\n"
     ]
    }
   ],
   "source": [
    "response = processor.batch_decode(outputs[:, inputs[\"input_ids\"].shape[-1]:])[0]\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f3f1be7-6960-474e-b19d-afd2efa56174",
   "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.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}