testdummy.py

import os
import multiprocessing
import logging
from fastapi import FastAPI
from pydantic import BaseModel
from huggingface_hub import login
from sentence_transformers import SentenceTransformer

# ──────────────── Logging 설정 ────────────────
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ──────────────── Hugging Face 로그인 ────────────────
HF_API_TOKEN = os.getenv("HF_API_TOKEN")

if multiprocessing.current_process().name == "MainProcess":
    if HF_API_TOKEN and HF_API_TOKEN.startswith("hf_"):
        logger.info("🔑 Hugging Face API 로그인 중 (MainProcess)...")
        login(token=HF_API_TOKEN)
    else:
        logger.warning("⚠️ HF_API_TOKEN이 없거나 잘못된 형식입니다.")

# ──────────────── FastAPI 인스턴스 ────────────────
app = FastAPI(title="FastAPI Multi-worker Example")

# ──────────────── 임베딩 모델 lazy-load ────────────────
embedding_model = None

def get_embedding_model():
    global embedding_model
    if embedding_model is None:
        logger.info("📦 임베딩 모델 로드 중...")
        embedding_model = SentenceTransformer("jhgan/ko-sroberta-multitask")
        logger.info("✅ 임베딩 모델 로드 완료!")
    return embedding_model

# ──────────────── Pydantic 모델 ────────────────
class Query(BaseModel):
    text: str

# ──────────────── 테스트 엔드포인트 ────────────────
@app.post("/encode")
def encode_text(query: Query):
    model = get_embedding_model()
    embedding = model.encode(query.text).tolist()
    return {"embedding": embedding}