gen_dis_explorer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from electra_model_tf2 import TFElectraGen, TFElectraDis\n",
    "from tokenizers.implementations import SentencePieceBPETokenizer\n",
    "from CocCocTokenizer import PyTokenizer\n",
    "import tensorflow as tf\n",
    "T = PyTokenizer(load_nontone_data=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer = SentencePieceBPETokenizer(\n",
    "    \"./vocab/vocab.json\",\n",
    "    \"./vocab/merges.txt\",\n",
    ")\n",
    "tokenizer.add_special_tokens([\"[PAD]\", \"[UNK]\", \"[CLS]\", \"[SEP]\", \"[MASK]\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def clean_text(txt_raw):\n",
    "    return ' '.join(T.word_tokenize(txt_raw, tokenize_option=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def encode_text(text):\n",
    "    text_encode = tokenizer.encode(text)\n",
    "    plain_text = ['[CLS]'] + text_encode.tokens + ['[SEP]']\n",
    "    indices = [tokenizer.token_to_id(\"[CLS]\")] + text_encode.ids + [tokenizer.token_to_id(\"[SEP]\")]\n",
    "    assert len(plain_text) == len(indices)\n",
    "    return plain_text, indices"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen_model = TFElectraGen.from_pretrained(\"./model_pretrained/gen/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "dis_model = TFElectraDis.from_pretrained(\"./model_pretrained/dis/\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[CLS] ▁trân_trọng ▁gửi ▁đến ▁quý ▁thành_viên ▁những ▁thông_tin ▁pháp_luật ▁nổi_bật ▁sau ▁đây [SEP]\n"
     ]
    }
   ],
   "source": [
    "text = \"trân trọng gửi đến quý thành viên những thông tin pháp luật nổi bật sau đây\"\n",
    "txt_plain, txt_indices = encode_text(clean_text(text))\n",
    "print(\" \".join(txt_plain))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Mask an item\n",
    "mask item at index 2. (word \"▁gửi\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[CLS] ▁trân_trọng [MASK] ▁đến ▁quý ▁thành_viên ▁những ▁thông_tin ▁pháp_luật ▁nổi_bật ▁sau ▁đây [SEP]\n",
      "[64002, 8806, 64004, 1168, 2673, 2545, 1125, 1705, 2211, 4973, 1288, 1433, 64003]\n"
     ]
    }
   ],
   "source": [
    "masked_index = 2\n",
    "txt_plain[masked_index] = '[MASK]'\n",
    "txt_indices[masked_index] = tokenizer.token_to_id('[MASK]')\n",
    "print(\" \".join(txt_plain))\n",
    "print(txt_indices)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Find masked item"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([1, 13, 64005])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gen_model_output = gen_model.get_output_generator_task(tf.constant([txt_indices]))\n",
    "gen_model_output.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Find top 50 words for the masked index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['▁quan_tâm', '▁sâu_sắc', '▁nhớ', '▁nhắc', '▁gửi', '▁tính', '▁cho', '▁liên_quan', '▁đề_cập', '▁mang', '▁tưởng_nhớ', '▁hướng', '▁biết', '▁thông_tin', '▁nhấn_mạnh', '▁chú_ý', '▁tốt_đẹp', '▁trân_trọng', '▁dành', '▁chân_thành', '▁cảm_ơn', '▁kỷ_niệm', '▁điện_mừng', '▁nói', '▁có', '▁đem', '▁giới_thiệu', '▁tin', '▁hơn', '▁gửi_gắm', '▁thông_điệp', '▁lời_chúc', '▁ý_kiến', '▁cảm_nhận', '▁ghi_nhận', '▁sự', '▁dẫn', '▁hình_ảnh', '▁nhận_thức', '▁sâu_đậm', '▁biết_ơn', '▁người', '▁tiếp', '▁đóng_góp', '▁đến', '▁thiết_thực', '▁,', '▁với', '▁nghĩ', '▁rõ']\n"
     ]
    }
   ],
   "source": [
    "print([tokenizer.id_to_token(item) for item in tf.math.top_k(gen_model_output[0, masked_index], 50).indices.numpy().tolist()])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Unmark item\n",
    "Choose a word and fill to mask index (word '▁quan_tâm' for example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[CLS] ▁trân_trọng ▁quan_tâm ▁đến ▁quý ▁thành_viên ▁những ▁thông_tin ▁pháp_luật ▁nổi_bật ▁sau ▁đây [SEP]\n",
      "[64002, 8806, 2768, 1168, 2673, 2545, 1125, 1705, 2211, 4973, 1288, 1433, 64003]\n"
     ]
    }
   ],
   "source": [
    "mark_token = '▁quan_tâm'\n",
    "txt_plain[masked_index] = mark_token\n",
    "txt_indices[masked_index] = tokenizer.token_to_id(mark_token)\n",
    "print(\" \".join(txt_plain))\n",
    "print(txt_indices)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Find replaced item"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(13,)\n"
     ]
    }
   ],
   "source": [
    "dis_model_output = dis_model.get_output_discriminator_task(tf.constant([txt_indices]))\n",
    "print(dis_model_output.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- True mean replaced item\n",
    "- False mean origin item"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[False, False, True, False, False, False, False, False, False, False, False, False, False]\n",
      "Replaced words: ['▁quan_tâm']\n"
     ]
    }
   ],
   "source": [
    "plain_result = list(tf.cast(((tf.math.sign(dis_model_output) + 1) / 2), tf.bool).numpy())\n",
    "print(plain_result)\n",
    "print(\"Replaced words: {}\".format([txt_plain[i] for i, value in enumerate(plain_result) if value]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Discriminator model extract features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['[CLS]', '▁trân_trọng', '▁gửi', '▁đến', '▁quý', '▁thành_viên', '▁những', '▁thông_tin', '▁pháp_luật', '▁nổi_bật', '▁sau', '▁đây', '[SEP]']\n",
      "[64002, 8806, 2320, 1168, 2673, 2545, 1125, 1705, 2211, 4973, 1288, 1433, 64003]\n"
     ]
    }
   ],
   "source": [
    "text = \"trân trọng gửi đến quý thành viên những thông tin pháp luật nổi bật sau đây\"\n",
    "txt_plain, txt_indices = encode_text(clean_text(text))\n",
    "print(txt_plain) \n",
    "print(txt_indices)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "dis_model_output = dis_model(tf.constant([txt_indices]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Discriminator output to detect replaced item"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([13])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dis_model_output[0].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Discriminator output features (12 layers + 1 output layer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "13"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(dis_model_output[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([1, 13, 256])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dis_model_output[1][-1].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}