Xolkin commited on
Commit
1859ba2
·
verified ·
1 Parent(s): ded136a

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +18 -15
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
- model_name = "microsoft/Phi-3-mini-4k-instruct"
 
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=150,
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=150, temperature=0.5, top_p=0.8):
27
- prompt = f"You are a medical chatbot. User says: '{message}'. Provide diagnosis and treatment in English."
 
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
- if "Diagnosis" in response:
41
- diagnosis = response.split("Diagnosis:")[1].split("Treatment:")[0].strip()
42
- treatment = response.split("Treatment:")[1].strip()
43
- response = f"Диагноз: {diagnosis}\nЛечение: {treatment}"
44
  except Exception as e:
45
  logger.error(f"Ошибка генерации ответа: {e}")
46
- return f"Ошибка генерации: {e}"
47
  return response
48
 
49
  demo = gr.Interface(
50
  fn=respond,
51
  inputs=[
52
  gr.Textbox(label="Ваше сообщение", placeholder="Опишите симптомы (например, 'Болит горло')..."),
53
- gr.Slider(minimum=50, maximum=300, value=150, step=10, label="Макс. токенов"),
54
  gr.Slider(minimum=0.1, maximum=1.0, value=0.5, label="Температура"),
55
- gr.Slider(minimum=0.1, maximum=1.0, value=0.8, label="Top-p")
56
  ],
57
  outputs="text",
58
- title="Медицинский чат-бот на базе Phi-3-mini-4k-instruct",
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__":