

import evaluate
from datasets import load_dataset
from transformers import AutoTokenizer, LongT5ForConditionalGeneration
dataset = load_dataset("scientific_papers", "pubmed", split="validation")
model = (
     LongT5ForConditionalGeneration.from_pretrained("Stancld/longt5-tglobal-large-16384-pubmed-3k_steps")
     .to("cuda")
     .half()
 )
tokenizer = AutoTokenizer.from_pretrained("Stancld/longt5-tglobal-large-16384-pubmed-3k_steps")
def generate_answers(batch):
     inputs_dict = tokenizer(
         batch["article"], max_length=16384, padding="max_length", truncation=True, return_tensors="pt"
     )
     input_ids = inputs_dict.input_ids.to("cuda")
     attention_mask = inputs_dict.attention_mask.to("cuda")
     output_ids = model.generate(input_ids, attention_mask=attention_mask, max_length=512, num_beams=2)
     batch["predicted_abstract"] = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
     return batch
result = dataset.map(generate_answer, batched=True, batch_size=2)
rouge = evaluate.load("rouge")
rouge.compute(predictions=result["predicted_abstract"], references=result["abstract"])