01-indobart-bertopic-00.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys\n",
    "import torch\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from datasets import load_dataset\n",
    "\n",
    "import nltk\n",
    "import evaluate\n",
    "from transformers import DataCollatorForSeq2Seq, AutoTokenizer\n",
    "from transformers import AutoModelForSeq2SeqLM, Seq2SeqTrainingArguments, Seq2SeqTrainer\n",
    "from transformers import pipeline\n",
    "from transformers import BertTokenizer, AutoModel\n",
    "\n",
    "from umap import UMAP\n",
    "from hdbscan import HDBSCAN\n",
    "from sentence_transformers import SentenceTransformer\n",
    "from sklearn.feature_extraction.text import CountVectorizer\n",
    "from nltk.corpus import stopwords\n",
    "\n",
    "from bertopic import BERTopic\n",
    "from bertopic.representation import KeyBERTInspired\n",
    "from bertopic.vectorizers import ClassTfidfTransformer\n",
    "\n",
    "from repo.indobenchmark.toolkit.tokenization_indonlg import IndoNLGTokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "using cuda\n"
     ]
    }
   ],
   "source": [
    "# check if cuda or mps available, if available, use one of them, otherwise use cpu\n",
    "\n",
    "device = torch.device(\"cpu\")\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda\")\n",
    "    print(\"using cuda\")\n",
    "elif torch.backends.mps.is_available():\n",
    "    device = torch.device(\"mps\")\n",
    "    # os.environ[\"PYTORCH_MPS_HIGH_WATERMARK_RATIO\"] = \"1\" # This is tracked as pytorch issue #98222\n",
    "    print(\"using mps\")\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "    print(\"using cpu\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify the directory and file path\n",
    "benc_directory = \"benc_result/01-indobart-bertopic/\"\n",
    "\n",
    "# Create the benc_directory if it doesn't exist\n",
    "os.makedirs(benc_directory, exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>document</th>\n",
       "      <th>id</th>\n",
       "      <th>summary</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Jakarta, CNN Indonesia - - Dokter Ryan Thamrin...</td>\n",
       "      <td>1501893029-lula-kamal-dokter-ryan-thamrin-saki...</td>\n",
       "      <td>Dokter Lula Kamal yang merupakan selebriti sek...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Selfie ialah salah satu tema terpanas di kalan...</td>\n",
       "      <td>1509072914-dua-smartphone-zenfone-baru-tawarka...</td>\n",
       "      <td>Asus memperkenalkan   ZenFone generasi keempat...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Jakarta, CNN Indonesia - - Dinas Pariwisata Pr...</td>\n",
       "      <td>1510613677-songsong-visit-2020-bengkulu-perkua...</td>\n",
       "      <td>Dinas Pariwisata Provinsi Bengkulu kembali men...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Merdeka.com - Indonesia Corruption Watch (ICW)...</td>\n",
       "      <td>1502706803-icw-ada-kejanggalan-atas-tewasnya-s...</td>\n",
       "      <td>Indonesia Corruption Watch (ICW) meminta Komis...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Merdeka.com - Presiden Joko Widodo (Jokowi) me...</td>\n",
       "      <td>1503039338-pembagian-sepeda-usai-upacara-penur...</td>\n",
       "      <td>Jokowi memimpin upacara penurunan bendera. Usa...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                            document  \\\n",
       "0  Jakarta, CNN Indonesia - - Dokter Ryan Thamrin...   \n",
       "1  Selfie ialah salah satu tema terpanas di kalan...   \n",
       "2  Jakarta, CNN Indonesia - - Dinas Pariwisata Pr...   \n",
       "3  Merdeka.com - Indonesia Corruption Watch (ICW)...   \n",
       "4  Merdeka.com - Presiden Joko Widodo (Jokowi) me...   \n",
       "\n",
       "                                                  id  \\\n",
       "0  1501893029-lula-kamal-dokter-ryan-thamrin-saki...   \n",
       "1  1509072914-dua-smartphone-zenfone-baru-tawarka...   \n",
       "2  1510613677-songsong-visit-2020-bengkulu-perkua...   \n",
       "3  1502706803-icw-ada-kejanggalan-atas-tewasnya-s...   \n",
       "4  1503039338-pembagian-sepeda-usai-upacara-penur...   \n",
       "\n",
       "                                             summary  \n",
       "0  Dokter Lula Kamal yang merupakan selebriti sek...  \n",
       "1  Asus memperkenalkan   ZenFone generasi keempat...  \n",
       "2  Dinas Pariwisata Provinsi Bengkulu kembali men...  \n",
       "3  Indonesia Corruption Watch (ICW) meminta Komis...  \n",
       "4  Jokowi memimpin upacara penurunan bendera. Usa...  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# ds = load_dataset('maryantocinn/indosum')\n",
    "ds = load_dataset(\"./repo/SEACrowd/indosum/indosum.py\")\n",
    "\n",
    "# pandas dataframe\n",
    "df_ds = {}\n",
    "for key in ds.keys():\n",
    "    df_ds[key] = ds[key].to_pandas()\n",
    "\n",
    "# show first 5 data from the dataset in pandas like table\n",
    "df_ds[\"train\"].head()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train dataset length:  14262\n",
      "Validation dataset length:  750\n",
      "Test dataset length:  3762\n"
     ]
    }
   ],
   "source": [
    "# check the length of the dataset\n",
    "print(\"Train dataset length: \", len(ds[\"train\"]))\n",
    "print(\"Validation dataset length: \", len(ds[\"validation\"]))\n",
    "print(\"Test dataset length: \", len(ds[\"test\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nltk.download(\"all\", quiet=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Topic Modelling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aad3f9c184d04a2abfff7a54cbf89789",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "modules.json:   0%|          | 0.00/229 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "630de9e821834789be9a774d9abfd25d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config_sentence_transformers.json:   0%|          | 0.00/176 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0a27232ee0684b53a34eeda3e1bfed78",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "README.md:   0%|          | 0.00/4.37k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2b36b35e14024b748a920db372832be5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "sentence_bert_config.json:   0%|          | 0.00/53.0 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "346e5c3916b64f56a583842cfe527ec4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config.json:   0%|          | 0.00/1.10k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f0c430209e4d4698a2b77a925af1bbeb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/498M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5f04345013ea40629432dfc470dc49ea",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/1.26k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cbe183b9a6ac4b3eaf29cb0d8e8dc768",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "vocab.txt:   0%|          | 0.00/229k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "08cc6fc4b5d44e4098c40778886dfcf5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer.json:   0%|          | 0.00/709k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a84080ef6ac54f86b16bde56f4ac354c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "special_tokens_map.json:   0%|          | 0.00/125 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6f4580a0bc084330b46c54bde264f7c7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "1_Pooling/config.json:   0%|          | 0.00/296 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "embedding_model = SentenceTransformer(\"LazarusNLP/all-indobert-base-v4\")\n",
    "\n",
    "stop_words = (\n",
    "    stopwords.words(\"indonesian\")\n",
    "    + stopwords.words(\"english\")\n",
    "    + [\"cnn\", \"dailysocial\", \"id\"]\n",
    ")\n",
    "vectorizer_model = CountVectorizer(stop_words=stop_words, token_pattern=\"[^\\W\\d_]+\")\n",
    "\n",
    "topic_model = BERTopic(\n",
    "    embedding_model=embedding_model,\n",
    "    vectorizer_model=vectorizer_model,\n",
    "    nr_topics=10,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-11-04 06:13:54,636 - BERTopic - Embedding - Transforming documents to embeddings.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d64cf9a28e644e3bb22b46a449d98e8b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Batches:   0%|          | 0/446 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-11-04 06:14:43,531 - BERTopic - Embedding - Completed ✓\n",
      "2024-11-04 06:14:43,533 - BERTopic - Dimensionality - Fitting the dimensionality reduction algorithm\n",
      "2024-11-04 06:15:28,358 - BERTopic - Dimensionality - Completed ✓\n",
      "2024-11-04 06:15:28,361 - BERTopic - Cluster - Start clustering the reduced embeddings\n",
      "2024-11-04 06:15:29,215 - BERTopic - Cluster - Completed ✓\n",
      "2024-11-04 06:15:29,216 - BERTopic - Representation - Extracting topics from clusters using representation models.\n",
      "2024-11-04 06:15:36,248 - BERTopic - Representation - Completed ✓\n",
      "2024-11-04 06:15:36,255 - BERTopic - Topic reduction - Reducing number of topics\n",
      "2024-11-04 06:15:41,550 - BERTopic - Topic reduction - Reduced number of topics from 163 to 10\n"
     ]
    }
   ],
   "source": [
    "topics, probs = topic_model.fit_transform(ds[\"train\"][\"document\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Topic</th>\n",
       "      <th>Count</th>\n",
       "      <th>Name</th>\n",
       "      <th>Representation</th>\n",
       "      <th>Representative_Docs</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>-1</td>\n",
       "      <td>4938</td>\n",
       "      <td>-1_indonesia_jakarta_orang_negara</td>\n",
       "      <td>[indonesia, jakarta, orang, negara, anak, sala...</td>\n",
       "      <td>[Jakarta, CNN Indonesia - - Komisi Pemberantas...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0</td>\n",
       "      <td>3207</td>\n",
       "      <td>0_pemain_gol_musim_laga</td>\n",
       "      <td>[pemain, gol, musim, laga, menit, tim, liga, p...</td>\n",
       "      <td>[Dua gol dari Isco berkontribusi membawa Spany...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>2295</td>\n",
       "      <td>1_jakarta_kpk_indonesia_jalan</td>\n",
       "      <td>[jakarta, kpk, indonesia, jalan, partai, ketua...</td>\n",
       "      <td>[Jakarta (ANTARA News) - Bimanesh Sutarjo, dok...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2</td>\n",
       "      <td>1175</td>\n",
       "      <td>2_film_lagu_album_jakarta</td>\n",
       "      <td>[film, lagu, album, jakarta, indonesia, orang,...</td>\n",
       "      <td>[Jakarta, CNN Indonesia - - Film horor Indones...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>3</td>\n",
       "      <td>1038</td>\n",
       "      <td>3_startup_teknologi_berita_inovasi</td>\n",
       "      <td>[startup, teknologi, berita, inovasi, fitur, i...</td>\n",
       "      <td>[Sebelum tips ini, kita sudah pernah memberika...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>4</td>\n",
       "      <td>620</td>\n",
       "      <td>4_trump_negara_presiden_israel</td>\n",
       "      <td>[trump, negara, presiden, israel, orang, ameri...</td>\n",
       "      <td>[Presiden Amerika Serikat Donald Trump secara ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>5</td>\n",
       "      <td>411</td>\n",
       "      <td>5_indonesia_makanan_pariwisata_kopi</td>\n",
       "      <td>[indonesia, makanan, pariwisata, kopi, wisata,...</td>\n",
       "      <td>[Jakarta, CNN Indonesia - - Festival Seni Buda...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>6</td>\n",
       "      <td>353</td>\n",
       "      <td>6_tubuh_kanker_makanan_penyakit</td>\n",
       "      <td>[tubuh, kanker, makanan, penyakit, kulit, pene...</td>\n",
       "      <td>[Merdeka.com - Musim hujan telah datang dan in...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>7</td>\n",
       "      <td>172</td>\n",
       "      <td>7_persen_rp_harga_pajak</td>\n",
       "      <td>[persen, rp, harga, pajak, indonesia, minyak, ...</td>\n",
       "      <td>[Jakarta, CNN Indonesia - - Menteri Koordinato...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>8</td>\n",
       "      <td>53</td>\n",
       "      <td>8_pendidikan_sekolah_guru_ui</td>\n",
       "      <td>[pendidikan, sekolah, guru, ui, muhadjir, perg...</td>\n",
       "      <td>[Jakarta, CNN Indonesia - - Menteri Pendidikan...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   Topic  Count                                 Name  \\\n",
       "0     -1   4938    -1_indonesia_jakarta_orang_negara   \n",
       "1      0   3207              0_pemain_gol_musim_laga   \n",
       "2      1   2295        1_jakarta_kpk_indonesia_jalan   \n",
       "3      2   1175            2_film_lagu_album_jakarta   \n",
       "4      3   1038   3_startup_teknologi_berita_inovasi   \n",
       "5      4    620       4_trump_negara_presiden_israel   \n",
       "6      5    411  5_indonesia_makanan_pariwisata_kopi   \n",
       "7      6    353      6_tubuh_kanker_makanan_penyakit   \n",
       "8      7    172              7_persen_rp_harga_pajak   \n",
       "9      8     53         8_pendidikan_sekolah_guru_ui   \n",
       "\n",
       "                                      Representation  \\\n",
       "0  [indonesia, jakarta, orang, negara, anak, sala...   \n",
       "1  [pemain, gol, musim, laga, menit, tim, liga, p...   \n",
       "2  [jakarta, kpk, indonesia, jalan, partai, ketua...   \n",
       "3  [film, lagu, album, jakarta, indonesia, orang,...   \n",
       "4  [startup, teknologi, berita, inovasi, fitur, i...   \n",
       "5  [trump, negara, presiden, israel, orang, ameri...   \n",
       "6  [indonesia, makanan, pariwisata, kopi, wisata,...   \n",
       "7  [tubuh, kanker, makanan, penyakit, kulit, pene...   \n",
       "8  [persen, rp, harga, pajak, indonesia, minyak, ...   \n",
       "9  [pendidikan, sekolah, guru, ui, muhadjir, perg...   \n",
       "\n",
       "                                 Representative_Docs  \n",
       "0  [Jakarta, CNN Indonesia - - Komisi Pemberantas...  \n",
       "1  [Dua gol dari Isco berkontribusi membawa Spany...  \n",
       "2  [Jakarta (ANTARA News) - Bimanesh Sutarjo, dok...  \n",
       "3  [Jakarta, CNN Indonesia - - Film horor Indones...  \n",
       "4  [Sebelum tips ini, kita sudah pernah memberika...  \n",
       "5  [Presiden Amerika Serikat Donald Trump secara ...  \n",
       "6  [Jakarta, CNN Indonesia - - Festival Seni Buda...  \n",
       "7  [Merdeka.com - Musim hujan telah datang dan in...  \n",
       "8  [Jakarta, CNN Indonesia - - Menteri Koordinato...  \n",
       "9  [Jakarta, CNN Indonesia - - Menteri Pendidikan...  "
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "topic_info = topic_model.get_topic_info()\n",
    "\n",
    "# save to excel\n",
    "topic_info.to_csv(benc_directory + \"topic_info.csv\")\n",
    "\n",
    "topic_info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>document</th>\n",
       "      <th>id</th>\n",
       "      <th>summary</th>\n",
       "      <th>Document</th>\n",
       "      <th>Topic</th>\n",
       "      <th>Name</th>\n",
       "      <th>Representation</th>\n",
       "      <th>Representative_Docs</th>\n",
       "      <th>Top_n_words</th>\n",
       "      <th>Probability</th>\n",
       "      <th>Representative_document</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Jakarta, CNN Indonesia - - Dokter Ryan Thamrin...</td>\n",
       "      <td>1501893029-lula-kamal-dokter-ryan-thamrin-saki...</td>\n",
       "      <td>Dokter Lula Kamal yang merupakan selebriti sek...</td>\n",
       "      <td>Jakarta, CNN Indonesia - - Dokter Ryan Thamrin...</td>\n",
       "      <td>-1</td>\n",
       "      <td>-1_indonesia_jakarta_orang_negara</td>\n",
       "      <td>[indonesia, jakarta, orang, negara, anak, sala...</td>\n",
       "      <td>[Jakarta, CNN Indonesia - - Komisi Pemberantas...</td>\n",
       "      <td>indonesia - jakarta - orang - negara - anak - ...</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Selfie ialah salah satu tema terpanas di kalan...</td>\n",
       "      <td>1509072914-dua-smartphone-zenfone-baru-tawarka...</td>\n",
       "      <td>Asus memperkenalkan   ZenFone generasi keempat...</td>\n",
       "      <td>Selfie ialah salah satu tema terpanas di kalan...</td>\n",
       "      <td>3</td>\n",
       "      <td>3_startup_teknologi_berita_inovasi</td>\n",
       "      <td>[startup, teknologi, berita, inovasi, fitur, i...</td>\n",
       "      <td>[Sebelum tips ini, kita sudah pernah memberika...</td>\n",
       "      <td>startup - teknologi - berita - inovasi - fitur...</td>\n",
       "      <td>0.502241</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Jakarta, CNN Indonesia - - Dinas Pariwisata Pr...</td>\n",
       "      <td>1510613677-songsong-visit-2020-bengkulu-perkua...</td>\n",
       "      <td>Dinas Pariwisata Provinsi Bengkulu kembali men...</td>\n",
       "      <td>Jakarta, CNN Indonesia - - Dinas Pariwisata Pr...</td>\n",
       "      <td>5</td>\n",
       "      <td>5_indonesia_makanan_pariwisata_kopi</td>\n",
       "      <td>[indonesia, makanan, pariwisata, kopi, wisata,...</td>\n",
       "      <td>[Jakarta, CNN Indonesia - - Festival Seni Buda...</td>\n",
       "      <td>indonesia - makanan - pariwisata - kopi - wisa...</td>\n",
       "      <td>0.931050</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Merdeka.com - Indonesia Corruption Watch (ICW)...</td>\n",
       "      <td>1502706803-icw-ada-kejanggalan-atas-tewasnya-s...</td>\n",
       "      <td>Indonesia Corruption Watch (ICW) meminta Komis...</td>\n",
       "      <td>Merdeka.com - Indonesia Corruption Watch (ICW)...</td>\n",
       "      <td>-1</td>\n",
       "      <td>-1_indonesia_jakarta_orang_negara</td>\n",
       "      <td>[indonesia, jakarta, orang, negara, anak, sala...</td>\n",
       "      <td>[Jakarta, CNN Indonesia - - Komisi Pemberantas...</td>\n",
       "      <td>indonesia - jakarta - orang - negara - anak - ...</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Merdeka.com - Presiden Joko Widodo (Jokowi) me...</td>\n",
       "      <td>1503039338-pembagian-sepeda-usai-upacara-penur...</td>\n",
       "      <td>Jokowi memimpin upacara penurunan bendera. Usa...</td>\n",
       "      <td>Merdeka.com - Presiden Joko Widodo (Jokowi) me...</td>\n",
       "      <td>1</td>\n",
       "      <td>1_jakarta_kpk_indonesia_jalan</td>\n",
       "      <td>[jakarta, kpk, indonesia, jalan, partai, ketua...</td>\n",
       "      <td>[Jakarta (ANTARA News) - Bimanesh Sutarjo, dok...</td>\n",
       "      <td>jakarta - kpk - indonesia - jalan - partai - k...</td>\n",
       "      <td>0.995418</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14257</th>\n",
       "      <td>Jakarta, CNN Indonesia - - Amerika Serikat dil...</td>\n",
       "      <td>1497645345-as-kirimkan-peluncur-rudal-ke-suria...</td>\n",
       "      <td>Amerika Serikat dilaporkan telah mengirimkan s...</td>\n",
       "      <td>Jakarta, CNN Indonesia - - Amerika Serikat dil...</td>\n",
       "      <td>4</td>\n",
       "      <td>4_trump_negara_presiden_israel</td>\n",
       "      <td>[trump, negara, presiden, israel, orang, ameri...</td>\n",
       "      <td>[Presiden Amerika Serikat Donald Trump secara ...</td>\n",
       "      <td>trump - negara - presiden - israel - orang - a...</td>\n",
       "      <td>0.403283</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14258</th>\n",
       "      <td>Bandung, CNN Indonesia - - Borneo FC berhasil ...</td>\n",
       "      <td>1495406700-borneo-bersyukur-tahan-persib-di-gbla</td>\n",
       "      <td>Borneo FC menahan imbang Persib Bandung pada l...</td>\n",
       "      <td>Bandung, CNN Indonesia - - Borneo FC berhasil ...</td>\n",
       "      <td>-1</td>\n",
       "      <td>-1_indonesia_jakarta_orang_negara</td>\n",
       "      <td>[indonesia, jakarta, orang, negara, anak, sala...</td>\n",
       "      <td>[Jakarta, CNN Indonesia - - Komisi Pemberantas...</td>\n",
       "      <td>indonesia - jakarta - orang - negara - anak - ...</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14259</th>\n",
       "      <td>JAKARTA (Pos Kota) – Komisi Pemberantasan Koru...</td>\n",
       "      <td>1513941815-mantan-dirjen-perhubungan-laut-sege...</td>\n",
       "      <td>Komisi Pemberantasan Korupsi (KPK) sudah melim...</td>\n",
       "      <td>JAKARTA (Pos Kota) – Komisi Pemberantasan Koru...</td>\n",
       "      <td>1</td>\n",
       "      <td>1_jakarta_kpk_indonesia_jalan</td>\n",
       "      <td>[jakarta, kpk, indonesia, jalan, partai, ketua...</td>\n",
       "      <td>[Jakarta (ANTARA News) - Bimanesh Sutarjo, dok...</td>\n",
       "      <td>jakarta - kpk - indonesia - jalan - partai - k...</td>\n",
       "      <td>0.781011</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14260</th>\n",
       "      <td>Merdeka.com - Sebuah kabar gembira datang bagi...</td>\n",
       "      <td>1496440800-rangking-fifa-indonesia-naik-dua-pe...</td>\n",
       "      <td>Kabar gembira datang bagi   sepakbola   Indone...</td>\n",
       "      <td>Merdeka.com - Sebuah kabar gembira datang bagi...</td>\n",
       "      <td>-1</td>\n",
       "      <td>-1_indonesia_jakarta_orang_negara</td>\n",
       "      <td>[indonesia, jakarta, orang, negara, anak, sala...</td>\n",
       "      <td>[Jakarta, CNN Indonesia - - Komisi Pemberantas...</td>\n",
       "      <td>indonesia - jakarta - orang - negara - anak - ...</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14261</th>\n",
       "      <td>Ada satu perangkat menarik lainnya yang tersaj...</td>\n",
       "      <td>1505421900-apple-tv-4k-usung-resolusi-super-be...</td>\n",
       "      <td>Ada satu perangkat menarik lainnya yang tersaj...</td>\n",
       "      <td>Ada satu perangkat menarik lainnya yang tersaj...</td>\n",
       "      <td>-1</td>\n",
       "      <td>-1_indonesia_jakarta_orang_negara</td>\n",
       "      <td>[indonesia, jakarta, orang, negara, anak, sala...</td>\n",
       "      <td>[Jakarta, CNN Indonesia - - Komisi Pemberantas...</td>\n",
       "      <td>indonesia - jakarta - orang - negara - anak - ...</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>14262 rows × 11 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                document  \\\n",
       "0      Jakarta, CNN Indonesia - - Dokter Ryan Thamrin...   \n",
       "1      Selfie ialah salah satu tema terpanas di kalan...   \n",
       "2      Jakarta, CNN Indonesia - - Dinas Pariwisata Pr...   \n",
       "3      Merdeka.com - Indonesia Corruption Watch (ICW)...   \n",
       "4      Merdeka.com - Presiden Joko Widodo (Jokowi) me...   \n",
       "...                                                  ...   \n",
       "14257  Jakarta, CNN Indonesia - - Amerika Serikat dil...   \n",
       "14258  Bandung, CNN Indonesia - - Borneo FC berhasil ...   \n",
       "14259  JAKARTA (Pos Kota) – Komisi Pemberantasan Koru...   \n",
       "14260  Merdeka.com - Sebuah kabar gembira datang bagi...   \n",
       "14261  Ada satu perangkat menarik lainnya yang tersaj...   \n",
       "\n",
       "                                                      id  \\\n",
       "0      1501893029-lula-kamal-dokter-ryan-thamrin-saki...   \n",
       "1      1509072914-dua-smartphone-zenfone-baru-tawarka...   \n",
       "2      1510613677-songsong-visit-2020-bengkulu-perkua...   \n",
       "3      1502706803-icw-ada-kejanggalan-atas-tewasnya-s...   \n",
       "4      1503039338-pembagian-sepeda-usai-upacara-penur...   \n",
       "...                                                  ...   \n",
       "14257  1497645345-as-kirimkan-peluncur-rudal-ke-suria...   \n",
       "14258   1495406700-borneo-bersyukur-tahan-persib-di-gbla   \n",
       "14259  1513941815-mantan-dirjen-perhubungan-laut-sege...   \n",
       "14260  1496440800-rangking-fifa-indonesia-naik-dua-pe...   \n",
       "14261  1505421900-apple-tv-4k-usung-resolusi-super-be...   \n",
       "\n",
       "                                                 summary  \\\n",
       "0      Dokter Lula Kamal yang merupakan selebriti sek...   \n",
       "1      Asus memperkenalkan   ZenFone generasi keempat...   \n",
       "2      Dinas Pariwisata Provinsi Bengkulu kembali men...   \n",
       "3      Indonesia Corruption Watch (ICW) meminta Komis...   \n",
       "4      Jokowi memimpin upacara penurunan bendera. Usa...   \n",
       "...                                                  ...   \n",
       "14257  Amerika Serikat dilaporkan telah mengirimkan s...   \n",
       "14258  Borneo FC menahan imbang Persib Bandung pada l...   \n",
       "14259  Komisi Pemberantasan Korupsi (KPK) sudah melim...   \n",
       "14260  Kabar gembira datang bagi   sepakbola   Indone...   \n",
       "14261  Ada satu perangkat menarik lainnya yang tersaj...   \n",
       "\n",
       "                                                Document  Topic  \\\n",
       "0      Jakarta, CNN Indonesia - - Dokter Ryan Thamrin...     -1   \n",
       "1      Selfie ialah salah satu tema terpanas di kalan...      3   \n",
       "2      Jakarta, CNN Indonesia - - Dinas Pariwisata Pr...      5   \n",
       "3      Merdeka.com - Indonesia Corruption Watch (ICW)...     -1   \n",
       "4      Merdeka.com - Presiden Joko Widodo (Jokowi) me...      1   \n",
       "...                                                  ...    ...   \n",
       "14257  Jakarta, CNN Indonesia - - Amerika Serikat dil...      4   \n",
       "14258  Bandung, CNN Indonesia - - Borneo FC berhasil ...     -1   \n",
       "14259  JAKARTA (Pos Kota) – Komisi Pemberantasan Koru...      1   \n",
       "14260  Merdeka.com - Sebuah kabar gembira datang bagi...     -1   \n",
       "14261  Ada satu perangkat menarik lainnya yang tersaj...     -1   \n",
       "\n",
       "                                      Name  \\\n",
       "0        -1_indonesia_jakarta_orang_negara   \n",
       "1       3_startup_teknologi_berita_inovasi   \n",
       "2      5_indonesia_makanan_pariwisata_kopi   \n",
       "3        -1_indonesia_jakarta_orang_negara   \n",
       "4            1_jakarta_kpk_indonesia_jalan   \n",
       "...                                    ...   \n",
       "14257       4_trump_negara_presiden_israel   \n",
       "14258    -1_indonesia_jakarta_orang_negara   \n",
       "14259        1_jakarta_kpk_indonesia_jalan   \n",
       "14260    -1_indonesia_jakarta_orang_negara   \n",
       "14261    -1_indonesia_jakarta_orang_negara   \n",
       "\n",
       "                                          Representation  \\\n",
       "0      [indonesia, jakarta, orang, negara, anak, sala...   \n",
       "1      [startup, teknologi, berita, inovasi, fitur, i...   \n",
       "2      [indonesia, makanan, pariwisata, kopi, wisata,...   \n",
       "3      [indonesia, jakarta, orang, negara, anak, sala...   \n",
       "4      [jakarta, kpk, indonesia, jalan, partai, ketua...   \n",
       "...                                                  ...   \n",
       "14257  [trump, negara, presiden, israel, orang, ameri...   \n",
       "14258  [indonesia, jakarta, orang, negara, anak, sala...   \n",
       "14259  [jakarta, kpk, indonesia, jalan, partai, ketua...   \n",
       "14260  [indonesia, jakarta, orang, negara, anak, sala...   \n",
       "14261  [indonesia, jakarta, orang, negara, anak, sala...   \n",
       "\n",
       "                                     Representative_Docs  \\\n",
       "0      [Jakarta, CNN Indonesia - - Komisi Pemberantas...   \n",
       "1      [Sebelum tips ini, kita sudah pernah memberika...   \n",
       "2      [Jakarta, CNN Indonesia - - Festival Seni Buda...   \n",
       "3      [Jakarta, CNN Indonesia - - Komisi Pemberantas...   \n",
       "4      [Jakarta (ANTARA News) - Bimanesh Sutarjo, dok...   \n",
       "...                                                  ...   \n",
       "14257  [Presiden Amerika Serikat Donald Trump secara ...   \n",
       "14258  [Jakarta, CNN Indonesia - - Komisi Pemberantas...   \n",
       "14259  [Jakarta (ANTARA News) - Bimanesh Sutarjo, dok...   \n",
       "14260  [Jakarta, CNN Indonesia - - Komisi Pemberantas...   \n",
       "14261  [Jakarta, CNN Indonesia - - Komisi Pemberantas...   \n",
       "\n",
       "                                             Top_n_words  Probability  \\\n",
       "0      indonesia - jakarta - orang - negara - anak - ...     0.000000   \n",
       "1      startup - teknologi - berita - inovasi - fitur...     0.502241   \n",
       "2      indonesia - makanan - pariwisata - kopi - wisa...     0.931050   \n",
       "3      indonesia - jakarta - orang - negara - anak - ...     0.000000   \n",
       "4      jakarta - kpk - indonesia - jalan - partai - k...     0.995418   \n",
       "...                                                  ...          ...   \n",
       "14257  trump - negara - presiden - israel - orang - a...     0.403283   \n",
       "14258  indonesia - jakarta - orang - negara - anak - ...     0.000000   \n",
       "14259  jakarta - kpk - indonesia - jalan - partai - k...     0.781011   \n",
       "14260  indonesia - jakarta - orang - negara - anak - ...     0.000000   \n",
       "14261  indonesia - jakarta - orang - negara - anak - ...     0.000000   \n",
       "\n",
       "       Representative_document  \n",
       "0                        False  \n",
       "1                        False  \n",
       "2                        False  \n",
       "3                        False  \n",
       "4                        False  \n",
       "...                        ...  \n",
       "14257                    False  \n",
       "14258                    False  \n",
       "14259                    False  \n",
       "14260                    False  \n",
       "14261                    False  \n",
       "\n",
       "[14262 rows x 11 columns]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "topic_document_info = topic_model.get_document_info(ds[\"train\"][\"document\"], df_ds[\"train\"])\n",
    "\n",
    "# save to excel\n",
    "topic_document_info.to_csv(benc_directory + \"topic_document_info.csv\")\n",
    "\n",
    "topic_document_info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c4e58c1eaea149d9801ca7387e3ce3e8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=8):   0%|          | 0/14262 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>document</th>\n",
       "      <th>id</th>\n",
       "      <th>summary</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>&lt;tag&gt; indonesia jakarta orang negara anak sala...</td>\n",
       "      <td>1501893029-lula-kamal-dokter-ryan-thamrin-saki...</td>\n",
       "      <td>Dokter Lula Kamal yang merupakan selebriti sek...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>&lt;tag&gt; startup teknologi berita inovasi fitur i...</td>\n",
       "      <td>1509072914-dua-smartphone-zenfone-baru-tawarka...</td>\n",
       "      <td>Asus memperkenalkan   ZenFone generasi keempat...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>&lt;tag&gt; indonesia makanan pariwisata kopi wisata...</td>\n",
       "      <td>1510613677-songsong-visit-2020-bengkulu-perkua...</td>\n",
       "      <td>Dinas Pariwisata Provinsi Bengkulu kembali men...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>&lt;tag&gt; indonesia jakarta orang negara anak sala...</td>\n",
       "      <td>1502706803-icw-ada-kejanggalan-atas-tewasnya-s...</td>\n",
       "      <td>Indonesia Corruption Watch (ICW) meminta Komis...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>&lt;tag&gt; jakarta kpk indonesia jalan partai ketua...</td>\n",
       "      <td>1503039338-pembagian-sepeda-usai-upacara-penur...</td>\n",
       "      <td>Jokowi memimpin upacara penurunan bendera. Usa...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                            document  \\\n",
       "0  <tag> indonesia jakarta orang negara anak sala...   \n",
       "1  <tag> startup teknologi berita inovasi fitur i...   \n",
       "2  <tag> indonesia makanan pariwisata kopi wisata...   \n",
       "3  <tag> indonesia jakarta orang negara anak sala...   \n",
       "4  <tag> jakarta kpk indonesia jalan partai ketua...   \n",
       "\n",
       "                                                  id  \\\n",
       "0  1501893029-lula-kamal-dokter-ryan-thamrin-saki...   \n",
       "1  1509072914-dua-smartphone-zenfone-baru-tawarka...   \n",
       "2  1510613677-songsong-visit-2020-bengkulu-perkua...   \n",
       "3  1502706803-icw-ada-kejanggalan-atas-tewasnya-s...   \n",
       "4  1503039338-pembagian-sepeda-usai-upacara-penur...   \n",
       "\n",
       "                                             summary  \n",
       "0  Dokter Lula Kamal yang merupakan selebriti sek...  \n",
       "1  Asus memperkenalkan   ZenFone generasi keempat...  \n",
       "2  Dinas Pariwisata Provinsi Bengkulu kembali men...  \n",
       "3  Indonesia Corruption Watch (ICW) meminta Komis...  \n",
       "4  Jokowi memimpin upacara penurunan bendera. Usa...  "
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def add_topic(example, idx):\n",
    "    # if already have <tag>, return the example\n",
    "    if \"<tag>\" in example[\"document\"]:\n",
    "        return example\n",
    "\n",
    "    curr_topic = \" \".join(topic_document_info[\"Representation\"].values[idx])\n",
    "    example[\"document\"] = f\"<tag> {curr_topic} <tag> {example['document']}\"\n",
    "    \n",
    "    return example\n",
    "\n",
    "# get the processor number and set the number of process\n",
    "ds[\"train\"] = ds[\"train\"].map(add_topic, with_indices=True, num_proc=os.cpu_count())\n",
    "\n",
    "df_ds[\"train\"] = ds[\"train\"].to_pandas()\n",
    "\n",
    "df_ds[\"train\"].head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "IndoNLGTokenizer(name_or_path='indobenchmark/indobart-v2', vocab_size=40004, model_max_length=1000000000000000019884624838656, is_fast=False, padding_side='right', truncation_side='right', special_tokens={'bos_token': '<s>', 'eos_token': '</s>', 'unk_token': '<unk>', 'sep_token': '</s>', 'pad_token': '<pad>', 'cls_token': '<s>', 'mask_token': '<mask>', 'additional_special_tokens': ['<tag>']}, clean_up_tokenization_spaces=False),  added_tokens_decoder={\n",
       "\t0: AddedToken(\"<s>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
       "\t1: AddedToken(\"<pad>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
       "\t2: AddedToken(\"</s>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
       "\t3: AddedToken(\"<unk>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
       "\t39942: AddedToken(\"<tag>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
       "\t40003: AddedToken(\"<mask>\", rstrip=False, lstrip=True, single_word=False, normalized=True, special=True),\n",
       "}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bart_model = AutoModelForSeq2SeqLM.from_pretrained(\"indobenchmark/indobart-v2\")\n",
    "indonlg_tokenizer = IndoNLGTokenizer.from_pretrained(\"indobenchmark/indobart-v2\")\n",
    "\n",
    "indonlg_tokenizer.add_special_tokens({\"additional_special_tokens\": [\"<tag>\"]})\n",
    "\n",
    "model = bart_model\n",
    "tokenizer = indonlg_tokenizer\n",
    "\n",
    "tokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "571471ef53f54107873e9981fb568b87",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/14262 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "811f0f71b401410c883f7e3fbf014bfd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/3762 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1b22c6ff266d47bd81e79952a8235cfd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/750 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using auto half precision backend\n"
     ]
    }
   ],
   "source": [
    "# Prepare and tokenize dataset\n",
    "def preprocess_function(examples):\n",
    "    model_inputs = tokenizer(examples[\"document\"], max_length=768, truncation=True)\n",
    "\n",
    "    labels = tokenizer(text_target=examples[\"summary\"], max_length=128, truncation=True)\n",
    "\n",
    "    model_inputs[\"labels\"] = labels[\"input_ids\"]\n",
    "    return model_inputs\n",
    "\n",
    "\n",
    "# Setup evaluation\n",
    "nltk.download(\"punkt_tab\", quiet=True)\n",
    "metric = evaluate.load(\"rouge\")\n",
    "\n",
    "\n",
    "def compute_metrics(eval_preds):\n",
    "    preds, labels = eval_preds\n",
    "\n",
    "    # decode preds and labels\n",
    "    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)\n",
    "    decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)\n",
    "    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)\n",
    "\n",
    "    # rougeLSum expects newline after each sentence\n",
    "    decoded_preds = [\n",
    "        \"\\n\".join(nltk.sent_tokenize(pred.strip())) for pred in decoded_preds\n",
    "    ]\n",
    "    decoded_labels = [\n",
    "        \"\\n\".join(nltk.sent_tokenize(label.strip())) for label in decoded_labels\n",
    "    ]\n",
    "\n",
    "    result = metric.compute(\n",
    "        predictions=decoded_preds, references=decoded_labels, use_stemmer=True\n",
    "    )\n",
    "    return result\n",
    "\n",
    "\n",
    "tokenized_ds = ds.map(preprocess_function, batched=True)\n",
    "\n",
    "# Load pretrained model and evaluate model after each epoch\n",
    "data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)\n",
    "\n",
    "per_device_batch_size = 8  # 8 for low hardware spec\n",
    "output_dir = \"./results/00-indobart\"\n",
    "\n",
    "training_args = Seq2SeqTrainingArguments(\n",
    "    output_dir=output_dir,\n",
    "    # overwrite_output_dir=True,\n",
    "    eval_strategy=\"epoch\",\n",
    "    learning_rate=3.75e-5,  # hf example: 2e-5\n",
    "    per_device_train_batch_size=per_device_batch_size,\n",
    "    per_device_eval_batch_size=per_device_batch_size,\n",
    "    weight_decay=0.01,\n",
    "    save_steps=1000,\n",
    "    save_total_limit=2,\n",
    "    num_train_epochs=3,  # hf example: 2\n",
    "    fp16=True,  # comment this if using mps/apple sillicon chip (not supported)\n",
    "    predict_with_generate=True,\n",
    "    generation_max_length=80,\n",
    "    log_level=\"info\",\n",
    "    logging_first_step=True,\n",
    "    resume_from_checkpoint=True,\n",
    ")\n",
    "\n",
    "trainer = Seq2SeqTrainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=tokenized_ds[\"train\"],\n",
    "    eval_dataset=tokenized_ds[\"validation\"],\n",
    "    processing_class=tokenizer,  # FutureWarning: `tokenizer` is deprecated and will be removed in version 5.0.0 for `Seq2SeqTrainer.__init__`.\n",
    "    data_collator=data_collator,\n",
    "    compute_metrics=compute_metrics,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The following columns in the training set don't have a corresponding argument in `MBartForConditionalGeneration.forward` and have been ignored: summary, id, document. If summary, id, document are not expected by `MBartForConditionalGeneration.forward`,  you can safely ignore this message.\n",
      "***** Running training *****\n",
      "  Num examples = 14,262\n",
      "  Num Epochs = 3\n",
      "  Instantaneous batch size per device = 8\n",
      "  Total train batch size (w. parallel, distributed & accumulation) = 8\n",
      "  Gradient Accumulation steps = 1\n",
      "  Total optimization steps = 5,349\n",
      "  Number of trainable parameters = 131,543,040\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='5349' max='5349' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [5349/5349 26:24, Epoch 3/3]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Rouge1</th>\n",
       "      <th>Rouge2</th>\n",
       "      <th>Rougel</th>\n",
       "      <th>Rougelsum</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>0.545100</td>\n",
       "      <td>0.533041</td>\n",
       "      <td>0.666536</td>\n",
       "      <td>0.583236</td>\n",
       "      <td>0.626986</td>\n",
       "      <td>0.656955</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>0.416700</td>\n",
       "      <td>0.513260</td>\n",
       "      <td>0.669243</td>\n",
       "      <td>0.586440</td>\n",
       "      <td>0.629475</td>\n",
       "      <td>0.660222</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>0.328600</td>\n",
       "      <td>0.522803</td>\n",
       "      <td>0.660575</td>\n",
       "      <td>0.576017</td>\n",
       "      <td>0.619246</td>\n",
       "      <td>0.651756</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to ./results/00-indobart/checkpoint-1000\n",
      "Configuration saved in ./results/00-indobart/checkpoint-1000/config.json\n",
      "Configuration saved in ./results/00-indobart/checkpoint-1000/generation_config.json\n",
      "Model weights saved in ./results/00-indobart/checkpoint-1000/model.safetensors\n",
      "tokenizer config file saved in ./results/00-indobart/checkpoint-1000/tokenizer_config.json\n",
      "Special tokens file saved in ./results/00-indobart/checkpoint-1000/special_tokens_map.json\n",
      "Deleting older checkpoint [results/00-indobart/checkpoint-1000] due to args.save_total_limit\n",
      "The following columns in the evaluation set don't have a corresponding argument in `MBartForConditionalGeneration.forward` and have been ignored: summary, id, document. If summary, id, document are not expected by `MBartForConditionalGeneration.forward`,  you can safely ignore this message.\n",
      "\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 750\n",
      "  Batch size = 8\n",
      "Saving model checkpoint to ./results/00-indobart/checkpoint-2000\n",
      "Configuration saved in ./results/00-indobart/checkpoint-2000/config.json\n",
      "Configuration saved in ./results/00-indobart/checkpoint-2000/generation_config.json\n",
      "Model weights saved in ./results/00-indobart/checkpoint-2000/model.safetensors\n",
      "tokenizer config file saved in ./results/00-indobart/checkpoint-2000/tokenizer_config.json\n",
      "Special tokens file saved in ./results/00-indobart/checkpoint-2000/special_tokens_map.json\n",
      "Deleting older checkpoint [results/00-indobart/checkpoint-2000] due to args.save_total_limit\n",
      "Saving model checkpoint to ./results/00-indobart/checkpoint-3000\n",
      "Configuration saved in ./results/00-indobart/checkpoint-3000/config.json\n",
      "Configuration saved in ./results/00-indobart/checkpoint-3000/generation_config.json\n",
      "Model weights saved in ./results/00-indobart/checkpoint-3000/model.safetensors\n",
      "tokenizer config file saved in ./results/00-indobart/checkpoint-3000/tokenizer_config.json\n",
      "Special tokens file saved in ./results/00-indobart/checkpoint-3000/special_tokens_map.json\n",
      "Deleting older checkpoint [results/00-indobart/checkpoint-3000] due to args.save_total_limit\n",
      "The following columns in the evaluation set don't have a corresponding argument in `MBartForConditionalGeneration.forward` and have been ignored: summary, id, document. If summary, id, document are not expected by `MBartForConditionalGeneration.forward`,  you can safely ignore this message.\n",
      "\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 750\n",
      "  Batch size = 8\n",
      "Saving model checkpoint to ./results/00-indobart/checkpoint-4000\n",
      "Configuration saved in ./results/00-indobart/checkpoint-4000/config.json\n",
      "Configuration saved in ./results/00-indobart/checkpoint-4000/generation_config.json\n",
      "Model weights saved in ./results/00-indobart/checkpoint-4000/model.safetensors\n",
      "tokenizer config file saved in ./results/00-indobart/checkpoint-4000/tokenizer_config.json\n",
      "Special tokens file saved in ./results/00-indobart/checkpoint-4000/special_tokens_map.json\n",
      "Deleting older checkpoint [results/00-indobart/checkpoint-4000] due to args.save_total_limit\n",
      "Saving model checkpoint to ./results/00-indobart/checkpoint-5000\n",
      "Configuration saved in ./results/00-indobart/checkpoint-5000/config.json\n",
      "Configuration saved in ./results/00-indobart/checkpoint-5000/generation_config.json\n",
      "Model weights saved in ./results/00-indobart/checkpoint-5000/model.safetensors\n",
      "tokenizer config file saved in ./results/00-indobart/checkpoint-5000/tokenizer_config.json\n",
      "Special tokens file saved in ./results/00-indobart/checkpoint-5000/special_tokens_map.json\n",
      "Saving model checkpoint to ./results/00-indobart/checkpoint-5349\n",
      "Configuration saved in ./results/00-indobart/checkpoint-5349/config.json\n",
      "Configuration saved in ./results/00-indobart/checkpoint-5349/generation_config.json\n",
      "Model weights saved in ./results/00-indobart/checkpoint-5349/model.safetensors\n",
      "tokenizer config file saved in ./results/00-indobart/checkpoint-5349/tokenizer_config.json\n",
      "Special tokens file saved in ./results/00-indobart/checkpoint-5349/special_tokens_map.json\n",
      "The following columns in the evaluation set don't have a corresponding argument in `MBartForConditionalGeneration.forward` and have been ignored: summary, id, document. If summary, id, document are not expected by `MBartForConditionalGeneration.forward`,  you can safely ignore this message.\n",
      "\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 750\n",
      "  Batch size = 8\n",
      "\n",
      "\n",
      "Training completed. Do not forget to share your model on huggingface.co/models =)\n",
      "\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=5349, training_loss=0.45139215052250414, metrics={'train_runtime': 1585.1852, 'train_samples_per_second': 26.991, 'train_steps_per_second': 3.374, 'total_flos': 1.5804686727512064e+16, 'train_loss': 0.45139215052250414, 'epoch': 3.0})"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Predict Test Data and evaluate the score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The following columns in the test set don't have a corresponding argument in `MBartForConditionalGeneration.forward` and have been ignored: summary, id, document. If summary, id, document are not expected by `MBartForConditionalGeneration.forward`,  you can safely ignore this message.\n",
      "\n",
      "***** Running Prediction *****\n",
      "  Num examples = 3762\n",
      "  Batch size = 8\n"
     ]
    },
    {
     "data": {
      "text/html": [],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ROUGE scores on the test set: {'rouge1': np.float64(0.6492739538349439), 'rouge2': np.float64(0.5628012293066681), 'rougeL': np.float64(0.6066622052994133), 'rougeLsum': np.float64(0.6391191735873272)}\n"
     ]
    }
   ],
   "source": [
    "# Generate predictions\n",
    "test_predictions = trainer.predict(tokenized_ds[\"test\"])\n",
    "\n",
    "# Get the predictions and labels from the result\n",
    "preds = test_predictions.predictions\n",
    "labels = test_predictions.label_ids\n",
    "\n",
    "# Evaluate using the compute_metrics function\n",
    "rouge_scores = compute_metrics((preds, labels))\n",
    "\n",
    "# Print the ROUGE scores\n",
    "print(\"ROUGE scores on the test set:\", rouge_scores)\n",
    "\n",
    "# save the results to a file 'rouge_scores.txt'\n",
    "with open(os.path.join(benc_directory, \"rouge_scores.txt\"), \"w\") as f:\n",
    "    f.write(str(rouge_scores))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pipeline summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "vscode": {
     "languageId": "ruby"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You seem to be using the pipelines sequentially on GPU. In order to maximize efficiency please use a dataset\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>document</th>\n",
       "      <th>summary</th>\n",
       "      <th>generated_summary</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Jakarta, CNN Indonesia - - Dilansir AFP, seora...</td>\n",
       "      <td>Eman Ahmed Abd El Aty memiliki berat badan men...</td>\n",
       "      <td>seorang warga mesir yang dipercaya sebagai wa...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Menteri Pertahanan Ryamizard Ryacudu menyambut...</td>\n",
       "      <td>Menteri Pertahanan Ryamizard Ryacudu menyambut...</td>\n",
       "      <td>pesawat tanpa awak itu rencananya digunakan u...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Jakarta, CNN Indonesia - - Meski sudah hampir ...</td>\n",
       "      <td>Rumah produksi film yang dibintangi Lindsay Lo...</td>\n",
       "      <td>meski sudah hampir 12 tahun berlalu, film mea...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Usai melaksanakan ibadah haji, Eggi Sudjana ak...</td>\n",
       "      <td>Eggi Sudjana akhirnya mendatangi kantor Baresk...</td>\n",
       "      <td>eggi sudjana akhirnya mendatangi kantor bares...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Banyak cara untuk memberikan pengajaran kepada...</td>\n",
       "      <td>Game permainan Kartu Muslim. Menggunakan basis...</td>\n",
       "      <td>permainan kartu tersebut memberikan pengajara...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                            document  \\\n",
       "0  Jakarta, CNN Indonesia - - Dilansir AFP, seora...   \n",
       "1  Menteri Pertahanan Ryamizard Ryacudu menyambut...   \n",
       "2  Jakarta, CNN Indonesia - - Meski sudah hampir ...   \n",
       "3  Usai melaksanakan ibadah haji, Eggi Sudjana ak...   \n",
       "4  Banyak cara untuk memberikan pengajaran kepada...   \n",
       "\n",
       "                                             summary  \\\n",
       "0  Eman Ahmed Abd El Aty memiliki berat badan men...   \n",
       "1  Menteri Pertahanan Ryamizard Ryacudu menyambut...   \n",
       "2  Rumah produksi film yang dibintangi Lindsay Lo...   \n",
       "3  Eggi Sudjana akhirnya mendatangi kantor Baresk...   \n",
       "4  Game permainan Kartu Muslim. Menggunakan basis...   \n",
       "\n",
       "                                   generated_summary  \n",
       "0   seorang warga mesir yang dipercaya sebagai wa...  \n",
       "1   pesawat tanpa awak itu rencananya digunakan u...  \n",
       "2   meski sudah hampir 12 tahun berlalu, film mea...  \n",
       "3   eggi sudjana akhirnya mendatangi kantor bares...  \n",
       "4   permainan kartu tersebut memberikan pengajara...  "
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# get device\n",
    "summarizer = pipeline(\"summarization\", model=model, tokenizer=tokenizer, device=device)\n",
    "\n",
    "# create table to show the result: document, summary, generated_summary\n",
    "df = pd.DataFrame(columns=['document', 'summary', 'generated_summary'])\n",
    "for i in range(100):\n",
    "    document = ds['test'][i]['document']\n",
    "    summary = ds['test'][i]['summary']\n",
    "    generated_summary = summarizer(document, min_length=5, max_length=80)\n",
    "    df = pd.concat([df, pd.DataFrame([[document, summary, generated_summary[0]['summary_text']]], columns=['document', 'summary', 'generated_summary'])], ignore_index=True)\n",
    "\n",
    "\n",
    "# Save the DataFrame to a CSV file\n",
    "df.to_csv(f'{benc_directory}/summarization_result.csv')\n",
    "df.to_json(f'{benc_directory}/summarization_result.json')\n",
    "\n",
    "df.head()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}