|
|
from PyQt6.QtCore import QThread, QMutex, QWaitCondition, pyqtSignal |
|
|
import time |
|
|
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 |
|
|
|
|
|
|
|
|
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, |
|
|
use_gpu_for_embeddings: bool = False, |
|
|
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)") |
|
|
if gpu_memory < 4: |
|
|
n_gpu_layers = min(n_gpu_layers, 10) |
|
|
elif gpu_memory < 8: |
|
|
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, |
|
|
} |
|
|
|
|
|
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, processing_type: str, |
|
|
max_context_tokens: int = 1500) -> Tuple[str, int]: |
|
|
"""Строит промпт с оптимизированным управлением токенами""" |
|
|
|
|
|
examples = self.retrieve(short_note) |
|
|
|
|
|
# Разные системные сообщения для жалоб и анамнеза |
|
|
if processing_type == 'complaints': |
|
|
system_msg = ( |
|
|
"На основе примеров напиши развёрнуто жалобы пациента, грамотно с медицинской точки зрения. " |
|
|
"Напиши жалобы в одно предложение, одной строкой. " |
|
|
"Не пиши вводных слов и фраз. Только жалобы пациента. " |
|
|
"Неуместно писать диагнозы и план лечения. " |
|
|
"Расшифруй все сокращения. " |
|
|
"Отвечай сразу без размышлений." |
|
|
) |
|
|
else: # history |
|
|
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} |
|
|
|
|
|
{"Жалобы" if processing_type == 'complaints' else "Анамнез"}: "{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 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, processing_type: str) -> str: |
|
|
"""Генерация с оптимизацией памяти""" |
|
|
prompt, prompt_tokens = self.build_prompt_with_token_management(short_note, processing_type) |
|
|
|
|
|
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 |
|
|
) |
|
|
|
|
|
result = output["choices"][0]["text"].strip() |
|
|
return result |
|
|
except Exception as e: |
|
|
print(f"❌ Ошибка генерации: {e}") |
|
|
return "" |
|
|
|
|
|
def __call__(self, short_note: str, processing_type: str) -> str: |
|
|
return self.generate(short_note, processing_type) |
|
|
|
|
|
|
|
|
class AIWorker(QThread): |
|
|
"""Worker для обработки данных нейросетью""" |
|
|
text_processed = pyqtSignal(str, str) # text, processing_type |
|
|
image_processed = pyqtSignal(str) |
|
|
processing_finished = pyqtSignal(str) # processing_type |
|
|
|
|
|
def __init__(self): |
|
|
super().__init__() |
|
|
self.mutex = QMutex() |
|
|
self.condition = QWaitCondition() |
|
|
self.text_to_process = None |
|
|
self.image_to_process = None |
|
|
self.processing_type = None # 'complaints', 'history', 'image' |
|
|
self._active = True |
|
|
|
|
|
# Инициализация MedicalRAG |
|
|
self.rag = None |
|
|
self._init_rag() |
|
|
|
|
|
def _init_rag(self): |
|
|
"""Инициализация MedicalRAG системы""" |
|
|
try: |
|
|
import psutil |
|
|
|
|
|
# Автоматическое определение режима низкой памяти |
|
|
total_memory = psutil.virtual_memory().total / 1024 ** 3 |
|
|
low_memory_mode = total_memory < 8 |
|
|
|
|
|
print(f"💾 Общая память: {total_memory:.1f} GB") |
|
|
print(f"🔧 Режим низкой памяти: {'Да' if low_memory_mode else 'Нет'}") |
|
|
|
|
|
# Путь к модели - нужно будет настроить под вашу систему |
|
|
model_path = "./models/YandexGPT-5-Lite-8B-instruct-Q4_K_M.gguf" |
|
|
|
|
|
# Если модель не найдена, используем имитацию |
|
|
if not os.path.exists(model_path): |
|
|
print("⚠️ Модель не найдена, используется имитация LLM") |
|
|
self.rag = None |
|
|
return |
|
|
|
|
|
self.rag = MedicalRAG( |
|
|
model_path=model_path, |
|
|
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 |
|
|
) |
|
|
print("✅ MedicalRAG инициализирован успешно!") |
|
|
|
|
|
except Exception as e: |
|
|
print(f"❌ Ошибка инициализации MedicalRAG: {e}") |
|
|
print("⚠️ Будет использоваться имитация LLM") |
|
|
self.rag = None |
|
|
|
|
|
def process_text(self, text, processing_type): |
|
|
"""Запускает обработку текста""" |
|
|
self.mutex.lock() |
|
|
self.text_to_process = text |
|
|
self.image_to_process = None |
|
|
self.processing_type = processing_type |
|
|
self.condition.wakeOne() |
|
|
self.mutex.unlock() |
|
|
|
|
|
def process_image(self, image_path): |
|
|
"""Запускает обработку изображения""" |
|
|
self.mutex.lock() |
|
|
self.image_to_process = image_path |
|
|
self.text_to_process = None |
|
|
self.processing_type = 'image' |
|
|
self.condition.wakeOne() |
|
|
self.mutex.unlock() |
|
|
|
|
|
def stop(self): |
|
|
"""Останавливает worker""" |
|
|
self._active = False |
|
|
self.condition.wakeOne() |
|
|
self.wait() |
|
|
|
|
|
def run(self): |
|
|
"""Основной цикл обработки""" |
|
|
while self._active: |
|
|
self.mutex.lock() |
|
|
if not self.text_to_process and not self.image_to_process: |
|
|
self.condition.wait(self.mutex) |
|
|
|
|
|
if not self._active: |
|
|
self.mutex.unlock() |
|
|
break |
|
|
|
|
|
if self.processing_type in ['complaints', 'history'] and self.text_to_process: |
|
|
text = self.text_to_process |
|
|
processing_type = self.processing_type |
|
|
self.text_to_process = None |
|
|
self.mutex.unlock() |
|
|
|
|
|
# Обработка текста с использованием LLM |
|
|
result = self.process_text_with_llm(text, processing_type) |
|
|
self.text_processed.emit(result, processing_type) |
|
|
|
|
|
elif self.processing_type == 'image' and self.image_to_process: |
|
|
image_path = self.image_to_process |
|
|
self.image_to_process = None |
|
|
self.mutex.unlock() |
|
|
|
|
|
# Обработка изображения |
|
|
result = self.simulate_image_processing(image_path) |
|
|
self.image_processed.emit(result) |
|
|
else: |
|
|
self.mutex.unlock() |
|
|
|
|
|
self.processing_finished.emit(self.processing_type) |
|
|
|
|
|
def process_text_with_llm(self, text, processing_type): |
|
|
"""Обработка текста с использованием LLM""" |
|
|
if not text.strip(): |
|
|
return "" |
|
|
|
|
|
# Если LLM не инициализирован, используем имитацию |
|
|
if self.rag is None: |
|
|
return self.simulate_text_processing(text, processing_type) |
|
|
|
|
|
try: |
|
|
print(f"🔍 Обработка {processing_type} с помощью LLM...") |
|
|
print(f"📥 Входной текст: {text}") |
|
|
|
|
|
# Обработка через LLM |
|
|
result = self.rag(text, processing_type) |
|
|
|
|
|
# Если LLM вернул пустой результат, используем исходный текст |
|
|
if not result or result.strip() == "": |
|
|
print("⚠️ LLM вернул пустой ответ, используется исходный текст") |
|
|
result = self._format_fallback_text(text, processing_type) |
|
|
else: |
|
|
print(f"📤 Результат LLM: {result}") |
|
|
|
|
|
return result |
|
|
|
|
|
except Exception as e: |
|
|
print(f"❌ Ошибка при обработке LLM: {e}") |
|
|
# В случае ошибки возвращаем форматированный исходный текст |
|
|
return self._format_fallback_text(text, processing_type) |
|
|
|
|
|
def _format_fallback_text(self, text, processing_type): |
|
|
"""Форматирует исходный текст как fallback""" |
|
|
if processing_type == 'complaints': |
|
|
return f"🤖 Анализ жалоб ИИ:\n\n{text}\n\nВывод: Требуется дополнительное обследование." |
|
|
else: |
|
|
return f"🤖 Анализ анамнеза ИИ:\n\n{text}\n\nВывод: Анамнез требует уточнения." |
|
|
|
|
|
def simulate_text_processing(self, text, processing_type): |
|
|
"""Имитация обработки текста (fallback)""" |
|
|
time.sleep(2) # Имитация задержки обработки |
|
|
|
|
|
if processing_type == 'complaints': |
|
|
return f"🤖 Анализ жалоб ИИ:\n\n{text}\n\nВывод: Выявлены характерные симптомы требующие дополнительного обследования." |
|
|
else: |
|
|
return f"🤖 Анализ анамнеза ИИ:\n\n{text}\n\nВывод: Анамнез указывает на хронический характер заболевания." |
|
|
|
|
|
def simulate_image_processing(self, image_path): |
|
|
"""Имитация обработки изображения нейросетью""" |
|
|
time.sleep(3) # Имитация задержки обработки |
|
|
|
|
|
filename = os.path.basename(image_path) |
|
|
return f"🖼️ Анализ изображения '{filename}':\n\nОбнаружены патологические изменения. Рекомендуется консультация специалиста." |
