Update app.py

app.py

@@ -2,43 +2,57 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-# Загружаем модель GPT-2 локально
-model_name = "gpt2"
+# Загружаем локальную модель distilgpt2 (более легкая, чем GPT-2)
+model_name = "distilgpt2"
 try:
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
     model = AutoModelForCausalLM.from_pretrained(model_name)
-    # Устанавливаем pad_token, если он не задан
+    # Устанавливаем pad_token, если не задан
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
+    model.eval()  # Режим оценки для оптимизации
 except Exception as e:
     print(f"Ошибка загрузки модели: {e}")
     exit(1)
 
-def respond(message, history, max_tokens=512, temperature=0.7, top_p=0.95):
+def respond(message, history, max_tokens=256, temperature=0.7, top_p=0.9):
     history = history or []
-    # Формируем историю чата
-    input_text = "\n".join([f"User: {msg[0]}\nAssistant: {msg[1]}" for msg in history] + [f"User: {message}"])
+    # Формируем входной текст с историей
+    input_text = ""
+    for user_msg, bot_msg in history:
+        input_text += f"User: {user_msg}\nAssistant: {bot_msg}\n"
+    input_text += f"User: {message}"
 
     # Токенизация
     try:
-        inputs = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=512, padding=True)
+        inputs = tokenizer(
+            input_text,
+            return_tensors="pt",
+            truncation=True,
+            max_length=512,
+            padding=True
+        )
     except Exception as e:
         return f"Ошибка токенизации: {e}", history
 
     # Генерация ответа
     try:
-        outputs = model.generate(
-            inputs["input_ids"],
-            max_length=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id,
-            no_repeat_ngram_size=2
-        )
+        with torch.no_grad():  # Отключаем градиенты для экономии памяти
+            outputs = model.generate(
+                inputs["input_ids"],
+                max_length=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                do_sample=True,
+                pad_token_id=tokenizer.eos_token_id,
+                no_repeat_ngram_size=2,
+                num_beams=2  # Добавляем beam search для лучшего качества
+            )
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # Удаляем входной текст из ответа
+        response = response[len(input_text):].strip()
     except Exception as e:
-        return f"Ошибка генерации: {e}", history
+        return f"Ошибка генерации ответа: {e}", history
 
     # Форматируем ответ
     formatted_response = format_response(response)
@@ -53,34 +67,70 @@ def format_response(response):
     return f"Предварительный диагноз: {diagnosis}\nОперация: {operation}\nЛечение: {treatment}"
 
 def extract_diagnosis(response):
-    return response.split(".")[0].strip() if "." in response else response.strip()
+    # Простое извлечение диагноза
+    sentences = response.split(".")
+    return sentences[0].strip() if sentences else response.strip()
 
 def extract_operation(response):
+    # Упрощенная логика: операция не требуется
     return "Не требуется"
 
 def extract_treatment(response):
-    return response.split(".")[-1].strip() if "." in response else "Не указано"
+    # Извлечение лечения
+    sentences = response.split(".")
+    return sentences[-1].strip() if len(sentences) > 1 else "Не указано"
 
-# Создаем 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")
+# Gradio интерфейс
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("## Медицинский чат-бот (на базе DistilGPT-2)")
+    chatbot = gr.Chatbot(label="История чата", height=400)
+    msg = gr.Textbox(
+        label="Ваше со��бщение",
+        placeholder="Опишите симптомы (например, 'Болит голова и температура')...",
+        lines=2
+    )
+    with gr.Row():
+        max_tokens = gr.Slider(
+            minimum=50,
+            maximum=512,
+            value=256,
+            step=10,
+            label="Макс. токенов"
+        )
+        temperature = gr.Slider(
+            minimum=0.1,
+            maximum=1.5,
+            value=0.7,
+            label="Температура"
+        )
+        top_p = gr.Slider(
+            minimum=0.1,
+            maximum=1.0,
+            value=0.9,
+            label="Top-p"
+        )
     clear = gr.Button("Очистить чат")
     state = gr.State(value=[])
 
     def submit_message(message, history, max_tokens, temperature, top_p):
+        if not message.strip():
+            return [], history, "Пожалуйста, введите сообщение."
         response, updated_history = respond(message, history, max_tokens, temperature, top_p)
         return [(message, response)], updated_history, ""
 
     def clear_chat():
         return [], [], ""
 
-    msg.submit(submit_message, [msg, state, max_tokens, temperature, top_p], [chatbot, state, msg])
-    clear.click(clear_chat, outputs=[chatbot, state, msg])
+    msg.submit(
+        fn=submit_message,
+        inputs=[msg, state, max_tokens, temperature, top_p],
+        outputs=[chatbot, state, msg],
+        queue=True
+    )
+    clear.click(
+        fn=clear_chat,
+        outputs=[chatbot, state, msg]
+    )
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)