geraskalnas/document-summarization

1"""2aggregate.py - module for aggregating text from multiple sources/multiple parts of a single source.3    Primary usage is through the BatchAggregator class.4 5How it works:61. We tell the language model to do it.72. The language model does it.83. Yaay!9"""10import logging11import pprint as pp12import time13 14import torch15from transformers import GenerationConfig, pipeline16 17from utils import compare_model_size18 19# Setting up logging20logging.basicConfig(21    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"22)23 24 25class BatchAggregator:26    """27    BatchAggregator is a class for aggregating text from multiple sources.28 29    Usage:30    >>> from aggregate import BatchAggregator31    >>> aggregator = BatchAggregator()32    >>> agg = aggregator.infer_aggregate(["This is a test", "This is another test"])33    >>> print(agg)34    """35 36    GENERIC_CONFIG = GenerationConfig(37        num_beams=8,38        early_stopping=True,39        do_sample=False,40        min_new_tokens=32,41        max_new_tokens=256,42        repetition_penalty=1.1,43        length_penalty=1.4,44        no_repeat_ngram_size=4,45        encoder_no_repeat_ngram_size=5,46    )47    CONFIGURED_MODELS = [48        "pszemraj/bart-large-mnli-dolly_hhrlhf-v1",49        "pszemraj/bart-base-instruct-dolly_hhrlhf",50        "pszemraj/flan-t5-large-instruct-dolly_hhrlhf",51        "pszemraj/flan-t5-base-instruct-dolly_hhrlhf",52    ]  # these have generation configs defined for this task in their model repos53 54    DEFAULT_INSTRUCTION = "Write a comprehensive yet concise summary that pulls together the main points of the following text:"55 56    def __init__(57        self,58        model_name: str = "pszemraj/bart-large-mnli-dolly_hhrlhf-v1",59        force_cpu: bool = False,60        **kwargs,61    ):62        """63        __init__ initializes the BatchAggregator class.64 65        :param str model_name: model name to use, default: "pszemraj/bart-large-mnli-dolly_hhrlhf-v1"66        :param bool force_cpu: force the model to run on CPU, default: False67        """68        self.device = None69        self.is_compiled = False70        self.model_name = None71        self.aggregator = None72        self.force_cpu = force_cpu73        self.logger = logging.getLogger(__name__)74        self.init_model(model_name)75 76    def init_model(self, model_name: str) -> None:77        """78        Initialize the model.79 80        :param model_name: The name of the model to use.81        """82        # Free up memory83        if torch.cuda.is_available():84            torch.cuda.empty_cache()85 86        self.logger.info(f"Setting model to {model_name}")87        self.model_name = model_name88        self.aggregator = self._create_pipeline(model_name)89        self._configure_model()90        # update the generation config with the specific tokenizer91        tokenizer_params = {92            "decoder_start_token_id": 093            if "t5" in model_name.lower()94            else self.aggregator.tokenizer.eos_token_id,95            "eos_token_id": 196            if "t5" in model_name.lower()97            else self.aggregator.tokenizer.eos_token_id,98            "pad_token_id": 099            if "t5" in model_name.lower()100            else self.aggregator.tokenizer.pad_token_id,101        }102        self.update_generation_config(**tokenizer_params)103 104    def _create_pipeline(105        self, model_name: str = "pszemraj/bart-large-mnli-dolly_hhrlhf-v1"106    ) -> pipeline:107        """108        _create_pipeline creates a pipeline for the model.109 110        :param str model_name: model name to use, default: "pszemraj/bart-large-mnli-dolly_hhrlhf-v1"111        :return pipeline: the pipeline for the model112 113        :raises Exception: if the pipeline cannot be created114        """115        self.device = 0 if torch.cuda.is_available() and not self.force_cpu else -1116        try:117            self.logger.info(118                f"Creating pipeline with model {model_name} on device {self.device}"119            )120            return pipeline(121                "text2text-generation",122                model_name,123                device=self.device,124                torch_dtype=torch.float32,125            )126        except Exception as e:127            self.logger.error(f"Failed to create pipeline: {e}")128            raise129 130    def _configure_model(self):131        """132        Configure the model for generation.133        """134        try:135            self.aggregator.model = torch.compile(self.aggregator.model)136            self.is_compiled = True137        except Exception as e:138            self.logger.warning(f"Could not compile model with Torch 2.0: {e}")139 140        if self.model_name not in self.CONFIGURED_MODELS:141            self.logger.info("Setting generation config to general defaults")142            self._set_default_generation_config()143        else:144            try:145                self.logger.info("Loading generation config from hub")146                self.aggregator.model.generation_config = (147                    GenerationConfig.from_pretrained(self.model_name)148                )149            except Exception as e:150                self.logger.warning(151                    f"Could not load generation config, using defaults: {e}"152                )153                self._set_default_generation_config()154 155        self.logger.info(self.aggregator.model.generation_config.to_json_string())156 157    def _set_default_generation_config(self):158        """159        Set the default generation configuration for the model.160        """161        self.aggregator.model.generation_config = self.GENERIC_CONFIG162 163        if (164            "large"165            or "xl" in self.model_name.lower()166            or compare_model_size(self.model_name, 500)167        ):168            upd = {"num_beams": 4}169            self.update_generation_config(**upd)170 171    def update_generation_config(self, **kwargs):172        """173        Update the generation configuration with the specified parameters.174 175        Args:176            **kwargs: The parameters to update in the generation configuration.177        """178        self.logger.info(f"Updating generation config with {pp.pformat(kwargs)}")179 180        self.aggregator.model.generation_config.update(**kwargs)181 182    def get_generation_config(self) -> dict:183        """184        Get the current generation configuration.185 186        Returns:187            dict: The current generation configuration.188        """189        return self.aggregator.model.generation_config.to_dict()190 191    def update_loglevel(self, level: str = "INFO"):192        """193        Update the log level.194 195        Args:196            level (str): The log level to set. Defaults to "INFO".197        """198        self.logger.setLevel(level)199 200    def infer_aggregate(201        self,202        text_list: list,203        instruction: str = DEFAULT_INSTRUCTION,204        **kwargs,205    ) -> str:206        f"""207        infer_aggregate - infers a consolidated summary from a list of texts.208 209        Args:210            text_list (list): The texts to summarize.211            instruction (str): The instruction for the summary. Defaults to {self.DEFAULT_INSTRUCTION}.212            **kwargs: Additional parameters to update in the generation configuration.213 214        Returns:215            The generated summary.216        """217        joined_text = "\n".join(text_list)218        prompt = f"{instruction}\n\n{joined_text}\n"219        if kwargs:220            self.update_generation_config(**kwargs)221        st = time.perf_counter()222        self.logger.info(f"inference on {len(text_list)} texts ...")223        result = self.aggregator(224            prompt,225            generation_config=self.aggregator.model.generation_config,226        )[0]["generated_text"]227        self.logger.info(f"Done. runtime:\t{round(time.perf_counter() - st, 2)}s")228        self.logger.info(229            f"Input tokens:\t{self.count_tokens(prompt)}. Output tokens:\t{self.count_tokens(result)}"230        )231        self.logger.debug(f"Generated text:\n{result}")232 233        return result234 235    def count_tokens(self, text: str) -> int:236        """count the number of tokens in a text"""237        return (238            len(self.aggregator.tokenizer.encode(text, truncation=False, padding=False))239            if text240            else 0241        )242