Spaces:
Sleeping
Sleeping
File size: 4,271 Bytes
4258647 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | # main.py
from __future__ import annotations
import asyncio
import time
from typing import List, Optional, Dict, Any
from fastapi import FastAPI, HTTPException, Depends
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field, validator
from get_embedding import EmbeddingFetcher
# -----------------------------
# Request / Response Schemas
# -----------------------------
class TextListRequest(BaseModel):
texts: List[str] = Field(..., description="List of strings to embed", min_items=1)
@validator("texts")
def non_empty_texts(cls, v: List[str]) -> List[str]:
if any((t is None or not isinstance(t, str) or t.strip() == "") for t in v):
raise ValueError("All items in 'texts' must be non-empty strings.")
return v
class EmbeddingResponse(BaseModel):
model_id: str
device: str
dims: int
count: int
elapsed_ms: float
embeddings: List[List[float]]
# -----------------------------
# App factory with lifespan
# -----------------------------
def create_app() -> FastAPI:
app = FastAPI(title="Embedding API", version="1.0.0")
# CORS: keep your original open policy (tighten in production)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Global container for services
app.state.container: Dict[str, Any] = {}
app.state.init_lock = asyncio.Lock()
@app.on_event("startup")
async def on_startup() -> None:
# Initialize the EmbeddingFetcher once, asynchronously
async with app.state.init_lock:
if "embedder" not in app.state.container:
fetcher = EmbeddingFetcher()
# Build models / download snapshots off the main thread
await fetcher.ensure_ready()
app.state.container["embedder"] = fetcher
@app.on_event("shutdown")
async def on_shutdown() -> None:
# Nothing special is required, but the hook is here for future cleanup
pass
# ---------------
# Dependencies
# ---------------
def get_embedder() -> EmbeddingFetcher:
fetcher: Optional[EmbeddingFetcher] = app.state.container.get("embedder")
if fetcher is None:
# Defensive: if a request sneaks in before startup finishes
raise HTTPException(status_code=503, detail="Service not ready. Try again shortly.")
return fetcher
# ---------------
# Routes
# ---------------
@app.get("/", tags=["meta"])
async def home():
return {"status": "ok", "message": "Embedding service is running."}
@app.get("/healthz", tags=["meta"])
async def healthz():
# Lightweight health; could add a test encode if you want deeper checks
return {"status": "healthy"}
@app.post("/get-embedding/", response_model=EmbeddingResponse, tags=["embedding"])
async def get_embedding(request: TextListRequest, embedder: EmbeddingFetcher = Depends(get_embedder)):
# Offload embedding to the service (async wrapper over blocking HF/Torch calls)
start = time.perf_counter()
try:
vectors = await embedder.embed(request.texts)
except ValueError as ve:
raise HTTPException(status_code=400, detail=str(ve)) from ve
except Exception as e:
raise HTTPException(status_code=500, detail=f"Embedding failed: {e}") from e
elapsed_ms = (time.perf_counter() - start) * 1000.0
dims = len(vectors[0]) if vectors and len(vectors[0]) else 0
return EmbeddingResponse(
model_id=embedder.model_id,
device=embedder.device_str,
dims=dims,
count=len(vectors),
elapsed_ms=round(elapsed_ms, 3),
embeddings=vectors,
)
return app
app = create_app()
# Optional: run via `python main.py` in development
if __name__ == "__main__":
import uvicorn
uvicorn.run(
"main:app",
host="0.0.0.0",
port=8000,
reload=True, # Turn off in production
workers=1, # Use a process manager (e.g., gunicorn) to scale
log_level="info",
)
|