import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

model_id = "mistralai/Mistral-7B-Instruct-v0.2"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=20)

def classify_email(text):
    prompt = f"""
Прочитай текст письма и классифицируй его по следующей схеме:
1 — человек хочет встретиться (например, предлагает даты или спрашивает когда удобно)
2 — человек отказывается, говорит "не интересно", "спасибо, не нужно"
3 — человек задаёт уточняющие вопросы, но не даёт прямого согласия или отказа

Ответь **только числом**: 1, 2 или 3.

Письмо:
{text}
"""
    result = pipe(prompt)[0]['generated_text']
    for token in ["1", "2", "3"]:
        if token in result:
            return token
    return "?"

iface = gr.Interface(fn=classify_email, inputs="text", outputs="text", title="Классификация писем")
iface.launch()