metaGenerator / app.py
Triok1's picture
Update app.py
a2a249e verified
raw
history blame
2.51 kB
from transformers import AutoTokenizer, T5ForConditionalGeneration
import torch
import gradio as gr
import re
import json
tokenizer = AutoTokenizer.from_pretrained("cointegrated/rut5-base-multitask", legacy=False)
model = T5ForConditionalGeneration.from_pretrained("cointegrated/rut5-base-multitask")
def smart_truncate(text, max_len):
if len(text) <= max_len:
return text
return text[:text[:max_len+1].rfind(' ')].strip()
def generate_meta(description):
# Упрощенный промт с примерами
prompt = f"""
Create a title and description for product page.
Product name: Fenix ARB-L18-4000U
Description: Аккумулятор 18650 с встроенным портом Type-C и защитой от перегрузок.
Output format:
{"title": "...", "description": "..."}
"""
inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=200,
num_beams=5,
early_stopping=True,
no_repeat_ngram_size=2
)
try:
result = tokenizer.decode(outputs[0], skip_special_tokens=True)
json_data = json.loads(re.search(r'\{.*\}', result, re.DOTALL).group())
# Принудительная постобработка
if "Fenix" not in json_data["title"]:
json_data["title"] = f"Аккумулятор Fenix {json_data['title']}"
return {
"title": smart_truncate(json_data["title"], 60),
"description": smart_truncate(json_data["description"], 160)
}
except:
# Фоллбэк
clean_text = re.sub(r'\s+', ' ', description)
return {
"title": smart_truncate(f"Аккумулятор Fenix {clean_text}", 60),
"description": smart_truncate(clean_text, 160)
}
# Интерфейс
with gr.Blocks() as app:
gr.Markdown("## Генератор метатегов (контроль длины)")
inp = gr.Textbox(label="Описание товара", lines=7)
btn = gr.Button("Сгенерировать")
with gr.Row():
out_title = gr.Textbox(label="Title (до 60 символов)", interactive=False)
out_desc = gr.Textbox(label="Description (до 160 символов)", lines=3, interactive=False)
btn.click(generate_meta, inputs=inp, outputs=[out_title, out_desc])
app.launch()