bot_innopolis / app.py
cody82's picture
Update app.py
4803a6e verified
raw
history blame
2.05 kB
import os
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
model_id = "sberbank-ai/rugpt3medium_based_on_gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
context = (
"Университет Иннополис был основан в 2012 году. "
"Это современный вуз в России, специализирующийся на IT и робототехнике, "
"расположенный в городе Иннополис, Татарстан.\n"
)
def respond(message, history=None):
prompt = f"Прочитай текст и ответь на вопрос:\n\n{context}\n\nВопрос: {message}\nОтвет:"
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
with torch.no_grad():
output_ids = model.generate(
input_ids,
max_new_tokens=100,
temperature=0.8,
top_p=0.9,
do_sample=True,
pad_token_id=tokenizer.eos_token_id
)
full_output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
if "Ответ:" in full_output:
answer = full_output.split("Ответ:")[-1].strip()
else:
answer = full_output[len(prompt):].strip()
return answer
### Gradio UI
chat = gr.ChatInterface(
fn=respond,
title="Бот об Университете Иннополис (на русском)",
chatbot=gr.Chatbot(label="Диалог", type="messages"), # Указан type="messages"
textbox=gr.Textbox(placeholder="Задай вопрос на русском...", label="Твой вопрос")
)
### ✅ API маршрут (Gradio 5+)
@gr.routes("/ask", methods=["POST"])
def ask_api(question: str):
return {"answer": respond(question)}
### Запуск
if __name__ == "__main__":
chat.launch()