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