alpertml/TopicModelingForSummarization

1### Imports2from transformers import PegasusForConditionalGeneration, PegasusTokenizer3from transformers import BartForConditionalGeneration, BartTokenizer4from transformers import T5ForConditionalGeneration, T5Tokenizer5from transformers import ProphetNetForConditionalGeneration, ProphetNetTokenizer6import torch7 8from config import config9 10### Classes and functions11 12##==========================================================================================================13class SummarizationUtilities:14    ##==========================================================================================================15    """16    Definition of attributes17    """18    model_name = None19    device = None20    tokenizer = None21    model = None22    ##==========================================================================================================23    """24    Function: __init__25    Arguments:26        - model_name27        - device28    """29    def __init__(self, model_name="google/pegasus-xsum", device=None, model_path=config.pegasus_model_path):30        self.model_name = model_name31        if device == None:32            self.device = self.detect_available_cuda_device()33        else:34            self.device = device35 36        self.tokenizer = PegasusTokenizer.from_pretrained(model_path)37        self.model = PegasusForConditionalGeneration.from_pretrained(model_path).to(device)38    ##=========================================================================================================39    """40    Function: detect_available_cuda_device41    Arguments: NA42    """43    def detect_available_cuda_device(self):44        self.device = "cuda" if torch.cuda.is_available() else "cpu"45    ##=========================================================================================================46    """47    Function: detect_available_cuda_device48    Arguments: NA49    """50    def tokenize(self, src_text, truncation = True, padding="longest", return_tensors="pt"):51        return self.tokenizer(src_text, truncation=truncation, padding=padding, return_tensors=return_tensors).to(self.device)52    ##=========================================================================================================53    """54    Function: generate55    Arguments: 56        - batch57    """58    def generate(self, batch):59        text_generated = self.model.generate(**batch)60        return text_generated61    ##=========================================================================================================62    """63    Function: decode_generated_text64    Arguments: 65        - batch66    """67    def decode_generated_text(self, generated_text, skip_special_tokens=True):68        return self.tokenizer.batch_decode(generated_text, skip_special_tokens=skip_special_tokens)69    ##=========================================================================================================70    """71    Function: get_summary72    Arguments: 73        - src_text74    """75    def get_summary(self, src_text):76        summary = None77 78        batch = self.tokenize(src_text)79        generated_text = self.generate(batch)80        target_text = self.decode_generated_text(generated_text)81        #print("target_text", target_text)82        summary = target_text83 84        return summary85 86    def summarize(self, src_text):87        summary = None88 89        batch = self.tokenize(src_text)90        generated_text = self.generate(batch)91        target_text = self.decode_generated_text(generated_text)92        #print("target_text", target_text)93        summary = target_text94 95        return summary96        97    ##=========================================================================================================98##==========================================================================================================99 100 101 102class BARTSummarizer:103    def __init__(self, device=None, model_path=config.bart_model_path):104        # https://stackoverflow.com/questions/66639722/why-does-huggingfaces-bart-summarizer-replicate-the-given-input-text105        self.device = device if device else torch.device("cuda" if torch.cuda.is_available() else "cpu")106        # self.tokenizer = BartTokenizer.from_pretrained("sshleifer/distilbart-xsum-6-6") #facebook/bart-large-cnn107        # self.model = BartForConditionalGeneration.from_pretrained("sshleifer/distilbart-xsum-6-6").to(self.device)108        self.tokenizer = BartTokenizer.from_pretrained(model_path)109        self.model = BartForConditionalGeneration.from_pretrained(model_path)110 111    def summarize(self, text):112        inputs = self.tokenizer([text], truncation=True, padding="longest", return_tensors="pt").to(self.device)113        summary_ids = self.model.generate(inputs["input_ids"], num_beams=4, max_length=200, early_stopping=True)114        summary = self.tokenizer.decode(summary_ids.squeeze(), skip_special_tokens=True)115        return summary116 117 118class T5Summarizer:119    def __init__(self, device=None, model_path=config.t5_model_path):120        self.device = device if device else torch.device("cuda" if torch.cuda.is_available() else "cpu")121        # self.tokenizer = T5Tokenizer.from_pretrained("t5-base")122        # self.model = T5ForConditionalGeneration.from_pretrained("t5-base").to(self.device)123        self.tokenizer = T5Tokenizer.from_pretrained(model_path)124        self.model = T5ForConditionalGeneration.from_pretrained(model_path).to(self.device)125 126    def summarize(self, text):127        inputs = self.tokenizer.encode_plus(text, return_tensors="pt", truncation=True, padding="longest").to(self.device)128        summary_ids = self.model.generate(inputs.input_ids)129        summary = self.tokenizer.decode(summary_ids.squeeze(), skip_special_tokens=True)130        return summary131 132 133class ProphetNetSummarizer:134    def __init__(self, device=None, model_path=config.prophetnet_model_path):135        self.device = device if device else torch.device("cuda" if torch.cuda.is_available() else "cpu")136        # self.tokenizer = ProphetNetTokenizer.from_pretrained("microsoft/prophetnet-large-uncased")137        # self.model = ProphetNetForConditionalGeneration.from_pretrained("microsoft/prophetnet-large-uncased").to(self.device)138        self.tokenizer = ProphetNetTokenizer.from_pretrained(model_path)139        self.model = ProphetNetForConditionalGeneration.from_pretrained(model_path).to(self.device)140 141    def summarize(self, text):142        inputs = self.tokenizer(text, return_tensors="pt", truncation=True, padding="longest").to(self.device)143        summary_ids = self.model.generate(inputs.input_ids)144        summary = self.tokenizer.decode(summary_ids.squeeze(), skip_special_tokens=True)145        return summary