| """ |
| Hugging Face Spaces (Docker) — Contrastive 질의 임베딩 API. |
| |
| 모델은 HF Hub에서 로드 (환경변수 HF_MODEL_ID, 예: <org>/<model-id>). |
| POST로 텍스트 받아 768-dim 벡터 반환. 백엔드 RAG_CONTRASTIVE_EMBEDDING_URL 연동용. |
| |
| 환경변수: |
| - HF_MODEL_ID: SentenceTransformer 모델 ID (필수) |
| - HF_TOKEN: 비공개 리포 접근용 토큰 (선택) |
| """ |
| import os |
| from typing import List, Union |
|
|
| 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 |
|
|
| model = None |
|
|
| app = FastAPI(title="Contrastive Embedding API") |
| app.add_middleware( |
| CORSMiddleware, |
| allow_origins=["*"], |
| allow_methods=["*"], |
| allow_headers=["*"], |
| ) |
|
|
|
|
| class EmbedRequest(BaseModel): |
| """백엔드 contrastive_retriever가 보내는 형식.""" |
| inputs: str |
|
|
|
|
| def load_model() -> None: |
| global model |
| from sentence_transformers import SentenceTransformer |
|
|
| print(f"[Embed] 모델 로딩: {MODEL_ID}") |
| model = SentenceTransformer(MODEL_ID, token=HF_TOKEN) |
| print("[Embed] 로드 완료.") |
|
|
|
|
| @app.on_event("startup") |
| def startup() -> None: |
| if MODEL_ID: |
| load_model() |
| else: |
| print("[Embed] HF_MODEL_ID 미설정. 설정 후 재시작.") |
|
|
|
|
| @app.post("/") |
| def embed(body: EmbedRequest) -> Union[List[List[float]], dict]: |
| """ |
| 질의 문자열을 768-dim 벡터로 변환. |
| 응답: [[float,...]] (백엔드가 기대하는 형식) 또는 {"embedding": [float,...]} |
| """ |
| global model |
| if not model: |
| return {"embedding": [], "error": "HF_MODEL_ID not set"} |
|
|
| text = (body.inputs or "").strip() |
| if not text: |
| return {"embedding": [], "error": "empty inputs"} |
|
|
| vec = model.encode([text], normalize_embeddings=True) |
| |
| return vec.tolist() |
|
|
|
|
| @app.get("/") |
| def root(): |
| """ |
| Hugging Face Spaces / 컨테이너 헬스체크용 루트 엔드포인트. |
| |
| - GET / → 200 OK |
| - GET /?logs=... → 200 OK (path는 / 그대로, query만 다름) |
| - UptimeRobot HTTP(S) 모니터의 HEAD / 도 FastAPI가 자동 처리 |
| """ |
| return { |
| "status": "ok", |
| "service": "contrastive-embedding", |
| "model_loaded": model is not None, |
| } |
|
|
|
|
| @app.get("/health") |
| def health(): |
| """단순 헬스 체크 엔드포인트.""" |
| return {"status": "ok", "model_loaded": model is not None} |
