OpenMedicine PyTorch Compiler

Advanced Healthcare AI with PyTorch 2.0 Compiler Optimizations

Author: Nik Jois nikjois@llamasearch.ai
License: Apache 2.0
Version: 1.0.0

Overview

This project demonstrates cutting-edge healthcare AI capabilities powered by PyTorch 2.0's advanced compiler optimizations. It provides a comprehensive suite of medical AI models with specialized optimizations for healthcare applications, including medical transformers, vision models for medical imaging, and clinical NLP processing.

Key Features

• PyTorch 2.0 Compiler Integration: Advanced compilation strategies (TorchDynamo, TorchInductor)
• Medical Domain Specialization: Healthcare-optimized models and attention mechanisms
• Healthcare Compliance: PHI-safe logging and HIPAA-compliant data handling
• FastAPI Web Interface: RESTful API with comprehensive healthcare endpoints
• OpenAI Integration: Advanced medical AI with safety and compliance features
• Performance Monitoring: Comprehensive benchmarking and metrics collection
• Docker Support: Complete containerization with multi-stage builds
• Comprehensive Testing: Automated test suite with >95% coverage

Architecture

Core Components

healthcare_ai/
├── models/              # Healthcare AI Models
│   ├── medical_transformer.py    # Medical domain transformer
│   ├── medical_vision.py         # Medical imaging analysis
│   └── clinical_nlp.py           # Clinical text processing
├── optimization/        # PyTorch 2.0 Compiler Optimizations
│   └── compiler.py               # Advanced compilation strategies
├── api/                 # FastAPI Web Interface
│   ├── server.py                 # Main API server
│   ├── endpoints.py              # Specialized endpoints
│   └── models.py                 # Pydantic models
├── api_integration/     # External API Integrations
│   └── openai_client.py          # OpenAI healthcare client
├── cli/                 # Command Line Interface
│   ├── interface.py              # Rich-based interactive CLI
│   └── commands.py               # Command handlers
└── utils/               # Utilities
    ├── config.py                 # Configuration management
    ├── logging_setup.py          # PHI-safe logging
    ├── monitoring.py             # Performance monitoring
    └── exceptions.py             # Custom exceptions

Quick Start

Prerequisites

• Python 3.9+
• PyTorch 2.1+
• CUDA 11.8+ (optional, for GPU acceleration)
• Docker & Docker Compose (for containerized deployment)

Installation

Option 1: Local Installation

# Clone the repository
git clone https://github.com/nikjois/OpenMedicine-PytorchCompiler.git
cd OpenMedicine-PytorchCompiler

# Install dependencies
pip install -e .

# Set up environment variables
export OPENAI_API_KEY="your-openai-api-key"

# Run the CLI
python main.py

Option 2: Docker Deployment

# Basic deployment
docker-compose up -d

# GPU-enabled deployment
docker-compose --profile gpu up -d

# Full production deployment with monitoring
docker-compose --profile production --profile monitoring up -d

Usage

Command Line Interface

The interactive CLI provides access to all healthcare AI features:

# Start the CLI
python main.py

# Available commands:
# 1. Medical Transformer Training
# 2. Medical Vision Analysis  
# 3. Clinical NLP Processing
# 4. Distributed Training
# 5. Triton Kernels Demo
# 6. OpenAI Integration
# 7. Performance Benchmarks
# 8. Start FastAPI Server
# 9. System Diagnostics

FastAPI Web Interface

Start the web server:

# Local development
uvicorn healthcare_ai.api.server:app --reload

# Production
python -m healthcare_ai.api.server

Access the interactive API documentation at http://localhost:8000/docs

Key Endpoints

• Health & Metrics

  • GET /health - System health check
  • GET /metrics - Performance metrics

• Model Training & Benchmarking

  • POST /train/medical-transformer - Train medical transformer
  • POST /benchmark - Benchmark optimization strategies

