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)