Update irpr/deps.py

irpr/deps.py

@@ -5,16 +5,24 @@ from typing import List, Dict, Optional
 import numpy as np
 from irpr.config import settings
 
-# 書き込み先確保
-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)
+# ===== 書き込み可能な場所にキャッシュを集約 =====
+BASE = settings.DATA_DIR or "./var"
+DEFAULT_CACHE = os.path.join(BASE, ".hf-cache")
+DEFAULT_HOME  = os.path.join(BASE, ".hf-home")
+
+# 既に設定が無ければ、ここで安全なデフォルトに設定
+os.environ.setdefault("HF_HOME", DEFAULT_HOME)
+os.environ.setdefault("TRANSFORMERS_CACHE", DEFAULT_CACHE)
+os.environ.setdefault("SENTENCE_TRANSFORMERS_HOME", DEFAULT_CACHE)
+os.environ.setdefault("HUGGINGFACE_HUB_CACHE", DEFAULT_CACHE)
+
+for d in [os.environ["HF_HOME"], os.environ["TRANSFORMERS_CACHE"],
+          os.environ["SENTENCE_TRANSFORMERS_HOME"], os.environ["HUGGINGFACE_HUB_CACHE"],
+          BASE, settings.CHROMA_PATH]:
+    try:
+        os.makedirs(d, exist_ok=True)
+    except Exception:
+        pass  # フォルダ作成に失敗しても致命ではない
 
 # 遅延ロード
 _EMB = None
@@ -27,7 +35,7 @@ def _get_embedder():
     global _EMB, _EMB_DIM
     if _EMB is None:
         from sentence_transformers import SentenceTransformer
-        _EMB = SentenceTransformer(settings.EMB_MODEL, cache_folder=os.environ.get("HF_HOME", "/data/.hf-cache"))
+        _EMB = SentenceTransformer(settings.EMB_MODEL, cache_folder=os.environ.get("HF_HOME", DEFAULT_CACHE))
         _EMB_DIM = _EMB.get_sentence_embedding_dimension()
     return _EMB
 
@@ -54,59 +62,4 @@ def add_to_index(records: List[Dict]):
     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():
-    """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 torch
-        name = settings.GEN_MODEL
-        _TOK = AutoTokenizer.from_pretrained(name, cache_dir=os.environ.get("HF_HOME", "/data/.hf-cache"))
-        _MODEL = AutoModelForCausalLM.from_pretrained(
-            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,
-        )
-        _GEN = pipeline("text-generation", model=_MODEL, tokenizer=_TOK)
-    return _GEN, _TOK
-
-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"]
-    return out[len(prompt):].strip()
+        rid = f"{doc_id}:{chunk_id}