Spaces:
Running
Running
""" | |
aggregate.py is a module for aggregating text from multiple sources, or multiple parts of a single source. | |
Primary usage is through the BatchAggregator class. | |
How it works: | |
1. We tell the language model to do it. | |
2. The language model does it. | |
3. Yaay! | |
""" | |
import pprint as pp | |
import logging | |
import time | |
import torch | |
from transformers import GenerationConfig, pipeline | |
from utils import compare_model_size | |
# Setting up logging | |
logging.basicConfig( | |
level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" | |
) | |
class BatchAggregator: | |
""" | |
BatchAggregator is a class for aggregating text from multiple sources. | |
Usage: | |
>>> from aggregate import BatchAggregator | |
>>> aggregator = BatchAggregator() | |
>>> aggregator.aggregate(["This is a test", "This is another test"]) | |
""" | |
GENERIC_CONFIG = GenerationConfig( | |
num_beams=8, | |
early_stopping=True, | |
do_sample=False, | |
min_new_tokens=32, | |
max_new_tokens=256, | |
repetition_penalty=1.1, | |
length_penalty=1.4, | |
no_repeat_ngram_size=4, | |
encoder_no_repeat_ngram_size=5, | |
) | |
CONFIGURED_MODELS = [ | |
"pszemraj/bart-large-mnli-dolly_hhrlhf-v1", | |
"pszemraj/bart-base-instruct-dolly_hhrlhf", | |
"pszemraj/flan-t5-large-instruct-dolly_hhrlhf", | |
"pszemraj/flan-t5-base-instruct-dolly_hhrlhf", | |
] # these have generation configs defined for this task in their model repos | |
DEFAULT_INSTRUCTION = "Write a comprehensive yet concise summary that pulls together the main points of the following text:" | |
def __init__( | |
self, model_name: str = "pszemraj/bart-large-mnli-dolly_hhrlhf-v1", **kwargs | |
): | |
""" | |
__init__ initializes the BatchAggregator class. | |
:param str model_name: model name to use, default: "pszemraj/bart-large-mnli-dolly_hhrlhf-v1" | |
""" | |
self.device = None | |
self.is_compiled = False | |
self.logger = logging.getLogger(__name__) | |
self.init_model(model_name) | |
def init_model(self, model_name: str) -> None: | |
""" | |
Initialize the model. | |
:param model_name: The name of the model to use. | |
""" | |
# Free up memory | |
if torch.cuda.is_available(): | |
torch.cuda.empty_cache() | |
self.logger.info(f"Setting model to {model_name}") | |
self.model_name = model_name | |
self.aggregator = self._create_pipeline(model_name) | |
self._configure_model() | |
# update the generation config with the specific tokenizer | |
tokenizer_params = { | |
"decoder_start_token_id": 0 | |
if "t5" in model_name.lower() | |
else self.aggregator.tokenizer.eos_token_id, | |
"eos_token_id": 1 | |
if "t5" in model_name.lower() | |
else self.aggregator.tokenizer.eos_token_id, | |
"pad_token_id": 0 | |
if "t5" in model_name.lower() | |
else self.aggregator.tokenizer.pad_token_id, | |
} | |
self.update_generation_config(**tokenizer_params) | |
def _create_pipeline( | |
self, model_name: str = "pszemraj/bart-large-mnli-dolly_hhrlhf-v1" | |
) -> pipeline: | |
""" | |
_create_pipeline creates a pipeline for the model. | |
:param str model_name: model name to use, default: "pszemraj/bart-large-mnli-dolly_hhrlhf-v1" | |
:return pipeline: the pipeline for the model | |
:raises Exception: if the pipeline cannot be created | |
""" | |
self.device = 0 if torch.cuda.is_available() else -1 | |
try: | |
self.logger.info( | |
f"Creating pipeline with model {model_name} on device {self.device}" | |
) | |
return pipeline( | |
"text2text-generation", | |
model_name, | |
device=self.device, | |
torch_dtype=torch.float32, | |
) | |
except Exception as e: | |
self.logger.error(f"Failed to create pipeline: {e}") | |
raise | |
def _configure_model(self): | |
""" | |
Configure the model for generation. | |
""" | |
try: | |
self.aggregator.model = torch.compile(self.aggregator.model) | |
self.is_compiled = True | |
except Exception as e: | |
self.logger.warning(f"Could not compile model with Torch 2.0: {e}") | |
if self.model_name not in self.CONFIGURED_MODELS: | |
self.logger.info("Setting generation config to general defaults") | |
self._set_default_generation_config() | |
else: | |
try: | |
self.logger.info("Loading generation config from hub") | |
self.aggregator.model.generation_config = ( | |
GenerationConfig.from_pretrained(self.model_name) | |
) | |
except Exception as e: | |
self.logger.warning( | |
f"Could not load generation config, using defaults: {e}" | |
) | |
self._set_default_generation_config() | |
self.logger.info(self.aggregator.model.generation_config.to_json_string()) | |
def _set_default_generation_config(self): | |
""" | |
Set the default generation configuration for the model. | |
""" | |
self.aggregator.model.generation_config = self.GENERIC_CONFIG | |
if ( | |
"large" | |
or "xl" in self.model_name.lower() | |
or compare_model_size(self.model_name, 500) | |
): | |
upd = {"num_beams": 4} | |
self.update_generation_config(**upd) | |
def update_generation_config(self, **kwargs): | |
""" | |
Update the generation configuration with the specified parameters. | |
Args: | |
**kwargs: The parameters to update in the generation configuration. | |
""" | |
self.logger.info(f"Updating generation config with {pp.pformat(kwargs)}") | |
self.aggregator.model.generation_config.update(**kwargs) | |
def update_loglevel(self, level: str = "INFO"): | |
""" | |
Update the log level. | |
Args: | |
level (str): The log level to set. Defaults to "INFO". | |
""" | |
self.logger.setLevel(level) | |
def infer_aggregate( | |
self, | |
text_list: list, | |
instruction: str = DEFAULT_INSTRUCTION, | |
**kwargs, | |
) -> str: | |
f""" | |
Generate a summary of the specified texts. | |
Args: | |
text_list (list): The texts to summarize. | |
instruction (str): The instruction for the summary. Defaults to {self.DEFAULT_INSTRUCTION}. | |
**kwargs: Additional parameters to update in the generation configuration. | |
Returns: | |
The generated summary. | |
""" | |
joined_text = "\n".join(text_list) | |
prompt = f"{instruction}\n\n{joined_text}\n" | |
if kwargs: | |
self.update_generation_config(**kwargs) | |
st = time.perf_counter() | |
self.logger.info(f"inference on {len(text_list)} texts ...") | |
result = self.aggregator( | |
prompt, | |
generation_config=self.aggregator.model.generation_config, | |
)[0]["generated_text"] | |
self.logger.info(f"Done. runtime:\t{round(time.perf_counter() - st, 2)}s") | |
self.logger.info( | |
f"Input tokens:\t{self.count_tokens(prompt)}. Output tokens:\t{self.count_tokens(result)}" | |
) | |
return result | |
def count_tokens(self, text: str) -> int: | |
"""count the number of tokens in a text""" | |
return ( | |
len(self.aggregator.tokenizer.encode(text, truncation=False, padding=False)) | |
if text | |
else 0 | |
) | |