Update app.py

app.py

@@ -16,7 +16,7 @@ try:
         model=model_name,
         device=-1,  # Используем CPU
         framework="pt",
-        max_length=100,  # Уменьшен для стабильности на CPU
+        max_length=100,  # Уменьшен для стабильности
         truncation=True,
         model_kwargs={"torch_dtype": torch.float32}
     )
@@ -26,8 +26,8 @@ except Exception as e:
     exit(1)
 
 def respond(message, max_tokens=100, temperature=0.5, top_p=0.7):
-    # Явный промпт для медицинского ответа на русском
-    prompt = f"Вы медицинский чат-бот. Пользователь говорит: '{message}'. Дайте краткий ответ с диагнозом и лечением на русском языке."
+    # Явный промпт с акцентом на медицинский ответ
+    prompt = f"Вы медицинский чат-бот. Пользователь говорит: '{message}'. Дайте краткий ответ только с диагнозом и лечением на русском языке в формате: Диагноз: [диагноз]. Лечение: [лечение]."
     try:
         logger.info(f"Генерация ответа для: {message}")
         outputs = generator(
@@ -40,13 +40,17 @@ def respond(message, max_tokens=100, temperature=0.5, top_p=0.7):
         )
         response = outputs[0]["generated_text"].replace(prompt, "").strip()
         logger.info(f"Ответ сгенерирован: {response}")
-        # Упрощаем ответ, если он содержит случайный текст
-        if "диагноз" not in response.lower() and "лечение" not in response.lower():
-            response = "Извините, не могу дать точный ответ. Проконсультируйтесь с врачом."
+
+        # Проверка и форматирование ответа
+        if "Диагноз:" in response and "Лечение:" in response:
+            return response
+        else:
+            # Если формат не соблюден, пытаемся извлечь диагноз и генерируем базовое лечение
+            diagnosis = response.split(".")[0].strip() if response else "Неизвестно"
+            return f"Диагноз: {diagnosis}. Лечение: Обратитесь к врачу для точной помощи."
     except Exception as e:
         logger.error(f"Ошибка генерации ответа: {e}")
-        response = f"Ошибка генерации: {e}. Попробуйте позже."
-    return response
+        return "Ошибка генерации. Проконсультируйтесь с врачом."
 
 demo = gr.Interface(
     fn=respond,
@@ -59,7 +63,7 @@ demo = gr.Interface(
     outputs="text",
     title="Медицинский чат-бот на базе RuGPT-3medium",
     theme=gr.themes.Soft(),
-    description="Введите симптомы, и чат-бот попробует предложить диагноз и лечение. Для точной помощи обратитесь к врачу."
+    description="Введите симптомы, и чат-бот предложит диагноз и лечение. Для точной помощи обратитесь к врачу."
 )
 
 if __name__ == "__main__":