Update app.py

app.py

@@ -1,29 +1,28 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline
 
-model_id = "mistralai/Mistral-7B-Instruct-v0.2"
+model_id = "HuggingFaceH4/zephyr-7b-beta"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
+model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype="auto")
 
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=20)
+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 — человек задаёт уточняющие вопросы, но не даёт прямого согласия или отказа
+    prompt = f"""Ты классификатор писем. Прочитай письмо и верни только одну цифру без пояснений.
 
-Ответь **только числом**: 1, 2 или 3.
+Схема классификации:
+1 — человек хочет встретиться (предлагает даты или спрашивает когда удобно)
+2 — человек отказывается, говорит "не интересно", "спасибо, не нужно"
+3 — человек задаёт уточняющие вопросы, но не даёт согласия или отказа
 
 Письмо:
 {text}
-"""
+
+Ответ:"""
     result = pipe(prompt)[0]['generated_text']
-    for token in ["1", "2", "3"]:
+    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()
+gr.Interface(fn=classify_email, inputs="text", outputs="text", title="Классификация писем").launch()