Update app.py

app.py

@@ -1,20 +1,29 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
-classifier = pipeline("zero-shot-classification", model="typeform/distilbert-base-uncased-mnli")
+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):
-    labels = ["клиент готов встретиться или хочется назначить время для звонка или встречи", "клиент отказывается от встречи, ему сейчас не интересно то, что мы предлагаем или у него нет времени", "клиент задает любые уточняющие вопросы и ему нужна помощь"]
-    result = classifier(text, candidate_labels=labels)
-    label = result["labels"][0]
-    if label == "клиент готов встретиться или хочется назначить время для звонка или встречи":
-        return "1"
-    elif label == "клиент отказывается от встречи, ему сейчас не интересно то, что мы предлагаем или у него нет времени":
-        return "2"
-    elif label == "клиент задает любые уточняющие вопросы и ему нужна помощь":
-        return "3"
-    else:
-        return "?"
+    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="Email Classifier")
+iface = gr.Interface(fn=classify_email, inputs="text", outputs="text", title="Классификация писем")
 iface.launch()