7.PRETRAIN_METHOD/7.4.1.gpt2_finetune_LM.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "import gluonnlp as nlp\n",
    "from gluonnlp.data import SentencepieceTokenizer\n",
    "from transformers import TFGPT2LMHeadModel\n",
    "\n",
    "from tensorflow.keras.preprocessing.sequence import pad_sequences\n",
    "\n",
    "from nltk.tokenize import sent_tokenize"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "아레 실행 커멘드는 gpt_ckpt 폴더가 있지 않은 경우에만 실행해주세요."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import wget\n",
    "import zipfile\n",
    "\n",
    "wget.download('https://github.com/NLP-kr/tensorflow-ml-nlp-tf2/releases/download/v1.0/gpt_ckpt.zip')\n",
    "\n",
    "with zipfile.ZipFile('gpt_ckpt.zip') as z:\n",
    "    z.extractall()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GPT2Model(tf.keras.Model):\n",
    "    def __init__(self, dir_path):\n",
    "        super(GPT2Model, self).__init__()\n",
    "        self.gpt2 = TFGPT2LMHeadModel.from_pretrained(dir_path)\n",
    "        \n",
    "    def call(self, inputs):\n",
    "        return self.gpt2(inputs)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BASE_MODEL_PATH = './gpt_ckpt'\n",
    "gpt_model = GPT2Model(BASE_MODEL_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 16\n",
    "NUM_EPOCHS = 10\n",
    "MAX_LEN = 30\n",
    "TOKENIZER_PATH = './gpt_ckpt/gpt2_kor_tokenizer.spiece'\n",
    "\n",
    "tokenizer = SentencepieceTokenizer(TOKENIZER_PATH)\n",
    "vocab = nlp.vocab.BERTVocab.from_sentencepiece(TOKENIZER_PATH,\n",
    "                                               mask_token=None,\n",
    "                                               sep_token=None,\n",
    "                                               cls_token=None,\n",
    "                                               unknown_token='<unk>',\n",
    "                                               padding_token='<pad>',\n",
    "                                               bos_token='<s>',\n",
    "                                               eos_token='</s>')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tf_top_k_top_p_filtering(logits, top_k=0, top_p=0.0, filter_value=-99999):\n",
    "    _logits = logits.numpy()\n",
    "    top_k = min(top_k, logits.shape[-1])  \n",
    "    if top_k > 0:\n",
    "        indices_to_remove = logits < tf.math.top_k(logits, top_k)[0][..., -1, None]\n",
    "        _logits[indices_to_remove] = filter_value\n",
    "\n",
    "    if top_p > 0.0:\n",
    "        sorted_logits = tf.sort(logits, direction='DESCENDING')\n",
    "        sorted_indices = tf.argsort(logits, direction='DESCENDING')\n",
    "        cumulative_probs = tf.math.cumsum(tf.nn.softmax(sorted_logits, axis=-1), axis=-1)\n",
    "\n",
    "        sorted_indices_to_remove = cumulative_probs > top_p\n",
    "        sorted_indices_to_remove = tf.concat([[False], sorted_indices_to_remove[..., :-1]], axis=0)\n",
    "        indices_to_remove = sorted_indices[sorted_indices_to_remove].numpy().tolist()\n",
    "        \n",
    "        _logits[indices_to_remove] = filter_value\n",
    "    return tf.constant([_logits])\n",
    "\n",
    "\n",
    "def generate_sent(seed_word, model, max_step=100, greedy=False, top_k=0, top_p=0.):\n",
    "    sent = seed_word\n",
    "    toked = tokenizer(sent)\n",
    "    \n",
    "    for _ in range(max_step):\n",
    "        input_ids = tf.constant([vocab[vocab.bos_token],]  + vocab[toked])[None, :] \n",
    "        outputs = model(input_ids)[:, -1, :]\n",
    "        if greedy:\n",
    "            gen = vocab.to_tokens(tf.argmax(outputs, axis=-1).numpy().tolist()[0])\n",
    "        else:\n",
    "            output_logit = tf_top_k_top_p_filtering(outputs[0], top_k=top_k, top_p=top_p)\n",
    "            gen = vocab.to_tokens(tf.random.categorical(output_logit, 1).numpy().tolist()[0])[0]\n",
    "        if gen == '</s>':\n",
    "            break\n",
    "        sent += gen.replace('▁', ' ')\n",
    "        toked = tokenizer(sent)\n",
    "\n",
    "    return sent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_sent('이때', gpt_model, greedy=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_sent('이때', gpt_model, top_k=0, top_p=0.95)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_IN_PATH = './data_in/KOR/'\n",
    "TRAIN_DATA_FILE = 'finetune_data.txt'\n",
    "\n",
    "sents = [s[:-1] for s in open(DATA_IN_PATH + TRAIN_DATA_FILE, encoding='utf-8').readlines()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_data = []\n",
    "output_data = []\n",
    "\n",
    "for s in sents:\n",
    "    tokens = [vocab[vocab.bos_token],]  + vocab[tokenizer(s)] + [vocab[vocab.eos_token],]\n",
    "    input_data.append(tokens[:-1])\n",
    "    output_data.append(tokens[1:])\n",
    "\n",
    "input_data = pad_sequences(input_data, MAX_LEN, value=vocab[vocab.padding_token])\n",
    "output_data = pad_sequences(output_data, MAX_LEN, value=vocab[vocab.padding_token])\n",
    "\n",
    "input_data = np.array(input_data, dtype=np.int64)\n",
    "output_data = np.array(output_data, dtype=np.int64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_object = tf.keras.losses.SparseCategoricalCrossentropy(\n",
    "    from_logits=True, reduction='none')\n",
    "\n",
    "train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='accuracy')\n",
    "\n",
    "def loss_function(real, pred):\n",
    "    mask = tf.math.logical_not(tf.math.equal(real, vocab[vocab.padding_token]))\n",
    "    loss_ = loss_object(real, pred)\n",
    "\n",
    "    mask = tf.cast(mask, dtype=loss_.dtype)\n",
    "    loss_ *= mask\n",
    "\n",
    "    return tf.reduce_mean(loss_)\n",
    "\n",
    "def accuracy_function(real, pred):\n",
    "    mask = tf.math.logical_not(tf.math.equal(real, vocab[vocab.padding_token]))\n",
    "    mask = tf.expand_dims(tf.cast(mask, dtype=pred.dtype), axis=-1)\n",
    "    pred *= mask    \n",
    "    acc = train_accuracy(real, pred)\n",
    "\n",
    "    return tf.reduce_mean(acc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gpt_model.compile(loss=loss_function,\n",
    "              optimizer=tf.keras.optimizers.Adam(1e-4),\n",
    "              metrics=[accuracy_function])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "history = gpt_model.fit(input_data, output_data, \n",
    "                    batch_size=BATCH_SIZE, epochs=NUM_EPOCHS,\n",
    "                    validation_split=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_OUT_PATH = './data_out'\n",
    "model_name = \"tf2_gpt2_finetuned_model\"\n",
    "\n",
    "save_path = os.path.join(DATA_OUT_PATH, model_name)\n",
    "\n",
    "if not os.path.exists(save_path):\n",
    "    os.makedirs(save_path)\n",
    "\n",
    "gpt_model.gpt2.save_pretrained(save_path)\n",
    "\n",
    "loaded_gpt_model = GPT2Model(save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_sent('이때', gpt_model, greedy=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_sent('이때', gpt_model, top_k=0, top_p=0.95)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}