feat: add summarizer

MultilingualSummarizer.ipynb

@@ -13,11 +13,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "vscode": {
-     "languageId": "python"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "import torch as pt\n",
@@ -30,6 +26,7 @@
     "from transformers import  AutoTokenizer\n",
     "import datasets\n",
     "\n",
+    "import re\n",
     "import logging\n",
     "logging.basicConfig(level=logging.INFO, format='%(asctime)s | %(name)s | %(levelname)s | %(message)s')\n",
     "\n",
@@ -56,10 +53,12 @@
     "    #\n",
     "    def __init__(self, model_name, language, inference_cfg=None, **kwargs):\n",
     "        logging.info(f\"Initializing multilingual summarizer {model_name}\")\n",
+    "        self.name = model_name.split('/')[-1]\n",
     "        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)\n",
-    "        self.dstTokenizer = DatasetTokenizer(model_name.split('/')[-1], model_name, language)\n",
+    "        self.dstTokenizer = DatasetTokenizer(self.name, model_name, language)\n",
     "        self.tokenizer  = self.dstTokenizer.get_tokenizer()\n",
     "        self.langid = self.dstTokenizer.get_langid()\n",
+    "        self.lang_token = self.dstTokenizer.get_lang_token()\n",
     "        self.inference_cfg = inference_cfg\n",
     "        self.enc_max_len = 512\n",
     "        self.language = language\n",
@@ -114,7 +113,8 @@
     "        summarizer = Summarizer(model = self.model, tokenizer = self.tokenizer,lcode=self.langid, batch_size = 8)\n",
     "        \n",
     "        #Summarize texts\n",
-    "        summarizer.summarize_dst(tok_dst,**self.inference_cfg)\n",
+    "        filter_fc = self._filter_final_summaries if self.name.startswith('mt5') else None\n",
+    "        summarizer.summarize_dst(tok_dst, filter_fc_batch = filter_fc,**self.inference_cfg)\n",
     "    \n",
     "        \n",
     "        scores = {}\n",
@@ -125,17 +125,16 @@
     "            \n",
     "        \n",
     "        return (summarizer.summarized_dst['summary'], scores)\n",
-    "    \n"
+    "    \n",
+    "    def _filter_final_summaries(self, batch, **kwargs):\n",
+    "        batch[\"summary\"] = [ re.sub(self.lang_token, '', tmp) for tmp in batch[\"summary\"]]\n",
+    "        return batch"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "vscode": {
-     "languageId": "python"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "## Configuration of summarization pipeline\n",
@@ -185,24 +184,36 @@
     "    ])\n",
     "    return cfg\n",
     "\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
     "cfg = summ_config()\n",
     "msummarizer = MultiSummarizer(**cfg)\n",
-    "ret = msummarizer(**cfg)\n"
+    "ret = msummarizer(**cfg)"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "vscode": {
-     "languageId": "python"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "ret = msummarizer(**cfg)\n",
     "print(ret)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
@@ -211,6 +222,18 @@
    "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.6.8"
+  },
   "orig_nbformat": 4
  },
  "nbformat": 4,