from transformers import pipeline

# Загрузка модели через pipeline
classifier = pipeline("text-classification", model="data-silence/news_classifier_ft")

# Словарь для преобразования меток
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})"