{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MljifiTVCT0_"
   },
   "source": [
    "# 推論用コード\n",
    "本コードはunslothで学習したqLoRAのアダプタを用いてELYZA-tasks-100-TVの出力を得るためのコードです。  \n",
    "Hugging Faceにアダプタをアップロードしてあることが前提となります。\n",
    "このコードはunslothライブラリを用いてモデルを読み込み、推論するためのコードとなります。\n",
    "このコードで生成されたjsonlファイルは課題の成果として提出可能なフォーマットになっております。\n",
    "\n",
    "※本コードはGoogle Colabでの動作を想定しており、Omnicampusでの動作を想定しておりません。\n",
    "Omnicampus向けのコードは別途用意しております。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "I5B5MOHuBy8b"
   },
   "outputs": [],
   "source": [
    "# 必要なライブラリをインストール\n",
    "%%capture\n",
    "!pip install unsloth\n",
    "!pip uninstall unsloth -y && pip install --upgrade --no-cache-dir \"unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git\"\n",
    "!pip install -U torch\n",
    "!pip install -U peft"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GM7SNRtACg9V"
   },
   "outputs": [],
   "source": [
    "# 必要なライブラリを読み込み\n",
    "from unsloth import FastLanguageModel\n",
    "from peft import PeftModel\n",
    "import torch\n",
    "import json\n",
    "from tqdm import tqdm\n",
    "import re"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JmdUATTVCtyr"
   },
   "outputs": [],
   "source": [
    "# ベースとなるモデルと学習したLoRAのアダプタ（Hugging FaceのIDを指定）。\n",
    "model_id = \"llm-jp/llm-jp-3-13b\"\n",
    "adapter_id = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hugging Face Token を指定。\n",
    "# 下記の URL から Hugging Face Token を取得できますので下記の HF_TOKEN に入れてください。\n",
    "# https://huggingface.co/settings/tokens  \n",
    "HF_TOKEN = \"\" #@param {type:\"string\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TB6Hzx-2B5g8"
   },
   "outputs": [],
   "source": [
    "# unslothのFastLanguageModelで元のモデルをロード。\n",
    "dtype = None # Noneにしておけば自動で設定\n",
    "load_in_4bit = True # 今回は13Bモデルを扱うためTrue\n",
    "\n",
    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
    "    model_name=model_id,\n",
    "    dtype=dtype,\n",
    "    load_in_4bit=load_in_4bit,\n",
    "    trust_remote_code=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 元のモデルにLoRAのアダプタを統合。\n",
    "model = PeftModel.from_pretrained(model, adapter_id, token = HF_TOKEN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fg_yURyiB8o6"
   },
   "outputs": [],
   "source": [
    "# タスクとなるデータの読み込み。\n",
    "# 事前にデータをアップロードしてください。\n",
    "datasets = []\n",
    "with open(\"./elyza-tasks-100-TV_0.jsonl\", \"r\") as f:\n",
    "    item = \"\"\n",
    "    for line in f:\n",
    "      line = line.strip()\n",
    "      item += line\n",
    "      if item.endswith(\"}\"):\n",
    "        datasets.append(json.loads(item))\n",
    "        item = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TwfZEra1CEJo"
   },
   "outputs": [],
   "source": [
    "# モデルを用いてタスクの推論。\n",
    "\n",
    "# 推論するためにモデルのモードを変更\n",
    "FastLanguageModel.for_inference(model)\n",
    "\n",
    "results = []\n",
    "for dt in tqdm(datasets):\n",
    "  input = dt[\"input\"]\n",
    "\n",
    "  prompt = f\"\"\"### 指示\\n{input}\\n### 回答\\n\"\"\"\n",
    "\n",
    "  inputs = tokenizer([prompt], return_tensors = \"pt\").to(model.device)\n",
    "\n",
    "  outputs = model.generate(**inputs, max_new_tokens = 512, use_cache = True, do_sample=False, repetition_penalty=1.2)\n",
    "  prediction = tokenizer.decode(outputs[0], skip_special_tokens=True).split('\\n### 回答')[-1]\n",
    "\n",
    "  results.append({\"task_id\": dt[\"task_id\"], \"input\": input, \"output\": prediction})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "voAPnXp5CKRL"
   },
   "outputs": [],
   "source": [
    "# 結果をjsonlで保存。\n",
    "\n",
    "# ここではadapter_idを元にファイル名を決定しているが、ファイル名は任意で問題なし。\n",
    "json_file_id = re.sub(\".*/\", \"\", adapter_id)\n",
    "with open(f\"/content/{json_file_id}_output.jsonl\", 'w', encoding='utf-8') as f:\n",
    "    for result in results:\n",
    "        json.dump(result, f, ensure_ascii=False)\n",
    "        f.write('\\n')"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}