rag_hf.py

"""
rag_hf.py — RAG module using Hugging Face Inference API for E5 embeddings.

Same algorithm as src/rag.py, but instead of loading the model locally,
embeddings are computed by calling the HF Inference API endpoint.

Model: intfloat/multilingual-e5-large
Uses: huggingface_hub.InferenceClient.feature_extraction()

Usage:
    from rag_hf import FAISSRetrieverHF
    retriever = FAISSRetrieverHF()
    retriever.load()
    results = retriever.retrieve("head pain, fever, stiff neck")
"""

import json
import logging
import os
import re
import time
from dataclasses import dataclass
from pathlib import Path

import faiss
import numpy as np
from huggingface_hub import InferenceClient

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
DEFAULT_CORPUS_PATH = (
    Path(__file__).parent / "data" / "corpus" / "corpus" / "protocols_corpus_enriched.jsonl"
)
FAISS_INDEX_PATH = (
    Path(__file__).parent / "data" / "faiss_index_hf.bin"
)
EMBEDDINGS_CACHE_PATH = (
    Path(__file__).parent / "data" / "e5_embeddings_cache_hf.npy"
)

HF_TOKEN = os.environ.get("HF_TOKEN", "")
HF_MODEL = "intfloat/multilingual-e5-large"

MAX_CHUNK_CHARS = 800
TOP_K = 5
# HF free-tier rate limit: keep batches small to avoid 429s
HF_BATCH_SIZE = 32
# Seconds to wait between batches
HF_BATCH_DELAY = 0.5
# Max retries on rate-limit / transient errors
HF_MAX_RETRIES = 5
HF_RETRY_DELAY = 10  # seconds

# Shared InferenceClient (initialised on first use)
_hf_client: InferenceClient | None = None


def _get_client() -> InferenceClient:
    global _hf_client
    if _hf_client is None:
        _hf_client = InferenceClient(model=HF_MODEL, token=HF_TOKEN)
    return _hf_client


def _reset_client() -> None:
    """Force re-creation of the shared client (e.g. after token change)."""
    global _hf_client
    _hf_client = None


@dataclass
class RetrievalResult:
    protocol_id: str
    title: str
    icd_codes: list[str]
    context: str
    score: float


@dataclass
class Protocol:
    id: str
    title: str
    icd_codes: list[str]
    text: str
    diagnostic_section: str = ""


# ---------------------------------------------------------------------------
# HF InferenceClient helpers
# ---------------------------------------------------------------------------

def _hf_embed(texts: list[str]) -> np.ndarray:
    """
    Embed a list of texts using InferenceClient.feature_extraction().
    Returns a float32 numpy array of shape (len(texts), dim), L2-normalised.
    Retries on transient errors.
    """
    client = _get_client()
    for attempt in range(1, HF_MAX_RETRIES + 1):
        try:
            result = client.feature_extraction(texts)
            # result is a list[list[float]] or numpy array
            arr = np.array(result, dtype=np.float32)
            # If the model returns a 3-D tensor (batch, seq, dim), mean-pool
            if arr.ndim == 3:
                arr = arr.mean(axis=1)
            # L2-normalise so inner-product == cosine similarity
            norms = np.linalg.norm(arr, axis=1, keepdims=True)
            norms = np.where(norms == 0, 1.0, norms)
            return arr / norms
        except Exception as e:
            msg = str(e)
            if attempt < HF_MAX_RETRIES:
                wait = HF_RETRY_DELAY
                logger.warning(
                    "HF embed error (attempt %d/%d): %s — retrying in %ds…",
                    attempt, HF_MAX_RETRIES, msg, wait,
                )
                time.sleep(wait)
            else:
                raise RuntimeError(f"HF embed failed after {HF_MAX_RETRIES} attempts: {e}") from e
    raise RuntimeError("HF API: exceeded max retries")


def _hf_embed_batched(texts: list[str], batch_size: int = HF_BATCH_SIZE) -> np.ndarray:
    """Embed a large list of texts in batches, with progress logging."""
    all_embeddings = []
    total = len(texts)
    for start in range(0, total, batch_size):
        batch = texts[start : start + batch_size]
        logger.info(
            "Embedding batch %d-%d / %d …",
            start + 1, min(start + batch_size, total), total,
        )
        emb = _hf_embed(batch)
        all_embeddings.append(emb)
        if start + batch_size < total:
            time.sleep(HF_BATCH_DELAY)
    return np.vstack(all_embeddings).astype(np.float32)


# ---------------------------------------------------------------------------
# FAISSRetrieverHF
# ---------------------------------------------------------------------------

