| from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
|
| import torch
|
|
|
|
|
|
|
|
|
| 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()
|
|
|
|
|
|
|
|
|
| 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) |