Spaces:
Running
Running
File size: 2,586 Bytes
a7b7a62 | 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 | from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer
from typing import List, Union, Optional
import uvicorn
app = FastAPI(
title="BGE Embedding API",
description="Embedding service using BAAI/bge-small-en-v1.5 - optimized for retrieval tasks",
version="2.0.0"
)
# Load BGE model once at startup - better for semantic search
model = SentenceTransformer('BAAI/bge-small-en-v1.5')
# BGE model instruction prefix for queries
QUERY_PREFIX = "Represent this sentence for searching relevant passages: "
class EmbeddingRequest(BaseModel):
texts: Union[str, List[str]]
normalize: bool = True
is_query: bool = False # Set to True for search queries, False for documents
class EmbeddingResponse(BaseModel):
embeddings: List[List[float]]
model: str = "BAAI/bge-small-en-v1.5"
dimension: int = 384
@app.get("/")
def root():
return {
"message": "BGE Embedding API - Optimized for Retrieval",
"model": "BAAI/bge-small-en-v1.5",
"dimension": 384,
"features": [
"Superior semantic search performance",
"Asymmetric query-document encoding",
"Instruction-based query prefix support"
],
"endpoints": {
"embed": "/embed (POST)",
"health": "/health (GET)"
}
}
@app.get("/health")
def health_check():
return {"status": "healthy", "model_loaded": True, "model": "BAAI/bge-small-en-v1.5"}
@app.post("/embed", response_model=EmbeddingResponse)
def create_embeddings(request: EmbeddingRequest):
try:
# Handle single string or list of strings
texts = [request.texts] if isinstance(request.texts, str) else request.texts
# Add query prefix if this is a search query (not a document)
if request.is_query:
texts = [QUERY_PREFIX + text for text in texts]
# Generate embeddings
embeddings = model.encode(
texts,
normalize_embeddings=request.normalize,
show_progress_bar=False
)
# Convert to list format
embeddings_list = embeddings.tolist()
return EmbeddingResponse(
embeddings=embeddings_list,
dimension=len(embeddings_list[0])
)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860) |