Initial codes commit

app.py

@@ -4,14 +4,16 @@ from fastapi import FastAPI, Request, Form
 from fastapi.responses import HTMLResponse
 from fastapi.templating import Jinja2Templates
 from api.endpoints import router
-from db.initializer import initialize
+from db.initializer import initialize_dbs
+from models.initializer import initialize_models
 from service.search import search
 
 templates = Jinja2Templates(directory="templates")
 
 @asynccontextmanager
-async def lifespan(_app: FastAPI):
-    initialize()
+async def lifespan(app: FastAPI):
+    initialize_dbs()
+    initialize_models(app)
     yield
 
 app = FastAPI(lifespan=lifespan)

config.py

@@ -5,7 +5,7 @@ import os
 HF_TOKEN = os.getenv("HF_TOKEN")
 
 # HF datasets repo info
-HF_REPO_ID = os.getenv("HF_REPO_ID", "m97j/pls-datasets")
+HF_DS_REPO_ID = os.getenv("HF_REPO_ID", "m97j/pls-datasets")
 HF_INDEX_FILE = os.getenv("HF_INDEX_FILE", "faiss/faiss_index_flat.faiss")
 HF_IDS_FILE = os.getenv("HF_IDS_FILE", "faiss/vector_ids.npy")
 
@@ -18,8 +18,16 @@ HF_CORPUS_SPLIT = os.getenv("HF_CORPUS_SPLIT", "train")
 MARKER_DIR = os.getenv("MARKER_DIR", "rag/state")
 CORPUS_READY_MARK = os.path.join(MARKER_DIR, ".corpus_ready")
 
-# Embedding / LLM model
-EMBED_MODEL = os.getenv("EMBED_MODEL", "intfloat/multilingual-e5-large")
+# Embedding model
+HF_MODEL_REPO_ID = os.getenv("HF_MODEL_REPO_ID", "m97j/pragmatic-search")
+EMBED_MODEL = os.getenv("EMBED_MODEL", "model_quantized.onnx")
+EMBED_DIR = os.getenv("EMBED_DIR", "embedder")
+
+# Reranking model
+RERANK_MODEL = os.getenv("RERANK_MODEL", "model_quantized.onnx")
+RERANK_DIR = os.getenv("RERANK_DIR", "reranker")
+
+# Retrieval settings
 TOP_K = int(os.getenv("TOP_K", "5"))
 
 

db/initializer.py

@@ -2,7 +2,7 @@
 import faiss
 import numpy as np
 from huggingface_hub import hf_hub_download
-from config import HF_REPO_ID, HF_INDEX_FILE, HF_IDS_FILE
+from config import HF_DS_REPO_ID, HF_INDEX_FILE, HF_IDS_FILE
 from modules.retriever import set_index
 from modules import corpus
 
@@ -10,8 +10,8 @@ _vector_ids = None
 
 def _load_index_in_memory():
     """HF Hub에서 인덱스/ID 매핑을 받아 메모리에 로드"""
-    index_path = hf_hub_download(repo_id=HF_REPO_ID, filename=HF_INDEX_FILE, repo_type="dataset")
-    ids_path = hf_hub_download(repo_id=HF_REPO_ID, filename=HF_IDS_FILE, repo_type="dataset")
+    index_path = hf_hub_download(repo_id=HF_DS_REPO_ID, filename=HF_INDEX_FILE, repo_type="dataset")
+    ids_path = hf_hub_download(repo_id=HF_DS_REPO_ID, filename=HF_IDS_FILE, repo_type="dataset")
     index = faiss.read_index(index_path)
     set_index(index)
     global _vector_ids
@@ -21,7 +21,7 @@ def get_vector_ids():
     global _vector_ids
     return _vector_ids
 
-def initialize():
+def initialize_dbs():
     # 1) 코퍼스 준비 (최초 1회만 다운로드)
     corpus.prepare_corpus()
     # 2) 인덱스/ID 매핑 메모리 로드

models/embedder.py

@@ -0,0 +1,26 @@
+# rag/models/embedder.py
+from typing import List
+import numpy as np
+import onnxruntime as ort
+from fastapi import Request
+
+def _l2_normalize(vec: np.ndarray) -> List[float]:
+    norm = np.linalg.norm(vec) or 1.0
+    return (vec / norm).tolist()
+
+def get_embedding(request: Request, text: str) -> List[float]:
+    """
+    request.app.state.embedder_sess : ONNX Runtime InferenceSession
+    request.app.state.embedder_tokenizer : 토크나이저
+    """
+    tokenizer = request.app.state.embedder_tokenizer
+    sess: ort.InferenceSession = request.app.state.embedder_sess
+
+    inputs = tokenizer(text, return_tensors="np", padding=True, truncation=True, max_length=256)
+    ort_inputs = {k: v for k, v in inputs.items()}
+    ort_outs = sess.run(None, ort_inputs)
+    # 일반적으로 첫 번째 출력이 [batch, dim] 임베딩
+    vec = ort_outs[0][0]
+    return _l2_normalize(vec)
+
+

models/initializer.py