class FAISSRetrieverHF:
    """
    FAISS-based retriever identical to FAISSRetriever in src/rag.py,
    but uses the Hugging Face Inference API for E5 embeddings instead
    of a locally loaded SentenceTransformer model.
    """

    def __init__(self, corpus_path: Path = DEFAULT_CORPUS_PATH):
        self.corpus_path = corpus_path
        self.protocols: list[Protocol] = []
        self.is_loaded = False
        self.index: faiss.IndexFlatIP | None = None

    def load(self) -> None:
        """Load corpus, then build or load the FAISS index."""
        if self.is_loaded:
            return

        logger.info("Loading clinical protocols from %s …", self.corpus_path)
        self._load_protocols()
        logger.info("Loaded %d protocols", len(self.protocols))

        if FAISS_INDEX_PATH.exists() and EMBEDDINGS_CACHE_PATH.exists():
            self._load_index()
        else:
            self._build_index()

        logger.info(
            "FAISS index ready: %d vectors, dimension %d",
            self.index.ntotal, self.index.d,
        )
        self.is_loaded = True

    def retrieve(self, query: str, top_k: int = TOP_K) -> list[RetrievalResult]:
        """Embed the query via HF API and search the FAISS index."""
        if not self.is_loaded:
            return []

        query_text = f"query: {query}"
        query_embedding = _hf_embed([query_text])  # shape (1, dim)

        scores, indices = self.index.search(query_embedding, top_k)

        if len(indices[0]) > 0:
            logger.info(
                "Retrieval: query='%.60s…', top_k=%d, best_score=%.4f",
                query, top_k, scores[0][0],
            )

        results = []
        for score, idx in zip(scores[0], indices[0]):
            if idx < 0:
                continue
            proto = self.protocols[idx]
            context = (
                proto.diagnostic_section[:MAX_CHUNK_CHARS]
                if proto.diagnostic_section
                else proto.text[:MAX_CHUNK_CHARS]
            )
            results.append(
                RetrievalResult(
                    protocol_id=proto.id,
                    title=proto.title,
                    icd_codes=proto.icd_codes,
                    context=context,
                    score=float(score),
                )
            )
        return results

    # ------------------------------------------------------------------
    # Internal helpers  (same logic as src/rag.py)
    # ------------------------------------------------------------------

    def _load_protocols(self) -> None:
        generic_titles = {
            "одобрен", "одобрено", "одобрена", "протокол", "протоколы",
            "рекомендовано", "рекомендован", "утверждено", "утвержден",
            "кп", "клинический протокол",
        }
        with open(self.corpus_path, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                obj = json.loads(line)
                text = obj.get("text", "")

                title = obj.get("title", "").strip().strip("«»\" \t\n")
                is_generic = (
                    not title
                    or len(title) < 5
                    or title.lower().strip() in generic_titles
                    or title.lower().startswith("одобрен")
                    or title.lower().startswith("рекомендован")
                    or title.lower().startswith("утвержден")
                )
                if is_generic:
                    title = self._extract_title_from_text(text)

                if not title:
                    source_file = obj.get("source_file", "")
                    if source_file:
                        clean = re.sub(r"\.(pdf|docx|txt)$", "", source_file, flags=re.I)
                        clean = clean.replace("_", " ").strip("«»\"' ").strip()
                        if clean and len(clean) > 3:
                            title = clean

                if not title:
                    title = "Protocol " + obj.get("protocol_id", "Unknown")

                icd_codes = obj.get("icd_codes", [])
                if not icd_codes:
                    icd_codes = self._extract_icd_codes_from_text(text)

                proto = Protocol(
                    id=obj.get("protocol_id", ""),
                    title=title,
                    icd_codes=icd_codes,
                    text=text,
                )
                proto.diagnostic_section = self._extract_diagnostic_section(text)
                self.protocols.append(proto)

    @staticmethod
    def _extract_title_from_text(text: str) -> str:
        patterns = [
            r"ПРОТОКОЛ\s+ДИАГНОСТИКИ\s+И\s+ЛЕЧЕНИ[ЯЮ]\s+(.+?)(?:\s*1\s*\.|\s*ВВОДНАЯ)",
            r"ПРОТОКОЛ\s+(?:ПО\s+)?ДИАГНОСТИК[ИЕ]\s+И\s+ЛЕЧЕНИ[ЯЮ]\s+(.+?)(?:\s*1\s*\.|\s*ВВОДНАЯ)",
            r"КЛИНИЧЕСКИЙ\s+ПРОТОКОЛ\s+(.+?)(?:\s*1\s*\.|\s*ВВОДНАЯ)",
        ]
        for pat in patterns:
            m = re.search(pat, text[:1500], re.IGNORECASE | re.DOTALL)
            if m:
                title = re.sub(r"\s+", " ", m.group(1)).strip()
                title = title.strip("«»\"' \t\n.:;,")
                if len(title) > 3:
                    return title
        return ""

    _CYRILLIC_TO_LATIN = {
        'А': 'A', 'В': 'B', 'С': 'C', 'D': 'D', 'Е': 'E',
        'F': 'F', 'G': 'G', 'Н': 'H', 'І': 'I', 'J': 'J',
        'К': 'K', 'L': 'L', 'М': 'M', 'N': 'N', 'О': 'O',
        'Р': 'P', 'Q': 'Q', 'R': 'R', 'S': 'S', 'Т': 'T',
        'У': 'Y', 'V': 'V', 'W': 'W', 'Х': 'X', 'Z': 'Z',
    }
    _VALID_ICD_PREFIXES = set('ABCDEFGHJKLMNOPQRSTVWXYZ')

    @staticmethod
    def _extract_icd_codes_from_text(text: str) -> list[str]:
        search_area = text[:8000]
        codes = re.findall(
            r"(?<!\d)([A-ZА-ЯЁ]\s?\d{2}(?:\.\d{1,2})?)(?!\d)",
            search_area,
        )
        cyr_map = FAISSRetrieverHF._CYRILLIC_TO_LATIN
        valid_prefixes = FAISSRetrieverHF._VALID_ICD_PREFIXES
        seen: set[str] = set()
        result = []
        for code in codes:
            normalized = code.replace(" ", "")
            if len(normalized) < 3:
                continue
            first_char = normalized[0]
            if first_char in cyr_map:
                normalized = cyr_map[first_char] + normalized[1:]
            if normalized[0] not in valid_prefixes:
                continue
            if normalized not in seen:
                seen.add(normalized)
                result.append(normalized)
        return result

    def _build_index(self) -> None:
        """Compute E5 embeddings via HF API and build FAISS index."""
        passage_texts = [
            f"passage: {p.title}. МКБ-10: {', '.join(p.icd_codes)}. {p.diagnostic_section[:500]}"
            for p in self.protocols
        ]
        logger.info(
            "Computing E5 embeddings for %d protocols via HF API (batch=%d)…",
            len(passage_texts), HF_BATCH_SIZE,
        )
        embeddings = _hf_embed_batched(passage_texts)
        logger.info("Embeddings computed: shape %s", embeddings.shape)

        dim = embeddings.shape[1]
        self.index = faiss.IndexFlatIP(dim)
        self.index.add(embeddings)

        faiss.write_index(self.index, str(FAISS_INDEX_PATH))
        np.save(str(EMBEDDINGS_CACHE_PATH), embeddings)
        logger.info("FAISS index saved  → %s", FAISS_INDEX_PATH)
        logger.info("Embeddings cached  → %s", EMBEDDINGS_CACHE_PATH)

    def _load_index(self) -> None:
        logger.info("Loading cached FAISS index from %s …", FAISS_INDEX_PATH)
        self.index = faiss.read_index(str(FAISS_INDEX_PATH))
        logger.info(
            "Cached FAISS index loaded: %d vectors, dim %d",
            self.index.ntotal, self.index.d,
        )

    @staticmethod
    def _extract_diagnostic_section(text: str) -> str:
        patterns = [
            r"(?:диагностические\s+критерии|критерии\s+диагностики)",
            r"(?:жалобы\s+и\s+анамнез|клиническ[аи][яе]\s+картин[аы])",
            r"(?:физикальн[оы][ей]\s+обследовани|осмотр)",
            r"(?:клинические\s+проявления|симптом[ыа])",
        ]
        best_start = -1
        for pat in patterns:
            m = re.search(pat, text, re.IGNORECASE)
            if m:
                best_start = max(0, m.start() - 50)
                break
        if best_start >= 0:
            return text[best_start : best_start + MAX_CHUNK_CHARS]
        return text[:MAX_CHUNK_CHARS]


# ---------------------------------------------------------------------------
# Format RAG context for the LLM prompt (same as src/rag.py)
# ---------------------------------------------------------------------------
def format_rag_context(results: list[RetrievalResult]) -> str:
    if not results:
        return ""
    lines = ["RELEVANT CLINICAL PROTOCOLS:"]
    for i, r in enumerate(results, 1):
        codes = ", ".join(r.icd_codes[:5])
        lines.append(
            f"\n--- Protocol {i}: {r.title} (ICD-10: {codes}) "
            f"[score={r.score:.3f}] ---\n"
            f"{r.context}"
        )
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Quick CLI test
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    import sys
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")

    retriever = FAISSRetrieverHF()
    print("Loading RAG (HF API) …")
    retriever.load()
    print("Ready.\n")

    query = " ".join(sys.argv[1:]) or "головная боль, высокая температура, ригидность шеи"
    print(f"Query: {query}\n")
    results = retriever.retrieve(query, top_k=5)
    for i, r in enumerate(results, 1):
        print(f"{i}. [{r.score:.4f}] {r.title}")
        print(f"   ICD-10: {', '.join(r.icd_codes)}")
    print()