import gradio as gr
from transformers import pipeline

def get_classifier():
    classifier = pipeline(
        "zero-shot-classification",
        model="joeddav/xlm-roberta-large-xnli",
        framework="pt"
    )
    return classifier

def classify_email(text):
    classifier = get_classifier()
    
    candidate_labels = [
        "назначение встречи",
        "отказ или отсутствие интереса",
        "уточняющие вопросы"
    ]
    
    result = classifier(
        text,
        candidate_labels,
        hypothesis_template="В этом письме клиент выражает {}",
        multi_label=False
    )
    
    # Преобразуем метки в категории 1,2,3
    label_map = {
        "назначение встречи": 1,
        "отказ или отсутствие интереса": 2,
        "уточняющие вопросы": 3
    }
    
    top_label = result["labels"][0]
    category = label_map[top_label]
    confidence = result["scores"][0]
    
    return {
        "category": category,
        "confidence": confidence,
        "label": top_label
    }

# Создаем интерфейс Gradio
def gradio_interface(text):
    result = classify_email(text)
    return (
        f"Категория: {result['category']}\n"
        f"Уверенность: {result['confidence']:.2f}\n"
        f"Метка: {result['label']}"
    )

# Создаем веб-интерфейс
iface = gr.Interface(
    fn=gradio_interface,
    inputs=gr.Textbox(lines=5, label="Введите текст письма"),
    outputs=gr.Textbox(label="Результат"),
    title="Классификация писем",
    description="Определение категории письма: 1 - назначение встречи, 2 - отказ, 3 - уточняющие вопросы"
)

if __name__ == "__main__":
    iface.launch()