File size: 1,207 Bytes
7ab3fbc 19533dc becbc70 7ab3fbc 19533dc 28a81fd 7ab3fbc 19533dc 51b4dcc 19533dc 7ab3fbc becbc70 28a81fd 7ab3fbc 19533dc 7ab3fbc 19533dc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 |
import gradio as gr
from transformers import pipeline
# Загружаем модель
classifier = pipeline("text-classification", model="sberbank-ai/ruRoberta-large")
# Категории в порядке, соответствующем меткам модели
CATEGORIES = ["Овощи", "Инструменты", "Коробки", "Прочее"] # Измените на ваши категории!
def classify(item: str, categories: list) -> str:
prompt = f"""
Товар: {item}
Категории: {", ".join(categories)}.
К какой категории относится товар? Ответь только названием категории.
"""
result = classifier(prompt, truncation=True)
label_index = int(result[0]['label'].split("_")[1])
return CATEGORIES[label_index]
# Интерфейс Gradio
iface = gr.Interface(
fn=classify,
inputs=[
gr.Textbox(label="Название товара Сбер"),
gr.Textbox(label="Категории (через запятую)", value="Овощи, Инструменты, Коробки, Прочее")
],
outputs=gr.Textbox(label="Категория")
)
iface.launch() |