# 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. **Клонируйте репозиторий**:
```bash
git clone <repository-url>
cd medical-rag
```

2. **Установите зависимости**:
```bash
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**](https://huggingface.co/yandex/YandexGPT-5-Lite-8B-instruct-GGUF/resolve/main/YandexGPT-5-Lite-8B-instruct-Q4_K_M.gguf?download=true)

Создайте папку и поместите туда модель:
```bash
mkdir models
# Поместите скачанную модель в папку models/
```

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

Создайте файл `rag_corpus.json` с медицинскими примерами в формате:
```json
[
  {
    "short": "краткое описание",
    "full": "развернутое описание жалоб"
  }
]
```

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

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

```python
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

```python
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
```

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

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

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

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