Update app.py

app.py

@@ -1,81 +1,103 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-# Используем модель GPT-2 для генерации текста
-model_name = "gpt2"  # Модель GPT-2
-model = AutoModelForCausalLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+# Загружаем модель GPT-2
+model_name = "gpt2"
+try:
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+except Exception as e:
+    print(f"Ошибка загрузки модели: {e}")
+    exit(1)
 
-def respond(message, history=None, system_message=None, max_tokens=512, temperature=0.7, top_p=0.95):
-    if history is None:
-        history = []  # Инициализируем пустой список, если history не передан
+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}"])
 
-    # Объединяем сообщения в историю
-    input_text = "\n".join([msg[1] for msg in history] + [message])
-
-    # Токенизация текста
-    inputs = tokenizer(input_text, return_tensors="pt", truncation=True, padding=True)
+    # Токенизация
+    inputs = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=512)
 
     # Генерация ответа
-    outputs = model.generate(
-        inputs["input_ids"],
-        max_length=max_tokens,
-        temperature=temperature,
-        top_p=top_p,
-        do_sample=True,
-    )
-
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    # Формируем ответ согласно шаблону
-    response = format_response(response)
-
-    # Добавляем текущий запрос и ответ в историю
-    history.append((message, response))
-
-    return response, 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
+        )
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    except Exception as e:
+        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)
 
-    formatted_response = f"Предварительный диагноз: {diagnosis}\nОперация: {operation}\nЛечение: {treatment}"
-
-    return formatted_response
+    return f"Предварительный диагноз: {diagnosis}\nОперация: {operation}\nЛечение: {treatment}"
 
 def extract_diagnosis(response):
-    # Простой способ извлечь диагноз (можно улучшить с помощью NLP)
-    diagnosis = response.split(".")[0]  # Пример: диагноз - первая часть ответа
-    return diagnosis.strip()
+    # Извлечение диагноза (упрощенно)
+    return response.split(".")[0].strip() if "." in response else response.strip()
 
 def extract_operation(response):
-    # Извлекаем название операции из ответа
-    operation = "Не требуется"  # Пример, что операция не требуется
-    return operation.strip()
+    # Упрощенная логика для операции
+    return "Не требуется"
 
 def extract_treatment(response):
-    # Извлекаем лечение (например, лечение как последняя часть ответа)
-    treatment = response.split(".")[-1]  # Пример: лечение - последняя часть
-    return treatment.strip()
+    # Извлечение лечения
+    return response.split(".")[-1].strip() if "." in response else "Не указано"
 
 # Интерфейс Gradio
-demo = gr.Interface(
-    fn=respond,
-    inputs=[
-        gr.Textbox(value="Здравствуйте. Отвечай кратко...", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max Tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, label="Temperature"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, label="Top-p"),
-        gr.State()  # Вход для состояния (history)
-    ],
-    outputs=[
-        gr.Textbox(label="Response"),  # Ответ от модели
-        gr.State()  # Выход для состояния (обновленная история)
-    ],
-    live=True  # Обновляем интерфейс в реальном времени
-)
+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="")
+
+        def clear_chat():
+            return [], [], ""
+
+        # Обработка отправки сообщения
+        msg.submit(
+            submit_message,
+            inputs=[msg, state, max_tokens, temperature, top_p],
+            outputs=[chatbot, state, msg]
+        )
+
+        # Очистка чата
+        clear.click(
+            clear_chat,
+            outputs=[chatbot, state, msg]
+        )
+
+    return demo
 
 if __name__ == "__main__":
-    demo.launch()
+    demo = create_interface()
+    demo.launch(server_name="0.0.0.0", server_port=7860)