Update irpr/deps.py

irpr/deps.py

@@ -1,9 +1,8 @@
 import os
 import pickle
 import numpy as np
-import torch
 
-# ---- キャッシュ/データディレクトリを先に用意（import 時でも安全） -----------------
+# ========= キャッシュ/データディレクトリ（import時は作成と権限確認だけ） =========
 CACHE_DIR = os.environ.get("HF_HOME") or "/tmp/hf-cache"
 for k in ["HF_HOME", "TRANSFORMERS_CACHE", "SENTENCE_TRANSFORMERS_HOME", "TORCH_HOME"]:
     os.environ.setdefault(k, CACHE_DIR)
@@ -14,38 +13,34 @@ from irpr.config import settings
 DATA_DIR = settings.DATA_DIR or "data"
 os.makedirs(DATA_DIR, exist_ok=True)
 
-# 永続化パス（NumPyベースの簡易インデックス）
 VEC_PATH = os.path.join(DATA_DIR, "vectors.npy")
 STORE_PATH = os.path.join(DATA_DIR, "store.pkl")
 
-# ---- 遅延ロード用のシングルトン ---------------------------------------------------
+# ========= グローバル状態（import時にモデルは作らない！） =========
 from typing import Optional, List, Dict
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-from sentence_transformers import SentenceTransformer
 
-_EMB: Optional[SentenceTransformer] = None
+_EMB = None           # SentenceTransformer
 _EMB_DIM: Optional[int] = None
+_TOK = None           # AutoTokenizer
+_MODEL = None         # AutoModelForCausalLM
+_GEN = None           # pipeline("text-generation")
 
-_TOK = None
-_MODEL = None
-_GEN = None
+_VECTORS: Optional[np.ndarray] = None  # shape [N, D] float32
+_STORE: Optional[List[Dict]] = None
 
-_VECTORS: Optional[np.ndarray] = None  # shape [N, D], float32, L2正規化済み
-_STORE: Optional[List[Dict]] = None    # メタ情報の配列
-
-# ---- 低レベルユーティリティ -------------------------------------------------------
+# ========= 永続インデックスの入出力 =========
 def _load_index():
-    """ベクトルとメタ情報を遅延ロード"""
-    global _VECTORS, _STORE, _EMB_DIM
+    """ベクトル/メタの遅延ロード（モデルは触らない）"""
+    global _VECTORS, _STORE
     if _VECTORS is None:
         if os.path.exists(VEC_PATH):
             try:
-                vec = np.load(VEC_PATH)
-                _VECTORS = vec.astype(np.float32, copy=False)
+                arr = np.load(VEC_PATH)
+                _VECTORS = arr.astype(np.float32, copy=False)
             except Exception:
-                _VECTORS = np.empty((0, _EMB_DIM or 0), dtype=np.float32)
+                _VECTORS = np.empty((0, 0), dtype=np.float32)
         else:
-            _VECTORS = np.empty((0, _EMB_DIM or 0), dtype=np.float32)
+            _VECTORS = np.empty((0, 0), dtype=np.float32)
     if _STORE is None:
         if os.path.exists(STORE_PATH):
             try:
@@ -61,31 +56,35 @@ def _save_index():
     global _VECTORS, _STORE
     if _VECTORS is None or _STORE is None:
         return
-    # np.save はディレクトリが必要
     os.makedirs(os.path.dirname(VEC_PATH), exist_ok=True)
     np.save(VEC_PATH, _VECTORS)
     with open(STORE_PATH, "wb") as f:
         pickle.dump(_STORE, f)
 
-# ---- モデル遅延ロード -------------------------------------------------------------
-def _get_emb_model() -> SentenceTransformer:
-    """SentenceTransformer を初回呼び出し時にロード"""
+# ========= モデル（初回呼び出し時にだけロード） =========
+def _get_emb_model():
+    """SentenceTransformer を初めて必要になったときだけロード"""
     global _EMB, _EMB_DIM
     if _EMB is None:
+        from sentence_transformers import SentenceTransformer  # ← ここで初めてimport
         model_name = settings.EMB_MODEL or "intfloat/multilingual-e5-base"
         _EMB = SentenceTransformer(model_name, cache_folder=CACHE_DIR)
         _EMB_DIM = _EMB.get_sentence_embedding_dimension()
