geraskalnas/document-summarization

1"""2summarize - a module for summarizing text using a model from the Hugging Face model hub3"""4import logging5import pprint as pp6 7logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(message)s")8 9import torch10from tqdm.auto import tqdm11from transformers import AutoModelForSeq2SeqLM, AutoTokenizer12 13from utils import validate_pytorch214 15 16def load_model_and_tokenizer(model_name: str) -> tuple:17    """18    load_model_and_tokenizer - load a model and tokenizer from a model name/ID on the hub19 20    :param str model_name: the model name/ID on the hub21    :return tuple: a tuple containing the model and tokenizer22    """23    device = "cuda" if torch.cuda.is_available() else "cpu"24    model = AutoModelForSeq2SeqLM.from_pretrained(25        model_name,26    ).to(device)27    model = model.eval()28 29    tokenizer = AutoTokenizer.from_pretrained(model_name)30 31    logging.info(f"Loaded model {model_name} to {device}")32 33    if validate_pytorch2():34        try:35            logging.info("Compiling model with Torch 2.0")36            model = torch.compile(model)37        except Exception as e:38            logging.warning(f"Could not compile model with Torch 2.0: {e}")39    else:40        logging.info("Torch 2.0 not detected, skipping compilation")41 42    return model, tokenizer43 44 45def summarize_and_score(46    ids, mask, model, tokenizer, is_general_attention_model=True, **kwargs47) -> tuple:48    """49    summarize_and_score - given a batch of ids and a mask, return a summary and a score for the summary50 51    Args:52        ids (): the batch of ids53        mask (): the attention mask for the batch54        model   (): the model to use for summarization55        tokenizer (): the tokenizer to use for summarization56        is_general_attention_model (bool, optional): whether the model is a general attention model. Defaults to True.57        **kwargs: any additional arguments to pass to the model58    Returns:59        tuple (str, float): the summary,  the score for the summary60    """61 62    ids = ids[None, :]63    mask = mask[None, :]64 65    input_ids = ids.to("cuda") if torch.cuda.is_available() else ids66    attention_mask = mask.to("cuda") if torch.cuda.is_available() else mask67 68    global_attention_mask = torch.zeros_like(attention_mask)69    # put global attention on <s> token70    global_attention_mask[:, 0] = 171 72    if is_general_attention_model:73        summary_pred_ids = model.generate(74            input_ids,75            attention_mask=attention_mask,76            output_scores=True,77            return_dict_in_generate=True,78            **kwargs,79        )80    else:81        summary_pred_ids = model.generate(82            input_ids,83            attention_mask=attention_mask,84            global_attention_mask=global_attention_mask,85            output_scores=True,86            return_dict_in_generate=True,87            **kwargs,88        )89    summary = tokenizer.batch_decode(90        summary_pred_ids.sequences,91        skip_special_tokens=True,92        remove_invalid_values=True,93    )94    score = round(summary_pred_ids.sequences_scores.cpu().numpy()[0], 4)95 96    return summary, score97 98 99def summarize_via_tokenbatches(100    input_text: str,101    model,102    tokenizer,103    batch_length=2048,104    batch_stride=16,105    min_batch_length=512,106    **kwargs,107) -> list:108    """109    summarize_via_tokenbatches - summarize a long string via batches of tokens110 111    Args:112        input_text (str): the text to summarize113        model (): the model to use for summarization114        tokenizer (): the tokenizer to use for summarization115        batch_length (int, optional): the length of each batch. Defaults to 2048.116        batch_stride (int, optional): the stride of each batch. Defaults to 16. The stride is the number of tokens that overlap between batches.117        min_batch_length (int, optional): the minimum length of each batch. Defaults to 512.118 119        **kwargs: any additional arguments to pass to the model for inference120    Returns:121        list: a list of dictionaries containing the input tokens, the summary, and the summary score122    """123 124    logger = logging.getLogger(__name__)125    # log all input parameters126    if batch_length < min_batch_length:127        logger.warning(128            f"batch_length must be at least {min_batch_length}. Setting batch_length to {min_batch_length}"129        )130        batch_length = min_batch_length131 132    logger.info(f"input parameters:\n{pp.pformat(kwargs)}")133    logger.info(f"batch_length: {batch_length}, batch_stride: {batch_stride}")134 135    encoded_input = tokenizer(136        input_text,137        padding="max_length",138        truncation=True,139        max_length=batch_length,140        stride=batch_stride,141        return_overflowing_tokens=True,142        add_special_tokens=False,143        return_tensors="pt",144    )145 146    in_id_arr, att_arr = encoded_input.input_ids, encoded_input.attention_mask147    gen_summaries = []148 149    pbar = tqdm(total=len(in_id_arr))150 151    for _id, _mask in zip(in_id_arr, att_arr):152        result, score = summarize_and_score(153            ids=_id,154            mask=_mask,155            model=model,156            tokenizer=tokenizer,157            **kwargs,158        )159        score = round(float(score), 4)160        _sum = {161            "input_tokens": _id,162            "summary": result,163            "summary_score": score,164        }165        gen_summaries.append(_sum)166        logger.debug(f"Score for batch: {score}. num chars: {len(repr(result))}")167        logger.debug(f"Summary:\n\t{result}")168        pbar.update()169 170    pbar.close()171 172    return gen_summaries173