Update app.py

app.py

@@ -1,42 +1,48 @@
 import torch
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-model_id = "google/flan-t5-base"  # можно flan-t5-large, если хватает памяти
+model_id = "openai-community/gpt2"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForSeq2SeqLM.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 и робототехнике, расположенный в городе Иннополис, Татарстан.
-"""
+system_message = "Ты — умный помощник по Университету Иннополис."
 
-def respond(message, history=None):
+def respond(user_message, history=None):
     if history is None:
         history = []
 
-    prompt = (
-        "Используя следующий контекст, ответь на вопрос четко и кратко.\n"
-        f"Контекст: {context}\n"
-        f"Вопрос: {message}\n"
-        "Ответ:"
-    )
+    # Формируем контекст для диалога
+    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=100,
+            max_new_tokens=150,
+            pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.eos_token_id,
             do_sample=False,
-            eos_token_id=tokenizer.eos_token_id
         )
-    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    history.append((message, answer))
-    return history
+    # Отсекание входной части prompt, оставляем только сгенерированный ответ
+    generated_text = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True).strip()
 
-iface = gr.ChatInterface(fn=respond, title="Innopolis Q&A")
-iface.launch()
+    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()