Browse Source

first commit

master
commit
f94df248d0
  1. 221
      .gitignore
  2. 1706
      rag_corpus.json
  3. 26
      readme.md
  4. 6
      requerements.txt
  5. 229
      yandex.py
  6. 232
      yandex_gpu.py
  7. 309
      yandex_gpu_lowPC.py

221
.gitignore vendored

@ -0,0 +1,221 @@
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
cover/
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
.pybuilder/
target/
# Jupyter Notebook
.ipynb_checkpoints
# IPython
profile_default/
ipython_config.py
# pyenv
# For a library or package, you might want to ignore these files since the code is
# intended to run in multiple environments; otherwise, check them in:
# .python-version
# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock
# poetry
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
# This is especially recommended for binary packages to ensure reproducibility, and is more
# commonly ignored for libraries.
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
#poetry.lock
# pdm
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
#pdm.lock
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
# in version control.
# https://pdm.fming.dev/#use-with-ide
.pdm.toml
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
__pypackages__/
# Celery stuff
celerybeat-schedule
celerybeat.pid
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
dmypy.json
# Pyre type checker
.pyre/
# pytype static type analyzer
.pytype/
# Cython debug symbols
cython_debug/
# PyCharm
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
# and can be added to the global gitignore or merged into this file. For a more nuclear
# option (not recommended) you can ignore the whole idea folder.
.idea/
# VS Code
.vscode/
*.code-workspace
# Data files
*.db
*.sqlite
*.sqlite3
# ChromaDB database
chroma_db/
# AI Models
models/
!models/.gitkeep
# Downloaded data
data/
!data/.gitkeep
# Temporary files
*.tmp
*.temp
*.log
# OS generated files
.DS_Store
.DS_Store?
._*
.Spotlight-V100
.Trashes
ehthumbs.db
Thumbs.db
# Windows
*.lnk
# Application specific
# rag_corpus.json
config.json
# LLM cache and temporary files
.cache/
*.gguf.tmp
# Logs
logs/
*.log
# Backup files
*.bak
*.backup
# Local configuration
local_config.py
settings_local.py
# Large files (добавьте сюда файлы больше 100MB если они есть)
# *.largefile
/scripts/

1706
rag_corpus.json

File diff suppressed because it is too large Load Diff

26
readme.md

