Mister2005/Cross-Encoder-Reranking-API
1
1"""2FinReg BGE Cross-Encoder Reranking API.3Hosted on Hugging Face Spaces (Mister2005/Cross-Encoder-Reranking-API).4 5Default Model: BAAI/bge-reranker-base (State-of-the-art Chinese & English Cross-Encoder)6Framework: FastAPI + SentenceTransformers / PyTorch CPU/GPU7"""8 9import os10import time11import torch12import uvicorn13from typing import List, Dict, Any, Optional14from fastapi import FastAPI, HTTPException15from fastapi.responses import HTMLResponse16from pydantic import BaseModel, Field17from sentence_transformers import CrossEncoder18 19MODEL_NAME = os.getenv("MODEL_NAME", "BAAI/bge-reranker-base")20DEVICE = "cuda" if torch.cuda.is_available() else "cpu"21 22app = FastAPI(23 title="FinReg BGE Reranker API",24 description="High-Precision Regulatory Document Re-Ranking powered by BAAI/bge-reranker-base.",25 version="2.0.0"26)27 28print(f"Loading CrossEncoder model '{MODEL_NAME}' on {DEVICE}...")29try:30 model = CrossEncoder(MODEL_NAME, max_length=512, device=DEVICE)31 print(f"Successfully initialized {MODEL_NAME}!")32except Exception as e:33 print(f"Error loading model: {e}")34 model = None35 36class RerankItem(BaseModel):37 rank: int38 original_index: int39 score: float40 document: str41 42class RerankRequest(BaseModel):43 query: str = Field(..., description="The search query or compliance question")44 documents: List[str] = Field(..., description="List of candidate text passages to re-rank")45 top_k: Optional[int] = Field(default=None, description="Number of top passages to return (defaults to all)")46 47class RerankResponse(BaseModel):48 query: str49 model: str50 total_evaluated: int51 latency_ms: float52 scores: List[float]53 ranked_indices: List[int]54 ranked_results: List[RerankItem]55 56@app.get("/", response_class=HTMLResponse)57def root_ui():58 return f"""59 <!DOCTYPE html>60 <html>61 <head>62 <title>FinReg BGE Reranker API</title>63 <meta charset="utf-8">64 <style>65 body {{ font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; max-width: 800px; margin: 40px auto; padding: 0 20px; color: #1e293b; background: #f8fafc; }}66 .card {{ background: white; padding: 30px; border-radius: 12px; box-shadow: 0 4px 6px -1px rgb(0 0 0 / 0.1); }}67 h1 {{ color: #0f172a; margin-top: 0; }}68 .badge {{ display: inline-block; padding: 4px 10px; border-radius: 9999px; font-size: 12px; font-weight: 600; background: #e0e7ff; color: #3730a3; }}69 .endpoint {{ background: #f1f5f9; padding: 12px; border-radius: 6px; font-family: monospace; margin: 12px 0; }}70 a.btn {{ display: inline-block; background: #2563eb; color: white; padding: 10px 18px; border-radius: 6px; text-decoration: none; font-weight: 500; margin-top: 15px; }}71 a.btn:hover {{ background: #1d4ed8; }}72 </style>73 </head>74 <body>75 <div class="card">76 <span class="badge">Active Microservice</span>77 <h1>⚖️ FinReg BGE Reranker API</h1>78 <p>Cloud cross-encoder service powering statutory retrieval re-ranking for Indian Regulatory Compliance.</p>79 80 <p><strong>Active Model:</strong> <code>{MODEL_NAME}</code> ({DEVICE.upper()})</p>81 <p><strong>Status:</strong> {'🟢 Online' if model else '🔴 Loading Error'}</p>82 83 <h3>API Endpoints:</h3>84 <div class="endpoint">POST /rerank (Standard JSON Payload)</div>85 <div class="endpoint">GET /health (Healthcheck)</div>86 <div class="endpoint">GET /docs (Interactive Swagger API Explorer)</div>87 88 <a class="btn" href="/docs">Open Interactive API Docs (Swagger) →</a>89 </div>90 </body>91 </html>92 """93 94@app.get("/health")95def health_check():96 return {97 "status": "healthy" if model else "unhealthy",98 "model": MODEL_NAME,99 "device": DEVICE,100 "model_loaded": model is not None101 }102 103@app.post("/rerank", response_model=RerankResponse)104def rerank_documents(request: RerankRequest):105 if not model:106 raise HTTPException(status_code=503, detail="Model is not loaded on server.")107 108 if not request.query.strip() or not request.documents:109 return RerankResponse(110 query=request.query,111 model=MODEL_NAME,112 total_evaluated=0,113 latency_ms=0.0,114 scores=[],115 ranked_indices=[],116 ranked_results=[]117 )118 119 start_time = time.time()120 try:121 # Create (query, doc) pairs122 pairs = [[request.query, doc] for doc in request.documents]123 124 # CrossEncoder scoring with sigmoid activation125 raw_scores = model.predict(pairs, convert_to_numpy=True, show_progress_bar=False)126 scores_list = [float(s) for s in raw_scores]127 128 # Sigmoid normalization: 1 / (1 + exp(-score))129 probs = [round(float(torch.sigmoid(torch.tensor(s)).item()), 4) for s in scores_list]130 131 # Rank pairs132 indexed = list(enumerate(probs))133 indexed.sort(key=lambda x: x[1], reverse=True)134 135 ranked_indices = [idx for idx, _ in indexed]136 137 top_k = request.top_k if request.top_k and request.top_k > 0 else len(request.documents)138 ranked_results = []139 for rank_num, (orig_idx, score) in enumerate(indexed[:top_k], 1):140 ranked_results.append(RerankItem(141 rank=rank_num,142 original_index=orig_idx,143 score=score,144 document=request.documents[orig_idx]145 ))146 147 latency = (time.time() - start_time) * 1000.0148 149 return RerankResponse(150 query=request.query,151 model=MODEL_NAME,152 total_evaluated=len(request.documents),153 latency_ms=round(latency, 2),154 scores=probs,155 ranked_indices=ranked_indices,156 ranked_results=ranked_results157 )158 except Exception as e:159 raise HTTPException(status_code=500, detail=str(e))160 161if __name__ == "__main__":162 uvicorn.run(app, host="0.0.0.0", port=7860)163