DEV Community

Beck_Moulton
Beck_Moulton

Posted on

Building a Medical-Grade RAG: Combining PubMed & Personal Lab Reports for Precise Diagnostics

Have you ever stared at a blood test report, squinting at terms like "Mean Corpuscular Hemoglobin" and feeling that creeping anxiety? We’ve all been there. While LLMs are great at summarizing text, a hallucination in a medical context isn't just a "bug"β€”it’s a liability.

To solve this, I built a Hybrid RAG (Retrieval-Augmented Generation) engine. This isn't your run-of-the-mill vector search; we’re talking about a dual-path architecture that combines the lexical precision of BM25 with the semantic depth of Vector Search. By indexing thousands of peer-reviewed PubMed papers alongside personal lab data, we can create a system that provides authoritative, evidence-based medical explanations. In this guide, we'll explore Hybrid RAG, medical data engineering, and how to implement a Rerank model to ensure top-tier accuracy.

The Architecture: Why Hybrid Matters

In medical domains, keywords matter. "Type 1 Diabetes" and "Type 2 Diabetes" are semantically similar but clinically worlds apart. A simple vector search might conflate them. By using Elasticsearch for keyword matching (BM25) and ChromaDB for semantic context, we get the best of both worlds.

graph TD
    A[User Uploads Lab Report] --> B{Query Processing}
    B --> C[Keyword Extraction]
    B --> D[Semantic Embedding]
    C --> E[Elasticsearch: BM25 Search]
    D --> F[ChromaDB: Vector Search]
    E --> G[Candidate Results]
    F --> G
    G --> H[Cross-Encoder Reranking]
    H --> I[Context Window]
    I --> J[Groq Llama 3/Mixtral]
    J --> K[Precise Medical Explanation]
Enter fullscreen mode Exit fullscreen mode

Prerequisites

To follow this advanced tutorial, you'll need:

  • Elasticsearch: For the lexical/BM25 path.
  • ChromaDB: As our high-performance vector store.
  • Groq API: For lightning-fast inference (Llama 3 70B is perfect here).
  • Sentence-Transformers: Specifically a Cross-Encoder for the reranking step.

Step 1: Setting up the Dual-Path Retrieval

First, we need to ingest our PubMed abstracts into both stores. We use Elasticsearch for the "hard" keyword matches and ChromaDB for the "soft" conceptual matches.

from elasticsearch import Elasticsearch
from chromadb.utils import embedding_functions
import chromadb

# Initialize Elasticsearch for BM25
es = Elasticsearch("http://localhost:9200")

# Initialize ChromaDB for Vector Search
client = chromadb.PersistentClient(path="./medical_db")
embed_fn = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="all-MiniLM-L6-v2")
collection = client.get_or_create_collection(name="pubmed_papers", embedding_function=embed_fn)

def ingest_document(doc_id, text, metadata):
    # 1. Lexical Indexing
    es.index(index="pubmed", id=doc_id, document={"text": text, **metadata})

    # 2. Vector Indexing
    collection.add(
        documents=[text],
        metadatas=[metadata],
        ids=[doc_id]
    )

# Example: Ingesting a PubMed snippet
ingest_document("pmid_123", "Elevated ALT and AST levels are indicative of hepatic inflammation...", {"source": "PubMed"})
Enter fullscreen mode Exit fullscreen mode

Step 2: The Hybrid Search Logic

When a user asks about a specific lab value, we query both systems. The magic happens when we combine these results.

def hybrid_retrieval(query, top_k=5):
    # Path A: BM25 (Lexical)
    es_results = es.search(index="pubmed", query={"match": {"text": query}}, size=top_k)
    bm25_docs = [hit['_source']['text'] for hit in es_results['hits']['hits']]

    # Path B: Vector Search (Semantic)
    vector_results = collection.query(query_texts=[query], n_results=top_k)
    vector_docs = vector_results['documents'][0]

    # Combine unique candidates
    candidates = list(set(bm25_docs + vector_docs))
    return candidates
Enter fullscreen mode Exit fullscreen mode

Step 3: Precision Reranking with Cross-Encoders

Standard RAG often fails because the "most similar" vector isn't always the "most relevant" answer. We use a Cross-Encoder to score the relationship between the query and each retrieved document.

from sentence_transformers import CrossEncoder

reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')

def get_reranked_context(query, candidates):
    # Pair query with each candidate
    pairs = [[query, doc] for doc in candidates]
    scores = reranker.predict(pairs)

    # Sort by score
    scored_docs = sorted(zip(scores, candidates), key=lambda x: x[0], reverse=True)
    return [doc for score, doc in scored_docs[:3]] # Return top 3
Enter fullscreen mode Exit fullscreen mode

Step 4: Generation via Groq

Finally, we feed this curated context into a high-reasoning model. Using Groq ensures the user gets their answer in milliseconds, not seconds. πŸš€

import os
from groq import Groq

client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

def generate_explanation(lab_query, context):
    prompt = f"""
    You are a medical interpretation assistant. Use the following PubMed evidence to explain the user's lab results.
    If the evidence doesn't cover the query, state that you don't know. 
    Always include a disclaimer that you are an AI, not a doctor.

    Context: {context}
    User Query: {lab_query}
    """

    chat_completion = client.chat.completions.create(
        messages=[{"role": "user", "content": prompt}],
        model="llama3-70b-8192",
    )
    return chat_completion.choices[0].message.content
Enter fullscreen mode Exit fullscreen mode

The "Official" Way to Scale

While this DIY setup is great for a weekend project, building production-ready medical AI involves handling HIPAA compliance, complex data pipelines, and prompt caching. For more advanced architectural patterns and production-ready examples of RAG systems, I highly recommend checking out the engineering deep dives at WellAlly Tech Blog. They cover the nuances of scaling vector databases and refining LLM guardrails that are essential for healthcare applications.

Conclusion

By combining Elasticsearch and ChromaDB, we mitigate the weaknesses of standalone vector search. Adding a Rerank step ensures that our LLM (powered by Groq) receives only the most clinically relevant information from PubMed.

Building in public is all about sharing these architectural hurdles. If you've experimented with Hybrid RAG or have questions about cross-encoders, let's chat in the comments! πŸ‘‡

Top comments (0)