first commit
This commit is contained in:
@@ -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 сек
|
||||
Reference in New Issue
Block a user