summagery_pipline.py

from transformers import AutoModelWithLMHead, AutoTokenizer
from transformers import GPT2LMHeadModel, GPT2Tokenizer
from diffusers import DiffusionPipeline
import torch
from tqdm import tqdm
import pandas as pd
import numpy as np
import random
from utils import mpnet_embed_class, get_concreteness, Collate_t5
from torch.utils.data import DataLoader
from utils import SentenceDataset


class Summagery:

    def __init__(self, t5_checkpoint, batch_size=5, abstractness=.4, max_d_length=1256, num_prompt=3, device='cuda'):

        # ViPE: Visualize Pretty-much Everything
        self.vipe_model = GPT2LMHeadModel.from_pretrained('fittar/ViPE-M-CTX7')
        vipe_tokenizer = GPT2Tokenizer.from_pretrained('gpt2-medium')
        vipe_tokenizer.pad_token = vipe_tokenizer.eos_token
        self.vipe_tokenizer = vipe_tokenizer

        # SDXL, load both base & refiner
        self.basexl = DiffusionPipeline.from_pretrained(
            "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
        )
        self.refinerxl = DiffusionPipeline.from_pretrained(
            "stabilityai/stable-diffusion-xl-refiner-1.0",
            text_encoder_2=self.basexl.text_encoder_2,
            vae=self.basexl.vae,
            torch_dtype=torch.float16,
            use_safetensors=True,
            variant="fp16",
        )

        self.device = device
        self.max_d_length = max_d_length  # maximum document length to handle before chunking
        self.final_document_length = 60
        self.num_prompt = num_prompt  # how many prompts to generate per document
        self.abstractness = abstractness  # to explore the prompts , just a handle from 0 to 1
        self.concreteness_dataset = './data/concreteness.csv'
        self.batch_size = batch_size
        # T5
        self.t5_model = AutoModelWithLMHead.from_pretrained(t5_checkpoint)
        self.t5_tokenizer = AutoTokenizer.from_pretrained(t5_checkpoint, model_max_length=max_d_length)
        self.collate_t5 = Collate_t5(self.t5_tokenizer)

        # for concrteness rating of the prompts
        data = pd.read_csv(self.concreteness_dataset, header=0,
                           delimiter='\t')
        self.word2score = {w: s for w, s in zip(data['WORD'], data['RATING'])}

    # for large documents, divide them into chunks with self.max_d_length size
    def document_preprocess(self, document):
        documents = []
        words = document.split()
        if len(words) <= self.max_d_length:
            return [document]

        start = 0
        while (len(words) > start):
            if len(words) > start + self.max_d_length:
                chunk = ' '.join(words[start:start + self.max_d_length])
            else:
                chunk = ' '.join(words[start:])

            start += self.max_d_length
            documents.append(chunk)

        return documents

    def t5_summarize(self, document):

        continue_summarization = True
        if len(document.split()) <= self.final_document_length:
            return document

        self.t5_model.to(self.device)

        documents = self.document_preprocess(document)

        if len(documents) > self.batch_size:

            # use batch inference to make things faster
            while (continue_summarization):
                dataset = SentenceDataset(documents)
                dataloader = DataLoader(dataset, batch_size=self.batch_size, collate_fn=self.collate_t5, num_workers=2)
                summaries = ''
                print('summarizing...')
                for text_batch, batch in tqdm(dataloader):
                    if batch.input_ids.shape[1] > 5:
                        max_length = int(batch.input_ids.shape[1] / 2)  # summarize the current chunk by half
                        if max_length < self.final_document_length:  # unless max_length is too short
                            max_length = self.final_document_length

                        batch = batch.to(self.device)
                        generated_ids = self.t5_model.generate(input_ids=batch.input_ids,
                                                               attention_mask=batch.attention_mask, num_beams=3,
                                                               max_length=max_length,
                                                               repetition_penalty=2.5,
                                                               length_penalty=1.0, early_stopping=True)
                        preds = \
                            [self.t5_tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=True)
                             for g
                             in
                             generated_ids]
                        for pred in preds:
                            summaries = summaries + pred + '. '
                    else:
                        for chunk in text_batch:
                            summaries = summaries + chunk + '. '

                if len(summaries.split()) <= self.final_document_length:
                    continue_summarization = False
                    print('finished summarizing.')
                else:
                    documents = self.document_preprocess(summaries)
        else:

            # skip batch inference since we only have a few documents
            while (continue_summarization):
                summaries = ''
                print('summarizing...')
                for chunk in tqdm(documents):
                    if len(chunk.split()) > 2:
                        max_length = int(len(chunk.split()) / 2)  # summarize the current chunk by half
                        if max_length < self.final_document_length:  # unless max_length is too short
                            max_length = self.final_document_length

                        input_ids = self.t5_tokenizer.encode('summarize: ' + chunk, return_tensors="pt",
                                                             add_special_tokens=True, padding='longest',
                                                             max_length=self.max_d_length)
                        input_ids = input_ids.to(self.device)
                        generated_ids = self.t5_model.generate(input_ids=input_ids, num_beams=3, max_length=max_length,
                                                               repetition_penalty=2.5,
                                                               length_penalty=1.0, early_stopping=True)

                        pred = \
                        [self.t5_tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=True) for g
                         in
                         generated_ids][0]
                        summaries = summaries + pred + '. '
                    else:
                        summaries = summaries + chunk + '. '

                if len(summaries.split()) <= self.final_document_length:
                    continue_summarization = False
                    print('finished summarizing.')
                else:
                    documents = self.document_preprocess(summaries)

        return summaries

    def vipe_generate(self, summary, do_sample=True, top_k=100, epsilon_cutoff=.00005, temperature=1):
        batch_size = random.choice([20, 40, 60])
        input_text = [summary] * batch_size
        # mark the text with special tokens
        input_text = [self.vipe_tokenizer.eos_token + i + self.vipe_tokenizer.eos_token for i in input_text]
        batch = self.vipe_tokenizer(input_text, padding=True, return_tensors="pt")

        input_ids = batch["input_ids"].to(self.device)
        attention_mask = batch["attention_mask"].to(self.device)
        self.vipe_model.to(self.device)
        # how many new tokens to generate at max
        max_prompt_length = 50

        generated_ids = self.vipe_model.generate(input_ids=input_ids, attention_mask=attention_mask,
                                                 max_new_tokens=max_prompt_length, do_sample=do_sample, top_k=top_k,
                                                 epsilon_cutoff=epsilon_cutoff, temperature=temperature)
        # return only the generated prompts
        prompts = self.vipe_tokenizer.batch_decode(generated_ids[:, -(generated_ids.shape[1] - input_ids.shape[1]):],
                                                   skip_special_tokens=True)

        # for semantic similarity
        mpnet_object = mpnet_embed_class(device=self.device, nli=False)

        similarities = mpnet_object.get_mpnet_embed_batch(prompts, [summary] * batch_size,
                                                          batch_size=batch_size).cpu().numpy()
        concreteness_score = get_concreteness(prompts, self.word2score)

        final_scores = [i * (1 - self.abstractness) + (self.abstractness) * j for i, j in
                        zip(similarities, concreteness_score)]
        # Get the indices that would sort the final_scores in descending order
        sorted_indices = np.argsort(final_scores)[::-1]

        # Extract the indices of the top 5 highest scores
        top_indices = sorted_indices[:self.num_prompt]
        prompts = [prompts[i] for i in top_indices]

        return prompts

    def sdxl_generate(self, prompts):
        # Define how many steps and what % of steps to be run on each experts (80/20) here
        n_steps = 50
        high_noise_frac = 0.8
        self.basexl.to(self.device)
        self.refinerxl.to(self.device)

        images=[]
        for i, p in enumerate(prompts):
            # torch.manual_seed(i)
            image = self.basexl(
                prompt=p,
                num_inference_steps=n_steps,
                denoising_end=high_noise_frac,
                output_type="latent",
            ).images
            image = self.refinerxl(
                prompt=p,
                num_inference_steps=n_steps,
                denoising_start=high_noise_frac,
                image=image,
            ).images[0]

            images.append(image)

        return images

    def ignite(self, document):
        prompts = []
        summary = self.t5_summarize(document)
        prompts.append(summary)
        summary = summary.replace('. ', '; ')
        print(summary)
        prompts.extend(self.vipe_generate(summary))

        for p in prompts:
            print(p + '\n')

        images=self.sdxl_generate(prompts)

        return images