Spaces:
Sleeping
Sleeping
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() | |