Update main.py

main.py

@@ -12,9 +12,9 @@ from typing import Dict, Any, List, Tuple, Optional
 
 import numpy as np
 import faiss
-from fastapi import FastAPI, HTTPException, Body
+from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse
+from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel
 
 import gradio as gr
@@ -30,18 +30,33 @@ if not LOG.handlers:
 LOG.setLevel(logging.INFO)
 
 # =============================================================================
-# CONFIG
+# CONFIG (via ENV)
 # =============================================================================
 PORT = int(os.getenv("PORT", "7860"))
 DATA_ROOT = os.getenv("DATA_ROOT", "/tmp/data")  # persistant dans le conteneur Space
 os.makedirs(DATA_ROOT, exist_ok=True)
 
-# Embeddings "provider"
-# - "dummy": hash → vecteurs 128D (léger pour un Space Free)
-# (Tu peux plus tard brancher HF Transformers si tu veux.)
+# Provider d'embeddings:
+#   - "dummy" : vecteurs aléatoires déterministes (très rapide)
+#   - "st"    : Sentence-Transformers (CPU-friendly, simple)
+#   - "hf"    : Transformers (AutoModel/AutoTokenizer, pooling manuel)
 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"))
+
+# 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
+_ST_MODEL = None
+_HF_TOKENIZER = None
+_HF_MODEL = None
+
 # =============================================================================
 # JOB STATE
 # =============================================================================
@@ -101,20 +116,93 @@ def _chunk_text(text: str, size: int = 200, overlap: int = 20) -> List[str]:
         i = j - overlap if (j - overlap) > i else j
     return chunks
 
+def _l2_normalize(x: np.ndarray) -> np.ndarray:
+    n = np.linalg.norm(x, axis=1, keepdims=True) + 1e-12
+    return x / n
+
+# ----------------------- PROVIDER: DUMMY --------------------------------------
 def _emb_dummy(texts: List[str], dim: int = EMB_DIM) -> np.ndarray:
-    # vecteurs déterministes à partir d’un hash
     vecs = np.zeros((len(texts), dim), dtype="float32")
     for i, t in enumerate(texts):
         h = hashlib.sha1((t or "").encode("utf-8")).digest()
         rng = np.random.default_rng(int.from_bytes(h[:8], "little", signed=False))
         v = rng.standard_normal(dim).astype("float32")
-        # normalisation
-        n = np.linalg.norm(v) + 1e-9
-        vecs[i] = v / n
+        vecs[i] = v / (np.linalg.norm(v) + 1e-9)
+    return vecs
+
+# ----------------- PROVIDER: Sentence-Transformers ----------------------------
+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}")
+    return _ST_MODEL
+
+def _emb_st(texts: List[str]) -> np.ndarray:
+    model = _get_st_model()
+    vecs = model.encode(
+        texts,
+        batch_size=max(1, EMB_BATCH),
+        convert_to_numpy=True,
+        normalize_embeddings=True,
+        show_progress_bar=False,
+    ).astype("float32")
     return vecs
 
+def _st_dim() -> int:
+    model = _get_st_model()
+    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])
+
+# ----------------------- PROVIDER: Transformers (HF) --------------------------
+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_MODEL.eval()
+        LOG.info(f"[hf] modèle chargé: {EMB_MODEL}")
+    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)
+    return summed / counts
+
+def _emb_hf(texts: List[str]) -> np.ndarray:
+    import torch
+    tok, mod = _get_hf_model()
+    all_vecs = []
+    bs = max(1, EMB_BATCH)
+    with torch.no_grad():
+        for i in range(0, len(texts), bs):
+            batch = texts[i:i+bs]
+            enc = tok(batch, padding=True, truncation=True, return_tensors="pt")
+            out = mod(**enc)
+            last = out.last_hidden_state  # (b, t, h)
+            pooled = _mean_pool(last.numpy(), enc["attention_mask"].numpy())  # numpy
+            all_vecs.append(pooled.astype("float32"))
+    vecs = np.concatenate(all_vecs, axis=0)
+    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))
+    except Exception:
+        return 768
+
+# ---------------------------- DATASET / FAISS ---------------------------------
 def _save_dataset(ds_dir: str, rows: List[Dict[str, Any]]):
-    # Format simple : on écrit un JSONL + un manifest JSON
     os.makedirs(ds_dir, exist_ok=True)
     data_path = os.path.join(ds_dir, "data.jsonl")
     with open(data_path, "w", encoding="utf-8") as f:
@@ -137,15 +225,14 @@ def _load_dataset(ds_dir: str) -> List[Dict[str, Any]]:
                 continue
     return out
 
-def _save_faiss(fx_dir: str, xb: np.ndarray):
+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é
-    # les embeddings _emb_dummy sont déjà 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:
-        json.dump({"dim": xb.shape[1], "count": int(index.ntotal), "provider": EMB_PROVIDER}, f)
+        json.dump(meta, f, ensure_ascii=False, indent=2)
 
 def _load_faiss(fx_dir: str) -> faiss.Index:
     idx_path = os.path.join(fx_dir, "emb.faiss")
@@ -163,7 +250,7 @@ def _tar_dir_to_bytes(dir_path: str) -> bytes:
 # =============================================================================
 # FASTAPI
 # =============================================================================
-fastapi_app = FastAPI(title="remote-indexer-min", version="1.0.0")
+fastapi_app = FastAPI(title="remote-indexer", version="2.0.0")
 fastapi_app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"],
@@ -183,12 +270,17 @@ class IndexRequest(BaseModel):
 
 @fastapi_app.get("/health")
 def health():
-    return {"ok": True, "service": "remote-indexer-min", "emb_provider": EMB_PROVIDER}
+    info = {
+        "ok": True,
+        "service": "remote-indexer",
+        "provider": EMB_PROVIDER,
+        "model": EMB_MODEL if EMB_PROVIDER != "dummy" else None
+    }
+    return info
 
 @fastapi_app.get("/")
 def root_redirect():
-    # petit état + invite à utiliser /ui
-    return {"ok": True, "service": "remote-indexer-min", "ui": "/ui"}
+    return {"ok": True, "service": "remote-indexer", "ui": "/ui"}
 
 @fastapi_app.post("/index")
 def index(req: IndexRequest):
@@ -196,7 +288,7 @@ def index(req: IndexRequest):
     st = JobState(job_id=job_id, project_id=req.project_id, stage="pending", messages=[])
     JOBS[job_id] = st
     _add_msg(st, f"Job {job_id} créé pour project {req.project_id}")
-    _add_msg(st, f"Index start project={req.project_id} files={len(req.files)} chunk_size={req.chunk_size} overlap={req.overlap} batch_size={req.batch_size} store_text={req.store_text}")
+    _add_msg(st, f"Index start project={req.project_id} files={len(req.files)} chunk_size={req.chunk_size} overlap={req.overlap} batch_size={req.batch_size} store_text={req.store_text} provider={EMB_PROVIDER} model={EMB_MODEL if EMB_PROVIDER!='dummy' else '-'}")
     try:
         base, ds_dir, fx_dir = _proj_dirs(req.project_id)
 
@@ -217,12 +309,17 @@ def index(req: IndexRequest):
         _set_stage(st, "embedding")
         if EMB_PROVIDER == "dummy":
             xb = _emb_dummy([r["text"] for r in rows], dim=EMB_DIM)
-        else:
-            # fallback sur dummy tant que pas d'autre provider
-            xb = _emb_dummy([r["text"] for r in rows], dim=EMB_DIM)
+            dim = xb.shape[1]
+        elif EMB_PROVIDER == "st":
+            xb = _emb_st([r["text"] for r in rows])
+            dim = xb.shape[1]
+        else:  # "hf"
+            xb = _emb_hf([r["text"] for r in rows])
+            dim = xb.shape[1]
+
         st.embedded = xb.shape[0]
         _add_msg(st, f"Embeddings {st.embedded}/{st.total_chunks}")
-        _add_msg(st, f"Embeddings dim={xb.shape[1]}")
+        _add_msg(st, f"Embeddings dim={dim}")
 
         # 3) Sauvegarde dataset (texte)
         _save_dataset(ds_dir, rows)
@@ -230,7 +327,13 @@ def index(req: IndexRequest):
 
         # 4) FAISS
         _set_stage(st, "indexing")
-        _save_faiss(fx_dir, xb)
+        faiss_meta = {
+            "dim": int(dim),
+            "count": int(xb.shape[0]),
+            "provider": EMB_PROVIDER,
+            "model": EMB_MODEL if EMB_PROVIDER != "dummy" else None
+        }
+        _save_faiss(fx_dir, xb, meta=faiss_meta)
         st.indexed = int(xb.shape[0])
         _add_msg(st, f"FAISS écrit sur {os.path.join(fx_dir, 'emb.faiss')}")
         _add_msg(st, f"OK — dataset+index prêts (projet={req.project_id})")
