Initial codes commit

Dockerfile

@@ -0,0 +1,12 @@
+FROM python:3.10-slim
+
+WORKDIR /app
+
+# requirements 설치
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# 앱 코드 복사
+COPY . .
+
+CMD ["sh", "-c", "uvicorn rag.app:app --host 0.0.0.0 --port ${PORT}"]

api/endpoints.py

@@ -0,0 +1,11 @@
+# rag/api/endpoints.py
+from fastapi import APIRouter
+from api.schemas import SearchRequest, SearchResponse
+from rag.service.search import search
+
+router = APIRouter()
+
+@router.post("/search", response_model=SearchResponse)
+def search_context(req: SearchRequest):
+    results = search(req.query)
+    return {"results": results}

api/schemas.py

@@ -0,0 +1,9 @@
+# rag/api/schemas.py
+from pydantic import BaseModel
+from typing import List, Dict, Any
+
+class SearchRequest(BaseModel):
+    query: str
+
+class SearchResponse(BaseModel):
+    results: List[Dict[str, Any]]

app.py

@@ -0,0 +1,29 @@
+# rag/app.py
+from contextlib import asynccontextmanager
+from fastapi import FastAPI, Request, Form
+from fastapi.responses import HTMLResponse
+from fastapi.templating import Jinja2Templates
+from api.endpoints import router
+from rag.db.initializer import initialize
+from rag.service.search import search
+
+templates = Jinja2Templates(directory="templates")
+
+@asynccontextmanager
+async def lifespan(_app: FastAPI):
+    initialize()
+    yield
+
+app = FastAPI(lifespan=lifespan)
+app.include_router(router)
+
+@app.get("/", response_class=HTMLResponse)
+def index(request: Request):
+    return templates.TemplateResponse("index.html", {"request": request, "results": None})
+
+@app.post("/demo", response_class=HTMLResponse)
+def demo(request: Request, query: str = Form(...)):
+    results = search(query)
+    return templates.TemplateResponse("index.html", {"request": request, "results": results, "query": query})
+
+

config.py

@@ -0,0 +1,25 @@
+# rag/config.py
+import os
+
+# Huggingface Hub token
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+# HF datasets repo info
+HF_REPO_ID = os.getenv("HF_REPO_ID", "m97j/pls-datasets")
+HF_INDEX_FILE = os.getenv("HF_INDEX_FILE", "faiss_index_flat.faiss")
+HF_IDS_FILE = os.getenv("HF_IDS_FILE", "vector_ids.npy")
+
+# Corpus dataset info
+HF_CORPUS_REPO = os.getenv("HF_CORPUS_REPO", "HuggingFaceFW/finewiki")
+HF_CORPUS_SUBSET = os.getenv("HF_CORPUS_SUBSET", "ko")
+HF_CORPUS_SPLIT = os.getenv("HF_CORPUS_SPLIT", "train")
+
+# Local paths
+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")
+TOP_K = int(os.getenv("TOP_K", "5"))
+
+

db/initializer.py

@@ -0,0 +1,32 @@
+# rag/db/initializer.py
+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 modules.utils import ensure_dir
+from modules.retriever import set_index
+from modules import corpus
+
+_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)
+    ids_path = hf_hub_download(repo_id=HF_REPO_ID, filename=HF_IDS_FILE)
+    index = faiss.read_index(index_path)
+    set_index(index)
+    global _vector_ids
+    _vector_ids = np.load(ids_path, allow_pickle=True)
+
+def get_vector_ids():
+    global _vector_ids
+    return _vector_ids
+
+def initialize():
+    # 1) 코퍼스 준비 (최초 1회만 다운로드)
+    corpus.prepare_corpus()
+    # 2) 인덱스/ID 매핑 메모리 로드
+    _load_index_in_memory()
+
+def force_update():
+    _load_index_in_memory()

modules/corpus.py

