data-silence commited on
Commit
0b5f5a1
1 Parent(s): ff19a01

Upload 2 files

Browse files
Files changed (2) hide show
  1. inference.py +34 -0
  2. widget.json +25 -0
inference.py ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import AutoTokenizer, AutoModelForSequenceClassification
2
+ import torch
3
+
4
+ # Загрузка модели и токенизатора
5
+ model_name = "data-silence/news_classifier_ft"
6
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
7
+ model = AutoModelForSequenceClassification.from_pretrained(model_name)
8
+
9
+ # Перевод модели в режим оценки и на нужное устройство
10
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
11
+ model = model.to(device)
12
+ model.eval()
13
+
14
+ # Словарь для маппинга индексов на категории
15
+ id2label = {
16
+ 0: 'climate', 1: 'conflicts', 2: 'culture', 3: 'economy', 4: 'gloss',
17
+ 5: 'health', 6: 'politics', 7: 'science', 8: 'society', 9: 'sports', 10: 'travel'
18
+ }
19
+
20
+ def predict(text):
21
+ # Токенизация входного текста
22
+ inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
23
+
24
+ # Получение предсказания
25
+ with torch.no_grad():
26
+ outputs = model(**inputs)
27
+
28
+ # Получение индекса предсказанной метки
29
+ predicted_label_id = outputs.logits.argmax(-1).item()
30
+
31
+ # Преобразование индекса в текстовую метку
32
+ predicted_label = id2label[predicted_label_id]
33
+
34
+ return {"label": predicted_label}
widget.json ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "task": "text-classification",
3
+ "inputs": [
4
+ {
5
+ "name": "text",
6
+ "type": "text",
7
+ "description": "Введите текст новости для классификации"
8
+ }
9
+ ],
10
+ "output": {
11
+ "type": "label",
12
+ "description": "Категория новости"
13
+ },
14
+ "examples": [
15
+ {
16
+ "text": "Власти отселили из районов боевых действий более 76 тыс. человек, заявил представитель МЧС России в межведомственном оперативном штабе Артем Шаров."
17
+ },
18
+ {
19
+ "text": "Ученые обнаружили новый вид динозавров в Аргентине."
20
+ },
21
+ {
22
+ "text": "Курс доллара упал на фоне новых экономических данных."
23
+ }
24
+ ]
25
+ }