@@ -265,10 +368,15 @@ def search(req: SearchRequest):
     if not rows:
         raise HTTPException(status_code=404, detail="dataset introuvable (index pas encore construit ?)")
 
-    # embeddings du query (dummy)
-    q = _emb_dummy([req.query], dim=EMB_DIM)[0:1, :]  # shape (1, d)
+    # Embedding de la requête avec le MÊME provider
+    if EMB_PROVIDER == "dummy":
+        q = _emb_dummy([req.query], dim=EMB_DIM)[0:1, :]
+    elif EMB_PROVIDER == "st":
+        q = _emb_st([req.query])[0:1, :]
+    else:
+        q = _emb_hf([req.query])[0:1, :]
 
-    # faiss
+    # FAISS
     index = _load_faiss(fx_dir)
     if index.d != q.shape[1]:
         raise HTTPException(status_code=500, detail=f"dim incompatibles: index.d={index.d} vs query={q.shape[1]}")
@@ -276,7 +384,6 @@ def search(req: SearchRequest):
     ids = ids[0].tolist()
     scores = scores[0].tolist()
 
-    # compose résultats
     out = []
     for idx, sc in zip(ids, scores):
         if idx < 0 or idx >= len(rows):
@@ -285,16 +392,14 @@ def search(req: SearchRequest):
         out.append({"path": r.get("path"), "text": r.get("text"), "score": float(sc)})
     return {"results": out}
 
-# ----------- ARTIFACTS EXPORT (ce qui manquait pour ton 404) -----------
+# ----------- ARTIFACTS EXPORT -----------
 @fastapi_app.get("/artifacts/{project_id}/dataset")
 def download_dataset(project_id: str):
     base, ds_dir, _ = _proj_dirs(project_id)
     if not os.path.isdir(ds_dir):
         raise HTTPException(status_code=404, detail="Dataset introuvable")
     buf = _tar_dir_to_bytes(ds_dir)
-    headers = {
-        "Content-Disposition": f'attachment; filename="{project_id}_dataset.tgz"'
-    }
+    headers = {"Content-Disposition": f'attachment; filename="{project_id}_dataset.tgz"'}
     return StreamingResponse(io.BytesIO(buf), media_type="application/gzip", headers=headers)
 
 @fastapi_app.get("/artifacts/{project_id}/faiss")
@@ -303,13 +408,11 @@ def download_faiss(project_id: str):
     if not os.path.isdir(fx_dir):
         raise HTTPException(status_code=404, detail="FAISS introuvable")
     buf = _tar_dir_to_bytes(fx_dir)
-    headers = {
-        "Content-Disposition": f'attachment; filename="{project_id}_faiss.tgz"'
-    }
+    headers = {"Content-Disposition": f'attachment; filename="{project_id}_faiss.tgz"'}
     return StreamingResponse(io.BytesIO(buf), media_type="application/gzip", headers=headers)
 
 # =============================================================================
-# GRADIO (UI facultative pour déclencher / tester rapidement)
+# GRADIO UI (facultatif)
 # =============================================================================
 def _ui_index(project_id: str, sample_text: str):
     files = [{"path": "sample.txt", "text": sample_text}]
@@ -331,8 +434,9 @@ def _ui_search(project_id: str, query: str, k: int):
     except Exception as e:
         return f"Erreur search: {e}"
 
-with gr.Blocks(title="Remote Indexer (FAISS mini)", analytics_enabled=False) as ui:
-    gr.Markdown("## Remote Indexer — demo UI (les vraies API sont sur `/index`, `/status/{job}`, `/search`, `/artifacts/...`).")
+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 '-'}'")
     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)
@@ -351,9 +455,9 @@ with gr.Blocks(title="Remote Indexer (FAISS mini)", analytics_enabled=False) as
 fastapi_app = gr.mount_gradio_app(fastapi_app, ui, path="/ui")
 
 # =============================================================================
-# MAIN (HF Space lancera ce module avec python -u main.py)
+# MAIN
 # =============================================================================
 if __name__ == "__main__":
     import uvicorn
     LOG.info("Démarrage Uvicorn sur 0.0.0.0:%s (UI_PATH=/ui)", PORT)
-    uvicorn.run(fastapi_app, host="0.0.0.0", port=PORT)
+    uvicorn.run(fastapi_app, host="0.0.0.0", port=PORT)