first commit
This commit is contained in:
+429
@@ -0,0 +1,429 @@
|
||||
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Обнаружены патологические изменения. Рекомендуется консультация специалиста."
|
||||
Reference in New Issue
Block a user