neural_news_lab_SPACE / new_sum.py
neural-news-lab's picture
Upload 3 files
c9d476f verified
Raw
History Blame
1.14 kB
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
# ======================
# MODEL SETUP
# ======================
MODEL_NAME = "cointegrated/rut5-base-multitask"
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME).to(device)
model.eval()
# ======================
# CORE FUNCTION
# ======================
def generate_summary(text: str) -> str:
if not text:
return ""
prompt = "summarize | " + text
inputs = tokenizer(
prompt,
return_tensors="pt",
truncation=True,
padding="max_length",
max_length=512
).to(device)
with torch.no_grad():
outputs = model.generate(
**inputs,
max_length=200,
min_length=30,
num_beams=3,
do_sample=False,
no_repeat_ngram_size=3,
repetition_penalty=1.2,
early_stopping=True
)
return tokenizer.decode(outputs[0], skip_special_tokens=True)