Update irpr/deps.py

irpr/deps.py

@@ -1,20 +1,107 @@
-from sentence_transformers import SentenceTransformer
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-import os, pickle, torch
+import os
+import pickle
 import numpy as np
+import torch
+
+# ---- キャッシュ/データディレクトリを先に用意（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)
+os.makedirs(CACHE_DIR, exist_ok=True)
 
 from irpr.config import settings
 
-# データディレクトリの存在を保証
-os.makedirs(settings.DATA_DIR, exist_ok=True)
+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")
+
+# ---- 遅延ロード用のシングルトン ---------------------------------------------------
+from typing import Optional, List, Dict
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from sentence_transformers import SentenceTransformer
+
+_EMB: Optional[SentenceTransformer] = None
+_EMB_DIM: Optional[int] = None
+
+_TOK = None
+_MODEL = None
+_GEN = None
+
+_VECTORS: Optional[np.ndarray] = None  # shape [N, D], float32, L2正規化済み
+_STORE: Optional[List[Dict]] = None    # メタ情報の配列
 
-# ---- Embeddings -------------------------------------------------------------
-emb_model = SentenceTransformer(settings.EMB_MODEL)
-EMB_DIM = emb_model.get_sentence_embedding_dimension()
+# ---- 低レベルユーティリティ -------------------------------------------------------
+def _load_index():
+    """ベクトルとメタ情報を遅延ロード"""
+    global _VECTORS, _STORE, _EMB_DIM
+    if _VECTORS is None:
+        if os.path.exists(VEC_PATH):
+            try:
+                vec = np.load(VEC_PATH)
+                _VECTORS = vec.astype(np.float32, copy=False)
+            except Exception:
+                _VECTORS = np.empty((0, _EMB_DIM or 0), dtype=np.float32)
+        else:
+            _VECTORS = np.empty((0, _EMB_DIM or 0), dtype=np.float32)
+    if _STORE is None:
+        if os.path.exists(STORE_PATH):
+            try:
+                with open(STORE_PATH, "rb") as f:
+                    s = pickle.load(f)
+                _STORE = s if isinstance(s, list) else []
+            except Exception:
+                _STORE = []
+        else:
+            _STORE = []
 
-def embed_texts(texts: list[str]) -> np.ndarray:
-    # すでに正規化付きで取得
-    v = emb_model.encode(
+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 を初回呼び出し時にロード"""
+    global _EMB, _EMB_DIM
+    if _EMB is None:
+        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))
+    return _EMB
+
+def _get_gen_pipeline():
+    """text-generation pipeline を初回呼び出し時にロード"""
+    global _TOK, _MODEL, _GEN
+    if _GEN is None:
+        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(
+            gen_name,
+            cache_dir=CACHE_DIR,
+            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+            device_map="auto",
+            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(
         texts,
         normalize_embeddings=True,
         convert_to_numpy=True,
@@ -22,91 +109,60 @@ def embed_texts(texts: list[str]) -> np.ndarray:
     )
     return v.astype(np.float32, copy=False)
 
-# ---- シンプル永続インデックス（NumPyのみ） ------------------------------------
-# ベクトル: data/vectors.npy（shape: [N, D], float32, 正規化済み）
-# メタ情報: data/store.pkl   （list[dict]）
-VEC_PATH = os.path.join(settings.DATA_DIR, "vectors.npy")
-STORE_PATH = os.path.join(settings.DATA_DIR, "store.pkl")
-
-# 初期化
-if os.path.exists(VEC_PATH):
-    vectors = np.load(VEC_PATH, mmap_mode=None)
-    # guard: 既存が空ファイルの可能性
-    if vectors.ndim != 2 or vectors.shape[1] != EMB_DIM:
-        vectors = np.empty((0, EMB_DIM), dtype=np.float32)
-else:
-    vectors = np.empty((0, EMB_DIM), dtype=np.float32)
-
-if os.path.exists(STORE_PATH):
-    with open(STORE_PATH, "rb") as f:
-        store = pickle.load(f)
-    if not isinstance(store, list):
-        store = []
-else:
-    store = []
-
-def _save_index():
-    np.save(VEC_PATH, vectors)
-    with open(STORE_PATH, "wb") as f:
-        pickle.dump(store, f)
+def _set_vectors(arr: np.ndarray):
+    """内部ベクトル配列を安全に差し替え"""
+    global _VECTORS
+    _VECTORS = arr.astype(np.float32, copy=False)
 
-def add_to_index(records: list[dict]):
+def add_to_index(records: List[Dict]):
     """
     records: [{"text":..., "source_url":..., "title":..., "doc_id":..., "chunk_id":...}]
     """
-    global vectors, store
     if not records:
         return
+    _load_index()
     vecs = embed_texts([r["text"] for r in records])  # [M, D]
-    if vectors.size == 0:
-        vectors = vecs
+    global _VECTORS, _STORE
+    if _VECTORS is None or _VECTORS.size == 0:
+        _set_vectors(vecs)
     else:
-        vectors = np.vstack([vectors, vecs])
-    store.extend(records)
+        _set_vectors(np.vstack([_VECTORS, vecs]))
+    _STORE.extend(records)
     _save_index()
 
 def search(query: str, top_k=8):
     """
-    内積（cos類似に相当：両方正規化済み）でトップKを返す簡易検索
+    内積ベース（両ベクトル正規化済み→cos 類似）で上位K件を返す
     """
-    if vectors.size == 0 or len(store) == 0:
+    _load_index()
+    if _VECTORS is None or _VECTORS.size == 0 or not _STORE:
         return []
 
-    qv = embed_texts([query])[0]           # [D]
-    sims = vectors @ qv                    # [N]  内積
-    if top_k > sims.shape[0]:
-        top_k = sims.shape[0]
-    idx = np.argpartition(-sims, top_k-1)[:top_k]
-    # スコア順に並べ替え
+    qv = embed_texts([query])[0]       # [D]
+    sims = _VECTORS @ qv               # [N]
+    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 rank, i in enumerate(idx.tolist()):
-        rec = store[i].copy()
+    for i in idx.tolist():
+        rec = _STORE[i].copy()
         rec["score"] = float(sims[i])
         hits.append(rec)
     return hits
 
-# ---- 生成モデル --------------------------------------------------------------
-tok = AutoTokenizer.from_pretrained(settings.GEN_MODEL)
-model = AutoModelForCausalLM.from_pretrained(
-    settings.GEN_MODEL,
-    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto",
-    low_cpu_mem_usage=True
-)
-gen = pipeline("text-generation", model=model, tokenizer=tok)
-
-def generate_chat(messages: list[dict], max_new_tokens=800, temperature=0.2):
+# ---- 生成ユーティリティ ----------------------------------------------------------
+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(
         prompt,
         do_sample=(temperature > 0.0),
         temperature=temperature,
-        max_new_tokens=max_new_tokens
+        max_new_tokens=max_new_tokens,
     )[0]["generated_text"]
     generated = out[len(prompt):].strip()
     return generated or out