NomClass / app.py
hh1199's picture
Update app.py
86bd5a4 verified
raw
history blame
1.8 kB
import gradio as gr
from transformers import pipeline
classifier = pipeline(
"zero-shot-classification",
model="cointegrated/rubert-tiny2",
device=-1 # Используем CPU для стабильности на бесплатных инстансах
)
def classify(item: str, categories: str) -> str:
# Формируем гипотезу для классификации
hypothesis_template = "Этот текст относится к категории {}."
categories_list = [c.strip() for c in categories.split(",")]
result = classifier(
item,
candidate_labels=categories_list,
hypothesis_template=hypothesis_template,
multi_label=False
)
# Возвращаем категорию с наибольшим скором
return f"{result['labels'][0]} (уверенность: {result['scores'][0]:.2f})"
iface = gr.Interface(
fn=classify,
inputs=[
gr.Textbox(label="Название товара", placeholder="Введите наименование товара"),
gr.Textbox(label="Категории (через запятую)",
value="Овощи, Инструменты, Техника, Упаковка")
],
outputs=gr.Textbox(label="Результат классификации"),
examples=[
["Молоток", "Овощи, Инструменты, Техника"],
["Свекла", "Фрукты, Овощи, Ягоды"],
["Ноутбук", "Техника, Мебель, Инструменты"]
],
title="Классификатор товаров",
description="Введите название товара и список категорий через запятую"
)
iface.launch(debug=True)