import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Загружаем модель и токенайзер
model_name = "cointegrated/rut5-base-summarization"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

def generate_meta_description(product_description):
    prompt = f"Сгенерируй meta description (до 160 символов) по следующему описанию товара: {product_description}"

    inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)
    summary_ids = model.generate(
        inputs["input_ids"],
        max_length=60,  # приблизительно ~160 символов на русском
        num_beams=4,
        no_repeat_ngram_size=2,
        early_stopping=True
    )
    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)

    # обрезаем аккуратно, чтобы не обрывать слова
    if len(summary) > 160:
        truncated = summary[:160]
        last_space = truncated.rfind(' ')
        summary = truncated[:last_space]
    
    return summary.strip()

iface = gr.Interface(
    fn=generate_meta_description,
    inputs=gr.Textbox(label="Описание товара", lines=5, placeholder="Например: Красивое мужское пальто из шерсти..."),
    outputs=gr.Textbox(label="Meta Description (до 160 символов)"),
    title="Meta Description генератор (русский)",
    description="Генерирует логичный и краткий meta description по описанию товара (до 160 символов, без обрезания слов)."
)

if __name__ == "__main__":
    iface.launch()