min-new-token 추가

app.py

@@ -2,13 +2,40 @@ import gradio as gr
 import torch
 import random
 import time
-from transformers import pipeline
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+
+
+model_name="heegyu/bluechat-v0"
+device="cuda:0" if torch.cuda.is_available() else 'cpu'
+model = AutoModelForCausalLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# generator = pipeline(
+#     'text-generation',
+#     model="heegyu/bluechat-v0",
+#     device="cuda:0" if torch.cuda.is_available() else 'cpu'
+# )
+
+def get_message(prompt, min_new_tokens=16, max_turn=4):
+    prompt = prompt.strip()
+    ids = tokenizer(prompt, return_tensors="pt").to(device)
+    min_length = ids['input_ids'].shape[1] + min_new_tokens
+
+    output = model.generate(
+        **ids,
+        no_repeat_ngram_size=3,
+        eos_token_id=2, # 375=\n 2=</s>, 0:open-end
+        max_new_tokens=128,
+        min_length=min_length,
+        do_sample=True,
+        top_p=0.7,
+        early_stopping=True
+    ) # [0]['generated_text']
+
+    output = tokenizer.decode(output.cpu()[0])
+    print(output)
+    return output[len(prompt):]
 
-generator = pipeline(
-    'text-generation',
-    model="heegyu/bluechat-v0",
-    device="cuda:0" if torch.cuda.is_available() else 'cpu'
-)
 
 def query(message, chat_history, max_turn=4):
     prompt = []
@@ -21,19 +48,7 @@ def query(message, chat_history, max_turn=4):
     prompt.append(f"<usr> {message}")
     prompt = "\n".join(prompt) + "\n<bot>"
 
-    output = generator(
-        prompt,
-        # repetition_penalty=1.3,
-        # no_repeat_ngram_size=2,
-        eos_token_id=2, # \n
-        max_new_tokens=128,
-        do_sample=True,
-        top_p=0.9,
-    )[0]['generated_text']
-
-    print(output)
-
-    response = output[len(prompt):]
+    response = get_message(prompt, 8)
     return response.strip()
 
 with gr.Blocks() as demo:

test.ipynb

@@ -2,161 +2,42 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 40,
    "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"
-    }
-   ],
+   "outputs": [],
    "source": [
     "import torch\n",
     "import random\n",
     "import time\n",
-    "from transformers import pipeline\n",
+    "from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer\n",
     "\n",
-    "generator = pipeline(\n",
-    "    'text-generation',\n",
-    "    model=\"heegyu/bluechat-v0\",\n",
-    "    device=\"cuda:0\" if torch.cuda.is_available() else 'cpu'\n",
-    ")"
+    "model_name=\"heegyu/bluechat-v0\"\n",
+    "device=\"cuda:0\" if torch.cuda.is_available() else 'cpu'\n",
+    "model = AutoModelForCausalLM.from_pretrained(model_name)\n",
+    "tokenizer = AutoTokenizer.from_pretrained(model_name)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 32,
+   "execution_count": 54,
    "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",
+    "def query(prompt, min_new_tokens=16, max_turn=4):\n",
+    "    ids = tokenizer(prompt.strip(), return_tensors=\"pt\").to(device)\n",
+    "    min_length = ids['input_ids'].shape[1] + min_new_tokens\n",
+    "    output = model.generate(\n",
+    "        **ids,\n",
+    "        no_repeat_ngram_size=3,\n",
+    "        eos_token_id=2, # 375=\\n 2=</s>, 0:open-end\n",
     "        max_new_tokens=128,\n",
+    "        min_length=min_length,\n",
     "        do_sample=True,\n",
     "        top_p=0.7,\n",
     "        early_stopping=True\n",
-    "    )[0]['generated_text']\n",
-    "\n",
+    "    ) # [0]['generated_text']\n",
+    "    output = tokenizer.decode(output.cpu()[0])\n",
     "    print(output)\n",
     "\n",
     "    # response = output[len(prompt):]\n",
@@ -165,19 +46,34 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 33,
+   "execution_count": 42,
    "metadata": {},
    "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n"
+     ]
+    },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
+      "\n",
       "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"
+      "1 : \n",
+      " music : music songs 수록곡을 즐겨들어요</s><bot> 앗 어떤 장르를 주로 들으시나요?</s>\n",
+      "1 : music songs 좋죠</s>\n",
+      "bot> 저도 요즘 들어 좋아하게 된 곡들 위주로 들어요 ㅎㅎ</s>\n",
+      "2 : music songs 어떤 노래들 자주 들어요?</s>\n",
+      "bot> 저 music songs someone이 제일 좋더라구요 ㅎㅎ</s>\n",
+      "1 : music songs는 어떤 곡들 주로 들어요?</s>\n",
+      "bot> 저 music songs는 주로 music songs를 많이 들어요 ㅎㅎ</s>\n"
      ]
     }
    ],
@@ -191,6 +87,91 @@
     "1 : \n",
     "\"\"\")"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 48,
+   "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": [
+      "<usr> 안녕하세요\n",
+      "<bot> 안녕하세요~ 저녁 드셨나요? ㅎㅎ? ㅎㅎ</s>\n"
+     ]
+    }
+   ],
+   "source": [
+    "query(\"\"\"\n",
+    "<usr> 안녕하세요\n",
+    "<bot>\n",
+    "\"\"\", 8)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 55,
+   "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": [
+      "<usr> 안녕하세요 식사 하셨나요?\n",
+      "<bot> 안녕하세요 네~ 점심 먹었어요 식사하셨나요?\n",
+      "네~ 뭐드셨나요?</s>\n"
+     ]
+    }
+   ],
+   "source": [
+    "query(\"\"\"\n",
+    "<usr> 안녕하세요 식사 하셨나요?\n",
+    "<bot>\n",
+    "\"\"\", 8)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 63,
+   "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": [
+      "<usr> 창업에 관심이 있나요?\n",
+      "<bot> 네! 근데 요즘 창업에 대한 관심이 많이 떨어지더라구요</s>\n"
+     ]
+    }
+   ],
+   "source": [
+    "query(\"\"\"\n",
+    "<usr> 창업에 관심이 있나요?\n",
+    "<bot>\n",
+    "\"\"\", 8)"
+   ]
   }
  ],
  "metadata": {