Abstract_to_title_arxiv.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0477f1c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1ea7c38a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-05T14:02:03.142888Z",
     "start_time": "2024-04-05T14:02:03.139166Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import json\n",
    "#import pandas as pd\n",
    "import re\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch.nn.utils.rnn import pad_sequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64671f92",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-05T13:47:18.115211Z",
     "start_time": "2024-04-05T13:47:18.059029Z"
    }
   },
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32e5a722",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-05T13:47:27.716681Z",
     "start_time": "2024-04-05T13:47:27.601110Z"
    }
   },
   "outputs": [],
   "source": [
    "gpu_info = !nvidia-smi\n",
    "gpu_info = '\\n'.join(gpu_info)\n",
    "if gpu_info.find('failed') >= 0:\n",
    "  print('Select the Runtime > \"Change runtime type\" menu to enable a GPU accelerator, ')\n",
    "  print('and then re-execute this cell.')\n",
    "else:\n",
    "  print(gpu_info)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93174e4f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-05T13:55:44.183978Z",
     "start_time": "2024-04-05T13:55:44.180250Z"
    }
   },
   "outputs": [],
   "source": [
    "def print_list(l, K=None):\n",
    "\tfor i, e in enumerate(l):\n",
    "\t\tif i == K:\n",
    "\t\t\tbreak\n",
    "\t\tprint(e)\n",
    "\tprint()\n",
    "\n",
    "def load_from_pickle(pickle_file):\n",
    "\twith open(pickle_file, \"rb\") as pickle_in:\n",
    "\t\treturn pickle.load(pickle_in)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea066862",
   "metadata": {},
   "source": [
    "# Data Loading and Exploration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9402b928",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-05T13:53:28.899229Z",
     "start_time": "2024-04-05T13:53:28.740016Z"
    }
   },
   "outputs": [],
   "source": [
    "cols = ['id', 'title', 'abstract', 'categories']\n",
    "data = []\n",
    "file_name = './arxiv-metadata-oai-snapshot.json'\n",
    "\n",
    "with open(file_name, encoding='latin-1') as f:\n",
    "    count = 0\n",
    "    for line in f:\n",
    "        doc = json.loads(line)\n",
    "        lst = [doc['id'], doc['title'], doc['abstract'], doc['categories']]\n",
    "        data.append(lst)\n",
    "        count += 1\n",
    "        if count >= 10000:\n",
    "            break\n",
    "        \n",
    "df = pd.DataFrame(data=data, columns=cols).sample(n=100, random_state=68)\n",
    "\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14cefe10",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-05T14:04:05.909496Z",
     "start_time": "2024-04-05T14:04:04.648061Z"
    }
   },
   "outputs": [],
   "source": [
    "pad_word = \"<pad>\"\n",
    "bos_word = \"<s>\"\n",
    "eos_word = \"</s>\"\n",
    "unk_word = \"<unk>\"\n",
    "pad_id = 0\n",
    "bos_id = 1\n",
    "eos_id = 2\n",
    "unk_id = 3\n",
    "\n",
    "def normalize_sentence(s):\n",
    "    s = re.sub(r\"([.!?])\", r\" \\1\", s)\n",
    "    s = re.sub(r\"[^a-zA-Z.!?]+\", r\" \", s)\n",
    "    s = re.sub(r\"\\s+\", r\" \", s).strip()\n",
    "    return s\n",
    "\n",
    "class Vocabulary:\n",
    "    def __init__(self):\n",
    "        self.word_to_id = {pad_word: pad_id, bos_word: bos_id, eos_word:eos_id, unk_word: unk_id}\n",
    "        self.word_count = {}\n",
    "        self.id_to_word = {pad_id: pad_word, bos_id: bos_word, eos_id: eos_word, unk_id: unk_word}\n",
    "        self.num_words = 4\n",
    "\n",
    "    def get_ids_from_sentence(self, sentence):\n",
    "        sentence = normalize_sentence(sentence)\n",
    "        sent_ids = [bos_id] + [self.word_to_id[word] if word in self.word_to_id \\\n",
    "                               else unk_id for word in sentence.split()] + \\\n",
    "                               [eos_id]\n",
    "        return sent_ids\n",
    "\n",
    "    def tokenized_sentence(self, sentence):\n",
    "        sent_ids = self.get_ids_from_sentence(sentence)\n",
    "        return [self.id_to_word[word_id] for word_id in sent_ids]\n",
    "\n",
    "    def decode_sentence_from_ids(self, sent_ids):\n",
    "        words = list()\n",
    "        for i, word_id in enumerate(sent_ids):\n",
    "            if word_id in [bos_id, eos_id, pad_id]:\n",
    "                # Skip these words\n",
    "                continue\n",
    "            else:\n",
    "                words.append(self.id_to_word[word_id])\n",
    "        return ' '.join(words)\n",
    "\n",
    "    def add_words_from_sentence(self, sentence):\n",
    "        sentence = normalize_sentence(sentence)\n",
    "        for word in sentence.split():\n",
    "            if word not in self.word_to_id:\n",
    "                # add this word to the vocabulary\n",
    "                self.word_to_id[word] = self.num_words\n",
    "                self.id_to_word[self.num_words] = word\n",
    "                self.word_count[word] = 1\n",
    "                self.num_words += 1\n",
    "            else:\n",
    "                # update the word count\n",
    "                self.word_count[word] += 1\n",
    "\n",
    "abstracts_only = []\n",
    "titles_only = []\n",
    "                \n",
    "vocab = Vocabulary()\n",
    "for ids, title, abstract, categories in data:\n",
    "    abstracts_only.append(abstract)\n",
    "    titles_only.append(title)\n",
    "    vocab.add_words_from_sentence(title)\n",
    "    vocab.add_words_from_sentence(abstract)\n",
    "print(f\"Total words in the vocabulary = {vocab.num_words}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85fa5355",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-05T13:55:46.730680Z",
     "start_time": "2024-04-05T13:55:46.716749Z"
    }
   },
   "outputs": [],
   "source": [
    "print_list(sorted(vocab.word_count.items(), key=lambda item: item[1], reverse=True), 30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f756730",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-05T13:57:05.191697Z",
     "start_time": "2024-04-05T13:57:05.185555Z"
    }
   },
   "outputs": [],
   "source": [
    "for ids, title, abstract, categories in data[:2]:\n",
    "    sentence = abstract\n",
    "    word_tokens = vocab.tokenized_sentence(sentence)\n",
    "\n",
    "    # Automatically adds bos_id and eos_id before and after sentence ids respectively\n",
    "    word_ids = vocab.get_ids_from_sentence(sentence)\n",
    "    print(sentence)\n",
    "    print(word_tokens)\n",
    "    print(word_ids)\n",
    "    print(vocab.decode_sentence_from_ids(word_ids))\n",
    "    print()\n",
    "\n",
    "word = \"the\"\n",
    "word_id = vocab.word_to_id[word]\n",
    "print(f\"Word = {word}\")\n",
    "print(f\"Word ID = {word_id}\")\n",
    "print(f\"Word decoded from ID = {vocab.decode_sentence_from_ids([word_id])}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4ed3cfc",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-05T14:05:55.031117Z",
     "start_time": "2024-04-05T14:05:55.021549Z"
    }
   },
   "outputs": [],
   "source": [
    "class ArXiv_dataset(Dataset):\n",
    "    \"\"\"ArXiv dataset consisting of Abstract, title pairs.\"\"\"\n",
    "\n",
    "    def __init__(self, abstracts, titles, vocab, device):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            conversations: list of tuple (src_string, tgt_string)\n",
    "                         - src_string: String of the source sentence\n",
    "                         - tgt_string: String of the target sentence\n",
    "            vocab: Vocabulary object that contains the mapping of\n",
    "                    words to indices\n",
    "            device: cpu or cuda\n",
    "        \"\"\"\n",
    "        self.abstract_title = list(zip(abstracts, titles))\n",
    "        self.vocab = vocab\n",
    "        self.device = device\n",
    "\n",
    "        def encode(src, tgt):\n",
    "            src_ids = self.vocab.get_ids_from_sentence(src)\n",
    "            tgt_ids = self.vocab.get_ids_from_sentence(tgt)\n",
    "            return (src_ids, tgt_ids)\n",
    "\n",
    "        # We will pre-tokenize the conversations and save in id lists for later use\n",
    "        self.tokenized_abstract_title = [encode(src, tgt) for src, tgt in self.abstract_title]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.abstract_title)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        if torch.is_tensor(idx):\n",
    "            idx = idx.tolist()\n",
    "\n",
    "        return {\"conv_ids\":self.tokenized_abstract_title[idx], \"conv\":self.abstract_title[idx]}\n",
    "\n",
    "def collate_fn(data):\n",
    "    \"\"\"Creates mini-batch tensors from the list of tuples (src_seq, trg_seq).\n",
    "    We should build a custom collate_fn rather than using default collate_fn,\n",
    "    because merging sequences (including padding) is not supported in default.\n",
    "    Seqeuences are padded to the maximum length of mini-batch sequences (dynamic padding).\n",
    "    Args:\n",
    "        data: list of dicts {\"conv_ids\":(src_ids, tgt_ids), \"conv\":(src_str, trg_str)}.\n",
    "            - src_ids: list of src piece ids; variable length.\n",
    "            - tgt_ids: list of tgt piece ids; variable length.\n",
    "            - src_str: String of src\n",
    "            - tgt_str: String of tgt\n",
    "    Returns: dict { \"conv_ids\":     (src_ids, tgt_ids),\n",
    "                    \"conv\":         (src_str, tgt_str),\n",
    "                    \"conv_tensors\": (src_seqs, tgt_seqs)}\n",
    "            src_seqs: torch tensor of shape (src_padded_length, batch_size).\n",
    "            trg_seqs: torch tensor of shape (tgt_padded_length, batch_size).\n",
    "            src_padded_length = length of the longest src sequence from src_ids\n",
    "            tgt_padded_length = length of the longest tgt sequence from tgt_ids\n",
    "\n",
    "    \"\"\"\n",
    "    # Sort conv_ids based on decreasing order of the src_lengths.\n",
    "    # This is required for efficient GPU computations.\n",
    "    src_ids = [torch.LongTensor(e[\"conv_ids\"][0]) for e in data]\n",
    "    tgt_ids = [torch.LongTensor(e[\"conv_ids\"][1]) for e in data]\n",
    "    src_str = [e[\"conv\"][0] for e in data]\n",
    "    tgt_str = [e[\"conv\"][1] for e in data]\n",
    "    data = list(zip(src_ids, tgt_ids, src_str, tgt_str))\n",
    "    data.sort(key=lambda x: len(x[0]), reverse=True)\n",
    "    src_ids, tgt_ids, src_str, tgt_str = zip(*data)\n",
    "\n",
    "    ### BEGIN YOUR CODE ###\n",
    "\n",
    "    # Pad the src_ids and tgt_ids using token pad_id to create src_seqs and tgt_seqs\n",
    "    src_seqs = pad_sequence(src_ids, batch_first=False, padding_value=pad_id)\n",
    "    tgt_seqs = pad_sequence(tgt_ids, batch_first=False, padding_value=pad_id)\n",
    "    ### END YOUR CODE ###\n",
    "\n",
    "    return {\"conv_ids\":(src_ids, tgt_ids), \"conv\":(src_str, tgt_str), \"conv_tensors\":(src_seqs.to(device), tgt_seqs.to(device))}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e22e2734",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-05T14:06:11.177003Z",
     "start_time": "2024-04-05T14:06:09.891108Z"
    }
   },
   "outputs": [],
   "source": [
    "# Create the DataLoader for all_conversations\n",
    "dataset = ArXiv_dataset(abstracts_only, titles_only, vocab, device)\n",
    "\n",
    "batch_size = 2\n",
    "\n",
    "data_loader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cebb0f4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-05T14:06:11.200114Z",
     "start_time": "2024-04-05T14:06:11.179037Z"
    }
   },
   "outputs": [],
   "source": [
    "# Test one batch of training data\n",
    "first_batch = next(iter(data_loader))\n",
    "print(f\"Testing first training batch of size {len(first_batch['conv'][0])}\")\n",
    "print(f\"List of source strings:\")\n",
    "print_list(first_batch[\"conv\"][0])\n",
    "print(f\"Tokenized source ids:\")\n",
    "print_list(first_batch[\"conv_ids\"][0])\n",
    "print(f\"Padded source ids as tensor (shape {first_batch['conv_tensors'][0].size()}):\")\n",
    "print(first_batch[\"conv_tensors\"][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e77015b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-05T14:06:42.834102Z",
     "start_time": "2024-04-05T14:06:42.824026Z"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(f\"Testing first training batch of size {len(first_batch['conv'][1])}\")\n",
    "print(f\"List of target strings:\")\n",
    "print_list(first_batch[\"conv\"][1])\n",
    "print(f\"Tokenized target ids:\")\n",
    "print_list(first_batch[\"conv_ids\"][1])\n",
    "print(f\"Padded target ids as tensor (shape {first_batch['conv_tensors'][1].size()}):\")\n",
    "print(first_batch[\"conv_tensors\"][1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4af76007",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "# Train_Test Split Before\n",
    "\n",
    "dataset = load_dataset(\"json\", data_files=\"./arxiv-metadata-oai-snapshot.json\")\n",
    "\n",
    "dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f641b1b8",
   "metadata": {},
   "source": [
    "# Model Definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6e121f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForSeq2SeqLM, AutoTokenizer\n",
    "\n",
    "model_name = 'facebook/bart-large'\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "model = AutoModelForSeq2SeqLM.from_pretrained(model_name)\n",
    "\n",
    "max_token_length = max(len(tokenizer.encode(abstract, truncation=True)) for abstract in abstracts_only)\n",
    "print(f\"The longest text is {max_token_length} tokens long.\")\n",
    "\n",
    "def get_feature(batch):\n",
    "    encodings = tokenizer(batch['text'], text_target=batch['keywords'],\n",
    "                        max_length=1024, truncation=True)\n",
    "\n",
    "    encodings = {'input_ids': encodings['input_ids'],\n",
    "               'attention_mask': encodings['attention_mask'],\n",
    "               'labels': encodings['labels']}\n",
    "\n",
    "    return encodings\n",
    "\n",
    "dataset_pt = dataset.map(get_feature, batched=True)\n",
    "dataset_pt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4e31359",
   "metadata": {},
   "outputs": [],
   "source": [
    "columns = ['input_ids', 'labels', 'attention_mask']\n",
    "dataset_pt.set_format(type='torch', columns=columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "362f3a86",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import DataCollatorForSeq2Seq\n",
    "\n",
    "data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "667e50d9",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bbf23e9",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-05T14:08:04.021654Z",
     "start_time": "2024-04-05T14:08:04.015179Z"
    }
   },
   "outputs": [],
   "source": [
    "from transformers import TrainingArguments, Trainer\n",
    "\n",
    "# we're using the Trainer API which abstracts away a lot of complexity\n",
    "training_args = TrainingArguments(\n",
    "    output_dir = 'bart_abs_title', # rename to what you want it to be called\n",
    "    num_train_epochs=3, # your choice\n",
    "    warmup_steps = 500,\n",
    "    per_device_train_batch_size=4, # keep a small batch size when working with a small GPU\n",
    "    per_device_eval_batch_size=4,\n",
    "    weight_decay = 0.01, # helps prevent overfitting\n",
    "    logging_steps = 10,\n",
    "    evaluation_strategy = 'steps',\n",
    "    eval_steps=50, # base this on the size of your dataset and number of training epochs\n",
    "    save_steps=1e6,\n",
    "    gradient_accumulation_steps=16 # running this on a small GPU\n",
    ")\n",
    "\n",
    "trainer = Trainer(model=model, args=training_args, tokenizer=tokenizer, data_collator=data_collator,\n",
    "                  train_dataset = dataset_pt['train'], eval_dataset = dataset_pt['validation'])\n",
    "\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0f2a15b",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.save_model('Abs_to_Title') # set the name you want"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3595c23",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import pipeline\n",
    "\n",
    "# test the model using Hugging Face's pipeline\n",
    "pipe = pipeline(model='Abs_to_Title')\n",
    "\n",
    "# test the first item in the test set to see how it does\n",
    "test_text = dataset['test'][0]['text']\n",
    "title = dataset['test'][0]['title']\n",
    "print(\"the text: \", text_test)\n",
    "print(\"generated title: \", pipe(test_text))\n",
    "print(\"actual title : \",keywords)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}