Update app.py

app.py

@@ -1,16 +1,18 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline
 import torch
 
-# Загружаем локальную модель distilgpt2 (более легкая, чем GPT-2)
+# Загружаем модель через pipeline (локально, но из Hugging Face Hub)
 model_name = "distilgpt2"
 try:
-    tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
-    # Устанавливаем pad_token, если не задан
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
-    model.eval()  # Режим оценки для оптимизации
+    generator = pipeline(
+        "text-generation",
+        model=model_name,
+        device=-1,  # -1 означает CPU, подходит для бесплатного Spaces
+        framework="pt",
+        max_length=512,
+        truncation=True
+    )
 except Exception as e:
     print(f"Ошибка загрузки модели: {e}")
     exit(1)
@@ -23,34 +25,19 @@ def respond(message, history, max_tokens=256, temperature=0.7, top_p=0.9):
         input_text += f"User: {user_msg}\nAssistant: {bot_msg}\n"
     input_text += f"User: {message}"
 
-    # Токенизация
+    # Генерация ответа через pipeline
     try:
-        inputs = tokenizer(
+        outputs = generator(
             input_text,
-            return_tensors="pt",
-            truncation=True,
-            max_length=512,
-            padding=True
+            max_length=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            do_sample=True,
+            no_repeat_ngram_size=2,
+            pad_token_id=generator.tokenizer.eos_token_id,
+            num_return_sequences=1
         )
-    except Exception as e:
-        return f"Ошибка токенизации: {e}", history
-
-    # Генерация ответа
-    try:
-        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()
+        response = outputs[0]["generated_text"][len(input_text):].strip()
     except Exception as e:
         return f"Ошибка генерации ответа: {e}", history
 
@@ -67,49 +54,33 @@ def format_response(response):
     return f"Предварительный диагноз: {diagnosis}\nОперация: {operation}\nЛечение: {treatment}"
 
 def extract_diagnosis(response):
-    # Простое извлечение диагноза
     sentences = response.split(".")
     return sentences[0].strip() if sentences else response.strip()
 
 def extract_operation(response):
-    # Упрощенная логика: операция не требуется
     return "Не требуется"
 
 def extract_treatment(response):
-    # Извлечение лечения
     sentences = response.split(".")
     return sentences[-1].strip() if len(sentences) > 1 else "Не указано"
 
 # 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
-    )
+    gr.Markdown("## Медицинский чат-бот на базе DistilGPT-2")
+    chatbot = gr.Chatbot(label="Чат", height=400)
     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"
+        msg = gr.Textbox(
+            label="Ваше сообщение",
+            placeholder="Опишите симптомы (например, 'Болит горло')...",
+            lines=2,
+            show_label=True
         )
-    clear = gr.Button("Очистить чат")
+        submit_btn = gr.Button("Отправить", variant="primary")
+    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_btn = gr.Button("Очистить чат", variant="secondary")
     state = gr.State(value=[])
 
     def submit_message(message, history, max_tokens, temperature, top_p):
@@ -121,13 +92,22 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     def clear_chat():
         return [], [], ""
 
+    # Кнопка "Отправить"
+    submit_btn.click(
+        fn=submit_message,
+        inputs=[msg, state, max_tokens, temperature, top_p],
+        outputs=[chatbot, state, msg],
+        queue=True
+    )
+    # Поддержка Enter
     msg.submit(
         fn=submit_message,
         inputs=[msg, state, max_tokens, temperature, top_p],
         outputs=[chatbot, state, msg],
         queue=True
     )
-    clear.click(
+    # Кнопка "Очистить"
+    clear_btn.click(
         fn=clear_chat,
         outputs=[chatbot, state, msg]
     )