@@ -0,0 +1,61 @@
+# rag/modules/corpus_store.py
+from typing import List, Dict, Any
+from datasets import load_dataset, DatasetDict, Dataset
+from config import HF_CORPUS_REPO, HF_CORPUS_SUBSET, HF_CORPUS_SPLIT, MARKER_DIR, CORPUS_READY_MARK
+from modules.utils import ensure_dir, exists, touch
+
+_datasets: Dict[str, Dataset] = {}
+
+def prepare_corpus():
+    """
+    최초 1회만 parquet split을 로컬에 다운로드.
+    이후에는 로컬 캐시 사용.
+    """
+    ensure_dir(MARKER_DIR)
+    if exists(CORPUS_READY_MARK):
+        return
+
+    subsets = HF_CORPUS_SUBSET.split(",")  # "ko,en" → ["ko","en"]
+    for subset in subsets:
+        load_dataset(HF_CORPUS_REPO, subset.strip(), split=HF_CORPUS_SPLIT)
+    touch(CORPUS_READY_MARK)
+
+def _get_datasets() -> Dict[str, Dataset]:
+    global _datasets
+    if not _datasets:
+        subsets = HF_CORPUS_SUBSET.split(",")
+        for subset in subsets:
+            _datasets[subset.strip()] = load_dataset(
+                HF_CORPUS_REPO, subset.strip(), split=HF_CORPUS_SPLIT
+            )
+    return _datasets
+
+def fetch_contexts_by_ids(ids: List[int]) -> List[Dict[str, Any]]:
+    if not ids:
+        return []
+
+    datasets = _get_datasets()
+    id_set = set(ids)
+    results: List[Dict[str, Any]] = []
+
+    # 모든 subset을 순회하며 page_id 매칭
+    for subset, ds in datasets.items():
+        # filter를 사용하면 전체 순회보다 빠름 (병렬 최적화)
+        rows = ds.filter(lambda r: r["page_id"] in id_set)
+
+        id_to_row = {r["page_id"]: r for r in rows}
+        for i in ids:
+            r = id_to_row.get(i)
+            if r:
+                results.append({
+                    "id": r["page_id"],
+                    "title": r.get("title", ""),
+                    "text": r.get("wikitext", ""),
+                    "url": r.get("url", ""),
+                    "metadata": {
+                        "date_modified": r.get("date_modified", ""),
+                        "in_language": r.get("in_language", ""),
+                        "wikidata_id": r.get("wikidata_id", "")
+                    }
+                })
+    return results

modules/embedder.py

@@ -0,0 +1,17 @@
+# 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]:
+    # feature_extraction은 항상 2차원 배열 반환: [batch_size, embedding_dim]
+    embedding_2d = _client.feature_extraction(text)
+    vec = embedding_2d[0]  # 첫 번째 행이 입력 문장의 벡터
+    return _l2_normalize(vec)

modules/reranker.py

@@ -0,0 +1,39 @@
+# rag/modules/reranker.py
+import os
+from typing import List, Dict
+from huggingface_hub import InferenceClient
+
+# 환경변수에서 모델명과 토큰 불러오기
+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]:
+    """
+    contexts: [{"id": ..., "text": ...}, ...]
+    반환: threshold 이상 점수만 포함된 reranked contexts
+    """
+    if not contexts:
+        return []
+
+    # reranker 입력: (query, passage) 쌍 리스트
+    pairs = [(query, ctx["text"]) for ctx in contexts]
+
+    # Inference API 호출 → 각 쌍에 대한 점수 반환
+    scores = _client.rerank(inputs=pairs)
+
+    # scores는 [{"score": float}, ...] 형태
+    for ctx, sc in zip(contexts, scores):
+        ctx["score"] = sc["score"]
+
+    # 점수 내림차순 정렬
+    reranked = sorted(contexts, key=lambda x: x["score"], reverse=True)
+
+    # threshold 이상만 필터링
+    reranked = [c for c in reranked if c["score"] >= THRESHOLD]
+
+    return reranked

modules/retriever.py

@@ -0,0 +1,20 @@
+# rag/modules/retriever.py
+import numpy as np
+from config import TOP_K
+
+_index = None  # in-memory FAISS index
+
+def set_index(index_obj):
+    global _index
+    _index = index_obj
+
+def has_index() -> bool:
+    return _index is not None
+
+def retrieve_ids(query_embedding: list[float]) -> list[int]:
+    if _index is None:
+        raise RuntimeError("FAISS index is not loaded in memory.")
+    q = np.array([query_embedding], dtype="float32")
+    _, idx = _index.search(q, TOP_K)
+    return [int(i) for i in idx[0]]
+

modules/utils.py

@@ -0,0 +1,12 @@
+# rag/modules/utils.py
+import os
+
+def ensure_dir(path: str):
+    os.makedirs(path, exist_ok=True)
+
+def touch(path: str):
+    with open(path, "a", encoding="utf-8") as f:
+        pass
+
+def exists(path: str) -> bool:
+    return os.path.exists(path)

requirements.txt

@@ -0,0 +1,7 @@
+fastapi==0.115.0
+uvicorn[standard]==0.30.0
+
+huggingface_hub==0.24.6
+faiss-cpu==1.8.0
+numpy==1.26.4
+pydantic==2.11.5

service/search.py

@@ -0,0 +1,11 @@
+from modules.embedder import get_embedding
+from modules.retriever import retrieve_ids
+from modules.corpus import fetch_contexts_by_ids
+from modules.reranker import rerank
+
+def search(query: str) -> list[dict]:
+    embedding = get_embedding(query)
+    ids = retrieve_ids(embedding)
+    contexts = fetch_contexts_by_ids(ids)
+    reranked = rerank(query, contexts)
+    return reranked

templates/index.html

@@ -0,0 +1,18 @@
+<!DOCTYPE html>
+<html>
+  <body>
+    <h1>RAG Search Demo</h1>
+    <form method="post" action="/demo">
+      <input type="text" name="query" placeholder="Enter query" required>
+      <button type="submit">Search</button>
+    </form>
+    {% if results %}
+      <h2>Results for: {{ query }}</h2>
+      <ul>
+        {% for r in results %}
+          <li>{{ r }}</li>
+        {% endfor %}
+      </ul>
+    {% endif %}
+  </body>
+</html>