Update app.py

app.py

@@ -1,54 +1,165 @@
 """
-Hugging Face Spaces에서 돌리는 Reranker API.
-모델은 HF Hub에서 로드 (환경변수 HF_MODEL_ID, 예: username/certweb-reranker).
-무료 CPU Basic(16GB RAM)에서 동작. 슬립 후 첫 요청은 콜드스타트 30초~1분 걸릴 수 있음.
+Hugging Face Spaces에서 돌리는 Reranker API (ONNX 최적화 버전).
+
+모델은 HF Hub에서 로드 (환경변수 HF_MODEL_ID, 예: cross-encoder/ms-marco-MiniLM-L-6-v2).
+ONNX Runtime으로 변환하여 CPU 추론 속도 2-3배 향상.
+
+환경변수:
+- HF_MODEL_ID: 모델 ID (필수)
+- HF_TOKEN: 비공개 리포 접근용 토큰 (선택)
+- USE_ONNX: "true"면 ONNX Runtime 사용 (기본: true)
+- ONNX_PROVIDER: "CPUExecutionProvider" 또는 "CUDAExecutionProvider" (기본: CPU)
 """
 import os
+import time
+from typing import List, Optional
 
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
 
 MODEL_ID = (os.environ.get("HF_MODEL_ID") or "").strip()
-HF_TOKEN = (os.environ.get("HF_TOKEN") or "").strip() or None  # 비공개 리포일 때만 설정
+HF_TOKEN = (os.environ.get("HF_TOKEN") or "").strip() or None
+USE_ONNX = os.environ.get("USE_ONNX", "true").lower() == "true"
+ONNX_PROVIDER = os.environ.get("ONNX_PROVIDER", "CPUExecutionProvider")
+
 model = None
+tokenizer = None
+is_onnx_mode = False
 
-app = FastAPI(title="Reranker API")
+app = FastAPI(title="Reranker API (ONNX)")
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
 
 
+class RerankerRequest(BaseModel):
+    query: str
+    passages: List[str]
+
+
+class RerankerResponse(BaseModel):
+    scores: List[float]
+    latency_ms: Optional[float] = None
+    mode: Optional[str] = None
+    error: Optional[str] = None
+
+
+def load_onnx_model():
+    """ONNX Runtime으로 모델 로드. 최초 실행 시 자동 변환."""
+    global model, tokenizer, is_onnx_mode
+    try:
+        from optimum.onnxruntime import ORTModelForSequenceClassification
+        from transformers import AutoTokenizer
+        
+        print(f"[ONNX] 모델 로딩 중: {MODEL_ID} (provider={ONNX_PROVIDER})")
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
+        model = ORTModelForSequenceClassification.from_pretrained(
+            MODEL_ID,
+            export=True,  # PyTorch → ONNX 자동 변환
+            provider=ONNX_PROVIDER,
+            token=HF_TOKEN,
+        )
+        is_onnx_mode = True
+        print("[ONNX] 로드 완료.")
+    except Exception as e:
+        print(f"[ONNX] 로드 실패, fallback to CrossEncoder: {e}")
+        load_crossencoder_model()
+
+
+def load_crossencoder_model():
+    """기존 CrossEncoder 모델 로드 (fallback)."""
+    global model, is_onnx_mode
+    from sentence_transformers import CrossEncoder
+    print(f"[CrossEncoder] 모델 로딩 중: {MODEL_ID}")
+    model = CrossEncoder(MODEL_ID, token=HF_TOKEN)
+    is_onnx_mode = False
+    print("[CrossEncoder] 로드 완료.")
+
+
 @app.on_event("startup")
-def load_model():
+def startup_load_model():
     global model
     if not MODEL_ID:
-        raise RuntimeError("Space 설정에서 HF_MODEL_ID 를 지정하세요 (예: username/certweb-reranker)")
-    from sentence_transformers import CrossEncoder
-    print(f"모델 로딩 중: {MODEL_ID}")
-    # 비공개 리포면 HF_TOKEN (Secrets) 설정 필요
-    model = CrossEncoder(MODEL_ID, token=HF_TOKEN)
-    print("로드 완료.")
+        raise RuntimeError("Space 설정에서 HF_MODEL_ID 를 지정하세요")
+    
+    if USE_ONNX:
+        load_onnx_model()
+    else:
+        load_crossencoder_model()
+
+
+def predict_onnx(query: str, passages: List[str]) -> List[float]:
+    """ONNX 모델로 batch predict."""
+    import torch
+    pairs = [[query, p] for p in passages]
+    inputs = tokenizer(
+        pairs,
+        padding=True,
+        truncation=True,
+        max_length=512,
+        return_tensors="pt",
+    )
+    outputs = model(**inputs)
+    # logits shape: (batch_size, num_labels) - 보통 1개 label (relevance score)
+    logits = outputs.logits
+    if logits.shape[-1] == 1:
+        scores = logits.squeeze(-1)
+    else:
+        # 2개 label이면 softmax 후 positive class 확률
+        scores = torch.softmax(logits, dim=-1)[:, 1]
+    return scores.tolist()
+
+
+def predict_crossencoder(query: str, passages: List[str]) -> List[float]:
+    """CrossEncoder 모델로 predict."""
+    pairs = [(query, p) for p in passages]
+    scores = model.predict(pairs)
+    if hasattr(scores, "tolist"):
+        scores = scores.tolist()
+    return [float(s) for s in scores]
 
 
 @app.get("/")
 def root():
-    """브라우저/App 탭 접속 시 안내. 실제 점수는 POST / 로 요청."""
+    """브라우저/App 탭 접속 시 안내."""
     return {
-        "message": "CertWeb Reranker API",
+        "message": "CertWeb Reranker API (ONNX Optimized)",
+        "model": MODEL_ID,
+        "mode": "onnx" if is_onnx_mode else "crossencoder",
         "usage": "POST / with JSON body: {\"query\": \"질문\", \"passages\": [\"문단1\", \"문단2\", ...]}",
-        "response": "{\"scores\": [0.9, 0.2, ...]}",
+        "response": "{\"scores\": [0.9, 0.2, ...], \"latency_ms\": 123.4, \"mode\": \"onnx\"}",
     }
 
 
-@app.post("/")
-def score(request: dict):
-    query = request.get("query") or ""
-    passages = request.get("passages") or []
+@app.post("/", response_model=RerankerResponse)
+def score(request: RerankerRequest):
+    query = request.query
+    passages = request.passages
+    
     if not query or not passages or model is None:
-        return {"scores": []}
+        return RerankerResponse(scores=[], error="Empty query/passages or model not loaded")
+    
     try:
-        pairs = [(query, p) for p in passages]
-        scores = model.predict(pairs)
-        if hasattr(scores, "tolist"):
-            scores = scores.tolist()
-        return {"scores": [float(s) for s in scores]}
+        start = time.perf_counter()
+        if is_onnx_mode:
+            scores = predict_onnx(query, passages)
+        else:
+            scores = predict_crossencoder(query, passages)
+        latency_ms = (time.perf_counter() - start) * 1000
+        
+        return RerankerResponse(
+            scores=scores,
+            latency_ms=round(latency_ms, 2),
+            mode="onnx" if is_onnx_mode else "crossencoder",
+        )
     except Exception as e:
-        return {"scores": [], "error": str(e)}
+        return RerankerResponse(scores=[], error=str(e))
+
+
+@app.get("/health")
+def health():
+    """Health check endpoint."""
+    return {
+        "status": "ok",
+        "model_loaded": model is not None,
+        "mode": "onnx" if is_onnx_mode else "crossencoder",
+    }