@@ -0,0 +1,37 @@
+# rag/models/initializer.py
+from transformers import AutoTokenizer
+import onnxruntime as ort
+from huggingface_hub import hf_hub_download
+from fastapi import FastAPI
+from config import HF_MODEL_REPO_ID, EMBED_MODEL, EMBED_DIR, RERANK_MODEL, RERANK_DIR
+
+def initialize_models(app: FastAPI):
+    # Embedder
+    embedder_tokenizer = AutoTokenizer.from_pretrained(
+        HF_MODEL_REPO_ID,
+        subfolder=EMBED_DIR   # 토크나이저 관련 파일이 embedder/ 안에 있으므로 지정
+    )
+    embedder_model_path = hf_hub_download(
+        repo_id=HF_MODEL_REPO_ID,
+        filename=EMBED_MODEL,
+        subfolder=EMBED_DIR
+    )
+    embedder_sess = ort.InferenceSession(embedder_model_path, providers=["CPUExecutionProvider"])
+
+    # Reranker
+    reranker_tokenizer = AutoTokenizer.from_pretrained(
+        HF_MODEL_REPO_ID,
+        subfolder=RERANK_DIR   # 토크나이저 관련 파일이 reranker/ 안에 있으므로 지정
+    )
+    reranker_model_path = hf_hub_download(
+        repo_id=HF_MODEL_REPO_ID,
+        filename=RERANK_MODEL,
+        subfolder=RERANK_DIR
+    )
+    reranker_sess = ort.InferenceSession(reranker_model_path, providers=["CPUExecutionProvider"])
+
+    # FastAPI app.state에 저장 → 전역 공유
+    app.state.embedder_tokenizer = embedder_tokenizer
+    app.state.embedder_sess = embedder_sess
+    app.state.reranker_tokenizer = reranker_tokenizer
+    app.state.reranker_sess = reranker_sess

{modules → models}/reranker.py

@@ -1,39 +1,32 @@
-# rag/modules/reranker.py
+# rag/models/reranker.py
 import os
 from typing import List, Dict
-from huggingface_hub import InferenceClient
+import onnxruntime as ort
+from fastapi import Request
 
-# 환경변수에서 모델명과 토큰 불러오기
-HF_TOKEN = os.getenv("HF_TOKEN")
-RERANK_MODEL = os.getenv("RERANK_MODEL", "BAAI/bge-reranker-large")
-
-_client = InferenceClient(model=RERANK_MODEL, token=HF_TOKEN)
-
-# threshold 값은 환경변수나 config에서 관리 가능
 THRESHOLD = float(os.getenv("RERANK_THRESHOLD", "0.3"))
 
-def rerank(query: str, contexts: List[Dict]) -> List[Dict]:
+def rerank(request: Request, query: str, contexts: List[Dict]) -> List[Dict]:
     """
+    request.app.state.reranker_sess : ONNX Runtime InferenceSession
+    request.app.state.reranker_tokenizer : 토크나이저
     contexts: [{"id": ..., "text": ...}, ...]
-    반환: threshold 이상 점수만 포함된 reranked contexts
     """
     if not contexts:
         return []
 
-    # reranker 입력: (query, passage) 쌍 리스트
-    pairs = [(query, ctx["text"]) for ctx in contexts]
+    tokenizer = request.app.state.reranker_tokenizer
+    sess: ort.InferenceSession = request.app.state.reranker_sess
 
-    # Inference API 호출 → 각 쌍에 대한 점수 반환
-    scores = _client.rerank(inputs=pairs)
+    pairs = [(query, ctx["text"]) for ctx in contexts]
+    inputs = tokenizer(pairs, return_tensors="np", padding=True, truncation=True, max_length=256)
+    ort_inputs = {k: v for k, v in inputs.items()}
+    scores = sess.run(None, ort_inputs)[0]  # [batch, 1] 형태라고 가정
+    scores = scores.squeeze(-1)
 
-    # scores는 [{"score": float}, ...] 형태
     for ctx, sc in zip(contexts, scores):
-        ctx["score"] = sc["score"]
+        ctx["score"] = float(sc)
 
-    # 점수 내림차순 정렬
     reranked = sorted(contexts, key=lambda x: x["score"], reverse=True)
-
-    # threshold 이상만 필터링
     reranked = [c for c in reranked if c["score"] >= THRESHOLD]
-
     return reranked

modules/embedder.py

@@ -1,23 +0,0 @@
-# rag/modules/embedder.py
-import math
-from typing import List
-from huggingface_hub import InferenceClient
-from config import EMBED_MODEL, HF_TOKEN
-
-# 모델과 토큰 지정
-_client = InferenceClient(model=EMBED_MODEL, token=HF_TOKEN)
-
-def _l2_normalize(vec: List[float]) -> List[float]:
-    norm = math.sqrt(sum(x * x for x in vec)) or 1.0
-    return [x / norm for x in vec]
-
-def get_embedding(text: str) -> List[float]:
-    # embeddings 태스크를 직접 지정
-    response = _client.post(
-        json={"inputs": text},
-        task="embeddings"
-    )
-    # 항상 [batch_size, embedding_dim] 형태 반환
-    vec = response[0]
-    return _l2_normalize(vec)
-

service/search.py

@@ -1,8 +1,8 @@
 # rag/service/search.py
-from modules.embedder import get_embedding
+from models.embedder import get_embedding
 from modules.retriever import retrieve_ids
 from modules.corpus import fetch_contexts_by_ids
-from modules.reranker import rerank
+from models.reranker import rerank
 
 def search(query: str) -> list[dict]:
     embedding = get_embedding(query)