Update main.py

main.py

@@ -33,26 +33,59 @@ LOG.setLevel(logging.INFO)
 # CONFIG (via ENV)
 # =============================================================================
 PORT = int(os.getenv("PORT", "7860"))
-DATA_ROOT = os.getenv("DATA_ROOT", "/tmp/data")  # persistant dans le conteneur Space
+DATA_ROOT = os.getenv("DATA_ROOT", "/tmp/data")  # stockage interne du Space
 os.makedirs(DATA_ROOT, exist_ok=True)
 
 # Provider d'embeddings:
 #   - "dummy" : vecteurs aléatoires déterministes (très rapide)
-#   - "st"    : Sentence-Transformers (CPU-friendly, simple)
-#   - "hf"    : Transformers (AutoModel/AutoTokenizer, pooling manuel)
+#   - "st"    : Sentence-Transformers (CPU-friendly)
+#   - "hf"    : Transformers pur (AutoModel/AutoTokenizer)
 EMB_PROVIDER = os.getenv("EMB_PROVIDER", "dummy").strip().lower()
-
-# Modèle embeddings (utilisé si provider != "dummy")
-# Reco rapide et multilingue (FR ok) : paraphrase-multilingual-MiniLM-L12-v2 (dim=384)
 EMB_MODEL = os.getenv("EMB_MODEL", "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2").strip()
-
-# Batch d'encodage
 EMB_BATCH = int(os.getenv("EMB_BATCH", "32"))
+EMB_DIM = int(os.getenv("EMB_DIM", "128"))  # utilisé pour dummy
 
-# Dimension par défaut (dummy) — pour st/hf on lit depuis le modèle
-EMB_DIM = int(os.getenv("EMB_DIM", "128"))
-
-# Cache global lazy
+# =============================================================================
+# CACHE DIRECTORIES (crucial pour éviter PermissionError: '/.cache')
+# =============================================================================
+def _setup_cache_dirs() -> Dict[str, str]:
+    # HOME peut être vide -> expanduser('~') => '/' -> '/.cache' -> Permission denied
+    os.environ.setdefault("HOME", "/home/user")
+
+    CACHE_ROOT = os.getenv("CACHE_ROOT", "/tmp/.cache").rstrip("/")
+    paths = {
+        "root": CACHE_ROOT,
+        "hf_home": f"{CACHE_ROOT}/huggingface",
+        "hf_hub": f"{CACHE_ROOT}/huggingface/hub",
+        "hf_tf": f"{CACHE_ROOT}/huggingface/transformers",
+        "torch": f"{CACHE_ROOT}/torch",
+        "st": f"{CACHE_ROOT}/sentence-transformers",
+        "mpl": f"{CACHE_ROOT}/matplotlib",
+    }
+    for p in paths.values():
+        try:
+            os.makedirs(p, exist_ok=True)
+        except Exception as e:
+            LOG.warning("Impossible de créer %s : %s", p, e)
+
+    # Variables standard HF/Transformers/Torch/ST
+    os.environ["HF_HOME"] = paths["hf_home"]
+    os.environ["HF_HUB_CACHE"] = paths["hf_hub"]
+    os.environ["TRANSFORMERS_CACHE"] = paths["hf_tf"]
+    os.environ["TORCH_HOME"] = paths["torch"]
+    os.environ["SENTENCE_TRANSFORMERS_HOME"] = paths["st"]
+    os.environ["MPLCONFIGDIR"] = paths["mpl"]  # évite les warnings matplotlib
+
+    # Qualité de vie
+    os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1")
+    os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+
+    LOG.info("Caches configurés: %s", json.dumps(paths, indent=2))
+    return paths
+
+CACHE_PATHS = _setup_cache_dirs()
+
+# Cache global lazy (pour les modèles)
 _ST_MODEL = None
 _HF_TOKENIZER = None
 _HF_MODEL = None
@@ -135,8 +168,8 @@ def _get_st_model():
     global _ST_MODEL
     if _ST_MODEL is None:
         from sentence_transformers import SentenceTransformer
-        _ST_MODEL = SentenceTransformer(EMB_MODEL)
-        LOG.info(f"[st] modèle chargé: {EMB_MODEL}")
+        _ST_MODEL = SentenceTransformer(EMB_MODEL, cache_folder=CACHE_PATHS["st"])
+        LOG.info("[st] modèle chargé: %s (cache=%s)", EMB_MODEL, CACHE_PATHS["st"])
     return _ST_MODEL
 
 def _emb_st(texts: List[str]) -> np.ndarray:
@@ -155,7 +188,6 @@ def _st_dim() -> int:
     try:
         return int(model.get_sentence_embedding_dimension())
     except Exception:
-        # fallback : encode une phrase et lit la shape
         v = model.encode(["dimension probe"], convert_to_numpy=True)
         return int(v.shape[1])
 
