hh1199 commited on
Commit
3ee972d
·
verified ·
1 Parent(s): 7dace81

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +15 -30
app.py CHANGED
@@ -1,38 +1,23 @@
1
  import gradio as gr
2
  from transformers import pipeline
3
 
4
- classifier = pipeline(
5
- "zero-shot-classification",
6
- model="cointegrated/rubert-tiny2",
7
- device=-1 # Используем CPU для стабильности на бесплатных инстансах
8
- )
9
 
10
- def classify(item: str, categories: str) -> str:
11
- # Формируем гипотезу для классификации
12
- hypothesis_template = "Этот текст относится к категории {}."
13
-
14
- categories_list = [c.strip() for c in categories.split(",")]
15
-
16
- result = classifier(
17
- item,
18
- candidate_labels=categories_list,
19
- hypothesis_template=hypothesis_template,
20
- multi_label=False
21
- )
22
-
23
- # Возвращаем категорию с наибольшим скором
24
- return f"{result['labels'][0]} (уверенность: {result['scores'][0]:.2f})"
25
 
26
- iface = gr.Interface(
27
  fn=classify,
28
  inputs=[
29
- gr.Textbox(label="Название товара", placeholder="Введите наименование товара"),
30
- gr.Textbox(label="Категории (через запятую)",
31
- value="Овощи, Инструменты, Техника, Упаковка")
32
  ],
33
- outputs=gr.Textbox(label="Результат классификации"),
34
- title="Классификатор товаров",
35
- description="Введите название товара и список категорий через запятую"
36
- )
37
-
38
- iface.launch(debug=True)
 
1
  import gradio as gr
2
  from transformers import pipeline
3
 
4
+ models = {
5
+ "ruBert-tiny2": "cointegrated/rubert-tiny2",
6
+ "ruRoberta-large": "sberbank-ai/ruRoberta-large",
7
+ "multilingual-e5": "intfloat/multilingual-e5-base"
8
+ }
9
 
10
+ def classify(model, text, labels):
11
+ classifier = pipeline("zero-shot-classification", model=models[model])
12
+ result = classifier(text, [l.strip() for l in labels.split(",")])
13
+ return result['labels'][0]
 
 
 
 
 
 
 
 
 
 
 
14
 
15
+ gr.Interface(
16
  fn=classify,
17
  inputs=[
18
+ gr.Dropdown(list(models.keys())),
19
+ gr.Textbox(),
20
+ gr.Textbox(value="Овощи, Инструменты, Техника")
21
  ],
22
+ outputs=gr.Textbox()
23
+ ).launch()