bot_innopolis / app.py
cody82's picture
Update app.py
d334b30 verified
raw
history blame
1.83 kB
import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
model_id = "openai-community/gpt2" # если нужен GPT-2 community версия
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 и робототехнике, расположенный в городе Иннополис, Татарстан.
"""
def respond(message, history=None):
if history is None:
history = []
# Формируем промпт — контекст + вопрос + "Ответ:"
prompt = (
f"Контекст: {context}\n"
f"Вопрос: {message}\n"
"Ответ:"
)
inputs = tokenizer(prompt, return_tensors="pt").to(device)
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=100,
do_sample=False,
eos_token_id=tokenizer.eos_token_id, # Может отсутствовать у GPT-2, можно убрать
pad_token_id=tokenizer.eos_token_id, # Чтобы избежать warning
)
# Декодируем срез outputs после длины входа — берем только сгенерированное продолжение
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
answer = generated_text[len(prompt):].strip()
history.append((message, answer))
return history
iface = gr.ChatInterface(fn=respond, title="Innopolis Q&A (GPT-2)")
iface.launch()