Initial codes commit

models/embedder.py

@@ -3,12 +3,20 @@ from typing import List
 import numpy as np
 import onnxruntime as ort
 from fastapi import Request
-from modules.utils import generate_position_ids
 
 def _l2_normalize(vec: np.ndarray) -> List[float]:
     norm = np.linalg.norm(vec) or 1.0
     return (vec / norm).tolist()
 
+def _generate_position_ids(input_ids: np.ndarray) -> np.ndarray:
+    """
+    input_ids: [batch_size, seq_len]
+    return: position_ids of shape [batch_size, seq_len] with int64 dtype
+    """
+    batch_size, seq_len = input_ids.shape
+    position_ids = np.arange(seq_len)[None, :].astype("int64")
+    return np.tile(position_ids, (batch_size, 1))
+
 def get_embedding(request: Request, text: str) -> List[float]:
     """
     request.app.state.embedder_sess : ONNX Runtime InferenceSession
@@ -19,7 +27,7 @@ def get_embedding(request: Request, text: str) -> List[float]:
 
     inputs = tokenizer(text, return_tensors="np", padding=True, truncation=True, max_length=256)
     input_ids = inputs["input_ids"]
-    inputs["position_ids"] = generate_position_ids(input_ids)
+    inputs["position_ids"] = _generate_position_ids(input_ids)
     ort_inputs = {k: v for k, v in inputs.items()}
     ort_outs = sess.run(None, ort_inputs)
     print([arr.shape for arr in ort_outs])

models/reranker.py

@@ -3,7 +3,6 @@ import os
 from typing import List, Dict
 import onnxruntime as ort
 from fastapi import Request
-from modules.utils import generate_position_ids
 
 THRESHOLD = float(os.getenv("RERANK_THRESHOLD", "0.3"))
 
@@ -21,8 +20,6 @@ def rerank(request: Request, query: str, contexts: List[Dict]) -> List[Dict]:
 
     pairs = [(query, ctx["text"]) for ctx in contexts]
     inputs = tokenizer(pairs, return_tensors="np", padding=True, truncation=True, max_length=256)
-    input_ids = inputs["input_ids"]
-    inputs["position_ids"] = generate_position_ids(input_ids)
     ort_inputs = {k: v for k, v in inputs.items()}
     scores = sess.run(None, ort_inputs)[0]  # [batch, 1] 형태
     scores = scores.squeeze(-1)

modules/utils.py

@@ -1,6 +1,5 @@
 # rag/modules/utils.py
 import os
-import numpy as np
 
 def ensure_dir(path: str):
     os.makedirs(path, exist_ok=True)
@@ -11,12 +10,3 @@ def touch(path: str):
 
 def exists(path: str) -> bool:
     return os.path.exists(path)
-
-def generate_position_ids(input_ids: np.ndarray) -> np.ndarray:
-    """
-    input_ids: [batch_size, seq_len]
-    return: position_ids of shape [batch_size, seq_len] with int64 dtype
-    """
-    batch_size, seq_len = input_ids.shape
-    position_ids = np.arange(seq_len)[None, :].astype("int64")
-    return np.tile(position_ids, (batch_size, 1))