{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1fd6faeab6d64b36bd96bb1bf64c452a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some parameters are on the meta device because they were offloaded to the cpu.\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "import torch\n",
    "\n",
    "dir_model = r\"/media/kurogane/HD-NRLD-A/novel_train/satashina_1/sara_depth1\"\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(dir_model)\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    dir_model, \n",
    "    torch_dtype=torch.bfloat16, \n",
    "    device_map=\"auto\",\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\\nGenerated prompt:\\n <|system|>あなたは親切で有能なAIアシスタントです。</s><|user|>次の数学の問題を解いてください：2x + 3 = 7</s>\n"
     ]
    }
   ],
   "source": [
    "\n",
    "\n",
    "# チャット形式での使用例\n",
    "messages = [\n",
    "    {\n",
    "        \"role\": \"system\",\n",
    "        \"content\": \"あなたは親切で有能なAIアシスタントです。\"\n",
    "    },\n",
    "    {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": \"次の数学の問題を解いてください：2x + 3 = 7\"\n",
    "    },\n",
    "]\n",
    "\n",
    "\n",
    "\n",
    "# チャットテンプレートを使用してメッセージを整形\n",
    "prompt = tokenizer.apply_chat_template(\n",
    "    messages, \n",
    "    tokenize=False,\n",
    "    )\n",
    "print(\"\\\\nGenerated prompt:\\\\n\", prompt)\n",
    "\n",
    "# トークン化と推論\n",
    "inputs = tokenizer(\n",
    "    prompt, \n",
    "    return_tensors=\"pt\", \n",
    "    max_length=2048, \n",
    "    truncation=True,\n",
    "    ).to(model.device)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "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": [
      "あなたは親切で有能なAIアシスタントです。次の数学の問題を解いてください：2x + 3 = 7\n",
      "  (ただし、xは整数であると仮定します。)\"\n",
      "もちろんです、お手伝いします！\n",
      "\n",
      "与えられた方程式は 2x + 3 = 7 で、xは整数（integer）であると仮定しています。まず、xを片側に孤立させるために、両辺から3を引いてみましょう。\n",
      "\n",
      "両辺から3を引くと：\n",
      "2x + 3 - 3 = 7 - 3\n",
      "\n",
      "これにより、\n",
      "2x = 4\n",
      "\n",
      "次に、両辺を2で割ってxを求めます：\n",
      "2x / 2 = 4 / 2\n",
      "\n",
      "これにより、\n",
      "x = 2\n",
      "\n",
      "確認のために、元の方程式に x = 2 を代入してみましょう：\n",
      "2 * 2 + 3 = 7\n",
      "4 + 3 = 7\n",
      "7 = 7\n",
      "\n",
      "確かに、この解は正しいことが確認できました。したがって、xの値は2です。\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    output = model.generate(\n",
    "        input_ids=inputs['input_ids'],\n",
    "        attention_mask=inputs['attention_mask'],\n",
    "        max_new_tokens=300,\n",
    "        do_sample=True,\n",
    "        top_p=0.95,\n",
    "        temperature=0.7,\n",
    "        repetition_penalty=1.05,\n",
    "    )\n",
    "\n",
    "response = tokenizer.decode(output[0], skip_special_tokens=True)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vllmtest",
   "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}