File size: 1,063 Bytes
a9082f6 |
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 |
from typing import List, Tuple
from transformers import AutoTokenizer, BartForConditionalGeneration, BartTokenizerFast
import torch
def summarize_text(tokenizer: BartTokenizerFast, model: BartForConditionalGeneration,
input_texts: List[str]):
inputs = tokenizer(input_texts, padding=True,
return_tensors='pt', truncation=True).to(1)
with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
summary_ids = model.generate(inputs["input_ids"])
summaries = tokenizer.batch_decode(summary_ids, skip_special_tokens=True,
clean_up_tokenization_spaces=False, batch_size=len(input_texts))
return summaries
def get_summarizer(model_id="ccdv/lsg-bart-base-4096-multinews") -> Tuple[BartTokenizerFast, BartForConditionalGeneration]:
tokenizer = BartTokenizerFast.from_pretrained(model_id)
model = BartForConditionalGeneration.from_pretrained(model_id).to(1)
model = torch.compile(model)
return tokenizer, model
# OpenAI reader
|