PolyakovK's picture
Update app.py
64f605f verified
raw
history blame
1.31 kB
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()