-        # 既存ベクトル配列のD次元を確定させる
+
+        # 既存ベクトル配列の次元整合
         _load_index()
-        if _VECTORS is not None and _VECTORS.size == 0:
-            # 空なら次元を確定
-            _set_vectors(np.empty((0, _EMB_DIM), dtype=np.float32))
+        global _VECTORS
+        if _VECTORS.size == 0 or (_VECTORS.ndim == 2 and _VECTORS.shape[1] != _EMB_DIM):
+            _VECTORS = np.empty((0, _EMB_DIM), dtype=np.float32)
     return _EMB
 
 def _get_gen_pipeline():
-    """text-generation pipeline を初回呼び出し時にロード"""
+    """text-generation pipeline を初回だけロード"""
     global _TOK, _MODEL, _GEN
     if _GEN is None:
+        from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline  # ← 遅延import
+        import torch  # ← 遅延import
+
         gen_name = settings.GEN_MODEL or "Qwen/Qwen2.5-3B-Instruct"
         _TOK = AutoTokenizer.from_pretrained(gen_name, cache_dir=CACHE_DIR)
         _MODEL = AutoModelForCausalLM.from_pretrained(
@@ -93,12 +92,12 @@ def _get_gen_pipeline():
             cache_dir=CACHE_DIR,
             torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
             device_map="auto",
-            low_cpu_mem_usage=True
+            low_cpu_mem_usage=True,
         )
         _GEN = pipeline("text-generation", model=_MODEL, tokenizer=_TOK)
     return _GEN, _TOK
 
-# ---- ベクトル化と検索 -------------------------------------------------------------
+# ========= 埋め込みと検索 =========
 def embed_texts(texts: List[str]) -> np.ndarray:
     emb = _get_emb_model()
     v = emb.encode(
@@ -109,11 +108,6 @@ def embed_texts(texts: List[str]) -> np.ndarray:
     )
     return v.astype(np.float32, copy=False)
 
-def _set_vectors(arr: np.ndarray):
-    """内部ベクトル配列を安全に差し替え"""
-    global _VECTORS
-    _VECTORS = arr.astype(np.float32, copy=False)
-
 def add_to_index(records: List[Dict]):
     """
     records: [{"text":..., "source_url":..., "title":..., "doc_id":..., "chunk_id":...}]
@@ -123,27 +117,22 @@ def add_to_index(records: List[Dict]):
     _load_index()
     vecs = embed_texts([r["text"] for r in records])  # [M, D]
     global _VECTORS, _STORE
-    if _VECTORS is None or _VECTORS.size == 0:
-        _set_vectors(vecs)
+    if _VECTORS.size == 0:
+        _VECTORS = vecs
     else:
-        _set_vectors(np.vstack([_VECTORS, vecs]))
+        _VECTORS = np.vstack([_VECTORS, vecs])
     _STORE.extend(records)
     _save_index()
 
 def search(query: str, top_k=8):
-    """
-    内積ベース（両ベクトル正規化済み→cos 類似）で上位K件を返す
-    """
     _load_index()
-    if _VECTORS is None or _VECTORS.size == 0 or not _STORE:
+    if _VECTORS.size == 0 or not _STORE:
         return []
-
-    qv = embed_texts([query])[0]       # [D]
-    sims = _VECTORS @ qv               # [N]
+    qv = embed_texts([query])[0]
+    sims = _VECTORS @ qv
     top_k = min(top_k, sims.shape[0])
     idx = np.argpartition(-sims, top_k - 1)[:top_k]
     idx = idx[np.argsort(-sims[idx])]
-
     hits = []
     for i in idx.tolist():
         rec = _STORE[i].copy()
@@ -151,11 +140,8 @@ def search(query: str, top_k=8):
         hits.append(rec)
     return hits
 
-# ---- 生成ユーティリティ ----------------------------------------------------------
+# ========= 生成ユーティリティ =========
 def generate_chat(messages: list[dict], max_new_tokens=800, temperature=0.2) -> str:
-    """
-    messages: [{"role":"system"/"user"/"assistant", "content":"..."}]
-    """
     gen, tok = _get_gen_pipeline()
     prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     out = gen(