Update app.py
Browse files
app.py
CHANGED
@@ -3,18 +3,20 @@ from transformers import pipeline
|
|
3 |
import torch
|
4 |
import logging
|
5 |
|
|
|
6 |
logging.basicConfig(level=logging.INFO)
|
7 |
logger = logging.getLogger(__name__)
|
8 |
|
9 |
-
|
|
|
10 |
try:
|
11 |
logger.info(f"Попытка загрузки модели {model_name}...")
|
12 |
generator = pipeline(
|
13 |
"text-generation",
|
14 |
model=model_name,
|
15 |
-
device=-1,
|
16 |
framework="pt",
|
17 |
-
max_length=
|
18 |
truncation=True,
|
19 |
model_kwargs={"torch_dtype": torch.float32}
|
20 |
)
|
@@ -23,8 +25,9 @@ except Exception as e:
|
|
23 |
logger.error(f"Ошибка загрузки модели: {e}")
|
24 |
exit(1)
|
25 |
|
26 |
-
def respond(message, max_tokens=
|
27 |
-
|
|
|
28 |
try:
|
29 |
logger.info(f"Генерация ответа для: {message}")
|
30 |
outputs = generator(
|
@@ -36,27 +39,27 @@ def respond(message, max_tokens=150, temperature=0.5, top_p=0.8):
|
|
36 |
num_return_sequences=1
|
37 |
)
|
38 |
response = outputs[0]["generated_text"].replace(prompt, "").strip()
|
39 |
-
|
40 |
-
|
41 |
-
|
42 |
-
|
43 |
-
response = f"Диагноз: {diagnosis}\nЛечение: {treatment}"
|
44 |
except Exception as e:
|
45 |
logger.error(f"Ошибка генерации ответа: {e}")
|
46 |
-
|
47 |
return response
|
48 |
|
49 |
demo = gr.Interface(
|
50 |
fn=respond,
|
51 |
inputs=[
|
52 |
gr.Textbox(label="Ваше сообщение", placeholder="Опишите симптомы (например, 'Болит горло')..."),
|
53 |
-
gr.Slider(minimum=50, maximum=
|
54 |
gr.Slider(minimum=0.1, maximum=1.0, value=0.5, label="Температура"),
|
55 |
-
gr.Slider(minimum=0.1, maximum=1.0, value=0.
|
56 |
],
|
57 |
outputs="text",
|
58 |
-
title="Медицинский чат-бот на базе
|
59 |
-
theme=gr.themes.Soft()
|
|
|
60 |
)
|
61 |
|
62 |
if __name__ == "__main__":
|
|
|
3 |
import torch
|
4 |
import logging
|
5 |
|
6 |
+
# Настройка логирования
|
7 |
logging.basicConfig(level=logging.INFO)
|
8 |
logger = logging.getLogger(__name__)
|
9 |
|
10 |
+
# Загружаем модель
|
11 |
+
model_name = "ai-forever/rugpt3medium_based_on_gpt2"
|
12 |
try:
|
13 |
logger.info(f"Попытка загрузки модели {model_name}...")
|
14 |
generator = pipeline(
|
15 |
"text-generation",
|
16 |
model=model_name,
|
17 |
+
device=-1, # Используем CPU
|
18 |
framework="pt",
|
19 |
+
max_length=100, # Уменьшен для стабильности на CPU
|
20 |
truncation=True,
|
21 |
model_kwargs={"torch_dtype": torch.float32}
|
22 |
)
|
|
|
25 |
logger.error(f"Ошибка загрузки модели: {e}")
|
26 |
exit(1)
|
27 |
|
28 |
+
def respond(message, max_tokens=100, temperature=0.5, top_p=0.7):
|
29 |
+
# Явный промпт для медицинского ответа на русском
|
30 |
+
prompt = f"Вы медицинский чат-бот. Пользователь говорит: '{message}'. Дайте краткий ответ с диагнозом и лечением на русском языке."
|
31 |
try:
|
32 |
logger.info(f"Генерация ответа для: {message}")
|
33 |
outputs = generator(
|
|
|
39 |
num_return_sequences=1
|
40 |
)
|
41 |
response = outputs[0]["generated_text"].replace(prompt, "").strip()
|
42 |
+
logger.info(f"Ответ сгенерирован: {response}")
|
43 |
+
# Упрощаем ответ, если он содержит случайный текст
|
44 |
+
if "диагноз" not in response.lower() and "лечение" not in response.lower():
|
45 |
+
response = "Извините, не могу дать точный ответ. Проконсультируйтесь с врачом."
|
|
|
46 |
except Exception as e:
|
47 |
logger.error(f"Ошибка генерации ответа: {e}")
|
48 |
+
response = f"Ошибка генерации: {e}. Попробуйте позже."
|
49 |
return response
|
50 |
|
51 |
demo = gr.Interface(
|
52 |
fn=respond,
|
53 |
inputs=[
|
54 |
gr.Textbox(label="Ваше сообщение", placeholder="Опишите симптомы (например, 'Болит горло')..."),
|
55 |
+
gr.Slider(minimum=50, maximum=200, value=100, step=10, label="Макс. токенов"),
|
56 |
gr.Slider(minimum=0.1, maximum=1.0, value=0.5, label="Температура"),
|
57 |
+
gr.Slider(minimum=0.1, maximum=1.0, value=0.7, label="Top-p")
|
58 |
],
|
59 |
outputs="text",
|
60 |
+
title="Медицинский чат-бот на базе RuGPT-3medium",
|
61 |
+
theme=gr.themes.Soft(),
|
62 |
+
description="Введите симптомы, и чат-бот попробует предложить диагноз и лечение. Для точной помощи обратитесь к врачу."
|
63 |
)
|
64 |
|
65 |
if __name__ == "__main__":
|