app/server_clinical_embedding.py

from typing import List, Optional
from fastapi import FastAPI, Query, UploadFile, File, HTTPException
from fastapi.responses import RedirectResponse
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
import uvicorn
import io
import csv
import os

from clinical_embedding import ModelManager

# Pydantic models for request/response
class EmbeddingRequest(BaseModel):
    sentences: List[str]
    pooling: str = 'cls'
    model: str = 'clinical_bert'

class EmbeddingResponse(BaseModel):
    embeddings: List[List[float]]
    shape: List[int]
    pooling: str
    model: str

# Initialize FastAPI app
app = FastAPI(
    title="Clinical Embedding API",
    description="API for generating embeddings using various models (ClinicalBERT, BERT, Word2Vec)",
    version="2.0.0"
)

# Add CORS middleware to allow web page access
app.add_middleware(
    CORSMiddleware,
    allow_origins=["https://santanche-clinical-embedding.hf.space"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Serve static files
app.mount("/app/static", StaticFiles(directory="static"), name="static")

# Initialize model manager (global instance)
model_manager = ModelManager()

@app.on_event("startup")
async def startup_event():
    """
    Load default model on startup.
    Other models will be loaded on demand (see ModelManager).
    """
    # Pre-load ClinicalBERT as it's the default
    model_manager.get_model('clinical_bert')

@app.get("/")
async def root():
    return RedirectResponse(url="/browser/")

@app.get("/browser/")
def get_browser():
    return FileResponse(os.path.join("static", "browser", "index.html"))

@app.get("/embeddings", response_model=EmbeddingResponse)
async def get_embeddings(
    sentences: List[str] = Query(..., description="List of sentences to embed"),
    pooling: str = Query('cls', description="Pooling strategy: mean, cls, or max"),
    model: str = Query('clinical_bert', description="Model to use: clinical_bert, bert, word2vec")
):
    """
    Generate embeddings for a list of sentences.
    Supports bracketed text for context-aware specific extraction.
    """
    # Validate pooling method
    if pooling not in ['mean', 'cls', 'max']:
        raise HTTPException(status_code=400, detail="Invalid pooling method. Choose from: mean, cls, max")
    
    try:
        embedder = model_manager.get_model(model)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    
    # Generate embeddings
    embeddings = embedder.get_embeddings(sentences, pooling=pooling)
    
    # Convert to list for JSON serialization
    embeddings_list = embeddings.tolist()
    
    return EmbeddingResponse(
        embeddings=embeddings_list,
        shape=list(embeddings.shape),
        pooling=pooling,
        model=model
    )

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "loaded_models": list(model_manager.models.keys())
    }

@app.post("/embeddings/batch")
async def post_embeddings_batch(request: EmbeddingRequest):
    """
    POST endpoint for batch embedding generation.
    """
    # Validate pooling method
    if request.pooling not in ['mean', 'cls', 'max']:
        raise HTTPException(status_code=400, detail="Invalid pooling method. Choose from: mean, cls, max")
    
    try:
        embedder = model_manager.get_model(request.model)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    
    # Generate embeddings
    embeddings = embedder.get_embeddings(request.sentences, pooling=request.pooling)
    
    # Convert to list for JSON serialization
    embeddings_list = embeddings.tolist()
    
    return EmbeddingResponse(
        embeddings=embeddings_list,
        shape=list(embeddings.shape),
        pooling=request.pooling,
        model=request.model
    )

@app.post("/embeddings/file")
async def upload_file_embeddings(
    file: UploadFile = File(...),
    pooling: str = Query('cls', description="Pooling strategy: mean, cls, or max"),
    model: str = Query('clinical_bert', description="Model to use: clinical_bert, bert, word2vec")
):
    """
    Upload a CSV file with terms and get embeddings back as CSV.
    """
    # Validate file type
    if not file.filename.endswith('.csv'):
        raise HTTPException(status_code=400, detail="File must be a CSV")
    
    # Validate pooling method
    if pooling not in ['mean', 'cls', 'max']:
        raise HTTPException(status_code=400, detail="Invalid pooling method. Choose from: mean, cls, max")
    
    try:
        embedder = model_manager.get_model(model)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    
    try:
        # Read CSV file
        contents = await file.read()
        csv_reader = csv.DictReader(io.StringIO(contents.decode('utf-8')))
        
        # Get column name (first column)
        fieldnames = csv_reader.fieldnames
        if not fieldnames or len(fieldnames) == 0:
            raise HTTPException(status_code=400, detail="CSV must have at least one column")
        
        column_name = fieldnames[0]
        
        # Extract terms
        terms = [row[column_name] for row in csv_reader if row[column_name].strip()]
        
        if not terms:
            raise HTTPException(status_code=400, detail="No terms found in CSV")
        
        # Generate embeddings
        embeddings = embedder.get_embeddings(terms, pooling=pooling)
        
        # Create output CSV
        output = io.StringIO()
        writer = csv.writer(output)
        
        # Write header (term + embedding dimensions)
        header = [column_name] + [f"dim_{i}" for i in range(embeddings.shape[1])]
        writer.writerow(header)
        
        # Write rows
        for term, embedding in zip(terms, embeddings):
            row = [term] + embedding.tolist()
            writer.writerow(row)
        
        # Prepare response
        output.seek(0)
        return StreamingResponse(
            io.BytesIO(output.getvalue().encode('utf-8')),
            media_type="text/csv",
            headers={"Content-Disposition": f"attachment; filename=embeddings_{file.filename}"}
        )
    
    except UnicodeDecodeError:
        raise HTTPException(status_code=400, detail="File must be UTF-8 encoded")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error processing file: {str(e)}")

if __name__ == "__main__":
    # Run the server
    uvicorn.run(
        "server_clinical_embedding:app",
        host="0.0.0.0",
        port=8000,
        reload=False
    )