Create app.py

app.py

@@ -0,0 +1,177 @@
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import JSONResponse
+import numpy as np
+from sentence_transformers import SentenceTransformer
+import asyncio
+import logging
+from typing import List
+from pydantic import BaseModel
+import time
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+app = FastAPI(title="Embedding API")
+
+# Load model on startup
+try:
+    model = SentenceTransformer("all-MiniLM-L12-v2")
+    logger.info("Model loaded successfully")
+except Exception as e:
+    logger.error(f"Failed to load model: {e}")
+    model = None
+
+# Request batching queue
+class EmbedRequest(BaseModel):
+    texts: List[str]
+
+class BatchItem:
+    def __init__(self, texts: List[str]):
+        self.texts = texts
+        self.future: asyncio.Future = asyncio.Future()
+
+pending_batch: List[BatchItem] = []
+batch_lock = asyncio.Lock()
+batch_event = asyncio.Event()
+
+BATCH_TIMEOUT = 0.05  # 50ms window to collect requests
+MAX_BATCH_SIZE = 32
+
+
+async def batch_processor():
+    """Continuously process batches of requests"""
+    global pending_batch
+    
+    while True:
+        try:
+            # Wait for requests or timeout
+            try:
+                await asyncio.wait_for(batch_event.wait(), timeout=BATCH_TIMEOUT)
+                batch_event.clear()
+            except asyncio.TimeoutError:
+                pass
+            
+            async with batch_lock:
+                if not pending_batch:
+                    continue
+                
+                # Take up to MAX_BATCH_SIZE items
+                batch_to_process = pending_batch[:MAX_BATCH_SIZE]
+                pending_batch = pending_batch[MAX_BATCH_SIZE:]
+                
+                if not batch_to_process:
+                    continue
+            
+            # Flatten all texts
+            all_texts = []
+            text_counts = []
+            for item in batch_to_process:
+                all_texts.extend(item.texts)
+                text_counts.append(len(item.texts))
+            
+            logger.info(f"Processing batch of {len(batch_to_process)} requests, {len(all_texts)} texts total")
+            
+            # Compute embeddings
+            start = time.time()
+            embeddings = model.encode(all_texts, convert_to_numpy=True)
+            elapsed = time.time() - start
+            logger.info(f"Embedding computed in {elapsed:.2f}s")
+            
+            # Split embeddings back to individual requests
+            idx = 0
+            for item, count in zip(batch_to_process, text_counts):
+                item_embeddings = embeddings[idx:idx+count].tolist()
+                item.future.set_result(item_embeddings)
+                idx += count
+        
+        except Exception as e:
+            logger.error(f"Error in batch processor: {e}")
+            async with batch_lock:
+                for item in pending_batch:
+                    if not item.future.done():
+                        item.future.set_exception(e)
+                pending_batch.clear()
+            await asyncio.sleep(1)
+
+
+@app.on_event("startup")
+async def startup():
+    """Start batch processor on app startup"""
+    if model is None:
+        raise RuntimeError("Model failed to load")
+    asyncio.create_task(batch_processor())
+    logger.info("Batch processor started")
+
+
+@app.get("/health")
+async def health():
+    """Health check endpoint"""
+    if model is None:
+        return JSONResponse({"status": "unhealthy", "reason": "model not loaded"}, status_code=503)
+    return {"status": "healthy"}
+
+
+@app.post("/embed")
+async def embed(request: EmbedRequest):
+    """
+    Embed texts. Automatically batches with concurrent requests.
+    
+    Example:
+    ```
+    POST /embed
+    {
+        "texts": ["hello world", "foo bar"]
+    }
+    ```
+    """
+    if model is None:
+        raise HTTPException(status_code=503, detail="Model not ready")
+    
+    if not request.texts:
+        raise HTTPException(status_code=400, detail="texts cannot be empty")
+    
+    if len(request.texts) > 512:
+        raise HTTPException(status_code=400, detail="Cannot embed more than 512 texts at once")
+    
+    # Create request item
+    batch_item = BatchItem(request.texts)
+    
+    async with batch_lock:
+        pending_batch.append(batch_item)
+        
+        # Signal processor if we hit batch size
+        if len(pending_batch) >= MAX_BATCH_SIZE:
+            batch_event.set()
+    
+    # Signal processor that there's work (in case it's idle)
+    batch_event.set()
+    
+    # Wait for result with timeout
+    try:
+        embeddings = await asyncio.wait_for(batch_item.future, timeout=30.0)
+        return {
+            "embeddings": embeddings,
+            "model": "all-MiniLM-L12-v2",
+            "count": len(request.texts)
+        }
+    except asyncio.TimeoutError:
+        raise HTTPException(status_code=504, detail="Embedding request timed out")
+
+
+@app.get("/")
+async def root():
+    """API info"""
+    return {
+        "name": "Embedding API",
+        "model": "all-MiniLM-L12-v2",
+        "endpoints": {
+            "POST /embed": "Embed texts",
+            "GET /health": "Health check"
+        }
+    }
+
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)