spaces demo

LICENSE

@@ -0,0 +1,202 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        https://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "{}"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2021 SummerTime
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       https://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+

SummerTime.egg-info/PKG-INFO

@@ -0,0 +1,124 @@
+Metadata-Version: 2.1
+Name: SummerTime
+Version: 0.1
+Summary: A summarization mode
+Home-page: https://github.com/LILYlab
+Author: Ansong Ni, Murori Mutuma, Zhangir Azerbayev, Yusen Zhang, Tao Yu, Dragomir Radev
+Author-email: ansong.ni@yale.edu, murorimutuma@gmail.com, zhangir.azerbayev@yale.edu
+License: UNKNOWN
+Description: # SummerTime
+        
+        A library to help users choose appropriate summarization tools based on their specific tasks or needs. Includes models, evaluation metrics, and datasets.
+        
+        
+        
+        ## Installation and setup
+        
+        #### Create and activate a new `conda` environment:
+        ```bash
+        conda create -n st python=3.7
+        conda activate st
+        ```
+        
+        #### `pip` dependencies for local demo:
+        ```bash
+        pip install -r requirements.txt
+        ```
+        
+        
+        
+        ## Quick Start
+        Imports model, initializes default model, and summarizes sample documents.
+        ```python
+        import model as st_model
+        
+        model = st_model.summarizer()
+        documents = [
+            """ PG&E stated it scheduled the blackouts in response to forecasts for high winds amid dry conditions. 
+            The aim is to reduce the risk of wildfires. Nearly 800 thousand customers were scheduled to be affected 
+            by the shutoffs which were expected to last through at least midday tomorrow."""
+        ]
+        model.summarize(documents)
+        
+        # ["California's largest electricity provider has turned off power to hundreds of thousands of customers."]
+        ```
+        
+        Also, please run `demo.ipynb` demo Jupyter notebook for more examples. To start demo Jupyter notebook on localhost:
+        ```bash
+        jupyter notebook demo.ipynb
+        ```
+        
+        
+        
+        ## Models
+        Import and initialization:
+        ```python
+        import model as st_model
+        
+        default_model = std_model.summarizer()
+        bart_model = std_model.bart_model.BartModel()
+        pegasus_model = std_model.pegasus_model.PegasusModel()
+        lexrank_model = std_model.lexrank_model.LexRankModel()
+        textrank_model = st_model.textrank_model.TextRankModel()
+        ```
+        
+        All models can be initialized with the following optional options:
+        ```python
+        def __init__(self,
+                 trained_domain: str=None,
+                 max_input_length: int=None,
+                 max_output_length: int=None,
+                 ):
+        ```
+        
+        All models implement the following methods:
+        ```python
+        def summarize(self,
+          corpus: Union[List[str], List[List[str]]],
+          queries: List[str]=None) -> List[str]:
+        
+        def show_capability(cls) -> None:
+        
+        def generate_basic_description(cls) -> str:
+        ```
+        
+        
+        
+        ## Evaluation
+        Import and initialization:
+        ```python
+        import eval as st_eval
+        
+        bert_eval = st_eval.bertscore()
+        bleu_eval = st_eval.bleu_eval()
+        rouge_eval = st_eval.rouge()
+        rougewe_eval = st_eval.rougewe()
+        ```
+        
+        All evaluation metrics can be initialized with the following optional arguments:
+        ```python
+        def __init__(self, metric_name):
+        ```
+        
+        All evaluation metric objects implement the following methods:
+        ```python
+        def evaluate(self, model, data):
+        
+        def get_dict(self, keys):
+        ```
+        
+        
+        ## Datasets
+        Import and initialization:
+        ```python
+        import dataset.stdatasets as st_data
+        ```
+        
+        ## Contributors
+        This repository is built by the [LILY Lab](https://yale-lily.github.io/) at Yale University, led by Prof. [Dragomir Radev](https://cpsc.yale.edu/people/dragomir-radev). The main contributors are [Ansong Ni](https://niansong1996.github.io), Zhangir Azerbayev, Troy Feng, Murori Mutuma and Yusen Zhang (Penn State). For comments and question, please open an issue.
+        
+Platform: UNKNOWN
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Description-Content-Type: text/markdown

SummerTime.egg-info/SOURCES.txt

@@ -0,0 +1,46 @@
+README.md
+setup.py
+summertime.py
+SummerTime.egg-info/PKG-INFO
+SummerTime.egg-info/SOURCES.txt
+SummerTime.egg-info/dependency_links.txt
+SummerTime.egg-info/top_level.txt
+dataset/__init__.py
+dataset/datasets_demo.py
+dataset/huggingface_datasets.py
+dataset/non_huggingface_datasets.py
+dataset/st_dataset.py
+evaluation/__init__.py
+evaluation/base_metric.py
+evaluation/bertscore_metric.py
+evaluation/bleu_metric.py
+evaluation/meteor_metric.py
+evaluation/rouge_metric.py
+evaluation/rougewe_metric.py
+evaluation/summeval_metric.py
+model/__init__.py
+model/base_model.py
+model/defaults.py
+model/dialogue/__init__.py
+model/dialogue/hmnet_model.py
+model/multi_doc/__init__.py
+model/multi_doc/base_multi_doc_model.py
+model/multi_doc/multi_doc_joint_model.py
+model/multi_doc/multi_doc_separate_model.py
+model/query_based/__init__.py
+model/query_based/base_query_based_model.py
+model/query_based/bm25_model.py
+model/query_based/tf_idf_model.py
+model/single_doc/__init__.py
+model/single_doc/bart_model.py
+model/single_doc/base_single_doc_model.py
+model/single_doc/lexrank_model.py
+model/single_doc/longformer_model.py
+model/single_doc/pegasus_model.py
+model/single_doc/textrank_model.py
+tests/__init__.py
+tests/dataset_test.py
+tests/demo_test.py
+tests/evaluation_test.py
+tests/integration_test.py
+tests/model_test.py

SummerTime.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

SummerTime.egg-info/top_level.txt

@@ -0,0 +1,4 @@
+dataset
+evaluation
+model
+tests

__init__.py

@@ -0,0 +1,3 @@
+import SummerTime.model
+import SummerTime.dataset.st_dataset as data
+import SummerTime.evaluation

build/scripts-3.9/summertime

@@ -0,0 +1,3 @@
+#!python
+
+print("welcome to Summer Time!")

dataset/__init__.py

@@ -0,0 +1,36 @@
+from dataset.dataset_loaders import (
+    CnndmDataset,
+    MultinewsDataset,
+    SamsumDataset,
+    XsumDataset,
+    PubmedqaDataset,
+    MlsumDataset,
+    ScisummnetDataset,
+    SummscreenDataset,
+    QMsumDataset,
+    ArxivDataset,
+)
+
+
+SUPPORTED_SUMM_DATASETS = [
+    CnndmDataset,
+    MultinewsDataset,
+    SamsumDataset,
+    XsumDataset,
+    PubmedqaDataset,
+    MlsumDataset,
+    ScisummnetDataset,
+    SummscreenDataset,
+    QMsumDataset,
+    ArxivDataset,
+]
+
+
+def list_all_datasets():
+    all_datasets = []
+    for ds in SUPPORTED_SUMM_DATASETS:
+        dataset_description = ds.generate_basic_description()
+
+        all_datasets.append((ds.dataset_name, dataset_description))
+
+    return all_datasets

dataset/dataset_loaders.py

@@ -0,0 +1,501 @@
+from os import path
+from tqdm import tqdm
+from typing import List, Generator, Optional, Union
+
+from datasets import Dataset
+
+from dataset.st_dataset import SummInstance, SummDataset
+
+
+# Set directory to load non_huggingface dataset scripts
+FILE_DIRECTORY_PATH = path.dirname(path.realpath(__file__))
+BASE_NONHUGGINGFACE_DATASETS_PATH = path.join(
+    FILE_DIRECTORY_PATH, "non_huggingface_datasets_builders"
+)
+
+
+# Huggingface Datasets
+
+
+class CnndmDataset(SummDataset):
+    """
+    The CNN/DM dataset
+    """
+
+    dataset_name = "CNN/DailyMail"
+
+    is_query_based = False
+    is_dialogue_based = False
+    is_multi_document = False
+
+    huggingface_dataset = True
+    huggingface_page = "https://huggingface.co/datasets/cnn_dailymail"
+
+    def __init__(self):
+        super().__init__(
+            dataset_args=(
+                "cnn_dailymail",
+                "3.0.0",
+            )
+        )
+
+    def _process_data(self, data: Dataset) -> Generator[SummInstance, None, None]:
+        """
+        Overrides the SummDataset '_process_data()' method
+        This method processes the data contained in the dataset
+            and puts each data instance into a SummInstance object
+        :param dataset: a train/validation/test dataset
+        :rtype: a generator yielding SummInstance objects
+        """
+        for instance in tqdm(data):
+            article: str = instance["article"]
+            highlights: str = instance["highlights"]
+            summ_instance = SummInstance(source=article, summary=highlights)
+
+            yield summ_instance
+
+
+class MultinewsDataset(SummDataset):
+    """
+    The Multi News dataset
+    """
+
+    dataset_name = "Multinews"
+
+    is_query_based = False
+    is_dialogue_based = False
+    is_multi_document = True
+
+    huggingface_dataset = True
+    huggingface_page = "https://huggingface.co/datasets/multi_news"
+
+    def __init__(self):
+        super().__init__(dataset_args=("multi_news",))
+
+    def _process_data(self, data: Dataset) -> Generator[SummInstance, None, None]:
+        """
+        Overrides the SummDataset '_process_data()' method
+        This method processes the data contained in the dataset
+            and puts each data instance into a SummInstance object
+        :param dataset: a train/validation/test dataset
+        :rtype: a generator yielding SummInstance objects
+        """
+        for instance in tqdm(data):
+            document: list = [
+                doc for doc in instance["document"].split("|||||") if doc
+            ]  # removes the empty string generated
+            # since each doc ends with the delimiting token '|||||'
+            # the final doc creates an empty string
+            summary: str = instance["summary"]
+            summ_instance = SummInstance(source=document, summary=summary)
+
+            yield summ_instance
+
+
+class SamsumDataset(SummDataset):
+    """
+    The SAMsum Dataset
+    """
+
+    dataset_name = "Samsum"
+
+    is_query_based = False
+    is_dialogue_based = True
+    is_multi_document = False
+
+    huggingface_dataset = True
+    huggingface_page = "https://huggingface.co/datasets/samsum"
+
+    def __init__(self):
+        super().__init__(dataset_args=("samsum",))
+
+    def _process_data(self, data: Dataset) -> Generator[SummInstance, None, None]:
+        """
+        Overrides the SummDataset '_process_data()' method
+        This method processes the data contained in the dataset
+            and puts each data instance into a SummInstance object
+        :param dataset: a train/validation/test dataset
+        :rtype: a generator yielding SummInstance objects
+        """
+        for instance in tqdm(data):
+            dialogue: List = instance["dialogue"].split(
+                "\r\n"
+            )  # split each dialogue into a list of strings such as
+            # ["speaker1 : utter..", "speaker2 : utter..."]
+            summary: str = instance["summary"]
+            summ_instance = SummInstance(source=dialogue, summary=summary)
+
+            yield summ_instance
+
+
+class XsumDataset(SummDataset):
+    """
+    The Xsum Dataset
+    """
+
+    dataset_name = "Xsum"
+
+    huggingface_dataset = True
+    huggingface_page = "https://huggingface.co/datasets/xsum"
+
+    is_query_based = False
+    is_dialogue_based = False
+    is_multi_document = False
+
+    def __init__(self):
+        super().__init__(dataset_args=("xsum",))
+
+    def _process_data(self, data: Dataset) -> Generator[SummInstance, None, None]:
+        """
+        Overrides the SummDataset '_process_data()' method
+        This method processes the data contained in the dataset
+            and puts each data instance into a SummInstance object
+        :param dataset: a train/validation/test dataset
+        :rtype: a generator yielding SummInstance objects
+        """
+        for instance in tqdm(data):
+            document: List = instance["document"]
+            summary: str = instance["summary"]
+            summ_instance = SummInstance(source=document, summary=summary)
+
+            yield summ_instance
+
+
+class PubmedqaDataset(SummDataset):
+    """
+    The Pubmed QA dataset
+    """
+
+    dataset_name = "Pubmedqa"
+
+    is_query_based = True
+    is_dialogue_based = False
+    is_multi_document = False
+
+    huggingface_dataset = True
+    huggingface_page = "https://huggingface.co/datasets/pubmed_qa"
+
+    def __init__(self, seed=None):
+        super().__init__(
+            dataset_args=(
+                "pubmed_qa",
+                "pqa_artificial",
+            )
+        )
+
+    def _process_data(self, data: Dataset) -> Generator[SummInstance, None, None]:
+        """
+        Overrides the SummDataset '_process_data()' method
+        This method processes the data contained in the dataset
+            and puts each data instance into a SummInstance object
+        :param dataset: a train/validation/test dataset
+        :rtype: a generator yielding SummInstance objects
+        """
+        for instance in tqdm(data):
+            context: str = " ".join(instance["context"]["contexts"])
+            answer: str = instance["long_answer"]
+            query: str = instance["question"]
+            summ_instance = SummInstance(source=context, summary=answer, query=query)
+
+            yield summ_instance
+
+
+class MlsumDataset(SummDataset):
+    """
+    The MLsum Dataset - A multi-lingual dataset featuring 5 languages
+    Includes 1.5 million news articles and their corresponding summaries
+
+    "de" - German
+    "es" - Spanish
+    "fr" - French
+    "ru" - Russian
+    "tu" - Turkish
+    """
+
+    dataset_name = "MlSum"
+
+    is_query_based = False
+    is_dialogue_based = False
+    is_multi_document = False
+
+    huggingface_dataset = True
+    huggingface_page = "https://huggingface.co/datasets/mlsum"
+    supported_languages = ["de", "es", "fr", "ru", "tu"]
+
+    mlsum_instantiation_guide = """The languages supported for the Mlsum Dataset are:
+                de - German
+                es - Spanish
+                fr - French
+                ru - Russian
+                tu - Turkish
+
+                Examples to instantiate the dataset:
+                1. Dataset with only one language
+                   dataset = MlsumDataset({language_token})
+                   dataset = MlsumDataset("es")
+                   dataset = MlsumDataset("tu")...
+
+                2. Dataset with a multiple languages
+                   dataset = MlsumDataset({list of language_token})
+                   dataset = MlsumDataset(["es","de"])
+                   dataset = MlsumDataset(["es","de", "tu"])...
+
+                3. Dataset with all supported languages (default)
+                   dataset = MlsumDataset(all)
+                   dataset = MlsumDataset()
+                """
+
+    def __init__(self, languages: Optional[Union[str, List[str]]] = "all"):
+        super().__init__(dataset_args=(languages,))
+
+    def _load_dataset_safe(self, languages: Optional[Union[str, List[str]]]):
+        """
+        Overrides the parent class method
+        Method loads multiple datasets of different languages provided in :param languages:
+            It then concatenates these datasets into one combined dataset
+        :rtype: datasetDict containing the combined dataset
+        :param languages: Optional, either a string or list of strings specifying the languages
+            to load
+        """
+        print(MlsumDataset.mlsum_instantiation_guide)
+
+        # Choose languages to download articles
+        if languages == "all":
+            selected_languages = MlsumDataset.supported_languages
+        elif isinstance(languages, list):
+            for language in languages:
+                assert self.is_supported(language)
+            selected_languages = languages
+        else:
+            assert self.is_supported(languages)
+            selected_languages = [languages]
+
+        # Concatenate selected languaeges into one dataset
+        language_datasets = []
+        for language in selected_languages:
+            dataset = super()._load_dataset_safe(
+                "mlsum",
+                language,
+            )
+
+            language_datasets.append(dataset)
+
+        mlsum_dataset = self._concatenate_dataset_dicts(language_datasets)
+
+        return mlsum_dataset
+
+    def _process_data(self, data: Dataset) -> Generator[SummInstance, None, None]:
+        """
+        Overrides the SummDataset '_process_data()' method
+        This method processes the data contained in the dataset
+            and puts each data instance into a SummInstance object
+        :param dataset: a train/validation/test dataset
+        :rtype: a generator yielding SummInstance objects
+        """
+        for instance in tqdm(data):
+            article: List = instance["text"]
+            summary: str = instance["summary"]
+            summ_instance = SummInstance(source=article, summary=summary)
+
+            yield summ_instance
+
+    def is_supported(self, language: str):
+        """
+        Checks whether the requested langues is supported
+        :param language: string containing the requested language
+        :rtype bool:
+        """
+        if language not in MlsumDataset.supported_languages:
+            print(MlsumDataset.mlsum_instantiation_guide)
+            raise ValueError(
+                f"The language(s): '{language}' entered is not supported. See above message for usage info"
+            )
+        else:
+            return True
+
+
+# Non-huggingface datasets
+
+
+class ScisummnetDataset(SummDataset):
+    """
+    The SciSummNet dataset. As a dataset not included by huggingface, we need to do manually download, set basic
+        information for the dataset
+    """
+
+    dataset_name = "ScisummNet"
+
+    version = "1.1.0"
+    description = (
+        "A summary of scientific papers should ideally incorporate the impact of the papers on the "
+        "research community reflected by citations. To facilitate research in citation-aware scientific "
+        "paper summarization (Scisumm), the CL-Scisumm shared task has been organized since 2014 for "
+        "papers in the computational linguistics and NLP domain."
+    )
+
+    is_dialogue_based = False
+    is_multi_document = False
+    is_query_based = False
+
+    huggingface_dataset = False
+    builder_script_path = path.join(
+        BASE_NONHUGGINGFACE_DATASETS_PATH, dataset_name.lower() + ".py"
+    )
+
+    def __init__(self, seed=None):
+        super().__init__()
+
+    def _process_data(self, data: Dataset) -> Generator[SummInstance, None, None]:
+        """
+        Overrides the SummDataset '_process_data()' method
+        This method processes the data contained in the dataset
+            and puts each data instance into a SummInstance object
+        :param dataset: a train/validation/test dataset
+        :rtype: a generator yielding SummInstance objects
+        """
+        for instance in tqdm(data):
+            docs: List = [
+                instance["document_xml"],
+                instance["citing_sentences_annotated.json"],
+            ]
+            summary: str = instance["summary"]
+            summ_instance = SummInstance(source=docs, summary=summary)
+
+            yield summ_instance
+
+
+class SummscreenDataset(SummDataset):
+    """
+    The SummScreen dataset. As a dataset not included by huggingface, we need to do manually download, set basic
+        information for the dataset
+    """
+
+    dataset_name = "Summscreen"
+
+    version = "1.1.0"
+    is_dialogue_based = True
+    is_multi_document = False
+    is_query_based = False
+
+    huggingface_dataset = False
+    builder_script_path = path.join(
+        BASE_NONHUGGINGFACE_DATASETS_PATH, dataset_name.lower() + ".py"
+    )
+
+    def __init__(self, seed=None):
+        super().__init__()
+
+    def _process_data(self, data: Dataset) -> Generator[SummInstance, None, None]:
+        """
+        Overrides the SummDataset '_process_data()' method
+        This method processes the data contained in the dataset
+            and puts each data instance into a SummInstance object
+        :param dataset: a train/validation/test dataset
+        :rtype: a generator yielding SummInstance objects
+        """
+        for instance in tqdm(data):
+            transcript: List = instance[
+                "transcript"
+            ]  # convert string into a list of string dialogues
+            recap: str = instance["recap"]
+            summ_instance = SummInstance(source=transcript, summary=recap)
+
+            yield summ_instance
+
+
+class QMsumDataset(SummDataset):
+    """
+    QMSum Dataset
+    """
+
+    dataset_name = "QMsum"
+    description = """
+    QMSum is a new human-annotated benchmark for query-based multi-domain meeting summarization task,
+    which consists of 1,808 query-summary pairs over 232 meetings in multiple domains.
+    """
+
+    is_dialogue_based = True
+    is_multi_document = False
+    is_query_based = True
+
+    huggingface_dataset = False
+    builder_script_path = path.join(
+        BASE_NONHUGGINGFACE_DATASETS_PATH, dataset_name.lower() + ".py"
+    )
+
+    def __init__(self):
+        super().__init__()
+
+    def _process_data(self, data: Dataset) -> Generator[SummInstance, None, None]:
+        """
+        Overrides the SummDataset '_process_data()' method
+        This method processes the data contained in the dataset
+            and puts each data instance into a SummInstance object
+        :param dataset: a train/validation/test dataset
+        :rtype: a generator yielding SummInstance objects
+        """
+        for instance in tqdm(data):
+            for query_set in (
+                instance["general_query_list"] + instance["specific_query_list"]
+            ):
+                meeting: List = [
+                    utterance["speaker"] + " : " + utterance["content"]
+                    for utterance in instance["meeting_transcripts"]
+                ]
+                query: str = query_set["query"]
+                summary: str = query_set["answer"]
+                summ_instance = SummInstance(
+                    source=meeting, summary=summary, query=query
+                )
+
+            yield summ_instance
+
+
+class ArxivDataset(SummDataset):
+    """
+    The Arxiv Dataset
+    """
+
+    dataset_name = "Arxiv_longsummarization"
+    description = """
+    A summarization dataset comprised of pairs of scientific papers.
+    The dataset provides a challenging testbed for abstractive summarization.
+    It contains papers and their abstracts.
+    """
+
+    is_dialogue_based = False
+    is_multi_document = False
+    is_query_based = False
+
+    huggingface_dataset = False
+    builder_script_path = path.join(
+        BASE_NONHUGGINGFACE_DATASETS_PATH, dataset_name.lower() + ".py"
+    )
+
+    def __init__(self):
+
+        print(
+            "*****************",
+            "***Attention***",
+            "This dataset is quite large (approx 5Gb and will need about 15 Gb for the extraction process",
+            "Cancel/interrupt the download if size and time constraints will not be met",
+            "*****************",
+            sep="\n",
+        )
+
+        super().__init__()
+
+    def _process_data(self, data: Dataset) -> Generator[SummInstance, None, None]:
+        """
+        Overrides the SummDataset '_process_data()' method
+        This method processes the data contained in the dataset
+            and puts each data instance into a SummInstance object
+        :param dataset: a train/validation/test dataset
+        :rtype: a generator yielding SummInstance objects
+        """
+        for instance in tqdm(data):
+            article: List = instance["article_text"]
+            abstract: str = " ".join(instance["abstract_text"])
+            summ_instance = SummInstance(source=article, summary=abstract)
+
+            yield summ_instance

dataset/non_huggingface_datasets_builders/arxiv_longsummarization.py

@@ -0,0 +1,104 @@
+import os
+import json
+import datasets
+
+
+"""Arxiv dataset."""
+
+
+_CITATION = """
+@article{Cohan_2018,
+   title={A Discourse-Aware Attention Model for Abstractive Summarization of
+            Long Documents},
+   url={http://dx.doi.org/10.18653/v1/n18-2097},
+   DOI={10.18653/v1/n18-2097},
+   journal={Proceedings of the 2018 Conference of the North American Chapter of
+          the Association for Computational Linguistics: Human Language
+          Technologies, Volume 2 (Short Papers)},
+   publisher={Association for Computational Linguistics},
+   author={Cohan, Arman and Dernoncourt, Franck and Kim, Doo Soon and Bui, Trung and Kim, Seokhwan and Chang, Walter and Goharian, Nazli},
+   year={2018}
+}
+"""
+
+_DESCRIPTION = """
+A summarization dataset comprised of pairs of scientific papers.
+The dataset provides a challenging testbed for abstractive summarization.
+It contains papers and their abstracts.
+"""
+
+_HOMEPAGE = "https://github.com/armancohan/long-summarization"
+
+_LICENSE = "Apache-2.0 License"
+
+_URL = "https://archive.org/download/armancohan-long-summarization-paper-code/arxiv-dataset.zip"
+
+
+class SummertimeArxiv(datasets.GeneratorBasedBuilder):
+    """Arxiv long summarization dataset."""
+
+    VERSION = datasets.Version("1.0.0")
+
+    BUILDER_CONFIGS = [
+        datasets.BuilderConfig(),
+    ]
+
+    def _info(self):
+        features = datasets.Features(
+            {
+                "article_id": datasets.Value("string"),
+                "article_text": [datasets.Value("string")],
+                "abstract_text": [datasets.Value("string")],
+            }
+        )
+        return datasets.DatasetInfo(
+            description=_DESCRIPTION,
+            features=features,
+            supervised_keys=None,
+            homepage=_HOMEPAGE,
+            license=_LICENSE,
+            citation=_CITATION,
+        )
+
+    def _split_generators(self, dl_manager):
+        """Returns SplitGenerators."""
+        my_urls = _URL
+        path = dl_manager.download_and_extract(my_urls)
+        path = os.path.join(path, "arxiv-dataset")
+
+        trainpath = os.path.join(path, "train.txt")
+        valpath = os.path.join(path, "val.txt")
+        testpath = os.path.join(path, "test.txt")
+
+        return [
+            datasets.SplitGenerator(
+                name=datasets.Split.TRAIN,
+                # These kwargs will be passed to _generate_examples
+                gen_kwargs={"filepath": trainpath, "split": "train"},
+            ),
+            datasets.SplitGenerator(
+                name=datasets.Split.VALIDATION,
+                # These kwargs will be passed to _generate_examples
+                gen_kwargs={"filepath": valpath, "split": "val"},
+            ),
+            datasets.SplitGenerator(
+                name=datasets.Split.TEST,
+                # These kwargs will be passed to _generate_examples
+                gen_kwargs={"filepath": testpath, "split": "test"},
+            ),
+        ]
+
+    def _generate_examples(self, filepath, split):
+        """Yields examples."""
+
+        with open(filepath, "r") as f:
+            for line in f:
+
+                instance = json.loads(line)
+
+                entry = {}
+                entry["article_id"] = instance["article_id"]
+                entry["article_text"] = instance["article_text"]
+                entry["abstract_text"] = instance["abstract_text"]
+
+                yield entry["article_id"], entry

dataset/non_huggingface_datasets_builders/qmsum.py

@@ -0,0 +1,119 @@
+import os
+import json
+import datasets
+
+
+"""QMsum dataset."""
+
+
+_CITATION = """
+@inproceedings{zhong2021qmsum,
+   title={{QMS}um: {A} {N}ew {B}enchmark for {Q}uery-based {M}ulti-domain {M}eeting {S}ummarization},
+   author={Zhong, Ming and Yin, Da and Yu, Tao and Zaidi, Ahmad and Mutuma, Mutethia and Jha, Rahul and Hassan Awadallah, Ahmed and Celikyilmaz, Asli and Liu, Yang and Qiu, Xipeng and Radev, Dragomir},
+   booktitle={North American Association for Computational Linguistics (NAACL)},
+   year={2021}
+}
+"""
+
+_DESCRIPTION = """
+QMSum is a new human-annotated benchmark for query-based multi-domain meeting summarization task, \
+which consists of 1,808 query-summary pairs over 232 meetings in multiple domains.
+"""
+
+_HOMEPAGE = "https://github.com/Yale-LILY/QMSum"
+
+_BASE_URL = "https://raw.githubusercontent.com/Yale-LILY/QMSum/main/data/ALL/jsonl"
+_URLs = {
+    "train": _BASE_URL + "/train.jsonl",
+    "val": _BASE_URL + "/val.jsonl",
+    "test": _BASE_URL + "/test.jsonl",
+}
+
+
+class SummertimeQmsum(datasets.GeneratorBasedBuilder):
+    """QMsum dataset."""
+
+    VERSION = datasets.Version("1.0.0")
+
+    BUILDER_CONFIGS = [
+        datasets.BuilderConfig(),
+    ]
+
+    def _info(self):
+        features = datasets.Features(
+            {
+                "entry_number": datasets.Value("string"),
+                "meeting_transcripts": [
+                    {
+                        "speaker": datasets.Value("string"),
+                        "content": datasets.Value("string"),
+                    }
+                ],
+                "general_query_list": [
+                    {
+                        "query": datasets.Value("string"),
+                        "answer": datasets.Value("string"),
+                    }
+                ],
+                "specific_query_list": [
+                    {
+                        "query": datasets.Value("string"),
+                        "answer": datasets.Value("string"),
+                        "relevant_text_span": [[datasets.Value("string")]],
+                    }
+                ],
+            }
+        )
+        return datasets.DatasetInfo(
+            description=_DESCRIPTION,
+            features=features,
+            supervised_keys=None,
+            homepage=_HOMEPAGE,
+            license=None,
+            citation=_CITATION,
+        )
+
+    def _split_generators(self, dl_manager):
+        """Returns SplitGenerators."""
+        my_urls = _URLs
+        downloaded_files = dl_manager.download_and_extract(my_urls)
+
+        trainpath = downloaded_files["train"]
+        valpath = downloaded_files["val"]
+        testpath = downloaded_files["test"]
+
+        return [
+            datasets.SplitGenerator(
+                name=datasets.Split.TRAIN,
+                # These kwargs will be passed to _generate_examples
+                gen_kwargs={"filepath": trainpath, "split": "train"},
+            ),
+            datasets.SplitGenerator(
+                name=datasets.Split.VALIDATION,
+                # These kwargs will be passed to _generate_examples
+                gen_kwargs={"filepath": valpath, "split": "val"},
+            ),
+            datasets.SplitGenerator(
+                name=datasets.Split.TEST,
+                # These kwargs will be passed to _generate_examples
+                gen_kwargs={"filepath": testpath, "split": "test"},
+            ),
+        ]
+
+    def _generate_examples(self, filepath, split):
+        """Yields examples."""
+
+        extraction_path = os.path.join(filepath)
+
+        with open(extraction_path) as f:
+            for i, line in enumerate(f):
+
+                instance = json.loads(line)
+
+                entry = {}
+                entry["entry_number"] = split + "_" + str(i)
+                entry["meeting_transcripts"] = instance["meeting_transcripts"]
+                entry["general_query_list"] = instance["general_query_list"]
+                entry["specific_query_list"] = instance["specific_query_list"]
+
+                yield entry["entry_number"], entry

dataset/non_huggingface_datasets_builders/scisummnet.py

@@ -0,0 +1,105 @@
+import os
+import datasets
+
+
+"""Scisummnet dataset."""
+
+
+_CITATION = """
+@InProceedings{yasunaga&al.19.scisumm,
+    title = {{ScisummNet}: A Large Annotated Corpus and Content-Impact Models for Scientific Paper Summarization with Citation Networks},
+    author = {Michihiro Yasunaga and Jungo Kasai and Rui Zhang and Alexander Fabbri and Irene Li and Dan Friedman and Dragomir Radev},
+    booktitle = {Proceedings of AAAI 2019},
+    year = {2019}
+}
+@InProceedings{yasunaga&al.17,
+  title = {Graph-based Neural Multi-Document Summarization},
+  author = {Yasunaga, Michihiro and Zhang, Rui and Meelu, Kshitijh and Pareek, Ayush and Srinivasan, Krishnan and Radev, Dragomir R.},
+  booktitle = {Proceedings of CoNLL 2017},
+  year = {2017}
+}
+"""
+
+_DESCRIPTION = """
+A summary of scientific papers should ideally incorporate the impact of the papers on the research community
+reflected by citations. To facilitate research in citation-aware scientific paper summarization (Scisumm),
+the CL-Scisumm shared task has been organized since 2014 for papers in the computational linguistics and NLP domain.
+"""
+
+_HOMEPAGE = "https://cs.stanford.edu/~myasu/projects/scisumm_net/"
+
+_LICENSE = "CC BY-SA 4.0"
+
+_URLs = "https://cs.stanford.edu/~myasu/projects/scisumm_net/scisummnet_release1.1__20190413.zip"
+
+
+class SummertimeScisummnet(datasets.GeneratorBasedBuilder):
+    """Scisummnet dataset."""
+
+    VERSION = datasets.Version("1.1.0")
+
+    BUILDER_CONFIGS = [
+        datasets.BuilderConfig(),
+    ]
+
+    def _info(self):
+        features = datasets.Features(
+            {
+                "entry_number": datasets.Value("string"),
+                "document_xml": datasets.Value("string"),
+                "citing_sentences_annotated.json": datasets.Value("string"),
+                "summary": datasets.Value("string"),
+            }
+        )
+        return datasets.DatasetInfo(
+            description=_DESCRIPTION,
+            features=features,
+            supervised_keys=None,
+            homepage=_HOMEPAGE,
+            license=_LICENSE,
+            citation=_CITATION,
+        )
+
+    def _split_generators(self, dl_manager):
+        """Returns SplitGenerators."""
+        my_urls = _URLs
+        path = dl_manager.download_and_extract(my_urls)
+        trainpath = os.path.join(
+            path, "scisummnet_release1.1__20190413", "top1000_complete"
+        )
+        return [
+            datasets.SplitGenerator(
+                name=datasets.Split.TRAIN,
+                # These kwargs will be passed to _generate_examples
+                gen_kwargs={"extraction_path": trainpath, "split": "train"},
+            )
+        ]
+
+    def _generate_examples(self, extraction_path, split):
+        """Yields examples."""
+
+        for folder in os.listdir(extraction_path):
+
+            entry = {}
+
+            entry["entry_number"] = folder
+
+            doc_xml_path = os.path.join(
+                extraction_path, folder, "Documents_xml", folder + ".xml"
+            )
+            with open(doc_xml_path, "r", encoding="utf-8") as f:
+                entry["document_xml"] = f.read()
+
+            cite_annot_path = os.path.join(
+                extraction_path, folder, "citing_sentences_annotated.json"
+            )
+            with open(cite_annot_path, "r", encoding="utf-8") as f:
+                entry["citing_sentences_annotated.json"] = f.read()
+
+            summary_path = os.path.join(
+                extraction_path, folder, "summary", folder + ".gold.txt"
+            )
+            with open(summary_path, "r", encoding="utf-8") as f:
+                entry["summary"] = f.read()
+
+            yield entry["entry_number"], entry

dataset/non_huggingface_datasets_builders/summscreen.py

@@ -0,0 +1,123 @@
+import os
+import json
+import datasets
+
+
+"""Summscreen dataset."""
+
+
+_CITATION = """
+@article{DBLP:journals/corr/abs-2104-07091,
+  author    = {Mingda Chen and
+               Zewei Chu and
+               Sam Wiseman and
+               Kevin Gimpel},
+  title     = {SummScreen: {A} Dataset for Abstractive Screenplay Summarization},
+  journal   = {CoRR},
+  volume    = {abs/2104.07091},
+  year      = {2021},
+  url       = {https://arxiv.org/abs/2104.07091},
+  archivePrefix = {arXiv},
+  eprint    = {2104.07091},
+  timestamp = {Mon, 19 Apr 2021 16:45:47 +0200},
+  biburl    = {https://dblp.org/rec/journals/corr/abs-2104-07091.bib},
+  bibsource = {dblp computer science bibliography, https://dblp.org}
+}
+"""
+
+_DESCRIPTION = """
+A summary of scientific papers should ideally incorporate the impact of the papers on the research community
+reflected by citations. To facilitate research in citation-aware scientific paper summarization (Scisumm),
+the CL-Scisumm shared task has been organized since 2014 for papers in the computational linguistics and NLP domain.
+"""
+
+_HOMEPAGE = "https://github.com/mingdachen/SummScreen"
+
+_LICENSE = "MIT Licencse"
+
+_URLs = "https://drive.google.com/uc?id=1BvdIllGBo9d2-bzXQRzWuJXB04XPVmfF"
+
+
+class SummertimeSummscreen(datasets.GeneratorBasedBuilder):
+    """Summscreen dataset."""
+
+    VERSION = datasets.Version("1.1.0")
+
+    BUILDER_CONFIGS = [
+        datasets.BuilderConfig(),
+    ]
+
+    def _info(self):
+        features = datasets.Features(
+            {
+                "entry_number": datasets.Value("string"),
+                "transcript": datasets.features.Sequence(datasets.Value("string")),
+                "recap": datasets.Value("string"),
+            }
+        )
+        return datasets.DatasetInfo(
+            description=_DESCRIPTION,
+            features=features,
+            supervised_keys=None,
+            homepage=_HOMEPAGE,
+            license=_LICENSE,
+            citation=_CITATION,
+        )
+
+    def _split_generators(self, dl_manager):
+        """Returns SplitGenerators."""
+        my_urls = _URLs
+        path = dl_manager.download_and_extract(my_urls)
+        path = os.path.join(path, "SummScreen")
+
+        trainpath_fd = os.path.join("ForeverDreaming", "fd_train.json")
+        trainpath_tms = os.path.join("TVMegaSite", "tms_train.json")
+        trainpaths = [trainpath_fd, trainpath_tms]
+
+        devpath_fd = os.path.join("ForeverDreaming", "fd_dev.json")
+        devpath_tms = os.path.join("TVMegaSite", "tms_dev.json")
+        devpaths = [devpath_fd, devpath_tms]
+
+        testpath_fd = os.path.join("ForeverDreaming", "fd_test.json")
+        testpath_tms = os.path.join("TVMegaSite", "tms_test.json")
+        testpaths = [testpath_fd, testpath_tms]
+
+        return [
+            datasets.SplitGenerator(
+                name=datasets.Split.TRAIN,
+                # These kwargs will be passed to _generate_examples
+                gen_kwargs={"filepaths": (path, trainpaths), "split": "train"},
+            ),
+            datasets.SplitGenerator(
+                name=datasets.Split.VALIDATION,
+                # These kwargs will be passed to _generate_examples
+                gen_kwargs={"filepaths": (path, devpaths), "split": "dev"},
+            ),
+            datasets.SplitGenerator(
+                name=datasets.Split.TEST,
+                # These kwargs will be passed to _generate_examples
+                gen_kwargs={"filepaths": (path, testpaths), "split": "test"},
+            ),
+        ]
+
+    def _generate_examples(self, filepaths, split):
+        """Yields examples."""
+
+        path, relative_filepaths = filepaths
+        for filepath in relative_filepaths:
+
+            extraction_path = os.path.join(path, filepath)
+
+            with open(extraction_path, "r") as f:
+                for line in f:
+                    processed_line = line.replace("@@ ", "")
+                    instance = json.loads(processed_line)
+
+                    entry = {}
+                    entry["entry_number"] = instance["filename"]
+                    entry["transcript"] = instance["Transcript"]
+                    entry["recap"] = instance["Recap"][
+                        0
+                    ]  # Recap is a single string in list
+
+                    yield entry["entry_number"], entry

dataset/st_dataset.py

@@ -0,0 +1,281 @@
+from abc import abstractmethod
+from pprint import pformat
+from time import sleep
+from typing import List, Tuple, Optional, Union, Generator
+
+from datasets import (
+    Dataset,
+    DatasetDict,
+    DatasetInfo,
+    concatenate_datasets,
+    load_dataset,
+)
+
+# Defualt values for retrying dataset download
+DEFAULT_NUMBER_OF_RETRIES_ALLOWED = 5
+DEFAULT_WAIT_SECONDS_BEFORE_RETRY = 5
+
+# Default value for creating missing val/test splits
+TEST_OR_VAL_SPLIT_RATIO = 0.1
+
+
+class SummInstance:
+    """
+    Basic instance for summarization tasks
+    """
+
+    def __init__(
+        self, source: Union[List[str], str], summary: str, query: Optional[str] = None
+    ):
+        """
+        Create a summarization instance
+        :rtype: object
+        :param source: either `List[str]` or `str`, depending on the dataset itself, string joining may needed to fit
+            into specific models. For example, for the same document, it could be simply `str` or `List[str]` for
+            a list of sentences in the same document
+        :param summary: a string summary that serves as ground truth
+        :param query: Optional, applies when a string query is present
+        """
+        self.source = source
+        self.summary = summary
+        self.query = query
+
+    def __repr__(self):
+        instance_dict = {"source": self.source, "summary": self.summary}
+        if self.query:
+            instance_dict["query"] = self.query
+
+        return str(instance_dict)
+
+    def __str__(self):
+        instance_dict = {"source": self.source, "summary": self.summary}
+        if self.query:
+            instance_dict["query"] = self.query
+
+        return pformat(instance_dict, indent=1)
+
+
+class SummDataset:
+    """
+    Dataset class for summarization, which takes into account of the following tasks:
+        * Single document summarization
+        * Multi-document/Dialogue summarization
+        * Query-based summarization
+    """
+
+    def __init__(
+        self, dataset_args: Optional[Tuple[str]] = None, splitseed: Optional[int] = None
+    ):
+        """Create dataset information from the huggingface Dataset class
+        :rtype: object
+        :param dataset_args: a tuple containing arguments to passed on to the 'load_dataset_safe' method.
+            Only required for datasets loaded from the Huggingface library.
+            The arguments for each dataset are different and comprise of a string or multiple strings
+        :param splitseed: a number to instantiate the random generator used to generate val/test splits
+            for the datasets without them
+        """
+
+        # Load dataset from huggingface, use default huggingface arguments
+        if self.huggingface_dataset:
+            dataset = self._load_dataset_safe(*dataset_args)
+        # Load non-huggingface dataset, use custom dataset builder
+        else:
+            dataset = self._load_dataset_safe(path=self.builder_script_path)
+
+        info_set = self._get_dataset_info(dataset)
+
+        # Ensure any dataset with a val or dev or validation split is standardised to validation split
+        if "val" in dataset:
+            dataset["validation"] = dataset["val"]
+            dataset.remove("val")
+        elif "dev" in dataset:
+            dataset["validation"] = dataset["dev"]
+            dataset.remove("dev")
+
+        # If no splits other other than training, generate them
+        assert (
+            "train" in dataset or "validation" in dataset or "test" in dataset
+        ), "At least one of train/validation test needs to be not empty!"
+
+        if not ("validation" in dataset or "test" in dataset):
+            dataset = self._generate_missing_val_test_splits(dataset, splitseed)
+
+        self.description = info_set.description
+        self.citation = info_set.citation
+        self.homepage = info_set.homepage
+
+        # Extract the dataset entries from folders and load into dataset
+        self._train_set = self._process_data(dataset["train"])
+        self._validation_set = self._process_data(
+            dataset["validation"]
+        )  # Some datasets have a validation split
+        self._test_set = self._process_data(dataset["test"])
+
+    @property
+    def train_set(self) -> Union[Generator[SummInstance, None, None], List]:
+        if self._train_set is not None:
+            return self._train_set
+        else:
+            print(
+                f"{self.dataset_name} does not contain a train set, empty list returned"
+            )
+            return list()
+
+    @property
+    def validation_set(self) -> Union[Generator[SummInstance, None, None], List]:
+        if self._validation_set is not None:
+            return self._validation_set
+        else:
+            print(
+                f"{self.dataset_name} does not contain a validation set, empty list returned"
+            )
+            return list()
+
+    @property
+    def test_set(self) -> Union[Generator[SummInstance, None, None], List]:
+        if self._test_set is not None:
+            return self._test_set
+        else:
+            print(
+                f"{self.dataset_name} does not contain a test set, empty list returned"
+            )
+            return list()
+
+    def _load_dataset_safe(self, *args, **kwargs) -> Dataset:
+        """
+        This method creates a wrapper around the huggingface 'load_dataset()' function for a more robust download function,
+            the original 'load_dataset()' function occassionally fails when it cannot reach a server especially after multiple requests.
+            This method tackles this problem by attempting the download multiple times with a wait time before each retry
+
+        The wrapper method passes all arguments and keyword arguments to the 'load_dataset' function with no alteration.
+        :rtype: Dataset
+        :param args: non-keyword arguments to passed on to the 'load_dataset' function
+        :param kwargs: keyword arguments to passed on to the 'load_dataset' function
+        """
+
+        tries = DEFAULT_NUMBER_OF_RETRIES_ALLOWED
+        wait_time = DEFAULT_WAIT_SECONDS_BEFORE_RETRY
+
+        for i in range(tries):
+            try:
+                dataset = load_dataset(*args, **kwargs)
+            except ConnectionError:
+                if i < tries - 1:  # i is zero indexed
+                    sleep(wait_time)
+                    continue
+                else:
+                    raise RuntimeError(
+                        "Wait for a minute and attempt downloading the dataset again. \
+                         The server hosting the dataset occassionally times out."
+                    )
+            break
+
+        return dataset
+
+    def _get_dataset_info(self, data_dict: DatasetDict) -> DatasetInfo:
+        """
+        Get the information set from the dataset
+        The information set contains: dataset name, description, version, citation and licence
+        :param data_dict: DatasetDict
+        :rtype: DatasetInfo
+        """
+        return data_dict["train"].info
+
+    @abstractmethod
+    def _process_data(self, dataset: Dataset) -> Generator[SummInstance, None, None]:
+        """
+        Abstract class method to process the data contained within each dataset.
+        Each dataset class processes it's own information differently due to the diversity in domains
+        This method processes the data contained in the dataset
+            and puts each data instance into a SummInstance object,
+            the SummInstance has the following properties [source, summary, query[optional]]
+        :param dataset: a train/validation/test dataset
+        :rtype: a generator yielding SummInstance objects
+        """
+        return
+
+    def _generate_missing_val_test_splits(
+        self, dataset_dict: DatasetDict, seed: int
+    ) -> DatasetDict:
+        """
+        Creating the train, val and test splits from a dataset
+            the generated sets are 'train: ~.80', 'validation: ~.10', and 'test: ~10' in size
+            the splits are randomized for each object unless a seed is provided for the random generator
+
+        :param dataset: Arrow Dataset with containing, usually the train set
+        :param seed: seed for the random generator to shuffle the dataset
+        :rtype: Arrow DatasetDict containing the three splits
+        """
+
+        # Return dataset if no train set available for splitting
+        if "train" not in dataset_dict:
+            if "validation" not in dataset_dict:
+                dataset_dict["validation"] = None
+            if "test" not in dataset_dict:
+                dataset_dict["test"] = None
+
+            return dataset_dict
+
+        # Create a 'test' split from 'train' if no 'test' set is available
+        if "test" not in dataset_dict:
+            dataset_traintest_split = dataset_dict["train"].train_test_split(
+                test_size=TEST_OR_VAL_SPLIT_RATIO, seed=seed
+            )
+            dataset_dict["train"] = dataset_traintest_split["train"]
+            dataset_dict["test"] = dataset_traintest_split["test"]
+
+        # Create a 'validation' split from the remaining 'train' set if no 'validation' set is available
+        if "validation" not in dataset_dict:
+            dataset_trainval_split = dataset_dict["train"].train_test_split(
+                test_size=TEST_OR_VAL_SPLIT_RATIO, seed=seed
+            )
+            dataset_dict["train"] = dataset_trainval_split["train"]
+            dataset_dict["validation"] = dataset_trainval_split["test"]
+
+        return dataset_dict
+
+    def _concatenate_dataset_dicts(
+        self, dataset_dicts: List[DatasetDict]
+    ) -> DatasetDict:
+        """
+        Concatenate two dataset dicts with similar splits and columns tinto one
+        :param dataset_dicts: A list of DatasetDicts
+        :rtype: DatasetDict containing the combined data
+        """
+
+        # Ensure all dataset dicts have the same splits
+        setsofsplits = set(tuple(dataset_dict.keys()) for dataset_dict in dataset_dicts)
+        if len(setsofsplits) > 1:
+            raise ValueError("Splits must match for all datasets")
+
+        # Concatenate all datasets into one according to the splits
+        temp_dict = {}
+        for split in setsofsplits.pop():
+            split_set = [dataset_dict[split] for dataset_dict in dataset_dicts]
+            temp_dict[split] = concatenate_datasets(split_set)
+
+        return DatasetDict(temp_dict)
+
+    @classmethod
+    def generate_basic_description(cls) -> str:
+        """
+        Automatically generate the basic description string based on the attributes
+        :rtype: string containing the description
+        :param cls: class object
+        """
+
+        basic_description = (
+            f": {cls.dataset_name} is a "
+            f"{'query-based ' if cls.is_query_based else ''}"
+            f"{'dialogue ' if cls.is_dialogue_based else ''}"
+            f"{'multi-document' if cls.is_multi_document else 'single-document'} "
+            f"summarization dataset."
+        )
+
+        return basic_description
+
+    def show_description(self):
+        """
+        Print the description of the dataset.
+        """
+        print(self.dataset_name, ":\n", self.description)

dependencies.txt

@@ -0,0 +1,11 @@
+Migrate information to documentation/pypi for first release. 
+
+Dependencies: 
+- lexrank
+- sentencepiece
+- torch
+- transformers
+
+# datasets
+- datasets
+- py7zr

download.py

@@ -0,0 +1,3 @@
+import nltk
+
+nltk.download("stopwords")

evaluation/__init__.py

@@ -0,0 +1,14 @@
+import site
+import os
+
+# needed so that rouge works
+package_path = site.getsitepackages()[0]
+os.environ["ROUGE_HOME"] = package_path + "/summ_eval/ROUGE-1.5.5/"
+
+from .rouge_metric import Rouge
+from .bertscore_metric import BertScore
+from .rougewe_metric import RougeWe
+from .bleu_metric import Bleu
+from .meteor_metric import Meteor
+
+SUPPORTED_EVALUATION_METRICS = [BertScore, Bleu, Rouge, RougeWe, Meteor]

evaluation/base_metric.py

@@ -0,0 +1,27 @@
+from typing import List, Tuple, Dict
+
+
+class SummMetric:
+    metric_name: str = None
+    range: Tuple[float, float] = None
+    higher_is_better: bool = None
+    requires_heavy_compute: bool = None
+
+    def evaluate(
+        self,
+        # TODO zhangir: integrate with dataset api
+        inputs: List[str],
+        targets: List[str],
+        keys: List[str],
+    ) -> Dict[str, float]:
+        """
+        All metrics should have this function.
+        :input: A list of summaries.
+        :target: A list of target summaries corresponding to each entry of input.
+        :keys: Which metrics to return,
+        e.g, ['rouge_1_f_score', 'rouge_2_f_score']
+        :return: A dictionary with keys metrics and values scores.
+        """
+        raise NotImplementedError(
+            "the base class for metrics shouldn't be instantiated!"
+        )

evaluation/bertscore_metric.py

@@ -0,0 +1,20 @@
+from summ_eval.bert_score_metric import BertScoreMetric
+from evaluation.summeval_metric import SummEvalMetric
+from typing import List, Dict
+
+
+class BertScore(SummEvalMetric):
+    metric_name = "bert score"
+    range = (0, 1)
+    higher_is_better = True
+    requires_heavy_compute = True
+
+    def __init__(self):
+        se_metric = BertScoreMetric()
+        super(BertScore, self).__init__(se_metric)
+
+    def evaluate(
+        self, inputs: List[str], targets: List[str], keys: List[str] = ["bert_score_f1"]
+    ) -> Dict[str, float]:
+        # TODO zhangir: update when datasets api is merged
+        return super(BertScore, self).evaluate(inputs, targets, keys)

evaluation/bleu_metric.py

@@ -0,0 +1,20 @@
+from summ_eval.bleu_metric import BleuMetric
+from evaluation.summeval_metric import SummEvalMetric
+from typing import List, Dict
+
+
+class Bleu(SummEvalMetric):
+    metric_name = "bleu"
+    range = (0, 100)
+    higher_is_better = True
+    requires_heavy_compute = False
+
+    def __init__(self):
+        se_metric = BleuMetric()
+        super(Bleu, self).__init__(se_metric)
+
+    def evaluate(
+        self, inputs: List[str], targets: List[str], keys: List[str] = ["bleu"]
+    ) -> Dict[str, float]:
+        # TODO zhangir: potentially update when dataset api is merged.
+        return super(Bleu, self).evaluate(inputs, targets, keys)

evaluation/meteor_metric.py

@@ -0,0 +1,31 @@
+from .base_metric import SummMetric
+from typing import List, Dict
+from nltk.translate import meteor_score as nltk_meteor
+import nltk
+import statistics
+
+
+class Meteor(SummMetric):
+    metric_name = "meteor"
+    range = (0, 1)
+    higher_is_better = True
+    requires_heavy_compute = False
+
+    def __init__(self):
+        nltk.download("wordnet")
+
+    def evaluate(
+        self, inputs: List[str], targets: List[str], keys=["meteor"]
+    ) -> Dict[str, float]:
+
+        for key in keys:
+            if key != "meteor":
+                raise KeyError(key, "is not a valid key")
+
+        meteor_scores = [
+            nltk_meteor.meteor_score([input], target)
+            for input, target in zip(inputs, targets)
+        ]
+        meteor_score = statistics.mean(meteor_scores)
+
+        return {key: meteor_score for key in keys}

evaluation/rouge_metric.py

@@ -0,0 +1,23 @@
+from summ_eval.rouge_metric import RougeMetric
+from evaluation.summeval_metric import SummEvalMetric
+from typing import List, Dict
+
+
+class Rouge(SummEvalMetric):
+    metric_name = "rouge"
+    range = (0, 1)
+    higher_is_better = True
+    requires_heavy_compute = False
+
+    def __init__(self):
+        se_metric = RougeMetric()
+        super(Rouge, self).__init__(se_metric)
+
+    def evaluate(
+        self,
+        inputs: List[str],
+        targets: List[str],
+        keys: List[str] = ["rouge_1_f_score", "rouge_2_f_score", "rouge_l_f_score"],
+    ) -> Dict[str, float]:
+        score_dict = self.se_metric.evaluate_batch(inputs, targets)
+        return {key: score_dict["rouge"][key] for key in keys}

evaluation/rougewe_metric.py

@@ -0,0 +1,24 @@
+from evaluation.summeval_metric import SummEvalMetric
+from typing import List, Dict
+
+import nltk
+
+
+class RougeWe(SummEvalMetric):
+    metric_name = "rougeWE"
+    range = (0, 1)
+    higher_is_better = True
+    requires_heavy_compute = True
+
+    def __init__(self):
+        from summ_eval.rouge_we_metric import RougeWeMetric
+
+        nltk.download("stopwords")
+        se_metric = RougeWeMetric()
+        super(RougeWe, self).__init__(se_metric)
+
+    def evaluate(
+        self, inputs: List[str], targets: List[str], keys: List[str] = ["rouge_we_3_f"]
+    ) -> Dict[str, float]:
+        # TODO zhangir: update when dataset api is merged.
+        return super(RougeWe, self).evaluate(inputs, targets, keys)

evaluation/summeval_metric.py

@@ -0,0 +1,18 @@
+from .base_metric import SummMetric
+from summ_eval.metric import Metric as SEMetric
+from typing import List, Dict
+
+
+class SummEvalMetric(SummMetric):
+    """
+    Generic class for a summarization metric whose backend is SummEval.
+    """
+
+    def __init__(self, se_metric: SEMetric):
+        self.se_metric = se_metric
+
+    def evaluate(
+        self, inputs: List[str], targets: List[str], keys: List[str]
+    ) -> Dict[str, float]:
+        score_dict = self.se_metric.evaluate_batch(inputs, targets)
+        return {key: score_dict[key] if key in score_dict else None for key in keys}

model/__init__.py

@@ -0,0 +1,34 @@
+from .single_doc import (
+    BartModel,
+    LexRankModel,
+    LongformerModel,
+    PegasusModel,
+    TextRankModel,
+)
+from .multi_doc import MultiDocJointModel, MultiDocSeparateModel
+from .dialogue import HMNetModel
+from .query_based import TFIDFSummModel, BM25SummModel
+from .defaults import summarizer
+
+SUPPORTED_SUMM_MODELS = [
+    BartModel,
+    LexRankModel,
+    LongformerModel,
+    PegasusModel,
+    TextRankModel,
+    MultiDocJointModel,
+    MultiDocSeparateModel,
+    HMNetModel,
+    TFIDFSummModel,
+    BM25SummModel,
+]
+
+
+def list_all_models():
+    all_model_tuples = []
+    for model_class in SUPPORTED_SUMM_MODELS:
+        model_description = model_class.generate_basic_description()
+
+        all_model_tuples.append((model_class, model_description))
+
+    return all_model_tuples

model/base_model.py

@@ -0,0 +1,81 @@
+from typing import List, Union
+
+
+class SummModel:
+    """
+    Base model class for SummerTime
+    """
+
+    # static variables
+    model_name = "None"
+    is_extractive = False
+    is_neural = False
+    is_query_based = False
+    is_dialogue_based = False
+    is_multi_document = False
+
+    def __init__(
+        self,
+        trained_domain: str = None,
+        max_input_length: int = None,
+        max_output_length: int = None,
+    ):
+        self.trained_domain = trained_domain
+        self.max_input_length = max_input_length
+        self.max_output_length = max_output_length
+
+    def summarize(
+        self, corpus: Union[List[str], List[List[str]]], queries: List[str] = None
+    ) -> List[str]:
+        """
+        All summarization models should have this function
+
+        :param corpus: each string in the list is a source document to be summarized; if the model is multi-document or
+            dialogue summarization model, then each instance contains a list of documents/utterances
+        :param queries: a list of queries if this is a query-based model
+        :return: a list of generated summaries
+        """
+        raise NotImplementedError(
+            "The base class for models shouldn't be instantiated!"
+        )
+
+    @classmethod
+    def assert_summ_input_type(
+        cls, corpus: Union[List[str], List[List[str]]], queries: Union[List[str], None]
+    ):
+        """
+        Verifies that type of input corpus or queries for summarization align with the model type.
+        """
+        raise NotImplementedError(
+            "The base class for models shouldn't be instantiated!"
+        )
+
+    @classmethod
+    def show_capability(cls) -> None:
+        """
+        Use concise language to show the strength and weakness for each model. Try not to use NLP terminologies
+        """
+        raise NotImplementedError(
+            "The base class for models shouldn't be instantiated!"
+        )
+
+    @classmethod
+    def generate_basic_description(cls) -> str:
+        """
+        Automatically generate the basic description string based on the attributes
+        """
+        extractive_abstractive = "extractive" if cls.is_extractive else "abstractive"
+        neural = "neural" if cls.is_neural else "non-neural"
+
+        basic_description = (
+            f"{cls.model_name} is a"
+            f"{'query-based' if cls.is_query_based else ''} "
+            f"{extractive_abstractive}, {neural} model for summarization."
+        )
+        if cls.is_multi_document or cls.is_dialogue_based:
+            basic_description += (
+                f"It can handle {'multi-document' if cls.is_multi_document else ''} "
+                f"{'dialogue' if cls.is_dialogue_based else ''} textual data."
+            )
+
+        return basic_description

model/defaults.py

@@ -0,0 +1,10 @@
+from .single_doc import PegasusModel
+
+
+class summarizer(PegasusModel):
+    def __init__(self, device="cpu"):
+        super(summarizer, self).__init__(device)
+
+    def show_capability(self):
+        print("Pegasus is the default singe-document summarization model.")
+        super(summarizer, self).show_capability()

model/dialogue/__init__.py

@@ -0,0 +1 @@
+from .hmnet_model import HMNetModel

model/dialogue/hmnet/ExampleRawData/meeting_summarization/AMI_proprec/test_ami.json

@@ -0,0 +1 @@
+[{"source": {"dataset": "../ExampleRawData/meeting_summarization/AMI_proprec/test/"}, "task": "meeting", "name": "ami"}]

model/dialogue/hmnet/ExampleRawData/meeting_summarization/role_dict_ext.json

@@ -0,0 +1 @@
+{}

model/dialogue/hmnet/config/dialogue.conf

@@ -0,0 +1,98 @@
+##################
+# Trainer settings
+##################
+
+MODEL	MeetingNet_Transformer
+TASK	HMNet
+CRITERION MLECriterion
+
+SEED    1033
+
+MAX_NUM_EPOCHS	20
+EVAL_PER_UPDATE_NUM	10
+UPDATES_PER_EPOCH 20
+
+# The actuall learning rate will be multiplied with the number of GPUs
+OPTIMIZER	RAdam
+START_LEARNING_RATE	1e-3
+LR_SCHEDULER	LnrWrmpInvSqRtDcyScheduler
+WARMUP_STEPS	16000
+WARMUP_INIT_LR  1e-4
+WARMUP_END_LR 1e-3
+
+# The actuall start learning rate equals START_LEARNING_RATE * GRADIENT_ACCUMULATE_STEP
+# Model will be updated after every MINI_BATCH * GRADIENT_ACCUMULATE_STEP samples
+GRADIENT_ACCUMULATE_STEP	5
+
+GRAD_CLIPPING    2
+
+##################
+# Task settings
+##################
+
+# This is the relative path to the directory where this conf file locates
+USE_REL_DATA_PATH
+TRAIN_FILE	../ExampleRawData/meeting_summarization/AMI_proprec/train_ami.json
+DEV_FILE	../ExampleRawData/meeting_summarization/AMI_proprec/valid_ami.json
+TEST_FILE	../ExampleRawData/meeting_summarization/AMI_proprec/test_ami.json
+ROLE_DICT_FILE  ../ExampleRawData/meeting_summarization/role_dict_ext.json
+
+MINI_BATCH	1
+MAX_PADDING_RATIO	1
+BATCH_READ_AHEAD	10
+DOC_SHUFFLE_BUF_SIZE	10
+SAMPLE_SHUFFLE_BUFFER_SIZE	10
+BATCH_SHUFFLE_BUFFER_SIZE	10
+
+MAX_TRANSCRIPT_WORD 8300
+#MAX_SENT_LEN 30
+MAX_SENT_LEN 12
+# MAX_SENT_NUM 300
+MAX_SENT_NUM 60
+
+##################
+# Model settings
+##################
+
+DROPOUT	0.1
+VOCAB_DIM	512
+ROLE_SIZE   32
+ROLE_DIM	16
+POS_DIM 16
+ENT_DIM 16
+
+USE_ROLE
+USE_POSENT
+
+USE_BOS_TOKEN
+USE_EOS_TOKEN
+
+TRANSFORMER_EMBED_DROPOUT	0.1
+TRANSFORMER_RESIDUAL_DROPOUT	0.1
+TRANSFORMER_ATTENTION_DROPOUT	0.1
+TRANSFORMER_LAYER	6
+TRANSFORMER_HEAD	8
+TRANSFORMER_POS_DISCOUNT	80
+
+PRE_TOKENIZER	TransfoXLTokenizer
+PRE_TOKENIZER_PATH  ../../../third_party/HMNet/ExampleInitModel/transfo-xl-wt103
+PYLEARN_MODEL ../../../third_party/HMNet/ExampleInitModel/AMI-finetuned
+# e.g. PYLEARN_MODEL conf_hmnet_AMI_conf~/run_1/11600
+
+##################
+# Tokenizer settings
+##################
+
+EXTRA_IDS	1000
+
+##################
+# Decoding settings
+##################
+
+BEAM_WIDTH      6
+EVAL_TOKENIZED
+EVAL_LOWERCASE
+# MAX_GEN_LENGTH  300
+MAX_GEN_LENGTH  60
+MIN_GEN_LENGTH  10
+NO_REPEAT_NGRAM_SIZE 3

model/dialogue/hmnet_model.py

@@ -0,0 +1,483 @@
+from model.base_model import SummModel
+import argparse
+import os
+import torch
+import gzip
+import json
+from model.third_party.HMNet.Models.Trainers.HMNetTrainer import HMNetTrainer
+from model.third_party.HMNet.Utils.Arguments import Arguments
+
+import spacy
+
+nlp = spacy.load("en_core_web_sm", disable=["parser"])
+# tagger = nlp.get_pipe('tagger')
+# ner = nlp.get_pipe('ner')
+# POS = {w: i for i, w in enumerate([''] + list(tagger.labels))}
+# ENT = {w: i for i, w in enumerate([''] + list(ner.move_names))}
+# These two dicts are adapted from SpaCy 2.3.1, since HMNet's embedding for POS and ENT is fixed
+POS = {
+    "": 0,
+    "$": 1,
+    "''": 2,
+    ",": 3,
+    "-LRB-": 4,
+    "-RRB-": 5,
+    ".": 6,
+    ":": 7,
+    "ADD": 8,
+    "AFX": 9,
+    "CC": 10,
+    "CD": 11,
+    "DT": 12,
+    "EX": 13,
+    "FW": 14,
+    "HYPH": 15,
+    "IN": 16,
+    "JJ": 17,
+    "JJR": 18,
+    "JJS": 19,
+    "LS": 20,
+    "MD": 21,
+    "NFP": 22,
+    "NN": 23,
+    "NNP": 24,
+    "NNPS": 25,
+    "NNS": 26,
+    "PDT": 27,
+    "POS": 28,
+    "PRP": 29,
+    "PRP$": 30,
+    "RB": 31,
+    "RBR": 32,
+    "RBS": 33,
+    "RP": 34,
+    "SYM": 35,
+    "TO": 36,
+    "UH": 37,
+    "VB": 38,
+    "VBD": 39,
+    "VBG": 40,
+    "VBN": 41,
+    "VBP": 42,
+    "VBZ": 43,
+    "WDT": 44,
+    "WP": 45,
+    "WP$": 46,
+    "WRB": 47,
+    "XX": 48,
+    "_SP": 49,
+    "``": 50,
+}
+ENT = {
+    "": 0,
+    "B-ORG": 1,
+    "B-DATE": 2,
+    "B-PERSON": 3,
+    "B-GPE": 4,
+    "B-MONEY": 5,
+    "B-CARDINAL": 6,
+    "B-NORP": 7,
+    "B-PERCENT": 8,
+    "B-WORK_OF_ART": 9,
+    "B-LOC": 10,
+    "B-TIME": 11,
+    "B-QUANTITY": 12,
+    "B-FAC": 13,
+    "B-EVENT": 14,
+    "B-ORDINAL": 15,
+    "B-PRODUCT": 16,
+    "B-LAW": 17,
+    "B-LANGUAGE": 18,
+    "I-ORG": 19,
+    "I-DATE": 20,
+    "I-PERSON": 21,
+    "I-GPE": 22,
+    "I-MONEY": 23,
+    "I-CARDINAL": 24,
+    "I-NORP": 25,
+    "I-PERCENT": 26,
+    "I-WORK_OF_ART": 27,
+    "I-LOC": 28,
+    "I-TIME": 29,
+    "I-QUANTITY": 30,
+    "I-FAC": 31,
+    "I-EVENT": 32,
+    "I-ORDINAL": 33,
+    "I-PRODUCT": 34,
+    "I-LAW": 35,
+    "I-LANGUAGE": 36,
+    "L-ORG": 37,
+    "L-DATE": 38,
+    "L-PERSON": 39,
+    "L-GPE": 40,
+    "L-MONEY": 41,
+    "L-CARDINAL": 42,
+    "L-NORP": 43,
+    "L-PERCENT": 44,
+    "L-WORK_OF_ART": 45,
+    "L-LOC": 46,
+    "L-TIME": 47,
+    "L-QUANTITY": 48,
+    "L-FAC": 49,
+    "L-EVENT": 50,
+    "L-ORDINAL": 51,
+    "L-PRODUCT": 52,
+    "L-LAW": 53,
+    "L-LANGUAGE": 54,
+    "U-ORG": 55,
+    "U-DATE": 56,
+    "U-PERSON": 57,
+    "U-GPE": 58,
+    "U-MONEY": 59,
+    "U-CARDINAL": 60,
+    "U-NORP": 61,
+    "U-PERCENT": 62,
+    "U-WORK_OF_ART": 63,
+    "U-LOC": 64,
+    "U-TIME": 65,
+    "U-QUANTITY": 66,
+    "U-FAC": 67,
+    "U-EVENT": 68,
+    "U-ORDINAL": 69,
+    "U-PRODUCT": 70,
+    "U-LAW": 71,
+    "U-LANGUAGE": 72,
+    "O": 73,
+}
+
+
+class HMNetModel(SummModel):
+    # static variables
+    model_name = "HMNET"
+    is_extractive = False
+    is_neural = True
+    is_dialogue_based = True
+
+    def __init__(
+        self,
+        min_gen_length: int = 10,
+        max_gen_length: int = 300,
+        beam_width: int = 6,
+        **kwargs,
+    ):
+        """
+        Create a summarization model with HMNet backbone. In the default setting, the inference speed will be
+        10s/sample (on one GPU), however, if one can tune these three parameters properly, e.g. min_gen_length=10,
+        max_gen_length=100, and beam_width=2, the inference speed will increase to 2s/sample (on one GPU).
+
+        Args:
+          min_gen_length (int): minimum generation length of the decoder
+          max_gen_length (int): maximum generation length of the decoder
+          beam_width (int): width of the beam when doing beam search in the decoding process
+          kwargs: the other valid parameters. The valid parameters can be found in
+              model/dialogue/hmnet/config/dialogue.conf . You can use either lower case or upper case for parameter
+              name. The valid parameter name is one of the following args, however, we do not encourage you to modify
+               them, since some unexpected, untested errors might be triggered:
+              ['MODEL', 'TASK', 'CRITERION', 'SEED', 'MAX_NUM_EPOCHS', 'EVAL_PER_UPDATE_NUM'
+              , 'UPDATES_PER_EPOCH', 'OPTIMIZER', 'START_LEARNING_RATE', 'LR_SCHEDULER', 'WARMUP_STEPS',
+              'WARMUP_INIT_LR', 'WARMUP_END_LR', 'GRADIENT_ACCUMULATE_STEP', 'GRAD_CLIPPING', 'USE_REL_DATA_PATH',
+              'TRAIN_FILE', 'DEV_FILE', 'TEST_FILE', 'ROLE_DICT_FILE', 'MINI_BATCH', 'MAX_PADDING_RATIO',
+              'BATCH_READ_AHEAD', 'DOC_SHUFFLE_BUF_SIZE', 'SAMPLE_SHUFFLE_BUFFER_SIZE', 'BATCH_SHUFFLE_BUFFER_SIZE',
+              'MAX_TRANSCRIPT_WORD', 'MAX_SENT_LEN', 'MAX_SENT_NUM', 'DROPOUT', 'VOCAB_DIM', 'ROLE_SIZE', 'ROLE_DIM',
+              'POS_DIM', 'ENT_DIM', 'USE_ROLE', 'USE_POSENT', 'USE_BOS_TOKEN', 'USE_EOS_TOKEN',
+              'TRANSFORMER_EMBED_DROPOUT', 'TRANSFORMER_RESIDUAL_DROPOUT', 'TRANSFORMER_ATTENTION_DROPOUT',
+              'TRANSFORMER_LAYER', 'TRANSFORMER_HEAD', 'TRANSFORMER_POS_DISCOUNT', 'PRE_TOKENIZER',
+              'PRE_TOKENIZER_PATH', 'PYLEARN_MODEL', 'EXTRA_IDS', 'BEAM_WIDTH', 'EVAL_TOKENIZED', 'EVAL_LOWERCASE',
+              'MAX_GEN_LENGTH', 'MIN_GEN_LENGTH', 'NO_REPEAT_NGRAM_SIZE']
+
+        Return an instance of HMNet model for dialogue summarization.
+        """
+        super(HMNetModel, self).__init__()
+        self.root_path = self._get_root()
+
+        # we leave the most influential params with prompt and the others as hidden kwargs
+        kwargs["MIN_GEN_LENGTH"] = min_gen_length
+        kwargs["MAX_GEN_LENGTH"] = max_gen_length
+        kwargs["BEAM_WIDTH"] = beam_width
+        self.opt = self._parse_args(kwargs)
+        self.model = HMNetTrainer(self.opt)
+
+    def _get_root(self):
+        root_path = os.getcwd()
+        while "model" not in os.listdir(root_path):
+            root_path = os.path.dirname(root_path)
+        root_path = os.path.join(root_path, "model/dialogue")
+        return root_path
+
+    def _parse_args(self, kwargs):
+        parser = argparse.ArgumentParser(
+            description="HMNet: Pretrain or fine-tune models for HMNet model."
+        )
+        parser.add_argument(
+            "--command", default="evaluate", help="Command: train/evaluate"
+        )
+        parser.add_argument(
+            "--conf_file",
+            default=os.path.join(self.root_path, "hmnet/config/dialogue.conf"),
+            help="Path to the BigLearn conf file.",
+        )
+        parser.add_argument(
+            "--PYLEARN_MODEL", help="Overrides this option from the conf file."
+        )
+        parser.add_argument(
+            "--master_port", help="Overrides this option default", default=None
+        )
+        parser.add_argument("--cluster", help="local, philly or aml", default="local")
+        parser.add_argument(
+            "--dist_init_path", help="Distributed init path for AML", default="./tmp"
+        )
+        parser.add_argument(
+            "--fp16",
+            action="store_true",
+            help="Whether to use 16-bit float precision instead of 32-bit",
+        )
+        parser.add_argument(
+            "--fp16_opt_level",
+            type=str,
+            default="O1",
+            help="For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
+            "See details at https://nvidia.github.io/apex/amp.html",
+        )
+        parser.add_argument("--no_cuda", action="store_true", help="Disable cuda.")
+        parser.add_argument(
+            "--config_overrides",
+            help="Override parameters on config, VAR=val;VAR=val;...",
+        )
+
+        cmdline_args = parser.parse_args()
+        command = cmdline_args.command
+        conf_file = cmdline_args.conf_file
+        conf_args = Arguments(conf_file)
+        opt = conf_args.readArguments()
+
+        if cmdline_args.config_overrides:
+            for config_override in cmdline_args.config_overrides.split(";"):
+                config_override = config_override.strip()
+                if config_override:
+                    var_val = config_override.split("=")
+                    assert (
+                        len(var_val) == 2
+                    ), f"Config override '{var_val}' does not have the form 'VAR=val'"
+                    conf_args.add_opt(opt, var_val[0], var_val[1], force_override=True)
+
+        opt["cuda"] = torch.cuda.is_available() and not cmdline_args.no_cuda
+        opt["confFile"] = conf_file
+        if "datadir" not in opt:
+            opt["datadir"] = os.path.dirname(
+                conf_file
+            )  # conf_file specifies where the data folder is
+        opt["basename"] = os.path.basename(
+            conf_file
+        )  # conf_file specifies where the name of save folder is
+        opt["command"] = command
+
+        # combine cmdline_args into opt dictionary
+        for key, val in cmdline_args.__dict__.items():
+            # if val is not None and key not in ['command', 'conf_file']:
+            if val is not None:
+                opt[key] = val
+
+        # combine kwargs into opt dictionary (we allow lower case)
+        for key, val in kwargs.items():
+            valid_keys = [x for x in opt.keys() if x.upper() == x]
+            if key.upper() not in valid_keys:
+                print("WARNING: {} is not a valid key in HMNet.".format(key))
+                print("The valid keys are:", valid_keys)
+                continue
+            if val is not None:
+                opt[key.upper()] = val
+
+        return opt
+
+    def summarize(self, corpus, queries=None):
+        print(f"HMNet model: processing document of {corpus.__len__()} samples")
+        # transform the original dataset to "dialogue" input
+        # we only use test set path for evaluation
+        data_folder = os.path.join(
+            os.path.dirname(self.opt["datadir"]),
+            "ExampleRawData/meeting_summarization/AMI_proprec/test",
+        )
+
+        self._create_datafolder(data_folder)
+        self._preprocess(corpus, data_folder)
+
+        # return self.model.eval()
+        results = self._evaluate()
+
+        return results
+
+    def _evaluate(self):
+        if self.opt["rank"] == 0:
+            self.model.log("-----------------------------------------------")
+            self.model.log("Evaluating model ... ")
+
+        self.model.set_up_model()
+
+        eval_dataset = "test"
+        batch_generator_eval = self.model.get_batch_generator(eval_dataset)
+        predictions = self._eval_batches(
+            self.model.module, batch_generator_eval, self.model.saveFolder, eval_dataset
+        )
+
+        return predictions
+
+    def _eval_batches(self, module, dev_batches, save_folder, label=""):
+        max_sent_len = int(self.opt["MAX_GEN_LENGTH"])
+
+        print("Decoding current model ... \nSaving folder is {}".format(save_folder))
+        print("Each sample will cost about 10 second.")
+        import time
+
+        start_time = time.time()
+        predictions = []  # prediction of tokens from model
+        if not isinstance(module.tokenizer, list):
+            decoder_tokenizer = module.tokenizer
+        elif len(module.tokenizer) == 1:
+            decoder_tokenizer = module.tokenizer[0]
+        elif len(module.tokenizer) == 2:
+            decoder_tokenizer = module.tokenizer[1]
+        else:
+            assert False, "len(module.tokenizer) > 2"
+
+        with torch.no_grad():
+            for j, dev_batch in enumerate(dev_batches):
+                for b in dev_batch:
+                    if torch.is_tensor(dev_batch[b]):
+                        dev_batch[b] = dev_batch[b].to(self.opt["device"])
+
+                beam_search_res = module(
+                    dev_batch, beam_search=True, max_sent_len=max_sent_len
+                )
+                pred = [
+                    [t[0] for t in x] if len(x) > 0 else [[]] for x in beam_search_res
+                ]
+                predictions.extend(
+                    [
+                        [
+                            self._convert_tokens_to_string(decoder_tokenizer, tt)
+                            for tt in t
+                        ]
+                        for t in pred
+                    ]
+                )
+
+                if (
+                    "DEBUG" in self.opt and j >= 10
+                ) or j >= self.model.task.evaluator.eval_batches_num:
+                    # in debug mode (decode first 10 batches) ortherwise decode first self.eval_batches_num bathes
+                    break
+
+        top1_predictions = [x[0] for x in predictions]
+
+        print("Total time for inference:", time.time() - start_time)
+        return top1_predictions
+
+    def _convert_tokens_to_string(self, tokenizer, tokens):
+        if "EVAL_TOKENIZED" in self.opt:
+            tokens = [t for t in tokens if t not in tokenizer.all_special_tokens]
+        if "EVAL_LOWERCASE" in self.opt:
+            tokens = [t.lower() for t in tokens]
+        if "EVAL_TOKENIZED" in self.opt:
+            return " ".join(tokens)
+        else:
+            return tokenizer.decode(
+                tokenizer.convert_tokens_to_ids(tokens), skip_special_tokens=True
+            )
+
+    def _preprocess(self, corpus, test_path):
+        samples = []
+        for i, sample in enumerate(corpus):
+            new_sample = {"id": i, "meeting": [], "summary": []}
+            if isinstance(sample, str):
+                raise RuntimeError(
+                    "Error: the input of HMNet should be dialogues, rather than documents."
+                )
+
+            # add all the turns one by one
+            for turn in sample:
+                turn = [x.strip() for x in turn.split(":")]
+                if len(turn) < 2:
+                    continue
+                tokenized_turn = nlp(turn[1])
+                # In case we can't find proper entity in move_names
+                ent_id = []
+                pos_id = []
+                for token in tokenized_turn:
+                    ent = (
+                        token.ent_iob_ + "-" + token.ent_type_
+                        if token.ent_iob_ != "O"
+                        else "O"
+                    )
+                    ent_id.append(ENT[ent] if ent in ENT else ENT[""])
+
+                    pos = token.tag_
+                    pos_id.append(POS[pos] if pos in POS else POS[""])
+
+                new_sample["meeting"].append(
+                    {
+                        "speaker": turn[0],
+                        "role": "",
+                        "utt": {
+                            "word": [str(token) for token in tokenized_turn],
+                            "pos_id": pos_id,
+                            "ent_id": ent_id,
+                        },
+                    }
+                )
+            new_sample["summary"].append(
+                "This is a dummy summary. HMNet will filter out the sample w/o summary!"
+            )
+            samples.append(new_sample)
+            # save to the gzip
+            file_path = os.path.join(test_path, "split_{}.jsonl.gz".format(i))
+            with gzip.open(file_path, "wt", encoding="utf-8") as file:
+                file.write(json.dumps(new_sample))
+
+    def _clean_datafolder(self, data_folder):
+        for name in os.listdir(data_folder):
+            name = os.path.join(data_folder, name)
+            if ".gz" in name:
+                os.remove(name)
+
+    def _create_datafolder(self, data_folder):
+        if os.path.exists(data_folder):
+            self._clean_datafolder(data_folder)
+        else:
+            os.makedirs(data_folder)
+        with open(
+            os.path.join(os.path.dirname(data_folder), "test_ami.json"),
+            "w",
+            encoding="utf-8",
+        ) as file:
+            json.dump(
+                [
+                    {
+                        "source": {
+                            "dataset": "../ExampleRawData/meeting_summarization/AMI_proprec/test/"
+                        },
+                        "task": "meeting",
+                        "name": "ami",
+                    }
+                ],
+                file,
+            )
+
+        with open(
+            os.path.join(
+                os.path.dirname(os.path.dirname(data_folder)), "role_dict_ext.json"
+            ),
+            "w",
+        ) as file:
+            json.dump({}, file)
+
+    @classmethod
+    def show_capability(cls) -> None:
+        basic_description = cls.generate_basic_description()
+        more_details = (
+            "A HMNet model finetuned on CNN-DM dataset for summarization.\n\n"
+            "Strengths:\n - High performance on dialogue summarization task.\n\n"
+            "Weaknesses:\n - Not suitable for datasets other than dialogues.\n\n"
+            "Initialization arguments:\n "
+            " - `corpus`: Unlabelled corpus of documents.\n"
+        )
+        print(f"{basic_description} \n {'#' * 20} \n {more_details}")

model/multi_doc/__init__.py

@@ -0,0 +1,2 @@
+from .multi_doc_joint_model import MultiDocJointModel
+from .multi_doc_separate_model import MultiDocSeparateModel

model/multi_doc/base_multi_doc_model.py

@@ -0,0 +1,40 @@
+from model.base_model import SummModel
+
+
+class MultiDocSummModel(SummModel):
+
+    is_multi_document = True
+
+    def __init__(
+        self,
+        trained_domain: str = None,
+        max_input_length: int = None,
+        max_output_length: int = None,
+    ):
+        super(MultiDocSummModel, self).__init__(
+            trained_domain=trained_domain,
+            max_input_length=max_input_length,
+            max_output_length=max_output_length,
+        )
+
+    @classmethod
+    def assert_summ_input_type(cls, corpus, query):
+        if not all(
+            [
+                isinstance(ins, list) and all([isinstance(doc, str) for doc in ins])
+                for ins in corpus
+            ]
+        ):
+            raise TypeError(
+                "Multi-document summarization models summarize instances of multiple documents (`List[List[str]]`)."
+            )
+
+        if query is not None:
+            if not isinstance(query, list):
+                raise TypeError(
+                    "Query-based single-document summarization requires query of `List[str]`."
+                )
+            if not all([isinstance(q, str) for q in query]):
+                raise TypeError(
+                    "Query-based single-document summarization requires query of `List[str]`."
+                )

model/multi_doc/multi_doc_joint_model.py

@@ -0,0 +1,51 @@
+from .base_multi_doc_model import MultiDocSummModel
+from model.base_model import SummModel
+from model.single_doc import TextRankModel
+from typing import Union, List
+
+
+class MultiDocJointModel(MultiDocSummModel):
+
+    model_name = "Multi-document joint"
+    is_multi_document = True
+
+    def __init__(self, model_backend: SummModel = TextRankModel, **kwargs):
+        super(MultiDocJointModel, self).__init__()
+        model = model_backend(**kwargs)
+        self.model = model
+
+    def summarize(
+        self,
+        corpus: Union[List[str], List[List[str]]],
+        query: Union[List[str], List[List[str]]] = None,
+    ) -> List[str]:
+        self.assert_summ_input_type(corpus, None)
+        joint_corpus = []
+        for instance in corpus:
+            joint_corpus.append(" ".join(instance))
+
+        summaries = self.model.summarize(joint_corpus)
+
+        return summaries
+
+    @classmethod
+    def generate_basic_description(cls) -> str:
+        basic_description = (
+            "MultiDocJointModel performs multi-document summarization by"
+            " first concatenating all documents,"
+            " and then performing single-document summarization on the concatenation."
+        )
+        return basic_description
+
+    @classmethod
+    def show_capability(cls):
+        basic_description = cls.generate_basic_description()
+        more_details = (
+            "A multi-document summarization model."
+            " Allows for custom model backend selection at initialization."
+            " Concatenates each document in corpus and returns single-document summarization of joint corpus.\n"
+            "Strengths: \n - Allows for control of backend model.\n"
+            "Weaknesses: \n - Assumes all documents are equally weighted.\n"
+            " - May fail to extract information from certain documents.\n"
+        )
+        print(f"{basic_description}\n{'#' * 20}\n{more_details}")

model/multi_doc/multi_doc_separate_model.py

@@ -0,0 +1,49 @@
+from .base_multi_doc_model import MultiDocSummModel
+from model.base_model import SummModel
+from model.single_doc import TextRankModel
+from typing import Union, List
+
+
+class MultiDocSeparateModel(MultiDocSummModel):
+
+    model_name = "Multi-document separate"
+    is_multi_document = True
+
+    def __init__(self, model_backend: SummModel = TextRankModel, **kwargs):
+        super(MultiDocSeparateModel, self).__init__()
+        model = model_backend(**kwargs)
+        self.model = model
+
+    def summarize(
+        self,
+        corpus: Union[List[str], List[List[str]]],
+        query: Union[List[str], List[List[str]]] = None,
+    ) -> List[str]:
+        self.assert_summ_input_type(corpus, None)
+        summaries = []
+        for instance in corpus:
+            instance_summaries = self.model.summarize(instance)
+            summaries.append(" ".join(instance_summaries))
+
+        return summaries
+
+    @classmethod
+    def generate_basic_description(cls) -> str:
+        basic_description = (
+            "MultiDocSeparateModel performs multi-document summarization by"
+            " first performing single-document summarization on each document,"
+            " and then concatenating the results."
+        )
+        return basic_description
+
+    @classmethod
+    def show_capability(cls):
+        basic_description = cls.generate_basic_description()
+        more_details = (
+            "A multi-document summarization model."
+            " Allows for custom model backend selection at initialization."
+            " Performs single-document summarization on each document in corpus and returns concatenated result.\n"
+            "Strengths: \n - Allows for control of backend model.\n"
+            "Weaknesses: \n - Assumes all documents are equally weighted.\n - May produce redundant information for similar documents.\n"
+        )
+        print(f"{basic_description}\n{'#' * 20}\n{more_details}")

model/query_based/__init__.py

@@ -0,0 +1,2 @@
+from .bm25_model import BM25SummModel
+from .tf_idf_model import TFIDFSummModel

model/query_based/base_query_based_model.py

@@ -0,0 +1,147 @@
+from model.base_model import SummModel
+from model.single_doc import TextRankModel
+from typing import List, Union
+
+from nltk import sent_tokenize, word_tokenize
+from nltk.corpus import stopwords
+from nltk.stem import PorterStemmer
+
+
+class QueryBasedSummModel(SummModel):
+
+    is_query_based = True
+
+    def __init__(
+        self,
+        trained_domain: str = None,
+        max_input_length: int = None,
+        max_output_length: int = None,
+        model_backend: SummModel = TextRankModel,
+        retrieval_ratio: float = 0.5,
+        preprocess: bool = True,
+        **kwargs,
+    ):
+        super(QueryBasedSummModel, self).__init__(
+            trained_domain=trained_domain,
+            max_input_length=max_input_length,
+            max_output_length=max_output_length,
+        )
+        self.model = model_backend(**kwargs)
+        self.retrieval_ratio = retrieval_ratio
+        self.preprocess = preprocess
+
+    def _retrieve(self, instance: List[str], query: List[str], n_best) -> List[str]:
+        raise NotImplementedError()
+
+    def summarize(
+        self,
+        corpus: Union[List[str], List[List[str]]],
+        queries: List[str] = None,
+    ) -> List[str]:
+        self.assert_summ_input_type(corpus, queries)
+
+        retrieval_output = []  # List[str]
+        for instance, query in zip(corpus, queries):
+            if isinstance(instance, str):
+                is_dialogue = False
+                instance = sent_tokenize(instance)
+            else:
+                is_dialogue = True
+            query = [query]
+
+            # instance & query now are List[str] for sure
+            if self.preprocess:
+                preprocessor = Preprocessor()
+                instance = preprocessor.preprocess(instance)
+                query = preprocessor.preprocess(query)
+
+            n_best = max(int(len(instance) * self.retrieval_ratio), 1)
+            top_n_sent = self._retrieve(instance, query, n_best)
+
+            if not is_dialogue:
+                top_n_sent = " ".join(top_n_sent)  # str
+            retrieval_output.append(top_n_sent)
+
+        summaries = self.model.summarize(
+            retrieval_output
+        )  # List[str] or List[List[str]]
+        return summaries
+
+    def generate_specific_description(self):
+        is_neural = self.model.is_neural & self.is_neural
+        is_extractive = self.model.is_extractive | self.is_extractive
+        model_name = "Pipeline with retriever: {}, summarizer: {}".format(
+            self.model_name, self.model.model_name
+        )
+
+        extractive_abstractive = "extractive" if is_extractive else "abstractive"
+        neural = "neural" if is_neural else "non-neural"
+
+        basic_description = (
+            f"{model_name} is a "
+            f"{'query-based' if self.is_query_based else ''} "
+            f"{extractive_abstractive}, {neural} model for summarization."
+        )
+
+        return basic_description
+
+    @classmethod
+    def assert_summ_input_type(cls, corpus, query):
+        if query is None:
+            raise TypeError(
+                "Query-based summarization models summarize instances of query-text pairs， however, query is missing."
+            )
+
+        if not isinstance(query, list):
+            raise TypeError(
+                "Query-based single-document summarization requires query of `List[str]`."
+            )
+        if not all([isinstance(q, str) for q in query]):
+            raise TypeError(
+                "Query-based single-document summarization requires query of `List[str]`."
+            )
+
+    @classmethod
+    def generate_basic_description(cls) -> str:
+        basic_description = (
+            "QueryBasedSummModel performs query-based summarization. Given a query-text pair,"
+            "the model will first extract the most relevant sentences in articles or turns in "
+            "dialogues, then use the single document summarization model to generate the summary"
+        )
+        return basic_description
+
+    @classmethod
+    def show_capability(cls):
+        basic_description = cls.generate_basic_description()
+        more_details = (
+            "A query-based summarization model."
+            " Allows for custom model backend selection at initialization."
+            " Retrieve relevant turns and then summarize the retrieved turns\n"
+            "Strengths: \n - Allows for control of backend model.\n"
+            "Weaknesses: \n - Heavily depends on the performance of both retriever and summarizer.\n"
+        )
+        print(f"{basic_description}\n{'#' * 20}\n{more_details}")
+
+
+class Preprocessor:
+    def __init__(self, remove_stopwords=True, lower_case=True, stem=False):
+        self.sw = stopwords.words("english")
+        self.stemmer = PorterStemmer()
+        self.remove_stopwords = remove_stopwords
+        self.lower_case = lower_case
+        self.stem = stem
+
+    def preprocess(self, corpus: List[str]) -> List[str]:
+        if self.lower_case:
+            corpus = [sent.lower() for sent in corpus]
+        tokenized_corpus = [word_tokenize(sent) for sent in corpus]
+        if self.remove_stopwords:
+            tokenized_corpus = [
+                [word for word in sent if word not in self.sw]
+                for sent in tokenized_corpus
+            ]
+        if self.stem:
+            tokenized_corpus = [
+                [self.stemmer.stem(word) for word in sent] for sent in tokenized_corpus
+            ]
+        return [" ".join(sent) for sent in tokenized_corpus]

model/query_based/bm25_model.py

@@ -0,0 +1,45 @@
+from .base_query_based_model import QueryBasedSummModel
+from model.base_model import SummModel
+from model.single_doc import TextRankModel
+from typing import List
+
+from gensim.summarization.bm25 import BM25
+from nltk import word_tokenize
+
+
+class BM25SummModel(QueryBasedSummModel):
+
+    # static variables
+    model_name = "BM25"
+    is_extractive = True  # only represents the retrieval part
+    is_neural = False  # only represents the retrieval part
+    is_query_based = True
+
+    def __init__(
+        self,
+        trained_domain: str = None,
+        max_input_length: int = None,
+        max_output_length: int = None,
+        model_backend: SummModel = TextRankModel,
+        retrieval_ratio: float = 0.5,
+        preprocess: bool = True,
+        **kwargs
+    ):
+        super(BM25SummModel, self).__init__(
+            trained_domain=trained_domain,
+            max_input_length=max_input_length,
+            max_output_length=max_output_length,
+            model_backend=model_backend,
+            retrieval_ratio=retrieval_ratio,
+            preprocess=preprocess,
+            **kwargs
+        )
+
+    def _retrieve(self, instance: List[str], query: List[str], n_best):
+        bm25 = BM25(word_tokenize(s) for s in instance)
+        scores = bm25.get_scores(query)
+        best_sent_ind = sorted(
+            range(len(scores)), key=lambda i: scores[i], reverse=True
+        )[:n_best]
+        top_n_sent = [instance[ind] for ind in sorted(best_sent_ind)]
+        return top_n_sent

model/query_based/tf_idf_model.py

@@ -0,0 +1,46 @@
+from .base_query_based_model import QueryBasedSummModel
+from model.base_model import SummModel
+from model.single_doc import TextRankModel
+from typing import List
+
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
+
+
+class TFIDFSummModel(QueryBasedSummModel):
+
+    # static variables
+    model_name = "TF-IDF"
+    is_extractive = True
+    is_neural = False
+    is_query_based = True
+
+    def __init__(
+        self,
+        trained_domain: str = None,
+        max_input_length: int = None,
+        max_output_length: int = None,
+        model_backend: SummModel = TextRankModel,
+        retrieval_ratio: float = 0.5,
+        preprocess: bool = True,
+        **kwargs
+    ):
+        super(TFIDFSummModel, self).__init__(
+            trained_domain=trained_domain,
+            max_input_length=max_input_length,
+            max_output_length=max_output_length,
+            model_backend=model_backend,
+            retrieval_ratio=retrieval_ratio,
+            preprocess=preprocess,
+            **kwargs
+        )
+        self.vectorizer = TfidfVectorizer()
+
+    def _retrieve(self, instance: List[str], query: List[str], n_best):
+        instance_vectors = self.vectorizer.fit_transform(instance)
+        query_vector = self.vectorizer.transform(query)
+
+        similarities = cosine_similarity(query_vector, instance_vectors).squeeze()
+        top_n_index = similarities.argsort()[::-1][0:n_best]
+        top_n_sent = [instance[ind] for ind in top_n_index]  # List[str]
+        return top_n_sent

model/single_doc/__init__.py

@@ -0,0 +1,5 @@
+from .bart_model import BartModel
+from .pegasus_model import PegasusModel
+from .lexrank_model import LexRankModel
+from .longformer_model import LongformerModel
+from .textrank_model import TextRankModel

model/single_doc/bart_model.py

@@ -0,0 +1,36 @@
+from transformers import BartForConditionalGeneration, BartTokenizer
+from .base_single_doc_model import SingleDocSummModel
+
+
+class BartModel(SingleDocSummModel):
+
+    # static variables
+    model_name = "BART"
+    is_extractive = False
+    is_neural = False
+
+    def __init__(self, device="cpu"):
+        super(BartModel, self).__init__()
+
+        self.device = device
+        model_name = "facebook/bart-large-cnn"
+        self.tokenizer = BartTokenizer.from_pretrained(model_name)
+        self.model = BartForConditionalGeneration.from_pretrained(model_name)
+
+    def summarize(self, corpus, queries=None):
+        self.assert_summ_input_type(corpus, queries)
+
+        batch = self.tokenizer(
+            corpus, truncation=True, padding="longest", return_tensors="pt"
+        ).to(self.device)
+        encoded_summaries = self.model.generate(**batch)
+        summaries = self.tokenizer.batch_decode(
+            encoded_summaries, skip_special_tokens=True
+        )
+
+        return summaries
+
+    @classmethod
+    def show_capability(cls) -> None:
+        # TODO zhangir: add the show capability function for BART
+        print(cls.generate_basic_description())

model/single_doc/base_single_doc_model.py

@@ -0,0 +1,36 @@
+from model.base_model import SummModel
+
+
+class SingleDocSummModel(SummModel):
+    def __init__(
+        self,
+        trained_domain: str = None,
+        max_input_length: int = None,
+        max_output_length: int = None,
+    ):
+        super(SingleDocSummModel, self).__init__(
+            trained_domain=trained_domain,
+            max_input_length=max_input_length,
+            max_output_length=max_output_length,
+        )
+
+    @classmethod
+    def assert_summ_input_type(cls, corpus, query):
+        if not isinstance(corpus, list):
+            raise TypeError(
+                "Single-document summarization requires corpus of `List[str]`."
+            )
+        if not all([isinstance(ins, str) for ins in corpus]):
+            raise TypeError(
+                "Single-document summarization requires corpus of `List[str]`."
+            )
+
+        if query is not None:
+            if not isinstance(query, list):
+                raise TypeError(
+                    "Query-based single-document summarization requires query of `List[str]`."
+                )
+            if not all([isinstance(q, str) for q in query]):
+                raise TypeError(
+                    "Query-based single-document summarization requires query of `List[str]`."
+                )

model/single_doc/lexrank_model.py

@@ -0,0 +1,50 @@
+from lexrank import STOPWORDS
+from lexrank import LexRank as LR
+import nltk
+
+from .base_single_doc_model import SingleDocSummModel
+
+
+class LexRankModel(SingleDocSummModel):
+    # static variables
+    model_name = "LexRank"
+    is_extractive = True
+    is_neural = False
+
+    def __init__(self, data, summary_length=2, threshold=0.1):
+        super(LexRankModel, self).__init__()
+
+        nltk.download("punkt", quiet=True)
+        corpus = [nltk.sent_tokenize(example) for example in data]
+        self.lxr = LR(corpus, stopwords=STOPWORDS["en"])
+        self.summary_length = summary_length
+        self.threshold = threshold
+
+    def summarize(self, corpus, queries=None):
+        self.assert_summ_input_type(corpus, queries)
+
+        documents = [nltk.sent_tokenize(document) for document in corpus]
+        summaries = [
+            " ".join(
+                self.lxr.get_summary(
+                    document, summary_size=self.summary_length, threshold=self.threshold
+                )
+            )
+            for document in documents
+        ]
+
+        return summaries
+
+    @classmethod
+    def show_capability(cls):
+        basic_description = cls.generate_basic_description()
+        more_details = (
+            "Works by using a graph-based method to identify the most salient sentences in the document. \n"
+            "Strengths: \n - Fast with low memory usage \n - Allows for control of summary length \n "
+            "Weaknesses: \n - Not as accurate as neural methods. \n "
+            "Initialization arguments: \n "
+            "- `corpus`: Unlabelled corpus of documents. ` \n "
+            "- `summary_length`: sentence length of summaries \n "
+            "- `threshold`: Level of salience required for sentence to be included in summary."
+        )
+        print(f"{basic_description} \n {'#'*20} \n {more_details}")

model/single_doc/longformer_model.py

@@ -0,0 +1,57 @@
+from transformers import LongformerTokenizer, EncoderDecoderModel
+from .base_single_doc_model import SingleDocSummModel
+
+
+class LongformerModel(SingleDocSummModel):
+
+    # static variables
+    model_name = "Longformer"
+    is_extractive = False
+    is_neural = True
+
+    def __init__(self):
+        super(LongformerModel, self).__init__()
+
+        self.model = EncoderDecoderModel.from_pretrained(
+            "patrickvonplaten/longformer2roberta-cnn_dailymail-fp16"
+        )
+        self.tokenizer = LongformerTokenizer.from_pretrained(
+            "allenai/longformer-base-4096"
+        )
+
+    def summarize(self, corpus, queries=None):
+        self.assert_summ_input_type(corpus, queries)
+
+        summaries = list(map(lambda doc: self.summarize_single(doc), corpus))
+
+        return summaries
+
+    def summarize_single(self, document):
+        # Tokenizes document and returns PyTorch torch.Tensor object with length attribute
+        tokenized_sequence = self.tokenizer(
+            document,
+            return_tensors="pt",
+            return_length=True,
+            truncation=True,
+            max_length=4096,
+        )
+        print(
+            f"Longformer model: processing document of {tokenized_sequence.length} tokens"
+        )
+        input_ids = tokenized_sequence.input_ids
+        # output_ids is tensor with one layer: output_ids[0] extracts tensor layer for decoding
+        output_ids = self.model.generate(input_ids)
+
+        return self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+    @classmethod
+    def show_capability(cls) -> None:
+        basic_description = cls.generate_basic_description()
+        more_details = (
+            "A Longformer2Roberta model finetuned on CNN-DM dataset for summarization.\n\n"
+            "Strengths:\n - Correctly handles longer (> 2000 tokens) corpus.\n\n"
+            "Weaknesses:\n - Less accurate on contexts outside training domain.\n\n"
+            "Initialization arguments:\n "
+            " - `corpus`: Unlabelled corpus of documents.\n"
+        )
+        print(f"{basic_description} \n {'#'*20} \n {more_details}")

model/single_doc/pegasus_model.py

@@ -0,0 +1,50 @@
+from transformers import PegasusForConditionalGeneration, PegasusTokenizer
+from .base_single_doc_model import SingleDocSummModel
+
+
+class PegasusModel(SingleDocSummModel):
+    # static variables
+    model_name = "Pegasus"
+    is_extractive = False
+    is_neural = True
+
+    def __init__(self, device="cpu"):
+        super(PegasusModel, self).__init__()
+
+        self.device = device
+        model_name = "google/pegasus-xsum"
+        print("init load pretrained tokenizer")
+        self.tokenizer = PegasusTokenizer.from_pretrained(model_name)
+        print("init load pretrained model with tokenizer on " + device)
+        # self.model = PegasusForConditionalGeneration.from_pretrained(model_name).to(device)
+        self.model = PegasusForConditionalGeneration.from_pretrained(model_name)
+
+    def summarize(self, corpus, queries=None):
+        self.assert_summ_input_type(corpus, queries)
+
+        print("batching")
+        # batch = self.tokenizer(corpus, truncation=True, padding='longest', return_tensors="pt").to(self.device)
+        batch = self.tokenizer(corpus, truncation=True, return_tensors="pt")
+        print("encoding batches")
+        # encoded_summaries = self.model.generate(**batch, max_length=40, max_time=120)
+        encoded_summaries = self.model.generate(batch["input_ids"], max_time=1024)
+        print("decoding batches")
+        # summaries = self.tokenizer.batch_decode(encoded_summaries, skip_special_tokens=True)
+        summaries = [self.tokenizer.decode(encoded_summaries[0])]
+
+        return summaries
+
+    @classmethod
+    def show_capability(cls):
+        basic_description = cls.generate_basic_description()
+        more_details = (
+            "Introduced in 2019, a large neural abstractive summarization model trained on web crawl and "
+            "news data.\n "
+            "Strengths: \n - High accuracy \n - Performs well on almost all kinds of non-literary written "
+            "text \n "
+            "Weaknesses: \n - High memory usage \n "
+            "Initialization arguments: \n "
+            "- `device = 'cpu'` specifies the device the model is stored on and uses for computation. "
+            "Use `device='gpu'` to run on an Nvidia GPU."
+        )
+        print(f"{basic_description} \n {'#'*20} \n {more_details}")

model/single_doc/textrank_model.py

@@ -0,0 +1,89 @@
+import spacy
+import pytextrank  # noqa: F401
+from math import sqrt
+from operator import itemgetter
+from .base_single_doc_model import SingleDocSummModel
+from typing import Union, List
+
+
+class TextRankModel(SingleDocSummModel):
+    # static variables
+    model_name = "TextRank"
+    is_extractive = True
+    is_neural = False
+
+    def __init__(self, num_sentences=1):
+        super(TextRankModel, self).__init__()
+
+        self.num_sentences = num_sentences
+        # load a spaCy model, depending on language, scale, etc.
+        self.nlp = spacy.load("en_core_web_sm")
+        self.nlp.add_pipe("textrank", last=True)
+
+    def summarize(
+        self, corpus: Union[List[str], List[List[str]]], queries: List[str] = None
+    ) -> List[str]:
+        self.assert_summ_input_type(corpus, queries)
+
+        return list(map(lambda x: " ".join(self.summarize_single(x)), corpus))
+
+    def summarize_single(self, corpus) -> List[str]:
+        # add PyTextRank to the spaCy pipeline
+        doc = self.nlp(corpus)
+        sent_bounds = [[s.start, s.end, set([])] for s in doc.sents]
+
+        limit_phrases = self.num_sentences
+        phrase_id = 0
+        unit_vector = []
+        for p in doc._.phrases:
+            unit_vector.append(p.rank)
+            for chunk in p.chunks:
+                for sent_start, sent_end, sent_vector in sent_bounds:
+                    if chunk.start >= sent_start and chunk.end <= sent_end:
+                        sent_vector.add(phrase_id)
+                        break
+            phrase_id += 1
+            if phrase_id == limit_phrases:
+                break
+
+        sum_ranks = sum(unit_vector)
+
+        unit_vector = [rank / sum_ranks for rank in unit_vector]
+
+        sent_rank = {}
+        sent_id = 0
+        for sent_start, sent_end, sent_vector in sent_bounds:
+            sum_sq = 0.0
+            for phrase_id in range(len(unit_vector)):
+                if phrase_id not in sent_vector:
+                    sum_sq += unit_vector[phrase_id] ** 2.0
+            sent_rank[sent_id] = sqrt(sum_sq)
+            sent_id += 1
+
+        sorted(sent_rank.items(), key=itemgetter(1))
+
+        sent_text = {}
+        sent_id = 0
+        limit_sentences = self.num_sentences
+        summary_sentences = []
+        for sent in doc.sents:
+            sent_text[sent_id] = sent.text
+            sent_id += 1
+        num_sent = 0
+        for sent_id, rank in sorted(sent_rank.items(), key=itemgetter(1)):
+            summary_sentences.append(sent_text[sent_id])
+            num_sent += 1
+            if num_sent == limit_sentences:
+                break
+
+        return summary_sentences
+
+    @classmethod
+    def show_capability(cls):
+        basic_description = cls.generate_basic_description()
+        more_details = (
+            "A graphbased ranking model for text processing. Extractive sentence summarization. \n "
+            "Strengths: \n - Fast with low memory usage \n - Allows for control of summary length \n "
+            "Weaknesses: \n - Not as accurate as neural methods."
+        )
+        print(f"{basic_description} \n {'#'*20} \n {more_details}")

model/third_party/HMNet/DataLoader/README.md

@@ -0,0 +1 @@
+This dataloader is adapted from Microsoft's [infinibatch](https://github.com/microsoft/infinibatch) implementation, which is a library of checkpointable iterators for randomized data loading of massive data sets in deep neural network training.

model/third_party/HMNet/DataLoader/__init__.py

@@ -0,0 +1 @@
+from .infinibatch.infinibatch import datasets, iterators