from transformers import pipeline

def get_classifier():
    classifier = pipeline(
        "zero-shot-classification",
        model="sberbank-ai/rugpt3small_based_on_gpt2",
        framework="pt"
    )
    return classifier

def classify_email(text):
    classifier = get_classifier()
    
    candidate_labels = [
        "Клиент хочет назначить встречу",
        "Клиент не заинтересован или отказывается",
        "Клиент задаёт уточняющие вопросы"
    ]
    
    result = classifier(
        text,
        candidate_labels,
        hypothesis_template="Это письмо о том, что {}."
    )
    
    # Получаем индекс наиболее вероятной метки (0, 1 или 2)
    label_index = result["labels"].index(result["labels"][0])
    
    # Возвращаем категорию (1, 2 или 3) и уверенность
    return {
        "category": label_index + 1,
        "confidence": result["scores"][label_index],
        "label": result["labels"][0]
    }

# Пример использования
if __name__ == "__main__":
    test_text = "Добрый день! Можно ли узнать подробнее о ваших услугах и ценах?"
    result = classify_email(test_text)
    print(f"Категория: {result['category']}")
    print(f"Уверенность: {result['confidence']:.2f}")
    print(f"Метка: {result['label']}")