| 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 |
|
|
|
|
| |
|
|