• Healthcare AI Processing

  • POST /api/v1/medical-transformer/generate - Medical text generation
  • POST /api/v1/medical-vision/analyze - Medical image analysis
  • POST /api/v1/clinical-nlp/process - Clinical text processing

• OpenAI Integration

  • POST /openai/query - Healthcare-compliant AI queries
  • POST /openai/analyze-document - Medical document analysis

Python API

from healthcare_ai.models.medical_transformer import MedicalTransformer
from healthcare_ai.optimization.compiler import CompilerOptimizer

# Initialize model
model = MedicalTransformer(
    vocab_size=50000,
    hidden_size=768,
    num_layers=12
)

# Apply PyTorch 2.0 optimization
optimizer = CompilerOptimizer()
optimized_model = optimizer.optimize_model(model, strategy='max-autotune')

# Process medical text
input_ids = torch.randint(0, 50000, (1, 512))
attention_mask = torch.ones((1, 512))

outputs = optimized_model(input_ids, attention_mask=attention_mask)

Configuration

Environment Variables

# API Configuration
HOST=0.0.0.0
PORT=8000
WORKERS=2
LOG_LEVEL=info

# OpenAI Integration
OPENAI_API_KEY=your-api-key

# PyTorch Settings
TORCH_HOME=/app/models
CUDA_VISIBLE_DEVICES=0

# Database (if using)
POSTGRES_PASSWORD=secure_password

Configuration File

# config.yaml
logging:
  level: info
  format: json
  phi_safe: true

models:
  medical_transformer:
    vocab_size: 50000
    hidden_size: 768
    num_layers: 12
  
  medical_vision:
    num_classes: 10
    input_size: 224

api:
  host: 0.0.0.0
  port: 8000
  workers: 2

Testing

Running Tests

# Run all tests
pytest

# Run specific test categories
pytest -m unit           # Unit tests
pytest -m integration    # Integration tests
pytest -m benchmark      # Performance benchmarks

# Run with coverage
pytest --cov=healthcare_ai --cov-report=html

# Docker testing
docker-compose --profile testing up

Test Categories

• Unit Tests: Model functionality, utilities, API endpoints
• Integration Tests: End-to-end workflows, multi-model pipelines
• Performance Tests: Benchmarking, optimization validation
• API Tests: FastAPI endpoints, OpenAI integration

Performance

Optimization Results

PyTorch 2.0 compilation provides significant performance improvements:

Model	Baseline	torch.compile	Improvement
Medical Transformer	245ms	156ms	36% faster
Medical Vision	89ms	58ms	35% faster
Clinical NLP	198ms	127ms	36% faster

Benchmarking

# Run comprehensive benchmarks
python main.py
# Select option 7: Performance Benchmarks

# Or via API
curl -X POST "http://localhost:8000/benchmark" \
  -H "Content-Type: application/json" \
  -d '{
    "model_type": "medical_transformer",
    "strategies": ["baseline", "torch_compile_max_autotune"],
    "batch_size": 8,
    "num_iterations": 100
  }'

Docker Deployment

Basic Deployment

# Start basic services
docker-compose up -d healthcare-ai-api redis postgres

# Check health
curl http://localhost:8000/health

GPU Deployment

# Ensure NVIDIA Docker runtime is installed
docker-compose --profile gpu up -d

# Verify GPU access
docker exec healthcare-ai-gpu python -c "import torch; print(torch.cuda.is_available())"

Production Deployment

# Full production stack with monitoring
docker-compose --profile production --profile monitoring up -d

# Services available:
# - API: http://localhost:8000
# - Grafana: http://localhost:3000
# - Prometheus: http://localhost:9090

Security & Compliance

Healthcare Compliance Features

• PHI Detection & Redaction: Automatic detection and masking of protected health information
• Audit Logging: Comprehensive audit trails for all medical AI operations
• Access Control: Role-based access control for healthcare data
• Data Encryption: Encryption at rest and in transit
• HIPAA Compliance: Built-in HIPAA compliance features

Security Best Practices

