Update app.py

app.py

@@ -1,30 +1,46 @@
-from transformers import pipeline
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
 
-# Загрузка модели через pipeline
-classifier = pipeline("text-classification", model="data-silence/news_classifier_ft")
+# Загрузка модели и токенизатора
+model_name = "your-username/your-model-name"  # Замените на путь к вашей модели на HF
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
-# Словарь для преобразования меток
-category_mapper = {
-    'LABEL_0': 'climate',
-    'LABEL_1': 'conflicts',
-    'LABEL_2': 'culture',
-    'LABEL_3': 'economy',
-    'LABEL_4': 'gloss',
-    'LABEL_5': 'health',
-    'LABEL_6': 'politics',
-    'LABEL_7': 'science',
-    'LABEL_8': 'society',
-    'LABEL_9': 'sports',
-    'LABEL_10': 'travel'
+# Перевод модели в режим оценки
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+model.eval()
+
+# Словарь для маппинга индексов на категории
+id2label = {
+    0: 'climate', 1: 'conflicts', 2: 'culture', 3: 'economy', 4: 'gloss',
+    5: 'health', 6: 'politics', 7: 'science', 8: 'society', 9: 'sports', 10: 'travel'
 }
 
-def classify(text):
-    result = classifier(text)
-    category = category_mapper[result[0]['label']]
-    score = result[0]['score']
-    return {"category": category, "confidence": score}
+def predict(text):
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    predicted_label_id = outputs.logits.argmax(-1).item()
+    predicted_label = id2label[predicted_label_id]
+    
+    # Получаем вероятности для всех классов
+    probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
+    probs_dict = {id2label[i]: float(prob) for i, prob in enumerate(probabilities[0])}
+    
+    return predicted_label, probs_dict
+
+# Создание интерфейса Gradio
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Textbox(lines=5, label="Введите текст новости"),
+    outputs=[
+        gr.Label(label="Предсказанная категория"),
+        gr.Label(label="Вероятности категорий")
+    ],
+    title="Классификатор новостей",
+    description="Введите текст новости, и модель предскажет её категорию."
+)
 
-# Для Gradio интерфейса
-def run_inference(text):
-    result = classify(text)
-    return f"Predicted category: {result['category']} (confidence: {result['confidence']:.2f})"
+iface.launch()