bart_large_cnn_abstract_summarization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e6830b4d-cd2e-4547-b2ba-7c0192ee941a",
   "metadata": {},
   "source": [
    "# BART Large CNN Abstract Text Summarization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54b6ae0b-6f2f-4451-a6d7-7b9022815d4c",
   "metadata": {},
   "source": [
    "Summarize text of abstracts from research papers on [arXiv](https://arxiv.org/) in PDF format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff4687ec-c508-4a07-bd02-14f7c7d7d626",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install required libraries\n",
    "!pip3 -q install torch transformers tokenizers optimum gradio pypdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69f2a116-525a-4d8c-a1c0-2f0a86fd3554",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import modules\n",
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForSeq2SeqLM\n",
    "from optimum.bettertransformer import BetterTransformer\n",
    "import gradio as gr\n",
    "import pypdf\n",
    "import re"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c18ee4ff-2ddb-4473-9eb4-c3abf364e975",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load model directly\n",
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "    \"facebook/bart-large-cnn\",\n",
    "    use_fast=True\n",
    ")\n",
    "\n",
    "model_hf = AutoModelForSeq2SeqLM.from_pretrained(\n",
    "    \"facebook/bart-large-cnn\",\n",
    "    torch_dtype=torch.bfloat16\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8d517d1-2dbd-4215-adab-e6b83bf572dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use BetterTransformer API from Optimum library\n",
    "model = BetterTransformer.transform(model_hf, keep_original_model=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b520109-8a3e-44b0-a88f-1995cb23ab4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define function to extract text from abstract in a PDF\n",
    "def extract_abstract(pdf_path):\n",
    "  with open(pdf_path, \"rb\") as f:\n",
    "    reader = pypdf.PdfReader(f)\n",
    "    text = reader.pages[0].extract_text(orientations=(0)) + reader.pages[1].extract_text(orientations=(0))\n",
    "    text = text.replace(\"\\n\", \"\")\n",
    "    \n",
    "    abstract_regex = re.compile(r\"Abstract|ABSTRACT\", re.IGNORECASE)\n",
    "    abstract_match = re.search(abstract_regex, text)\n",
    "    if not abstract_match:\n",
    "        return \"\"\n",
    "    \n",
    "    abstract_start = abstract_match.start() + 8\n",
    "    introduction_regex = re.compile(r\"Introduction|ntroduction|INTRODUCTION|NTRODUCTION\")\n",
    "    abstract_end = re.search(introduction_regex, text[abstract_start:]).start() + abstract_start\n",
    "\n",
    "    if abstract_start != -1 and abstract_end != -1:\n",
    "        return text[abstract_start:abstract_end]\n",
    "    else:\n",
    "        return \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d263c62-2018-461d-9bf6-2edc87dabde9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define function for Gradio application\n",
    "def summarize_abstract(pdf_path):\n",
    "    abstract_text = extract_abstract(pdf_path)\n",
    "    \n",
    "    inputs = tokenizer(\n",
    "        abstract_text,\n",
    "        truncation=True,\n",
    "        max_length=130,\n",
    "        return_tensors=\"pt\"\n",
    "    )\n",
    "    \n",
    "    prediction = model.generate(**inputs)\n",
    "    prediction = tokenizer.batch_decode(prediction, skip_special_tokens=True)\n",
    "    \n",
    "    return prediction[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a863999e-23a7-4e42-8219-c139aadd89ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the Gradio application\n",
    "demo = gr.Interface(\n",
    "    fn=summarize_abstract,\n",
    "    inputs=[gr.File(label=\"PDF path\")],\n",
    "    outputs=[gr.Textbox(label=\"Abstract summary\")],\n",
    "    description=\"\"\"\n",
    "    # BART Large CNN Abstract Summarization\n",
    "    [Code](https://github.com/darylalim/bart-large-cnn-abstract-summarization)\n",
    "    \"\"\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b949d9b-140f-44b4-ad4e-d884685eea43",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start queue with default settings\n",
    "demo.queue()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ef17835-cda7-4879-a3f7-be2c776dae51",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the Gradio application\n",
    "demo.launch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a413fac8-1703-451b-b22d-d659bcd88dc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Close the Gradio application\n",
    "demo.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efda6a60-9492-40fe-a30f-286bb8b4c11b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}