134 lines
4.8 KiB
Python
134 lines
4.8 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Indexation incrémentale des fiches avec BGE‑M3 (FlagEmbedding) + FAISS
|
||
---------------------------------------------------------------------
|
||
• Parcourt récursivement le dossier Fiches (ext .md .MD .markdown .txt)
|
||
• Découpe chaque fichier en blocs de ~800 tokens (chevauchement 100)
|
||
• Encode uniquement les passages provenant de fichiers **nouveaux ou modifiés**
|
||
depuis la dernière indexation (basé sur l'horodatage mtime).
|
||
• Les embeddings sont ajoutés à l'index existant sans toucher aux anciens.
|
||
• Écrit/Met à jour : corpus.idx (vecteurs) + corpus.meta.json (métadonnées)
|
||
|
||
➡ Robustesse améliorée : le script détecte tous les formats de retour possibles
|
||
de `BGEM3FlagModel.encode` (ndarray ou dict avec clef `embedding`, etc.).
|
||
"""
|
||
|
||
from pathlib import Path
|
||
import json
|
||
import re
|
||
import time
|
||
from datetime import datetime
|
||
import faiss
|
||
import numpy as np
|
||
from FlagEmbedding import BGEM3FlagModel
|
||
|
||
# --- Paramètres -------------------------------------------------------------
|
||
ROOT = Path("Fiches") # répertoire local des fiches
|
||
MODEL_NAME = "BAAI/bge-m3" # embedder multilingue, MIT
|
||
CHUNK = 800 # taille cible d'un bloc (≈600 mots)
|
||
OVERLAP = 100 # chevauchement pour la cohérence
|
||
INDEX_FILE = "corpus.idx"
|
||
META_FILE = "corpus.meta.json"
|
||
EXTENSIONS = ["*.md", "*.MD", "*.markdown", "*.txt"]
|
||
BATCH = 128 # plus grand batch : encode plus vite
|
||
|
||
# --- Fonctions utilitaires --------------------------------------------------
|
||
|
||
def split(text: str, chunk_size: int = CHUNK, overlap: int = OVERLAP):
|
||
"""Découpe un texte en morceaux de chunk_size mots avec overlap mots."""
|
||
sentences = re.split(r"(?<=[\.!?])\s+", text)
|
||
chunks, buf = [], []
|
||
for s in sentences:
|
||
buf.append(s)
|
||
if len(" ".join(buf)) > chunk_size:
|
||
chunks.append(" ".join(buf))
|
||
buf = buf[-overlap:]
|
||
if buf:
|
||
chunks.append(" ".join(buf))
|
||
return chunks
|
||
|
||
|
||
def gather_files(root: Path):
|
||
for pattern in EXTENSIONS:
|
||
yield from root.rglob(pattern)
|
||
|
||
# --- Chargement éventuel de l'index existant -------------------------------
|
||
|
||
def load_existing():
|
||
if not Path(INDEX_FILE).exists():
|
||
return None, [], 0 # pas d'index, pas de meta
|
||
|
||
index = faiss.read_index(INDEX_FILE)
|
||
meta = json.load(open(META_FILE, encoding="utf-8"))
|
||
return index, meta, len(meta)
|
||
|
||
# --- Pipeline principal -----------------------------------------------------
|
||
|
||
def main():
|
||
index, meta, existing = load_existing()
|
||
|
||
# mapping chemin relatif ➜ mtime stocké
|
||
meta_mtime = {m["file"]: m.get("mtime", 0) for m in meta}
|
||
|
||
new_docs, new_meta = [], []
|
||
files_scanned, files_updated = 0, 0
|
||
|
||
for fp in gather_files(ROOT):
|
||
files_scanned += 1
|
||
rel = fp.relative_to(ROOT).as_posix()
|
||
mtime = int(fp.stat().st_mtime)
|
||
if meta_mtime.get(rel) == mtime:
|
||
continue # inchangé
|
||
files_updated += 1
|
||
text = fp.read_text(encoding="utf-8", errors="ignore")
|
||
for i, chunk in enumerate(split(text)):
|
||
new_docs.append(chunk)
|
||
new_meta.append({"file": rel, "part": i, "mtime": mtime})
|
||
|
||
if not new_docs:
|
||
print("Aucun fichier nouveau ou modifié. Index à jour ✔︎")
|
||
return
|
||
|
||
print(
|
||
f"Nouveaux passages : {len(new_docs)} issus de {files_updated} fiches ("\
|
||
f"{files_scanned} fiches scannées). Génération des embeddings…"
|
||
)
|
||
|
||
model = BGEM3FlagModel(MODEL_NAME, device="cpu")
|
||
emb_out = model.encode(new_docs, batch_size=BATCH)
|
||
|
||
# --- Gestion robuste du format de sortie --------------------------------
|
||
if isinstance(emb_out, np.ndarray):
|
||
emb = emb_out
|
||
elif isinstance(emb_out, dict):
|
||
for key in ("embedding", "embeddings", "sentence_embeds", "sentence_embedding"):
|
||
if key in emb_out:
|
||
emb = np.asarray(emb_out[key])
|
||
break
|
||
else:
|
||
emb = np.asarray(next(iter(emb_out.values())))
|
||
else: # liste de tableaux ?
|
||
emb = np.asarray(emb_out)
|
||
|
||
emb = emb.astype("float32")
|
||
emb /= np.linalg.norm(emb, axis=1, keepdims=True) + 1e-12
|
||
|
||
if index is None:
|
||
index = faiss.IndexFlatIP(emb.shape[1])
|
||
index.add(emb)
|
||
|
||
# Mettre à jour les métadonnées (anciens + nouveaux)
|
||
meta.extend(new_meta)
|
||
faiss.write_index(index, INDEX_FILE)
|
||
json.dump(meta, open(META_FILE, "w", encoding="utf-8"), ensure_ascii=False, indent=2)
|
||
|
||
print(
|
||
f"Index mis à jour : {len(meta)} vecteurs au total ("\
|
||
f"+{len(new_docs)}). Dernière maj : {datetime.now().isoformat(timespec='seconds')}"
|
||
)
|
||
|
||
if __name__ == "__main__":
|
||
t0 = time.time()
|
||
main()
|
||
print(f"Terminé en {time.time() - t0:.1f} s")
|