Skip to content

EMBEDDINGS_FIXED.md

Latest commit

 

History

History
165 lines (137 loc) · 6.03 KB

EMBEDDINGS_FIXED.md

File metadata and controls

165 lines (137 loc) · 6.03 KB

Medical Image Search Embeddings - FIXED! ✅

Issue

Medical image search was returning 0 results because all 50 images in the database had mock embeddings (arrays of zeros) instead of real NV-CLIP embeddings.

Root Cause

The NVCLIPEmbeddings class (src/embeddings/nvclip_embeddings.py) was hardcoded to use the NVIDIA Cloud API (https://integrate.api.nvidia.com/v1). However:

  • AWS deployment has a local NV-CLIP NIM container running on port 8002
  • The embedder was trying to connect to the cloud API and failing silently
  • On failure, it fell back to mock embeddings (all zeros)

Solution Implemented

Code Fix

Modified src/embeddings/nvclip_embeddings.py:

  • Added base_url parameter to __init__()
  • Added support for NVCLIP_BASE_URL environment variable
  • Made API key optional for local NIM deployments
  • Added auto-detection of local vs cloud based on hostname
  • Added clear messaging showing which endpoint is being used

Key Changes (lines 21-62)

def __init__(
    self,
    api_key: str = None,
    base_url: str = None  # NEW: Configurable base URL
):
    # Get base URL from parameter, env var, or default
    self.base_url = base_url or os.getenv('NVCLIP_BASE_URL', 'https://integrate.api.nvidia.com/v1')

    # API key: required for cloud, optional for local NIM
    self.api_key = api_key or os.getenv('NVIDIA_API_KEY')
    is_local_nim = 'localhost' in self.base_url or '127.0.0.1' in self.base_url

    if not self.api_key and not is_local_nim:
        raise ValueError(
            "NVIDIA API key required for cloud API. Set NVIDIA_API_KEY env var or pass api_key parameter.\n"
            "Get your key at: https://build.nvidia.com/nvidia/nvclip\n"
            "For local NIM, set NVCLIP_BASE_URL=http://localhost:8002/v1"
        )

    # Initialize OpenAI client
    self.client = OpenAI(
        api_key=self.api_key or "dummy",  # Local NIM doesn't validate key
        base_url=self.base_url
    )

Deployment Status

AWS (Production) ✅ WORKING

  • Status: 50 images with real NV-CLIP embeddings
  • Endpoint: Local NIM at http://localhost:8002/v1
  • Configuration: NVCLIP_BASE_URL=http://localhost:8002/v1
  • Verification: 
    magnitude=0.9998 ✅ REAL
magnitude=0.9998 ✅ REAL
magnitude=0.9998 ✅ REAL
  • Performance: 38.4 img/sec ingestion speed
  • Search: Ready for semantic image search

Local (Development) ✅ FIXED

  • Status: 50 images with real embeddings + 5 memories with real embeddings
  • Fix Applied: SSH tunnel established to AWS NIM on port 8002
  • Tunnel Command: 
    ssh -f -N -L 8002:localhost:8002 -i ~/.ssh/medical-graphrag-key.pem ubuntu@3.84.250.46
  • Memories Fixed: Ran scripts/fix_memory_embeddings.py with tunnel active

Verification Commands

Check Embeddings Quality

python3 -c "
from src.db.connection import get_connection
conn = get_connection()
cursor = conn.cursor()
cursor.execute('SELECT TOP 5 ImageID, VECTOR_DOT_PRODUCT(Vector, Vector) as magnitude FROM VectorSearch.MIMICCXRImages')
print('Embeddings check:')
for row in cursor.fetchall():
    status = '✅ REAL' if row[1] > 0 else '❌ MOCK'
    print(f'  {row[0][:30]}: magnitude={row[1]:.4f} {status}')
"

Test NV-CLIP API

# Local NIM (via SSH tunnel)
curl -X POST http://localhost:8002/v1/embeddings \
  -H 'Content-Type: application/json' \
  -d '{"input": ["chest X-ray showing pneumonia"], "model": "nvidia/nvclip"}'

# Cloud API (requires NVIDIA_API_KEY)
curl -X POST https://integrate.api.nvidia.com/v1/embeddings \
  -H "Authorization: Bearer $NVIDIA_API_KEY" \
  -H 'Content-Type: application/json' \
  -d '{"input": ["chest X-ray showing pneumonia"], "model": "nvidia/nvclip"}'

Usage

Run Ingestion with Local NIM

# AWS (direct access)
cd medical-graphrag
source venv/bin/activate
export NVCLIP_BASE_URL="http://localhost:8002/v1"
python ingest_mimic_images.py tests/fixtures/sample_medical_images --limit 50

Run Streamlit with NV-CLIP

cd mcp-server
NVCLIP_BASE_URL="http://localhost:8002/v1" streamlit run streamlit_app.py --server.port 8501

Test Medical Image Search

In the Streamlit UI (http://localhost:8501), try:

  • "Show me chest X-rays of pneumonia"
  • "Find chest X-rays showing cardiomegaly"
  • "Search for lateral view chest X-rays"

Files Modified

  • src/embeddings/nvclip_embeddings.py - Added configurable base_url support

Next Steps

  1. ✅ AWS deployment has real embeddings and working search
  2. ✅ SSH tunnel established for local development
  3. ✅ Local images re-ingested with real embeddings
  4. ✅ Memory search fixed with real embeddings
  5. ⏳ Test multi-modal search (text → image results)

Success Criteria Met ✅

  • ✅ Real NV-CLIP embeddings (magnitude ~1.0, not 0.0)
  • ✅ AWS has 50 images ready for search
  • ✅ Local has 50 images + 5 memories with real embeddings
  • ✅ Code supports both local NIM and cloud API
  • ✅ Environment variable configuration working
  • ✅ Clear error messages and endpoint detection
  • ✅ SSH tunnel working for local development
  • ✅ Memory search operational

Status: ALL EMBEDDINGS FIXED! Both medical image search AND memory search are now operational.

Session State UI Fix (2025-11-22)

Memory Search Button Issue

After fixing embeddings, the memory search button in Streamlit sidebar was not displaying results.

Root Cause: Streamlit button state doesn't persist across re-renders. The condition if st.button(...) and search_query meant results were calculated but immediately lost when the page re-rendered.

Solution: Modified mcp-server/streamlit_app.py (lines 984-1014):

  • Added st.session_state.memory_search_results to persist results
  • Separated button click logic (stores to session state) from results display (reads from session state)
  • Results now persist after button click

Verification:

  • Streamlit restarted with fix applied
  • SSH tunnel to AWS NIM still active on port 8002
  • App accessible at http://localhost:8501
  • Memory search now fully functional: search → results display → persist across interactions