bot_innopolis / app.py
cody82's picture
Update app.py
c440868 verified
raw
history blame
1.68 kB
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr
model_id = "cody82/unitrip"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
system_message = "Ты — умный помощник по Университету Иннополис."
def respond(user_message, history):
if history is None:
history = []
# Формируем полный контекст (если нужно)
prompt = system_message + "\n"
for user_text, bot_text in history:
prompt += f"User: {user_text}\nAssistant: {bot_text}\n"
prompt += f"User: {user_message}\nAssistant:"
inputs = tokenizer(prompt, return_tensors="pt").to(device)
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=150,
pad_token_id=tokenizer.eos_token_id,
eos_token_id=tokenizer.eos_token_id,
do_sample=False,
)
generated_text = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True).strip()
history.append((user_message, generated_text))
return history, history
with gr.Blocks() as demo:
chatbot = gr.Chatbot()
message = gr.Textbox(placeholder="Введите вопрос...")
state = gr.State([]) # История сообщений
message.submit(respond, inputs=[message, state], outputs=[chatbot, state])
message.submit(lambda: "", None, message) # Очистить поле ввода после отправки
demo.launch(share=True)