from transformers import pipeline
from flask import Flask, request, jsonify

app = Flask(__name__)

# Загружаем zero-shot классификатор
classifier = pipeline("zero-shot-classification", model="typeform/distilbert-base-uncased-mnli")

# Гипотезы
LABELS = {
    "wants_meeting": "Клиент хочет назначить встречу или обсудить время",
    "not_interested": "Клиент не заинтересован во встрече или у него нет времени",
    "asking_questions": "Клиент задает уточняющие вопросы по теме"
}

@app.route("/analyze", methods=["POST"])
def analyze():
    data = request.get_json()
    emails = data.get("emails", [])

    if not emails or not isinstance(emails, list):
        return jsonify({"error": "Field 'emails' must be a non-empty list"}), 400

    results = []
    for email in emails:
        prediction = classifier(email, list(LABELS.values()), multi_label=True)
        scored_labels = dict(zip(prediction["labels"], prediction["scores"]))

        # Вывод в формате {label: score}
        result = {
            "text": email,
            "intents": {
                key: round(scored_labels[label], 4) for key, label in LABELS.items()
            }
        }
        results.append(result)

    return jsonify(results)

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)