@ -0,0 +1,26 @@
# Medical RAG System with YandexGPT
Система для автоматического преобразования кратких медицинских записей в развернутые формулировки жалоб пациентов с использованием RAG (Retrieval-Augmented Generation) и модели YandexGPT.
## 🚀 Особенности
- **RAG-архитектура**: Поиск релевантных медицинских примеров из базы знаний
- **Гибкое использование GPU**: Автоматическое определение и использование GPU для ускорения работы
- **Управление токенами**: Интеллектуальное ограничение длины контекста
- **Русскоязычная оптимизация**: Специально настроена для работы с медицинскими текстами на русском языке
- **Хранение состояния**: Сохранение векторной базы данных между сессиями
## 📋 Требования
### Аппаратные требования
- **Минимально**: CPU с 8+ GB RAM
- **Рекомендуется**: GPU с 8+ GB VRAM (NVIDIA)
- **Память**: 10+ GB свободного места
### Программные требования
- Python 3.8+
- PyTorch (с поддержкой CUDA при наличии GPU)
- Библиотеки: `chromadb`, `llama-cpp-python`, `sentence-transformers`, `tiktoken`
### Модель
[YandexGPT-5-Lite-8B-instruct-Q4_K_M](https://huggingface.co/yandex/YandexGPT-5-Lite-8B-instruct-GGUF/resolve/main/YandexGPT-5-Lite-8B-instruct-Q4_K_M.gguf?download=true)

6
requerements.txt

@ -0,0 +1,6 @@
torch>=2.0.0
chromadb>=0.4.0
llama-cpp-python>=0.2.0
sentence-transformers>=2.2.0
tiktoken>=0.4.0
numpy>=1.21.0

229
yandex.py

@ -0,0 +1,229 @@
# medical_rag.py
import os
import json
import tiktoken
from typing import List, Tuple
import chromadb
from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction
from llama_cpp import Llama
class MedicalRAG:
def __init__(
self,
model_path: str,
corpus_path: str = "rag_corpus.json",
db_path: str = "./chroma_db",
embedding_model_name: str = "cointegrated/rubert-tiny2",
top_k: int = 3,
n_ctx: int = 4096, # Увеличил контекст для YandexGPT
n_threads: int = 4,
token_multiplier: int = 5 # Множитель токенов для ответа
):
self.corpus_path = corpus_path
self.top_k = top_k
self.token_multiplier = token_multiplier
# === Инициализация токенизатора ===
print("Инициализация токенизатора...")
try:
self.encoding = tiktoken.get_encoding("cl100k_base") # Совместим с большинством моделей
except:
print(" Не удалось загрузить токенизатор, используется базовый подсчет символов")
self.encoding = None
# === Эмбеддинги ===
print("Загрузка эмбеддинг-модели...")
self.embedding_function = SentenceTransformerEmbeddingFunction(
model_name=embedding_model_name,
device="cpu"
)
# === ChromaDB ===
print("Инициализация ChromaDB...")
self.client = chromadb.PersistentClient(path=db_path)
self.collection = self.client.get_or_create_collection(
name="medical_anamnesis",
embedding_function=self.embedding_function,
metadata={"hnsw:space": "cosine"}
)
if self.collection.count() == 0:
print("Коллекция пуста. Загрузка данных...")
self._load_corpus()
else:
print(f"Коллекция уже содержит {self.collection.count()} записей.")
# === LLM (YandexGPT) ===
if not os.path.exists(model_path):
raise FileNotFoundError(
f"Модель не найдена: {model_path}\n"
"Скачайте YandexGPT-5-Lite-8B-instruct-Q4_K_M.gguf и поместите в ./models/"
)
print("Загрузка языковой модели (YandexGPT)...")
self.llm = Llama(
model_path=model_path,
n_ctx=n_ctx,
n_threads=n_threads,
verbose=False
)
print("✅ Система готова к работе!")
def _load_corpus(self):
with open(self.corpus_path, "r", encoding="utf-8") as f:
data = json.load(f)
self.collection.add(
documents=[item["full"] for item in data],
metadatas=[{"short": item["short"]} for item in data],
ids=[f"id_{i}" for i in range(len(data))]
)
print(f"✅ Загружено {len(data)} записей.")
def count_tokens(self, text: str) -> int:
"""Подсчет токенов в тексте"""
if self.encoding:
return len(self.encoding.encode(text))
else:
# Фолбэк: приблизительный подсчет (1 токен ≈ 4 символа)
return len(text) // 4
def build_prompt_with_token_management(self, short_note: str, max_context_tokens: int = 3000) -> Tuple[str, int]:
"""Строит промпт с управлением токенами"""
examples = self.retrieve(short_note)
# Системное сообщение
system_msg = (
"На основе приведённых клинических примеров напиши развёрнуто жалобы пациента, грамотно с медицинской точки зрения. "
"Напиши жалобы в одно предложение, одной строкой. "
"Не пиши вводных слов и фраз. Только жалобы пациента. "
"Неуместно писать диагнозы и план лечения. "
"Расшифруй все сокращения. "
"Отвечай сразу без размышлений."
)
# Считаем токены системного сообщения и короткой заметки
system_tokens = self.count_tokens(system_msg)
note_tokens = self.count_tokens(short_note)
# Доступное количество токенов для примеров
available_tokens = max_context_tokens - system_tokens - note_tokens - 100 # Запас на форматирование
# Отбираем примеры, которые помещаются в контекст
selected_examples = []
current_tokens = 0
for example in examples:
example_tokens = self.count_tokens(example)
if current_tokens + example_tokens <= available_tokens:
selected_examples.append(example)
current_tokens += example_tokens
else:
print(" RAG: Token limit exceeded.")
break
context = "\n\n".join([f"Пример: {ex}" for ex in selected_examples])
user_msg = f"""Примеры развёрнутых описаний:
{context}
Жалобы пациента: "{short_note}"
"""
# 🔑 Формат промпта для YandexGPT
prompt = (
f"<|im_start|>system\n{system_msg}<|im_end|>\n"
f"<|im_start|>user\n{user_msg}<|im_end|>\n"
"<|im_start|>assistant\n"
)
prompt_tokens = self.count_tokens(prompt)
return prompt, prompt_tokens
def retrieve(self, query: str, n: int = None) -> List[str]:
n = n or self.top_k
results = self.collection.query(query_texts=[query], n_results=n)
return results["documents"][0]
def generate(self, short_note: str) -> str:
prompt, prompt_tokens = self.build_prompt_with_token_management(short_note)
# Вычисляем max_tokens на основе длины промпта
available_tokens = 4096 - prompt_tokens - 50 # Запас
max_tokens = min(prompt_tokens * self.token_multiplier, available_tokens)
print(f"📊 Токены: промпт={prompt_tokens}, макс.ответ={max_tokens}")
output = self.llm(
prompt,
max_tokens=max_tokens,
temperature=0.1,
stop=["<|im_end|>"],
echo=False
)
result = output["choices"][0]["text"].strip()
return result
def __call__(self, short_note: str) -> str:
return self.generate(short_note)
# === Отключаем телеметрию Chroma ===
os.environ["CHROMA_TELEMETRY"] = "false"
import time
# === Запуск ===
if __name__ == "__main__":
rag = MedicalRAG(
model_path="./models/YandexGPT-5-Lite-8B-instruct-Q4_K_M.gguf",
n_ctx=8192 # Увеличиваем контекст для YandexGPT
)
# Промты для тестирования
test_notes = [
"Кашель сухой, температура 38",
"А.д. 140, заложенность ушей, частичная потеря слуха",
"а.д. 140/80, т.36.6, выделения из левого уха гнойные",
"ушные палочки 5 лет, снижение слуха 2 года"
]
for note in test_notes:
print(f"\n📥 Кратко: {note}")
t1 = time.time()
result = rag(note)
elapsed_time = time.time() - t1
print(f"⏱ Время выполнения: {elapsed_time:.2f} сек")
if result:
print(f"📤 Развёрнуто:\n{result}")
else:
print("❌ Пустой ответ от модели.")
print("" * 60)
"""
📥 Кратко: Кашель сухой, температура 38
📊 Токены: промпт=392, макс.ответ=1960
📤 Развёрнуто:
Пациент жалуется на сухой кашель и повышение температуры до 38°C.
📥 Кратко: А.д. 140, заложенность ушей, частичная потеря слуха
📊 Токены: промпт=379, макс.ответ=1895
📤 Развёрнуто:
Пациент жалуется на повышенное артериальное давление 140 мм рт. ст., заложенность ушей и частичную потерю слуха.
📥 Кратко: а.д. 140/80, т.36.6, ОСГО л. уха
📊 Токены: промпт=405, макс.ответ=2025
Пустой ответ от модели.
📥 Кратко: ушные палочки 5 лет, снижение слуха 2 года
📊 Токены: промпт=430, макс.ответ=2150
📤 Развёрнуто:
Пациент жалуется на использование ушных палочек в течение пяти лет и снижение слуха, которое наблюдается в течение двух лет.
"""
# YandexGPT-5-Lite-8B-instruct-Q4_K_M 7-30 сек

232
yandex_gpu.py

@ -0,0 +1,232 @@
# medical_rag.py
import os
import json
import tiktoken
from typing import List, Tuple
import chromadb
from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction
from llama_cpp import Llama
import torch # Добавляем импорт torch
class MedicalRAG:
def __init__(
self,
model_path: str,
corpus_path: str = "rag_corpus.json",
db_path: str = "./chroma_db",
embedding_model_name: str = "cointegrated/rubert-tiny2",
top_k: int = 3,
n_ctx: int = 4096,
n_threads: int = 4,
token_multiplier: int = 5,
n_gpu_layers: int = -1, # Автоматическое определение слоев для GPU
main_gpu: int = 0, # Основной GPU
tensor_split: List[float] = None, # Разделение тензоров между GPU
use_gpu_for_embeddings: bool = True # Использовать GPU для эмбеддингов
):
self.corpus_path = corpus_path
self.top_k = top_k
self.token_multiplier = token_multiplier
# === Проверка доступности GPU ===
self.has_gpu = torch.cuda.is_available()
if self.has_gpu:
print(f"✅ GPU доступен: {torch.cuda.get_device_name()}")
print(f"✅ Количество GPU: {torch.cuda.device_count()}")
print(f"✅ Память GPU: {torch.cuda.get_device_properties(0).total_memory / 1024 ** 3:.1f} GB")
else:
print(" GPU не доступен, используется CPU")
# === Инициализация токенизатора ===
print("Инициализация токенизатора...")
try:
self.encoding = tiktoken.get_encoding("cl100k_base")
except:
print(" Не удалось загрузить токенизатор, используется базовый подсчет символов")
self.encoding = None
# === Эмбеддинги с GPU ===
print("Загрузка эмбеддинг-модели...")
device = "cuda" if self.has_gpu and use_gpu_for_embeddings else "cpu"
print(f"Эмбеддинг-модель будет использовать: {device.upper()}")
self.embedding_function = SentenceTransformerEmbeddingFunction(
model_name=embedding_model_name,
device=device
)
# === ChromaDB ===
print("Инициализация ChromaDB...")
self.client = chromadb.PersistentClient(path=db_path)
self.collection = self.client.get_or_create_collection(
name="medical_anamnesis",
embedding_function=self.embedding_function,
metadata={"hnsw:space": "cosine"}
)
if self.collection.count() == 0:
print("Коллекция пуста. Загрузка данных...")
self._load_corpus()
else:
print(f"Коллекция уже содержит {self.collection.count()} записей.")
# === LLM (YandexGPT) с GPU ===
if not os.path.exists(model_path):
raise FileNotFoundError(
f"Модель не найдена: {model_path}\n"
"Скачайте YandexGPT-5-Lite-8B-instruct-Q4_K_M.gguf и поместите в ./models/"
)
print("Загрузка языковой модели (YandexGPT)...")
# Параметры для GPU
gpu_params = {}
if self.has_gpu:
gpu_params.update({
"n_gpu_layers": n_gpu_layers, # -1 = все слои на GPU
"main_gpu": main_gpu,
"tensor_split": tensor_split,
"low_vram": False, # Отключаем для лучшей производительности, если достаточно памяти
"flash_attn": True # Включаем flash attention для ускорения
})
print(f"Используется GPU с {n_gpu_layers} слоями")
else:
print("Используется CPU")
self.llm = Llama(
model_path=model_path,
n_ctx=n_ctx,
n_threads=n_threads,
verbose=False,
**gpu_params
)
print("✅ Система готова к работе!")
def _load_corpus(self):
with open(self.corpus_path, "r", encoding="utf-8") as f:
data = json.load(f)
self.collection.add(
documents=[item["full"] for item in data],
metadatas=[{"short": item["short"]} for item in data],
ids=[f"id_{i}" for i in range(len(data))]
)
print(f"✅ Загружено {len(data)} записей.")
def count_tokens(self, text: str) -> int:
"""Подсчет токенов в тексте"""
if self.encoding:
return len(self.encoding.encode(text))
else:
return len(text) // 4
def build_prompt_with_token_management(self, short_note: str, max_context_tokens: int = 3000) -> Tuple[str, int]:
"""Строит промпт с управлением токенами"""
examples = self.retrieve(short_note)
system_msg = (
"На основе приведённых клинических примеров напиши развёрнуто жалобы пациента, грамотно с медицинской точки зрения. "
"Напиши жалобы в одно предложение, одной строкой. "
"Не пиши вводных слов и фраз. Только жалобы пациента. "
"Неуместно писать диагнозы и план лечения. "
"Расшифруй все сокращения. "
"Отвечай сразу без размышлений."
)
system_tokens = self.count_tokens(system_msg)
note_tokens = self.count_tokens(short_note)
available_tokens = max_context_tokens - system_tokens - note_tokens - 100
selected_examples = []
current_tokens = 0
for example in examples:
example_tokens = self.count_tokens(example)
if current_tokens + example_tokens <= available_tokens:
selected_examples.append(example)
current_tokens += example_tokens
else:
print(" RAG: Token limit exceeded.")
break
context = "\n\n".join([f"Пример: {ex}" for ex in selected_examples])
user_msg = f"""Примеры развёрнутых описаний:
{context}
Жалобы пациента: "{short_note}"
"""
prompt = (
f"<|im_start|>system\n{system_msg}<|im_end|>\n"
f"<|im_start|>user\n{user_msg}<|im_end|>\n"
"<|im_start|>assistant\n"
)
prompt_tokens = self.count_tokens(prompt)
return prompt, prompt_tokens
def retrieve(self, query: str, n: int = None) -> List[str]:
n = n or self.top_k
results = self.collection.query(query_texts=[query], n_results=n)
return results["documents"][0]
def generate(self, short_note: str) -> str:
prompt, prompt_tokens = self.build_prompt_with_token_management(short_note)
available_tokens = 4096 - prompt_tokens - 50
max_tokens = min(prompt_tokens * self.token_multiplier, available_tokens)
print(f"📊 Токены: промпт={prompt_tokens}, макс.ответ={max_tokens}")
print(f" Устройство: {'GPU' if self.has_gpu else 'CPU'}")
output = self.llm(
prompt,
max_tokens=max_tokens,
temperature=0.1,
stop=["<|im_end|>"],
echo=False
)
result = output["choices"][0]["text"].strip()
return result
def __call__(self, short_note: str) -> str:
return self.generate(short_note)
# === Отключаем телеметрию Chroma ===
os.environ["CHROMA_TELEMETRY"] = "false"
import time
# === Запуск ===
if __name__ == "__main__":
rag = MedicalRAG(
model_path="./models/YandexGPT-5-Lite-8B-instruct-Q4_K_M.gguf",
n_ctx=8192,
n_gpu_layers=35, # Количество слоев для GPU (можно настроить)
use_gpu_for_embeddings=True
)
# Промты для тестирования
test_notes = [
"Кашель сухой, температура 38",
"А.д. 140, заложенность ушей, частичная потеря слуха",
"а.д. 140/80, т.36.6, ОСГО л. уха",
"ушные палочки 5 лет, снижение слуха 2 года"
]
for note in test_notes:
print(f"\n📥 Кратко: {note}")
t1 = time.time()
result = rag(note)
elapsed_time = time.time() - t1
print(f"⏱ Время выполнения: {elapsed_time:.2f} сек")
if result:
print(f"📤 Развёрнуто:\n{result}")
else:
print("❌ Пустой ответ от модели.")
print("" * 60)

309
yandex_gpu_lowPC.py

@ -0,0 +1,309 @@
import os
import json
import tiktoken
from typing import List, Tuple, Optional
import chromadb
from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction
from llama_cpp import Llama
import torch
# Отключаем телеметрию Chroma в самом начале
os.environ["CHROMA_TELEMETRY"] = "false"
class MedicalRAG:
def __init__(
self,
model_path: str,
corpus_path: str = "rag_corpus.json",
db_path: str = "./chroma_db",
embedding_model_name: str = "cointegrated/rubert-tiny2",
top_k: int = 3,
n_ctx: int = 2048, # Уменьшено для слабых компьютеров
n_threads: int = None, # Автоматическое определение
token_multiplier: int = 3, # Уменьшено
n_gpu_layers: int = 0, # По умолчанию CPU для совместимости
use_gpu_for_embeddings: bool = False, # По умолчанию CPU
low_memory: bool = True # Режим низкой памяти
):
self.corpus_path = corpus_path
self.top_k = top_k
self.token_multiplier = token_multiplier
self.low_memory = low_memory
# === Автоматическое определение потоков ===
if n_threads is None:
import multiprocessing
n_threads = max(1, multiprocessing.cpu_count() - 1)
# === Проверка доступности GPU (с оптимизацией) ===
self.has_gpu = torch.cuda.is_available() and not low_memory
if self.has_gpu:
gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1024 ** 3
print(f"✅ GPU доступен: {torch.cuda.get_device_name()} ({gpu_memory:.1f} GB)")
# Автоматическая настройка слоев GPU в зависимости от памяти
if gpu_memory < 4: # Маломощные GPU
n_gpu_layers = min(n_gpu_layers, 10)
elif gpu_memory < 8: # Средние GPU
n_gpu_layers = min(n_gpu_layers, 20)
else:
print(" Используется CPU режим")
n_gpu_layers = 0
use_gpu_for_embeddings = False
# === Оптимизированная инициализация токенизатора ===
print("Инициализация токенизатора...")
self.encoding = None
try:
self.encoding = tiktoken.get_encoding("cl100k_base")
except Exception as e:
print(f" Токенизатор не загружен: {e}")
# === Эмбеддинги с оптимизацией памяти ===
print("Загрузка эмбеддинг-модели...")
device = "cuda" if (self.has_gpu and use_gpu_for_embeddings) else "cpu"
# Параметры для экономии памяти
model_kwargs = {}
if low_memory:
model_kwargs = {
'device': device,
'model_kwargs': {'torch_dtype': torch.float16} # Половина точности
}
self.embedding_function = SentenceTransformerEmbeddingFunction(
model_name=embedding_model_name,
**model_kwargs
)
# === ChromaDB с оптимизацией ===
print("Инициализация ChromaDB...")
self.client = chromadb.PersistentClient(path=db_path)
# Упрощенные настройки для экономии памяти
collection_metadata = {"hnsw:space": "cosine"}
if low_memory:
collection_metadata.update({
"hnsw:construction_ef": 100, # Меньше использование памяти
"hnsw:M": 16, # Меньше связей в графе
})
self.collection = self.client.get_or_create_collection(
name="medical_anamnesis",
embedding_function=self.embedding_function,
metadata=collection_metadata
)
if self.collection.count() == 0:
print("Загрузка данных в коллекцию...")
self._load_corpus()
else:
print(f"Коллекция содержит {self.collection.count()} записей")
# === Оптимизированная загрузка LLM ===
if not os.path.exists(model_path):
raise FileNotFoundError(
f"Модель не найдена: {model_path}\n"
"Для слабых компьютеров рекомендуется использовать меньшие модели."
)
print("Загрузка языковой модели...")
# Базовые параметры для всех устройств
llm_params = {
"model_path": model_path,
"n_ctx": n_ctx,
"n_threads": n_threads,
"verbose": False,
"low_vram": low_memory, # Всегда включаем для совместимости
"use_mlock": not low_memory, # Блокировка памяти только если достаточно RAM
}
# Параметры только для GPU
if self.has_gpu and n_gpu_layers > 0:
llm_params.update({
"n_gpu_layers": n_gpu_layers,
"main_gpu": 0,
"tensor_split": None, # Упрощаем для совместимости
})
print(f"Используется GPU с {n_gpu_layers} слоями")
else:
print("Используется CPU")
try:
self.llm = Llama(**llm_params)
print("✅ Система готова к работе!")
except Exception as e:
print(f"❌ Ошибка загрузки модели: {e}")
raise
def _load_corpus(self):
"""Загрузка корпуса с обработкой ошибок"""
try:
with open(self.corpus_path, "r", encoding="utf-8") as f:
data = json.load(f)
# Пакетная обработка для больших корпусов
batch_size = 50 if self.low_memory else 100
for i in range(0, len(data), batch_size):
batch = data[i:i + batch_size]
self.collection.add(
documents=[item["full"] for item in batch],
metadatas=[{"short": item["short"]} for item in batch],
ids=[f"id_{i + j}" for j in range(len(batch))]
)
print(f"Загружено {min(i + batch_size, len(data))}/{len(data)} записей")
except Exception as e:
print(f"❌ Ошибка загрузки корпуса: {e}")
raise
def count_tokens(self, text: str) -> int:
"""Оптимизированный подсчет токенов"""
if self.encoding:
return len(self.encoding.encode(text))
else:
# Упрощенный подсчет для совместимости
return len(text.split()) # Приблизительно по словам
def build_prompt_with_token_management(self, short_note: str, max_context_tokens: int = 1500) -> Tuple[str, int]:
"""Строит промпт с оптимизированным управлением токенами"""
examples = self.retrieve(short_note)
system_msg = (
"На основе примеров напиши развёрнуто жалобы пациента, грамотно с медицинской точки зрения. "
"Напиши жалобы в одно предложение, одной строкой. "
"Не пиши вводных слов и фраз. Только жалобы пациента. "
"Неуместно писать диагнозы и план лечения. "
"Расшифруй все сокращения. "
"Отвечай сразу без размышлений."
)
system_tokens = self.count_tokens(system_msg)
note_tokens = self.count_tokens(short_note)
# Более консервативный расчет доступных токенов
available_tokens = max_context_tokens - system_tokens - note_tokens - 150
selected_examples = []
current_tokens = 0
for example in examples:
example_tokens = self.count_tokens(example)
if current_tokens + example_tokens <= available_tokens:
selected_examples.append(example)
current_tokens += example_tokens
else:
if self.low_memory:
break # Быстрый выход в режиме низкой памяти
elif len(selected_examples) > 0:
break # Сохраняем хотя бы один пример
# Упрощенный контекст
context = "\n".join([f"Пример: {ex}" for ex in selected_examples])
user_msg = f"""Примеры:
{context}
Жалобы: "{short_note}"
"""
import pprint
pprint.pprint(system_msg + " " + user_msg)
prompt = (
f"<|im_start|>system\n{system_msg}<|im_end|>\n"
f"<|im_start|>user\n{user_msg}<|im_end|>\n"
"<|im_start|>assistant\n"
)
prompt_tokens = self.count_tokens(prompt)
return prompt, prompt_tokens
def retrieve(self, query: str, n: int = None) -> List[str]:
"""Оптимизированный поиск"""
n = n or min(self.top_k, 3) # Ограничиваем для слабых компьютеров
try:
results = self.collection.query(
query_texts=[query],
n_results=n
)
return results["documents"][0]
except Exception as e:
print(f" Ошибка поиска: {e}")
return []
def generate(self, short_note: str) -> str:
"""Генерация с оптимизацией памяти"""
prompt, prompt_tokens = self.build_prompt_with_token_management(short_note)
# Более консервативный расчет максимальных токенов
available_tokens = 2048 - prompt_tokens - 30 # Запас уменьшен
max_tokens = min(prompt_tokens * self.token_multiplier, available_tokens, 512) # Жесткий лимит
print(f"📊 Токены: промпт={prompt_tokens}, ответ={max_tokens}")
print(f" Устройство: {'GPU' if self.has_gpu else 'CPU'}")
try:
output = self.llm(
prompt,
max_tokens=max_tokens,
temperature=0.1,
stop=["<|im_end|>"],
echo=False,
stream=False # Отключаем streaming для стабильности
)
result = output["choices"][0]["text"].strip()
return result
except Exception as e:
print(f"❌ Ошибка генерации: {e}")
return ""
def __call__(self, short_note: str) -> str:
return self.generate(short_note)
# === Упрощенный запуск ===
if __name__ == "__main__":
import time
# Автоматическое определение режима низкой памяти
import psutil
total_memory = psutil.virtual_memory().total / 1024 ** 3
low_memory_mode = total_memory < 8 # Меньше 8GB RAM
print(f"💾 Общая память: {total_memory:.1f} GB")
print(f"🔧 Режим низкой памяти: {'Да' if low_memory_mode else 'Нет'}")
rag = MedicalRAG(
model_path="./models/YandexGPT-5-Lite-8B-instruct-Q4_K_M.gguf",
n_ctx=2048, # Уменьшенный контекст
n_gpu_layers=10 if not low_memory_mode else 0, # Адаптивное количество слоев
use_gpu_for_embeddings=not low_memory_mode,
low_memory=low_memory_mode
)
# Промты для тестирования
test_notes = [
"Кашель сухой, температура 38",
"А.д. 140, заложенность ушей",
"а.д. 140/80, т.36.6",
"снижение слуха 2 года"
]
for note in test_notes:
print(f"\n📥 Кратко: {note}")
start_time = time.time()
result = rag(note)
elapsed_time = time.time() - start_time
print(f"⏱ Время: {elapsed_time:.2f} сек")
if result:
print(f"📤 Развёрнуто: {result}")
else:
print("❌ Пустой ответ")
print("" * 50)
Loading…
Cancel
Save