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

model_id = "HuggingFaceH4/zephyr-7b-beta"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype="auto")

pipe = TextGenerationPipeline(model=model, tokenizer=tokenizer, max_new_tokens=64, do_sample=True, temperature=0.2)

def classify_email(text):
    prompt = f"""Ты классификатор писем. Прочитай письмо и верни только одну цифру без пояснений.

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

Письмо:
{text}

Ответ:"""
    result = pipe(prompt)[0]['generated_text']
    for token in ['1', '2', '3']:
        if token in result:
            return token
    return "?"

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