data-silence's picture
Upload 2 files
79fe0fe verified
raw
history blame
944 Bytes
from transformers import pipeline
# Загрузка модели через pipeline
classifier = pipeline("text-classification", model="data-silence/news_classifier")
# Словарь для преобразования меток
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'
}
def classify(text):
result = classifier(text)
category = category_mapper[result[0]['label']]
score = result[0]['score']
return {"category": category, "confidence": score}
# Для Gradio интерфейса
def run_inference(text):
result = classify(text)
return f"Predicted category: {result['category']} (confidence: {result['confidence']:.2f})"