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

⚙️ Установка

1. Клонируйте репозиторий:

git clone <repository-url>
cd medical-rag

2. Установите зависимости:

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install chromadb llama-cpp-python sentence-transformers tiktoken

3. Скачайте модель YandexGPT:

Загрузите модель по ссыке:
YandexGPT-5-Lite-8B-instruct-Q4_K_M.gguf

Создайте папку и поместите туда модель:

mkdir models
# Поместите скачанную модель в папку models/

4. Подготовьте корпус данных:

Создайте файл rag_corpus.json с медицинскими примерами в формате:

[
  {
    "short": "краткое описание",
    "full": "развернутое описание жалоб"
  }
]

🎯 Использование

Базовый запуск

from yandex_gpu import MedicalRAG

# Инициализация системы
rag = MedicalRAG(
    model_path="./models/YandexGPT-5-Lite-8B-instruct-Q4_K_M.gguf",
    n_ctx=8192,
    n_gpu_layers=35,
    use_gpu_for_embeddings=True
)

# Преобразование краткой записи
result = rag("Кашель сухой, температура 38")
print(result)

Параметры инициализации

Параметр	Тип	По умолчанию	Описание
model_path	str	-	Путь к файлу модели GGUF
corpus_path	str	"rag_corpus.json"	Путь к JSON-файлу с корпусом
db_path	str	"./chroma_db"	Путь к базе данных ChromaDB
embedding_model_name	str	"cointegrated/rubert-tiny2"	Модель для эмбеддингов
top_k	int	3	Количество извлекаемых примеров
n_ctx	int	4096	Размер контекста
n_gpu_layers	int	-1	Количество слоев на GPU (-1 = все)
use_gpu_for_embeddings	bool	True	Использовать GPU для эмбеддингов

📊 Примеры работы

Вход	Выход
"Кашель сухой, температура 38"	"Пациент предъявляет жалобы на сухой кашель, повышение температуры тела до 38 градусов."
"А.д. 140, заложенность ушей"	"Пациент жалуется на повышение артериального давления до 140 мм рт.ст., заложенность ушей и частичное снижение слуха."
"ушные палочки 5 лет, снижение слуха 2 года"	"Пациент отмечает использование ушных палочек в течение 5 лет и постепенное снижение слуха на протяжении последних 2 лет."

🛠 Технические детали

Архитектура системы

1. Векторный поиск: ChromaDB с русскоязычными эмбеддингами
2. Языковая модель: YandexGPT-5-Lite-8B с оптимизацией для GPU
3. Управление контекстом: Автоматическое ограничение токенов
4. Шаблон промптов: Структурированные промпты для согласованных ответов

Производительность

• На CPU: ~10-30 секунд на запрос
• На GPU: ~2-8 секунд на запрос
• Память: ~6-8 GB для модели + ~1-2 GB для эмбеддингов

🔧 Настройка GPU

rag = MedicalRAG(
    model_path="./models/YandexGPT-5-Lite-8B-instruct-Q4_K_M.gguf",
    n_gpu_layers=35,      # Количество слоев на GPU
    main_gpu=0,           # Основной GPU
    tensor_split=None,    # Разделение между несколькими GPU
    use_gpu_for_embeddings=True  # Ускорение эмбеддингов на GPU
)

📁 Структура проекта

medical-rag/
├── yandex_gpu.py          # Основной класс MedicalRAG
├── rag_corpus.json        # Медицинский корпус
├── chroma_db/            # Векторная база данных
├── models/
│   └── YandexGPT-5-Lite-8B-instruct-Q4_K_M.gguf
└── requirements.txt

🚀 Быстрый старт

# Запуск тестового примера
python yandex_gpu.py

⚠️ Примечания

• Модель автоматически загружается при первом запуске
• Векторная база сохраняется между запусками
• Для лучшей производительности рекомендуется использовать GPU
• Система оптимизирована для русскоязычных медицинских текстов