Update irpr/deps.py

irpr/deps.py

@@ -1,95 +1,103 @@
+# irpr/deps.py  --- Chromadb版（faiss不使用・LLMなしでも動く）
+from __future__ import annotations
 import os
-import pickle
+from typing import List, Dict, Optional
 import numpy as np
-
-# ========= キャッシュ/データディレクトリ（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
 
-DATA_DIR = settings.DATA_DIR or "data"
-os.makedirs(DATA_DIR, exist_ok=True)
-
-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
-
-_EMB = None           # SentenceTransformer
+# 書き込み先確保
+for d in [
+    os.environ.get("HF_HOME", "/data/.hf-home"),
+    os.environ.get("TRANSFORMERS_CACHE", "/data/.hf-cache"),
+    os.environ.get("SENTENCE_TRANSFORMERS_HOME", "/data/.hf-cache"),
+    os.environ.get("HUGGINGFACE_HUB_CACHE", "/data/.hf-cache"),
+    settings.DATA_DIR,
+    settings.CHROMA_PATH,
+]:
+    if d: os.makedirs(d, exist_ok=True)
+
+# 遅延ロード
+_EMB = None
 _EMB_DIM: Optional[int] = None
-_TOK = None           # AutoTokenizer
-_MODEL = None         # AutoModelForCausalLM
-_GEN = None           # pipeline("text-generation")
+_CHROMA_COLLECTION = None
+_GEN = None
+_TOK = None
 
-_VECTORS: Optional[np.ndarray] = None  # shape [N, D] float32
-_STORE: Optional[List[Dict]] = None
-
-# ========= 永続インデックスの入出力 =========
-def _load_index():
-    """ベクトル/メタの遅延ロード（モデルは触らない）"""
-    global _VECTORS, _STORE
-    if _VECTORS is None:
-        if os.path.exists(VEC_PATH):
-            try:
-                arr = np.load(VEC_PATH)
-                _VECTORS = arr.astype(np.float32, copy=False)
-            except Exception:
-                _VECTORS = np.empty((0, 0), dtype=np.float32)
-        else:
-            _VECTORS = np.empty((0, 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 _save_index():
-    global _VECTORS, _STORE
-    if _VECTORS is None or _STORE is None:
-        return
-    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 を初めて必要になったときだけロード"""
+def _get_embedder():
     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)
+        from sentence_transformers import SentenceTransformer
+        _EMB = SentenceTransformer(settings.EMB_MODEL, cache_folder=os.environ.get("HF_HOME", "/data/.hf-cache"))
         _EMB_DIM = _EMB.get_sentence_embedding_dimension()
-
-        # 既存ベクトル配列の次元整合
-        _load_index()
-        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 embed_texts(texts: List[str]) -> np.ndarray:
+    emb = _get_embedder()
+    arr = emb.encode(texts, batch_size=16, normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=False)
+    return arr.astype(np.float32, copy=False)
+
+def _get_chroma():
+    global _CHROMA_COLLECTION
+    if _CHROMA_COLLECTION is None:
+        import chromadb
+        from chromadb.config import Settings as CS
+        client = chromadb.PersistentClient(path=settings.CHROMA_PATH, settings=CS(allow_reset=True))
+        _CHROMA_COLLECTION = client.get_or_create_collection(name="irpr_docs")
+    return _CHROMA_COLLECTION
+
+def add_to_index(records: List[Dict]):
+    if not records: return
+    col = _get_chroma()
+    texts = [r["text"] for r in records]
+    embs = embed_texts(texts)
+    ids, metas = [], []
+    for r in records:
+        doc_id = r.get("doc_id") or "doc"
+        chunk_id = r.get("chunk_id") or ""
+        rid = f"{doc_id}:{chunk_id}" if chunk_id else doc_id
+        ids.append(rid)
+        metas.append({
+            "source_url": r.get("source_url"),
+            "title": r.get("title"),
+            "doc_id": doc_id,
+            "chunk_id": chunk_id,
+        })
+    col.add(ids=ids, documents=texts, embeddings=embs, metadatas=metas)
+
+def search(query: str, top_k=8) -> List[Dict]:
+    col = _get_chroma()
+    q_emb = embed_texts([query])
+    res = col.query(query_embeddings=q_emb, n_results=top_k, include=["documents","metadatas","distances","ids"])
+    docs  = res.get("documents", [[]])[0]
+    metas = res.get("metadatas", [[]])[0]
+    dists = res.get("distances", [[]])[0]
+    out: List[Dict] = []
+    for text, meta, dist in zip(docs, metas, dists):
+        score = 1.0 - float(dist)/2.0 if dist is not None else None
+        out.append({
+            "text": text,
+            "source_url": (meta or {}).get("source_url"),
+            "title": (meta or {}).get("title"),
+            "doc_id": (meta or {}).get("doc_id"),
+            "chunk_id": (meta or {}).get("chunk_id"),
+            "score": score,
+        })
+    return out
+
+# ==== 生成（任意） ====
 def _get_gen_pipeline():
-    """text-generation pipeline を初回だけロード"""
-    global _TOK, _MODEL, _GEN
+    """GEN_MODEL が空なら LLM 無効の合図として例外を投げる"""
+    if not settings.GEN_MODEL:
+        raise RuntimeError("GEN_MODEL is empty (LLM disabled).")
+    global _GEN, _TOK
     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)
+        from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+        import torch
+        name = settings.GEN_MODEL
+        _TOK = AutoTokenizer.from_pretrained(name, cache_dir=os.environ.get("HF_HOME", "/data/.hf-cache"))
         _MODEL = AutoModelForCausalLM.from_pretrained(
-            gen_name,
-            cache_dir=CACHE_DIR,
+            name,
+            cache_dir=os.environ.get("HF_HOME", "/data/.hf-cache"),
             torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
             device_map="auto",
             low_cpu_mem_usage=True,
@@ -97,58 +105,8 @@ def _get_gen_pipeline():
         _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,
-        show_progress_bar=False,
-    )
-    return v.astype(np.float32, copy=False)
-
-def add_to_index(records: List[Dict]):
-    """
-    records: [{"text":..., "source_url":..., "title":..., "doc_id":..., "chunk_id":...}]
-    """
-    if not records:
-        return
-    _load_index()
-    vecs = embed_texts([r["text"] for r in records])  # [M, D]
-    global _VECTORS, _STORE
-    if _VECTORS.size == 0:
-        _VECTORS = vecs
-    else:
-        _VECTORS = np.vstack([_VECTORS, vecs])
-    _STORE.extend(records)
-    _save_index()
-
-def search(query: str, top_k=8):
-    _load_index()
-    if _VECTORS.size == 0 or not _STORE:
-        return []
-    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()
-        rec["score"] = float(sims[i])
-        hits.append(rec)
-    return hits
-
-# ========= 生成ユーティリティ =========
-def generate_chat(messages: list[dict], max_new_tokens=800, temperature=0.2) -> str:
+def generate_chat(messages: List[Dict], max_new_tokens=600, temperature=0.2) -> str:
     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,
-    )[0]["generated_text"]
-    generated = out[len(prompt):].strip()
-    return generated or out
+    out = gen(prompt, do_sample=(temperature>0.0), temperature=temperature, max_new_tokens=max_new_tokens)[0]["generated_text"]
+    return out[len(prompt):].strip()