KoAlpaca-Polyglot_6B.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"MODEL = 'beomi/KoAlpaca-Polyglot-5.8B'\"\"\"\n",
    "\n",
    "\n",
    "import torch\n",
    "from transformers import pipeline, AutoModelForCausalLM\n",
    "\n",
    "MODEL = 'beomi/KoAlpaca-Polyglot-5.8B'\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    MODEL,\n",
    "    torch_dtype=torch.float16,\n",
    "    low_cpu_mem_usage=True,\n",
    ").to(device=f\"cuda\", non_blocking=True)\n",
    "model.eval()\n",
    "\n",
    "pipe = pipeline(\n",
    "    'text-generation', \n",
    "    model=model,\n",
    "    tokenizer=MODEL,\n",
    "    device=0\n",
    ")\n",
    "\n",
    "def ask(x, context='', is_input_full=False):\n",
    "    ans = pipe(\n",
    "        f\"### 질문: {x}\\n\\n### 맥락: {context}\\n\\n### 답변:\" if context else f\"### 질문: {x}\\n\\n### 답변:\", \n",
    "        do_sample=True, \n",
    "        max_new_tokens=512,\n",
    "        temperature=0.7,\n",
    "        top_p=0.9,\n",
    "        return_full_text=False,\n",
    "        eos_token_id=2,\n",
    "    )\n",
    "    print(ans[0]['generated_text'])\n",
    "\n",
    "ask(\"딥러닝이 뭐야?\")\n"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}