{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Oc7WGlw38EOg"
   },
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9epjFWXJ8F5w"
   },
   "outputs": [],
   "source": [
    "MODEL_ID = \"Qwen/Qwen3-4B\"\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    MODEL_ID,\n",
    "    torch_dtype=torch.float16,\n",
    "    device_map=\"cuda\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kPaTF3zA8Hm0"
   },
   "outputs": [],
   "source": [
    "messages = [\n",
    "    {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": \"Write a limerick about being frugal.\",\n",
    "    },\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pb2NNXar8MRH"
   },
   "outputs": [],
   "source": [
    "text = tokenizer.apply_chat_template(\n",
    "    messages,\n",
    "    tokenize=False,\n",
    "    add_generation_prompt=True,\n",
    "    enable_thinking=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PmiXhpLt8MUR"
   },
   "outputs": [],
   "source": [
    "inputs = tokenizer(text, return_tensors=\"pt\").to(\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TzPH__c28MXC"
   },
   "outputs": [],
   "source": [
    "outputs = model.generate(\n",
    "    **inputs,\n",
    "    max_new_tokens=200,\n",
    "    temperature=0.5,\n",
    "    do_sample=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "u_-eCHPR8Hpl",
    "outputId": "250b2954-098d-495c-f18e-911e6e61327b"
   },
   "outputs": [],
   "source": [
    "response = tokenizer.decode(\n",
    "    outputs[0][inputs[\"input_ids\"].shape[-1] :],\n",
    "    skip_special_tokens=True,\n",
    ")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NwX3beTv8Hun"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}