Production-ready Medical Coding API with Phi-3 support

Dockerfile

@@ -0,0 +1,58 @@
+FROM python:3.10-slim
+
+# Set working directory
+WORKDIR /app
+
+# Set environment variables
+ENV PYTHONUNBUFFERED=1 \
+    PYTHONDONTWRITEBYTECODE=1 \
+    PIP_NO_CACHE_DIR=1 \
+    PIP_DISABLE_PIP_VERSION_CHECK=1 \
+    TRANSFORMERS_CACHE=/app/.cache/transformers \
+    HF_HOME=/app/.cache/huggingface \
+    DEBIAN_FRONTEND=noninteractive
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    git \
+    git-lfs \
+    build-essential \
+    curl \
+    ca-certificates \
+    && git lfs install \
+    && rm -rf /var/lib/apt/lists/* \
+    && apt-get clean
+
+# Upgrade pip and install build tools
+RUN pip install --no-cache-dir --upgrade \
+    pip==24.0 \
+    setuptools==69.5.1 \
+    wheel==0.43.0
+
+# Copy requirements first for better Docker caching
+COPY requirements.txt .
+
+# Install Python dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy application code
+COPY app/ ./app/
+
+# Create necessary directories with proper permissions
+RUN mkdir -p /app/offload /app/.cache/transformers /app/.cache/huggingface && \
+    chmod -R 777 /app/offload /app/.cache
+
+# Expose port 7860 (HuggingFace Spaces standard)
+EXPOSE 7860
+
+# Health check - more lenient for model loading
+HEALTHCHECK --interval=30s --timeout=20s --start-period=300s --retries=5 \
+    CMD curl -f http://localhost:7860/health || exit 1
+
+# Run the application with increased timeouts
+CMD ["uvicorn", "app.api:app", \
+    "--host", "0.0.0.0", \
+    "--port", "7860", \
+    "--timeout-keep-alive", "300", \
+    "--workers", "1", \
+    "--log-level", "info"]

README.md

@@ -0,0 +1,114 @@
+---
+title: Medical Coding API
+emoji: 🏥
+colorFrom: blue
+colorTo: green
+sdk: docker
+pinned: false
+license: mit
+app_port: 7860
+tags:
+  - medical
+  - healthcare
+  - icd-10
+  - cpt
+  - phi-3
+  - fastapi
+---
+
+# 🏥 Medical Coding API
+
+AI-powered API for extracting **ICD-10** and **CPT codes** from clinical provider notes using Microsoft Phi-3.
+
+## 🚀 Features
+
+- ✅ Extract ICD-10 diagnosis codes
+- ✅ Extract CPT procedure codes
+- ✅ Supports notes up to 10,000 characters (~2,500 words)
+- ✅ JSON output format
+- ✅ GPU-accelerated inference (when available)
+- ✅ Automatic text truncation
+- ✅ Production-ready with error handling
+
+## 📡 API Endpoints
+
+### POST `/predict`
+
+Extract medical codes from clinical note.
+
+**Request:**
+
+```json
+{
+  "note": "Your clinical note here..."
+}
+```
+
+**Response:**
+
+```json
+{
+  "result": {
+    "icd10_codes": ["I10", "E11.9"],
+    "cpt_codes": ["99213"]
+  },
+  "raw_output": "...",
+  "note_length": 250,
+  "truncated": false,
+  "processing_time": 3.45
+}
+```
+
+### GET `/health`
+
+Check API health status.
+
+### GET `/docs`
+
+Interactive API documentation (Swagger UI).
+
+## 🧪 Usage Examples
+
+### cURL
+
+```bash
+curl -X POST "https://YOUR-SPACE.hf.space/predict" \
+  -H "Content-Type: application/json" \
+  -d '{"note": "Patient with HTN, BP 160/95. Prescribed lisinopril."}'
+```
+
+### Python
+
+```python
+import requests
+
+response = requests.post(
+    "https://YOUR-SPACE.hf.space/predict",
+    json={"note": "Patient with diabetes, HbA1c 8.2. Started metformin."}
+)
+print(response.json())
+```
+
+### JavaScript
+
+```javascript
+fetch("https://YOUR-SPACE.hf.space/predict", {
+  method: "POST",
+  headers: { "Content-Type": "application/json" },
+  body: JSON.stringify({ note: "Clinical note here..." }),
+})
+  .then((res) => res.json())
+  .then((data) => console.log(data));
+```
+
+## ⚙️ Technical Details
+
+- **Model:** RayyanAhmed9477/med-coding (Phi-3 based)
+- **Framework:** FastAPI + Transformers
+- **Deployment:** HuggingFace Spaces (Docker)
+- **First Request:** 30-60 seconds (model loading)
+- **Subsequent Requests:** 2-10 seconds
+
+## 📝 License
+
+MIT License

app/__init__.py

@@ -0,0 +1,3 @@
+"""Medical Coding API - Extract ICD-10 and CPT codes from clinical notes."""
+
+__version__ = "1.0.0"

app/api.py

@@ -0,0 +1,282 @@
+# api.py
+import re
+import json
+import gc
+import time
+from typing import Optional
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel, Field
+from .model_loader import load_model_and_tokenizer
+from .prompt_template import PROMPT_TEMPLATE
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+app = FastAPI(
+    title="Medical Coding API",
+    description="Extract ICD-10 and CPT codes from clinical notes using AI",
+    version="1.0.0",
+    docs_url="/docs",
+    redoc_url="/redoc"
+)
+
+class NoteRequest(BaseModel):
+    note: str = Field(
+        ..., 
+        min_length=10,
+        max_length=50000,
+        description="Clinical provider note (10-50,000 characters)"
+    )
+    
+    class Config:
+        json_schema_extra = {
+            "example": {
+                "note": "Patient presents with essential hypertension. BP 160/95. Prescribed lisinopril 10mg daily. Office visit for established patient."
+            }
+        }
+
+class CodingResponse(BaseModel):
+    result: dict = Field(..., description="Extracted ICD-10 and CPT codes")
+    raw_output: str = Field(..., description="Raw model output")
+    note_length: int = Field(..., description="Length of input note in characters")
+    truncated: bool = Field(..., description="Whether note was truncated")
+    processing_time: float = Field(..., description="Time taken to process in seconds")
+
+# Global variables for lazy loading
+_gen_pipeline = None
+_tokenizer = None
+_model_load_time = None
+
+def get_model():
+    """Lazy load model on first request with error handling."""
+    global _gen_pipeline, _tokenizer, _model_load_time
+    
+    if _gen_pipeline is None:
+        logger.info("🔄 Loading model for the first time...")
+        start_time = time.time()
+        
+        try:
+            _gen_pipeline, _tokenizer = load_model_and_tokenizer()
+            _model_load_time = time.time() - start_time
+            logger.info(f"✅ Model loaded in {_model_load_time:.2f} seconds")
+        except Exception as e:
+            logger.error(f"❌ Failed to load model: {str(e)}")
+            raise HTTPException(
+                status_code=503,
+                detail=f"Model loading failed: {str(e)}. Please try again in a few moments."
+            )
+    
+    return _gen_pipeline, _tokenizer
+
+def extract_json_from_text(text: str) -> Optional[str]:
+    """Extract JSON object from text using brace counting."""
+    start_idx = text.find('{')
+    if start_idx == -1:
+        return None
+    
+    brace_count = 0
+    for i in range(start_idx, len(text)):
+        if text[i] == '{':
+            brace_count += 1
+        elif text[i] == '}':
+            brace_count -= 1
+            if brace_count == 0:
+                return text[start_idx:i+1]
+    return None
+
+def truncate_note(note: str, max_chars: int = 10000) -> str:
+    """Truncate note to prevent token limit issues."""
+    if len(note) <= max_chars:
+        return note
+    
+    logger.warning(f"Note truncated from {len(note)} to {max_chars} characters")
+    return note[:max_chars]
+
+# ===== ENDPOINTS =====
+
+@app.get("/")
+async def root():
+    """Root endpoint with API information."""
+    return {
+        "name": "Medical Coding API",
+        "version": "1.0.0",
+        "description": "Extract ICD-10 and CPT codes from clinical notes",
+        "model": "RayyanAhmed9477/med-coding (Phi-3 based)",
+        "endpoints": {
+            "/predict": "POST - Extract medical codes from clinical note",
+            "/health": "GET - Check API health status",
+            "/docs": "GET - Interactive API documentation",
+            "/metrics": "GET - API usage metrics"
+        },
+        "usage": {
+            "endpoint": "/predict",
+            "method": "POST",
+            "body": {"note": "Your clinical note here (10-50,000 chars)"},
+            "max_note_length": "50,000 characters (~10,000 words)"
+        }
+    }
+
+@app.get("/health")
+async def health_check():
+    """Health check endpoint."""
+    return {
+        "status": "healthy",
+        "model": "RayyanAhmed9477/med-coding",
+        "model_loaded": _gen_pipeline is not None,
+        "model_load_time": f"{_model_load_time:.2f}s" if _model_load_time else "not loaded yet"
+    }
+
+@app.get("/metrics")
+async def metrics():
+    """Get API usage metrics."""
+    return {
+        "model_loaded": _gen_pipeline is not None,
+        "model_load_time_seconds": _model_load_time,
+        "status": "operational"
+    }
+
+@app.post("/predict", response_model=CodingResponse)
+async def predict(request: NoteRequest):
+    """
+    Extract ICD-10 and CPT codes from clinical notes.
+    
+    **Input:** Clinical note (10-50,000 characters)
+    
+    **Output:** JSON with extracted codes:
+    - icd10_codes: List of ICD-10 diagnosis codes
+    - cpt_codes: List of CPT procedure codes
+    
+    **Note:** First request may take 30-60 seconds as model loads into memory.
+    Subsequent requests will be faster (2-10 seconds).
+    """
+    start_time = time.time()
+    
+    try:
+        # Validate input
+        note = request.note.strip()
+        if not note:
+            raise HTTPException(status_code=400, detail="Empty note provided")
+        
+        # Load model (lazy loading)
+        logger.info(f"📝 Processing note ({len(note)} characters)")
+        gen_pipeline, tokenizer = get_model()
+        
+        # Truncate if needed
+        original_length = len(note)
+        note_truncated = truncate_note(note, max_chars=10000)
+        
+        # Build prompt
+        prompt = PROMPT_TEMPLATE.format(note=note_truncated)
+        logger.info(f"🔮 Generating prediction (prompt length: {len(prompt)} chars)")
+        
+        # Generate prediction
+        outputs = gen_pipeline(
+            prompt,
+            max_new_tokens=600,
+            do_sample=False,
+            num_return_sequences=1,
+            pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+            temperature=0.1,
+            top_p=0.95,
+            return_full_text=False
+        )
+        
+        # Extract generated text
+        if isinstance(outputs, list) and len(outputs) > 0:
+            text = outputs[0].get("generated_text", "")
+        elif isinstance(outputs, dict):
+            text = outputs.get("generated_text", "")
+        else:
+            text = str(outputs)
+        
+        logger.info(f"📤 Model output length: {len(text)} characters")
+        
+        # Remove prompt if present
+        if prompt in text:
+            text = text.replace(prompt, "").strip()
+        
+        # Extract JSON
+        json_str = extract_json_from_text(text)
+        
+        if json_str is None:
+            logger.error(f"No JSON found in output: {text[:500]}")
+            raise HTTPException(
+                status_code=500,
+                detail={
+                    "error": "No valid JSON found in model output",
+                    "raw_output_preview": text[:300],
+                    "suggestion": "Model may need fine-tuning or prompt adjustment"
+                }
+            )
+        
+        # Parse JSON
+        try:
+            parsed = json.loads(json_str)
+        except json.JSONDecodeError as e:
+            logger.error(f"JSON parse error: {str(e)}")
+            raise HTTPException(
+                status_code=500,
+                detail={
+                    "error": f"Invalid JSON format: {str(e)}",
+                    "json_preview": json_str[:300]
+                }
+            )
+        
+        # Validate response structure
+        if not isinstance(parsed, dict):
+            raise HTTPException(
+                status_code=500,
+                detail="Model output is not a valid JSON object"
+            )
+        
+        # Clean up memory
+        gc.collect()
+        
+        processing_time = time.time() - start_time
+        logger.info(f"✅ Prediction completed in {processing_time:.2f} seconds")
+        
+        return CodingResponse(
+            result=parsed,
+            raw_output=text,
+            note_length=original_length,
+            truncated=original_length > 10000,
+            processing_time=round(processing_time, 2)
+        )
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"❌ Prediction failed: {str(e)}", exc_info=True)
+        gc.collect()
+        raise HTTPException(
+            status_code=500,
+            detail=f"Prediction failed: {str(e)}"
+        )
+
+@app.exception_handler(Exception)
+async def global_exception_handler(request: Request, exc: Exception):
+    """Global exception handler for unhandled errors."""
+    logger.error(f"Unhandled exception: {str(exc)}", exc_info=True)
+    return JSONResponse(
+        status_code=500,
+        content={
+            "detail": "Internal server error",
+            "error": str(exc),
+            "path": str(request.url)
+        }
+    )
+
+# Startup event
+@app.on_event("startup")
+async def startup_event():
+    """Log startup information."""
+    logger.info("=" * 60)
+    logger.info("🚀 Medical Coding API Starting...")
+    logger.info("=" * 60)
+    logger.info("⏳ Model will be loaded on first /predict request")
+    logger.info("📚 API Documentation: /docs")
+    logger.info("=" * 60)

app/model_loader.py

@@ -0,0 +1,133 @@
+# model_loader.py
+import os
+import sys
+from transformers import (
+    AutoTokenizer, 
+    AutoModelForCausalLM, 
+    AutoConfig,
+    pipeline
+)
+import torch
+import warnings
+
+warnings.filterwarnings("ignore")
+
+MODEL_NAME = "RayyanAhmed9477/med-coding"
+
+def load_model_and_tokenizer():
+    """
+    Loads Phi-3 model with comprehensive error handling and fallbacks.
+    Supports both CPU and GPU with automatic detection.
+    """
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"🔧 Using device: {device}")
+    print(f"🔧 PyTorch version: {torch.__version__}")
+    print(f"🔧 Transformers version: {sys.modules['transformers'].__version__}")
+    
+    # Get HuggingFace token from environment
+    hf_token = os.getenv("HF_TOKEN")
+    
+    try:
+        # ===== STEP 1: Load Tokenizer =====
+        print(f"📥 Loading tokenizer: {MODEL_NAME}")
+        tokenizer = AutoTokenizer.from_pretrained(
+            MODEL_NAME,
+            trust_remote_code=True,  # Critical for Phi-3
+            token=hf_token,
+            use_fast=True,
+            legacy=False
+        )
+        
+        # Configure tokenizer
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        if not hasattr(tokenizer, 'padding_side') or tokenizer.padding_side is None:
+            tokenizer.padding_side = "left"
+        
+        print("✅ Tokenizer loaded successfully")
+        
+        # ===== STEP 2: Load Configuration with trust_remote_code =====
+        print(f"📥 Loading model configuration: {MODEL_NAME}")
+        config = AutoConfig.from_pretrained(
+            MODEL_NAME,
+            trust_remote_code=True,  # Critical for Phi-3
+            token=hf_token
+        )
+        print(f"✅ Config loaded: {config.model_type}")
+        
+        # ===== STEP 3: Load Model =====
+        print(f"📥 Loading model: {MODEL_NAME}")
+        print("⏳ This may take 2-5 minutes on first load...")
+        
+        if device == "cuda":
+            # GPU Configuration
+            print("🎮 Using GPU with bfloat16 precision")
+            model = AutoModelForCausalLM.from_pretrained(
+                MODEL_NAME,
+                config=config,
+                trust_remote_code=True,
+                torch_dtype=torch.bfloat16,
+                device_map="auto",
+                token=hf_token,
+                low_cpu_mem_usage=True,
+                attn_implementation="eager"  # More stable than flash attention
+            )
+        else:
+            # CPU Configuration - optimized for stability
+            print("💻 Using CPU with float32 precision")
+            model = AutoModelForCausalLM.from_pretrained(
+                MODEL_NAME,
+                config=config,
+                trust_remote_code=True,
+                torch_dtype=torch.float32,
+                device_map={"": "cpu"},
+                token=hf_token,
+                low_cpu_mem_usage=True,
+                offload_folder="offload",
+                attn_implementation="eager"
+            )
+        
+        # Set model to evaluation mode
+        model.eval()
+        
+        # Disable gradients to save memory
+        for param in model.parameters():
+            param.requires_grad = False
+        
+        print("✅ Model loaded successfully!")
+        
+        # ===== STEP 4: Create Pipeline =====
+        print("🔧 Creating text generation pipeline...")
+        gen_pipeline = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            device=0 if device == "cuda" else -1,
+            torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
+            framework="pt"
+        )
+        
+        print("✅ Pipeline created successfully!")
+        print("=" * 60)
+        print("🎉 MODEL READY FOR INFERENCE")
+        print("=" * 60)
+        
+        return gen_pipeline, tokenizer
+        
+    except Exception as e:
+        print(f"❌ Error during model loading: {str(e)}")
+        print("\n🔍 Diagnostic Information:")
+        print(f"   - Model: {MODEL_NAME}")
+        print(f"   - Device: {device}")
+        print(f"   - Token available: {hf_token is not None}")
+        
+        import traceback
+        traceback.print_exc()
+        
+        raise RuntimeError(
+            f"Failed to load model {MODEL_NAME}. "
+            "Please check: "
+            "1) Internet connection, "
+            "2) HuggingFace token (if model is private), "
+            "3) Transformers version (requires >=4.36.0 for Phi-3)"
+        ) from e

app/prompt_template.py

@@ -0,0 +1,28 @@
+# prompt_template.py
+PROMPT_TEMPLATE = """<|system|>
+You are an expert medical coding assistant specialized in extracting ICD-10 diagnosis codes and CPT procedure codes from clinical notes.
+
+Your task:
+1. Analyze the clinical note carefully
+2. Extract all relevant ICD-10 codes (diagnosis codes)
+3. Extract all relevant CPT codes (procedure/service codes)
+4. Return ONLY valid medical codes found in the note
+5. Format your response as JSON with this exact structure:
+
+{{
+  "icd10_codes": ["code1", "code2"],
+  "cpt_codes": ["code1", "code2"]
+}}
+
+Rules:
+- Only include codes explicitly mentioned or clearly implied in the note
+- Use standard ICD-10 and CPT code formats
+- If no codes found, return empty arrays: {{"icd10_codes": [], "cpt_codes": []}}
+- Do not include explanations, only the JSON object
+<|end|>
+<|user|>
+Clinical Note:
+{note}
+<|end|>
+<|assistant|>
+"""

requirements.txt

@@ -0,0 +1,21 @@
+# Web Framework
+fastapi==0.109.2
+uvicorn[standard]==0.27.1
+python-multipart==0.0.9
+
+# Machine Learning - CRITICAL VERSIONS FOR PHI-3
+transformers==4.41.2
+torch==2.2.2
+accelerate==0.30.1
+safetensors==0.4.3
+sentencepiece==0.2.0
+
+# Utilities
+pydantic==2.7.1
+pydantic-settings==2.2.1
+python-dotenv==1.0.1
+protobuf==4.25.3
+einops==0.8.0
+
+# Monitoring
+psutil==5.9.8