from transformers import AutoTokenizer, T5ForConditionalGeneration
import torch
import gradio as gr
import re

# Загрузка модели и токенизатора
tokenizer = AutoTokenizer.from_pretrained("cointegrated/rut5-base-multitask")
model = T5ForConditionalGeneration.from_pretrained("cointegrated/rut5-base-multitask")

def generate_meta_tags(description):
    description = description.strip()

    # Формируем чёткий и понятный промт
    prompt = """
Описание товара:
{description}

Задача: создать SEO мета-теги для этого товара.
Формат вывода:
{"title": "SEO заголовок до 60 символов", "description": "SEO описание до 160 символов"}
""".format(description=description)

    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=200,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
            do_sample=True,
            temperature=0.7
        )

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

    # Проверяем, является ли вывод JSON
    if generated_text.startswith("{") and generated_text.endswith("}"):
        # Если JSON валидный — возвращаем как есть
        return generated_text.replace('\n', ' ').replace('“', '"').replace('”', '"')
    else:
        # Если не получилось — формируем безопасный вариант
        title = description[:60]  # Берём начало описания как title
        description_short = description[:160]

        return '{"title": "%s", "description": "%s"}' % (
            title.replace('"', "").replace("\\", ""),
            description_short.replace('"', "").replace("\\", "")
        )

# Интерфейс Gradio
demo = gr.Interface(
    fn=generate_meta_tags,
    inputs=gr.Textbox(label="Введите описание товара", lines=10),
    outputs=gr.Textbox(label="Сгенерированные мета-теги"),
    title="SEO Meta Tag Generator",
    description="Генерация meta-title и description на основе описания товара"
)

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