@@ -164,17 +196,16 @@ def _get_hf_model():
     global _HF_TOKENIZER, _HF_MODEL
     if _HF_MODEL is None or _HF_TOKENIZER is None:
         from transformers import AutoTokenizer, AutoModel
-        _HF_TOKENIZER = AutoTokenizer.from_pretrained(EMB_MODEL)
-        _HF_MODEL = AutoModel.from_pretrained(EMB_MODEL)
+        _HF_TOKENIZER = AutoTokenizer.from_pretrained(EMB_MODEL, cache_dir=CACHE_PATHS["hf_tf"])
+        _HF_MODEL = AutoModel.from_pretrained(EMB_MODEL, cache_dir=CACHE_PATHS["hf_tf"])
         _HF_MODEL.eval()
-        LOG.info(f"[hf] modèle chargé: {EMB_MODEL}")
+        LOG.info("[hf] modèle chargé: %s (cache=%s)", EMB_MODEL, CACHE_PATHS["hf_tf"])
     return _HF_TOKENIZER, _HF_MODEL
 
 def _mean_pool(last_hidden_state: "np.ndarray", attention_mask: "np.ndarray") -> "np.ndarray":
-    # mean pooling masquée
-    mask = attention_mask[..., None].astype(last_hidden_state.dtype)  # (b, t, 1)
-    summed = (last_hidden_state * mask).sum(axis=1)                   # (b, h)
-    counts = mask.sum(axis=1).clip(min=1e-9)                          # (b, 1)
+    mask = attention_mask[..., None].astype(last_hidden_state.dtype)
+    summed = (last_hidden_state * mask).sum(axis=1)
+    counts = mask.sum(axis=1).clip(min=1e-9)
     return summed / counts
 
 def _emb_hf(texts: List[str]) -> np.ndarray:
@@ -194,7 +225,6 @@ def _emb_hf(texts: List[str]) -> np.ndarray:
     return _l2_normalize(vecs)
 
 def _hf_dim() -> int:
-    # essaie de lire hidden_size
     try:
         _, mod = _get_hf_model()
         return int(getattr(mod.config, "hidden_size", 768))
@@ -228,7 +258,7 @@ def _load_dataset(ds_dir: str) -> List[Dict[str, Any]]:
 def _save_faiss(fx_dir: str, xb: np.ndarray, meta: Dict[str, Any]):
     os.makedirs(fx_dir, exist_ok=True)
     idx_path = os.path.join(fx_dir, "emb.faiss")
-    index = faiss.IndexFlatIP(xb.shape[1])  # cosine ~ inner product si normalisé
+    index = faiss.IndexFlatIP(xb.shape[1])  # cosine ~ inner product si embeddings normalisés
     index.add(xb)
     faiss.write_index(index, idx_path)
     with open(os.path.join(fx_dir, "meta.json"), "w", encoding="utf-8") as f:
@@ -250,7 +280,7 @@ def _tar_dir_to_bytes(dir_path: str) -> bytes:
 # =============================================================================
 # FASTAPI
 # =============================================================================
-fastapi_app = FastAPI(title="remote-indexer", version="2.0.0")
+fastapi_app = FastAPI(title="remote-indexer", version="2.1.0")
 fastapi_app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"],
@@ -274,7 +304,8 @@ def health():
         "ok": True,
         "service": "remote-indexer",
         "provider": EMB_PROVIDER,
-        "model": EMB_MODEL if EMB_PROVIDER != "dummy" else None
+        "model": EMB_MODEL if EMB_PROVIDER != "dummy" else None,
+        "cache_root": os.getenv("CACHE_ROOT", "/tmp/.cache"),
     }
     return info
 
@@ -436,7 +467,7 @@ def _ui_search(project_id: str, query: str, k: int):
 
 with gr.Blocks(title="Remote Indexer (FAISS)", analytics_enabled=False) as ui:
     gr.Markdown("## Remote Indexer — demo UI (API: `/index`, `/status/{job}`, `/search`, `/artifacts/...`).")
-    gr.Markdown(f"**Provider**: `{EMB_PROVIDER}` — **Model**: `{EMB_MODEL if EMB_PROVIDER!='dummy' else '-'}'")
+    gr.Markdown(f"**Provider**: `{EMB_PROVIDER}` — **Model**: `{EMB_MODEL if EMB_PROVIDER!='dummy' else '-'}` — **Cache**: `{os.getenv('CACHE_ROOT', '/tmp/.cache')}`")
     with gr.Tab("Index"):
         pid = gr.Textbox(label="Project ID", value="DEEPWEB")
         sample = gr.Textbox(label="Texte d’exemple", value="Alpha bravo charlie delta echo foxtrot.", lines=4)