File size: 3,686 Bytes
6ef31de |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
import torch
from transformers import T5ForConditionalGeneration, T5TokenizerFast
class TextSummarizer:
def __init__(self, device='cuda'):
self._load_model(
model_type="t5",
model_dir=
"./pretrained_models/flan-t5-large-finetuned-openai-summarize_from_feedback",
device=device)
def _load_model(self,
model_type: str = "t5",
model_dir: str = "outputs",
device: str = 'cuda'):
"""
loads a checkpoint for inferencing/prediction
Args:
model_type (str, optional): "t5" or "mt5". Defaults to "t5".
model_dir (str, optional): path to model directory.
Defaults to "outputs".
device (str, optional): device to run. Defaults to "cuda".
"""
if model_type == "t5":
self.model = T5ForConditionalGeneration.from_pretrained(
f"{model_dir}")
self.tokenizer = T5TokenizerFast.from_pretrained(f"{model_dir}")
else:
raise NotImplementedError(
f"model_type {model_type} not implemented")
self.device = torch.device(device)
self.model = self.model.to(self.device)
def _predict(
self,
source_text: str,
max_length: int = 512,
num_return_sequences: int = 1,
num_beams: int = 2,
top_k: int = 50,
top_p: float = 0.95,
do_sample: bool = True,
repetition_penalty: float = 2.5,
length_penalty: float = 1.0,
early_stopping: bool = True,
skip_special_tokens: bool = True,
clean_up_tokenization_spaces: bool = True,
):
"""
generates prediction for T5/MT5 model
Args:
source_text (str): any text for generating predictions
max_length (int, optional): max token length of prediction.
Defaults to 512.
num_return_sequences (int, optional): number of predictions to be
returned. Defaults to 1.
num_beams (int, optional): number of beams. Defaults to 2.
top_k (int, optional): Defaults to 50.
top_p (float, optional): Defaults to 0.95.
do_sample (bool, optional): Defaults to True.
repetition_penalty (float, optional): Defaults to 2.5.
length_penalty (float, optional): Defaults to 1.0.
early_stopping (bool, optional): Defaults to True.
skip_special_tokens (bool, optional): Defaults to True.
clean_up_tokenization_spaces (bool, optional): Defaults to True.
Returns:
list[str]: returns predictions
"""
input_ids = self.tokenizer.encode(
source_text, return_tensors="pt", add_special_tokens=True)
input_ids = input_ids.to(self.device)
generated_ids = self.model.generate(
input_ids=input_ids,
num_beams=num_beams,
max_length=max_length,
repetition_penalty=repetition_penalty,
length_penalty=length_penalty,
early_stopping=early_stopping,
top_p=top_p,
top_k=top_k,
num_return_sequences=num_return_sequences,
)
preds = [
self.tokenizer.decode(
g,
skip_special_tokens=skip_special_tokens,
clean_up_tokenization_spaces=clean_up_tokenization_spaces,
) for g in generated_ids
]
return preds
def __call__(self, source_text):
generated_text = self._predict(source_text=source_text)
return generated_text
|