{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/anaconda3/lib/python3.9/site-packages/huggingface_hub/utils/_hf_folder.py:92: UserWarning: A token has been found in `/Users/casa/.huggingface/token`. This is the old path where tokens were stored. The new location is `/Users/casa/.cache/huggingface/token` which is configurable using `HF_HOME` environment variable. Your token has been copied to this new location. You can now safely delete the old token file manually or use `huggingface-cli logout`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e42b34cf3f07417592f26316fea86e1a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)lve/main/config.json:   0%|          | 0.00/944 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4f89d76d6b7e4cf59a9dd631bd739221",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading pytorch_model.bin:   0%|          | 0.00/1.66G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a690f8b53a204d489f4d53a937068ac6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)neration_config.json:   0%|          | 0.00/111 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "14302bef459f485a998d908b131f43ec",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)okenizer_config.json:   0%|          | 0.00/771 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "33826da838e1402581f62fafd3657b90",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)olve/main/vocab.json:   0%|          | 0.00/1.27M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3ebc87d16a79449998bcb21e33d2ec0b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)olve/main/merges.txt:   0%|          | 0.00/925k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d70c4a2755d04e0d995686f9425b49f8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)/main/tokenizer.json:   0%|          | 0.00/3.07M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cd341cbb7ff445daa312695cc9be1a13",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)cial_tokens_map.json:   0%|          | 0.00/96.0 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "import random\n",
    "import time\n",
    "from transformers import pipeline\n",
    "\n",
    "generator = pipeline(\n",
    "    'text-generation',\n",
    "    model=\"heegyu/bluechat-v0\",\n",
    "    device=\"cuda:0\" if torch.cuda.is_available() else 'cpu'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def query(prompt, max_turn=4):\n",
    "    output = generator(\n",
    "        prompt.strip(),\n",
    "        # no_repeat_ngram_size=2,\n",
    "        eos_token_id=0, # 375=\\n 2=</s>, 0:open-end\n",
    "        max_new_tokens=128,\n",
    "        do_sample=True,\n",
    "        top_p=0.7,\n",
    "        early_stopping=True\n",
    "    )[0]['generated_text']\n",
    "\n",
    "    print(output)\n",
    "\n",
    "    # response = output[len(prompt):]\n",
    "    # return response.strip()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 : 안녕하세요</s>\n",
      "1 : 반가워요</s>\n",
      "0 : 요즘 좋아하는 음악 있으신가요?</s>\n",
      "1 : 최근에 들어서인지 너무 많이 들어요</s>\n",
      "0 : 음 주로 어떤거요?</s>\n",
      "1 : 최근에 들어올린 음악은 무엇인가요?0 : 네 키키 제가 좋아하는 곡은 바로 아이유에요1 : 아 아이유 노래 정말 좋네요0 : 아이유 노래 참 좋아요1 : 아이유 노래 진짜 좋아요0 : 아 진짜 아이유 노래 잘부르세요1 : 네 아이유 노래 좋아요0 : 아이유 노래 진짜 좋죠1 : 아 진짜 좋네요0 : 아이유 노래는 참 좋아요1 : 아이유 노래 정말 좋아요0 : 아이유 노래 정말 좋아요1 : 아이유 노래 정말 좋아요0 : 아이유 노래 진짜 좋아요1 : 아이유 노래 정말 좋아요0 : 아 진짜 좋아요1 : 아 진짜 좋아요0 : 아이유 노래\n"
     ]
    }
   ],
   "source": [
    "query(\"\"\"\n",
    "0 : 안녕하세요</s>\n",
    "1 : 반가워요</s>\n",
    "0 : 요즘 좋아하는 음악 있으신가요?</s>\n",
    "1 : 최근에 들어서인지 너무 많이 들어요</s>\n",
    "0 : 음 주로 어떤거요?</s>\n",
    "1 : \n",
    "\"\"\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.9.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}