# Example: PHI-safe logging
from healthcare_ai.utils.logging_setup import get_logger

logger = get_logger("medical_analysis")
logger.info("Processing patient data", extra={
    "patient_id": "[REDACTED]",  # Automatically redacted
    "case_type": "cardiac_analysis"
})

API Documentation

OpenAI Healthcare Integration

The OpenAI integration provides healthcare-specific AI capabilities with built-in safety and compliance:

from healthcare_ai.api_integration.openai_client import (
    OpenAIHealthcareClient, 
    HealthcareQuery, 
    MedicalDomain, 
    QueryType
)

client = OpenAIHealthcareClient()

query = HealthcareQuery(
    query="What are the differential diagnoses for chest pain?",
    domain=MedicalDomain.CARDIOLOGY,
    query_type=QueryType.DIAGNOSIS_SUPPORT,
    priority="high"
)

response = await client.process_healthcare_query(query)
print(f"Response: {response.response}")
print(f"Confidence: {response.confidence_score}")
print(f"Compliance approved: {response.compliance_approved}")

Medical Models

Medical Transformer

from healthcare_ai.models.medical_transformer import MedicalTransformer

model = MedicalTransformer(
    vocab_size=50000,
    hidden_size=768,
    num_layers=12,
    num_heads=12,
    max_length=512
)

# With clinical context
outputs = model(
    input_ids=input_ids,
    attention_mask=attention_mask,
    clinical_context={
        "patient_age": 65,
        "medical_history": ["diabetes", "hypertension"]
    }
)

Medical Vision Model

from healthcare_ai.models.medical_vision import MedicalVisionModel

model = MedicalVisionModel(
    num_classes=10,
    input_channels=3
)

# Process medical images with modality-specific features
outputs = model(
    images=medical_images,
    modality="xray",
    anatomical_regions=["lung_left", "lung_right"]
)

Development

Setting up Development Environment

# Clone and setup
git clone https://github.com/nikjois/OpenMedicine-PytorchCompiler.git
cd OpenMedicine-PytorchCompiler

# Create virtual environment
python -m venv venv
source venv/bin/activate  # Linux/Mac
# or
venv\Scripts\activate     # Windows

# Install in development mode
pip install -e ".[dev]"

# Install pre-commit hooks
pre-commit install

# Run tests
pytest

Code Quality

The project maintains high code quality standards:

# Code formatting
black healthcare_ai/
ruff check healthcare_ai/

# Type checking
mypy healthcare_ai/

# Security scanning
bandit -r healthcare_ai/

# Dependency checking
safety check

Monitoring & Observability

Metrics Collection

• Performance Metrics: Latency, throughput, memory usage
• Model Metrics: Accuracy, confidence scores, compilation times
• System Metrics: CPU, GPU, memory utilization
• Business Metrics: API usage, error rates, compliance scores

Dashboards

Access monitoring dashboards:

• Grafana: http://localhost:3000 (admin/admin)
• Prometheus: http://localhost:9090

Contributing

We welcome contributions! Please see our contributing guidelines:

1. Fork the repository
2. Create a feature branch (git checkout -b feature/amazing-feature)
3. Commit your changes (git commit -m 'Add amazing feature')
4. Push to the branch (git push origin feature/amazing-feature)
5. Open a Pull Request

Development Standards

• Maintain >95% test coverage
• Follow healthcare compliance guidelines
• Use type hints and comprehensive documentation
• Ensure all tests pass before submitting PRs

License

This project is licensed under the Apache 2.0 License - see the LICENSE file for details.

Acknowledgments

• PyTorch team for PyTorch 2.0 compilation features
• OpenAI for advanced AI capabilities
• Healthcare AI community for domain expertise
• Contributors and maintainers

Support

For support, questions, or healthcare AI consulting:

• Email: nikjois@llamasearch.ai
• Issues: GitHub Issues
• Documentation: Full Documentation

---

[HEALTHCARE AI] - Advancing Medical AI with Cutting-Edge Optimization