from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoConfig
import torch

MODEL_NAME = "IlyaGusev/rut5_base_sum_gazeta"
device = "cuda" if torch.cuda.is_available() else "cpu"

config = AutoConfig.from_pretrained(MODEL_NAME)
config.tie_word_embeddings = False  

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME, config=config).to(device)
model.eval()


def generate_summary(text: str) -> str:
    if not text:
        return ""

    # чуть лучше для T5
    prompt = "summarize: " + text

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        truncation=True,
        padding="longest",
        max_length=512
    ).to(device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_length=150,
            min_length=30,
            num_beams=4,
            do_sample=False,
            repetition_penalty=2.0,
            no_repeat_ngram_size=3,
            early_stopping=True
        )

    summary = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # 🔥 ВАЖНО: защита от мусорных токенов
    if "<0x" in summary or len(summary.strip()) < 10:
        return "Model output invalid or unstable. Try different input."

    return summary