Update app.py

app.py

@@ -2,22 +2,28 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-# Загружаем модель GPT-2
+# Загружаем модель GPT-2 локально
 model_name = "gpt2"
 try:
-    model = AutoModelForCausalLM.from_pretrained(model_name)
     tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    # Устанавливаем pad_token, если он не задан
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
 except Exception as e:
     print(f"Ошибка загрузки модели: {e}")
     exit(1)
 
 def respond(message, history, max_tokens=512, temperature=0.7, top_p=0.95):
-    # Формируем историю чата
     history = history or []
+    # Формируем историю чата
     input_text = "\n".join([f"User: {msg[0]}\nAssistant: {msg[1]}" for msg in history] + [f"User: {message}"])
 
     # Токенизация
-    inputs = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=512)
+    try:
+        inputs = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=512, padding=True)
+    except Exception as e:
+        return f"Ошибка токенизации: {e}", history
 
     # Генерация ответа
     try:
@@ -32,72 +38,49 @@ def respond(message, history, max_tokens=512, temperature=0.7, top_p=0.95):
         )
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     except Exception as e:
-        return f"Ошибка генерации ответа: {e}", history
+        return f"Ошибка генерации: {e}", history
 
     # Форматируем ответ
     formatted_response = format_response(response)
-
-    # Обновляем историю
     history.append((message, formatted_response))
 
     return formatted_response, history
 
 def format_response(response):
-    # Упрощенное форматирование ответа
     diagnosis = extract_diagnosis(response)
     operation = extract_operation(response)
     treatment = extract_treatment(response)
-
     return f"Предварительный диагноз: {diagnosis}\nОперация: {operation}\nЛечение: {treatment}"
 
 def extract_diagnosis(response):
-    # Извлечение диагноза (упрощенно)
     return response.split(".")[0].strip() if "." in response else response.strip()
 
 def extract_operation(response):
-    # Упрощенная логика для операции
     return "Не требуется"
 
 def extract_treatment(response):
-    # Извлечение лечения
     return response.split(".")[-1].strip() if "." in response else "Не указано"
 
-# Интерфейс Gradio
-def create_interface():
-    with gr.Blocks() as demo:
-        gr.Markdown("## Медицинский чат-бот на базе GPT-2")
-        chatbot = gr.Chatbot(label="Чат")
-        msg = gr.Textbox(label="Ваше сообщение", placeholder="Опишите симптомы...")
-        max_tokens = gr.Slider(minimum=50, maximum=1024, value=512, step=1, label="Максимальная длина ответа")
-        temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, label="Температура")
-        top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, label="Top-p")
-        clear = gr.Button("Очистить чат")
-
-        # Состояние для истории
-        state = gr.State(value=[])
-
-        def submit_message(message, history, max_tokens, temperature, top_p):
-            response, updated_history = respond(message, history, max_tokens, temperature, top_p)
-            return response, updated_history, gr.update(value="")
+# Создаем Gradio интерфейс
+with gr.Blocks() as demo:
+    gr.Markdown("## Медицинский чат-бот на базе GPT-2")
+    chatbot = gr.Chatbot(label="Чат")
+    msg = gr.Textbox(label="Ваше сообщение", placeholder="Опишите симптомы...")
+    max_tokens = gr.Slider(minimum=50, maximum=1024, value=512, step=1, label="Максимальная длина ответа")
+    temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, label="Температура")
+    top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, label="Top-p")
+    clear = gr.Button("Очистить чат")
+    state = gr.State(value=[])
 
-        def clear_chat():
-            return [], [], ""
+    def submit_message(message, history, max_tokens, temperature, top_p):
+        response, updated_history = respond(message, history, max_tokens, temperature, top_p)
+        return [(message, response)], updated_history, ""
 
-        # Обработка отправки сообщения
-        msg.submit(
-            submit_message,
-            inputs=[msg, state, max_tokens, temperature, top_p],
-            outputs=[chatbot, state, msg]
-        )
-
-        # Очистка чата
-        clear.click(
-            clear_chat,
-            outputs=[chatbot, state, msg]
-        )
+    def clear_chat():
+        return [], [], ""
 
-    return demo
+    msg.submit(submit_message, [msg, state, max_tokens, temperature, top_p], [chatbot, state, msg])
+    clear.click(clear_chat, outputs=[chatbot, state, msg])
 
 if __name__ == "__main__":
-    demo = create_interface()
     demo.launch(server_name="0.0.0.0", server_port=7860)