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