From ae6356ded44a7002b7d5951c5536098d7b216675 Mon Sep 17 00:00:00 2001 From: Corback Date: Wed, 17 Jun 2026 13:14:21 +0000 Subject: [PATCH] =?UTF-8?q?Add=20embed=5Ffallout=20v3=20=E2=80=94=20struct?= =?UTF-8?q?ural=20chunking=20+=20multi-collection=20routing?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pipelines/embed_fallout.py | 328 +++++++++++++++++++++++++++++++++++++ 1 file changed, 328 insertions(+) create mode 100644 pipelines/embed_fallout.py diff --git a/pipelines/embed_fallout.py b/pipelines/embed_fallout.py new file mode 100644 index 0000000..0fd4e7d --- /dev/null +++ b/pipelines/embed_fallout.py @@ -0,0 +1,328 @@ +""" +embed_fallout.py — Ingestion PDFs Fallout JdR → ChromaDB (v3) + +Fusion : + - Chunking structurel par paragraphes (fallout_ingest_v2, 15/06) + - Routing multi-collections selon README priorités (16/06) + +Collections : + fallout_lore_canon ← lore_post_guerre/ (canonique, priorité max) + fallout_regles ← regles/core|supplements|fiches_de_jeu/ + fallout_lore_contexte ← lore_pre_guerre/ + aventures/ + ambiance/radio/ + +Chemins container : PDF_DIR=/pdfs/jdr STATE=/pdfs/fallout_embed_state.json +Env vars : CHROMA_URL, OLLAMA_URL +""" +import json, os, re, subprocess, time, urllib.request +from pathlib import Path + +CHROMA_URL = os.getenv("CHROMA_URL", "http://chromadb:8000") +CHROMA_BASE = CHROMA_URL + "/api/v2/tenants/default_tenant/databases/default_database" +OLLAMA_URL = os.getenv("OLLAMA_URL", "http://172.18.0.1:11434") +EMBED_MODEL = "nomic-embed-text" + +PDF_DIR = Path(os.getenv("PDF_DIR", "/pdfs/jdr")) +STATE_FILE = Path(os.getenv("STATE_FILE", "/pdfs/fallout_embed_state.json")) +BATCH_SIZE = 50 + +ROUTING = [ + ("lore_post_guerre", "fallout_lore_canon", "lore_post_guerre", 450, 80, True), + ("regles/core", "fallout_regles", "regles_core", 350, 60, False), + ("regles/supplements", "fallout_regles", "regles_supplement", 350, 60, False), + ("regles/fiches_de_jeu", "fallout_regles", "fiches_jeu", 350, 60, False), + ("lore_pre_guerre", "fallout_lore_contexte", "lore_pre_guerre", 450, 80, False), + ("aventures", "fallout_lore_contexte", "aventures", 450, 80, False), + ("ambiance/radio", "fallout_lore_contexte", "ambiance_radio", 200, 40, False), +] + +SKIP_DIRS = { + "assets_graphiques", "battlemaps", "fonts", "photoshop_psd", + "fonds_de_page", "images_ref", "croquis", "maps", "logos", + "couvertures", ".claude", "archive", +} +SKIP_SLUGS = {"archive__meta_data__compressed"} +MIN_WORDS = 40 + + +# ── State ───────────────────────────────────────────────────────────────────── + +def load_state(): + try: + return json.loads(STATE_FILE.read_text()) + except Exception: + return {"done": [], "failed": []} + +def save_state(s): + STATE_FILE.write_text(json.dumps(s, indent=2)) + + +# ── ChromaDB ────────────────────────────────────────────────────────────────── + +def _http(method, url, data=None, timeout=30): + body = json.dumps(data).encode() if data is not None else None + req = urllib.request.Request(url, data=body, method=method, + headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=timeout) as r: + return json.loads(r.read()) + +def get_or_create_collection(name): + try: + r = _http("POST", f"{CHROMA_BASE}/collections", + {"name": name, "metadata": {"hnsw:space": "cosine"}}) + return r["id"] + except urllib.error.HTTPError as e: + if e.code == 409: + cols = _http("GET", f"{CHROMA_BASE}/collections") + return next(c["id"] for c in cols if c["name"] == name) + raise + +def upsert_batch(col_id, ids, embeddings, documents, metadatas): + _http("POST", f"{CHROMA_BASE}/collections/{col_id}/upsert", { + "ids": ids, "embeddings": embeddings, + "documents": documents, "metadatas": metadatas, + }, timeout=60) + + +# ── Embeddings ──────────────────────────────────────────────────────────────── + +def embed(text): + safe = " ".join(text.split()[:280]) + r = _http("POST", f"{OLLAMA_URL}/api/embeddings", + {"model": EMBED_MODEL, "prompt": safe}, timeout=60) + return r["embedding"] + + +# ── PDF extraction ──────────────────────────────────────────────────────────── + +def extract_text(pdf_path): + """pdftotext -layout en priorité, fallback pypdf, fallback fitz.""" + text = "" + try: + r = subprocess.run( + ["pdftotext", "-layout", str(pdf_path), "-"], + capture_output=True, timeout=180, + ) + if r.returncode == 0: + text = r.stdout.decode("utf-8", errors="replace") + except (subprocess.TimeoutExpired, FileNotFoundError): + pass + + if len(text.strip()) < 200: + try: + from pypdf import PdfReader + reader = PdfReader(str(pdf_path)) + text = "\n".join(p.extract_text() or "" for p in reader.pages) + except Exception: + try: + import fitz + doc = fitz.open(str(pdf_path)) + text = "\n".join(p.get_text() for p in doc) + doc.close() + except Exception as e: + raise RuntimeError(f"Aucune méthode d'extraction n'a fonctionné : {e}") + + return text.strip() + + +# ── Chunking structurel ─────────────────────────────────────────────────────── + +def is_section_header(line): + l = line.strip() + if not l or len(l) > 80: + return False + if l.isupper() and len(l.split()) <= 8: + return True + if re.match(r'^(CHAPITRE|CHAPTER|PARTIE|PART|SECTION|\d+[\.\)])\s+\S', l, re.I): + return True + return False + +def chunk_structural(text, chunk_words=450, overlap_words=80): + raw_blocks = re.split(r'\n{2,}', text) + blocks = [] + for b in raw_blocks: + b = b.strip() + if not b: + continue + if len(b.split()) > 800: + sub = [s.strip() for s in b.split('\n') if s.strip()] + blocks.extend(sub) + else: + blocks.append(b) + + chunks = [] + current_blocks, current_words = [], 0 + + for block in blocks: + bwords = len(block.split()) + if bwords < MIN_WORDS and not is_section_header(block): + current_blocks.append(block) + current_words += bwords + continue + + if current_words >= chunk_words: + chunk = "\n\n".join(current_blocks).strip() + if len(chunk.split()) >= MIN_WORDS: + chunks.append(chunk) + overlap_b, overlap_w = [], 0 + for ob in reversed(current_blocks): + ow = len(ob.split()) + if overlap_w + ow <= overlap_words: + overlap_b.insert(0, ob) + overlap_w += ow + else: + break + current_blocks = overlap_b + [block] + current_words = overlap_w + bwords + else: + current_blocks.append(block) + current_words += bwords + + if current_blocks: + chunk = "\n\n".join(current_blocks).strip() + if len(chunk.split()) >= MIN_WORDS: + chunks.append(chunk) + + return chunks + + +# ── Routing & métadonnées ───────────────────────────────────────────────────── + +def route_pdf(rel_path): + rel = rel_path.replace("\\", "/") + for part in rel.split("/")[:-1]: + if part.lower() in SKIP_DIRS: + return None + rel_low = rel.lower() + for prefix, col, cat, cw, ov, canon in ROUTING: + if rel_low.startswith(prefix.lower()): + return col, cat, cw, ov, canon + return None + +def detect_langue(filename): + n = filename.lower() + if re.search(r'(_fr[_\-\.])|(_fr$)|(manuel.*officiel)|(bestiaire)|(armes)|' + r'(special.*comp)|(regles_maison)|(frequence)|(pizza)', n): + return "fr" + if re.search(r'(_en[_\-\.])|(_en$)|(quickstart)|(booklet)|(toolkit)|(settlers)|' + r'(enclave)|(rust_devil)|(winter_of_atom)|(cascadia)|(hunted)|' + r'(showdown)|(fully_operational)|(last_boat)', n): + return "en" + return "fr" + +def detect_version(filename): + m = re.search(r'v(\d+[\.\d]*)', filename.lower()) + return m.group(1) if m else "" + + +# ── Main ────────────────────────────────────────────────────────────────────── + +def main(): + print("=== embed_fallout v3 — chunking structurel + routing multi-collections ===\n") + + all_pdfs = [] + for root, dirs, files in os.walk(PDF_DIR): + dirs[:] = [d for d in dirs if d.lower() not in SKIP_DIRS] + for f in sorted(files): + if f.lower().endswith(".pdf"): + all_pdfs.append(Path(root) / f) + all_pdfs.sort() + + to_process, skipped = [], [] + for pdf_path in all_pdfs: + rel = str(pdf_path.relative_to(PDF_DIR)) + slug = rel.replace(os.sep, "__").replace(" ", "_").lower()[:-4] + r = route_pdf(rel) + if r is None or slug in SKIP_SLUGS: + skipped.append(rel) + continue + col, cat, cw, ov, canon = r + to_process.append({ + "path": pdf_path, "rel": rel, "slug": slug, + "col": col, "cw": cw, "ov": ov, + "meta": { + "source": pdf_path.name, + "categorie": cat, + "langue": detect_langue(pdf_path.name), + "canonique": "true" if canon else "false", + "version": detect_version(pdf_path.name), + }, + }) + + print(f"PDFs trouvés : {len(all_pdfs)}") + print(f" A ingérer : {len(to_process)}") + print(f" Ignorés : {len(skipped)}") + for s in skipped: + print(f" [SKIP] {s}") + + col_ids = {} + print() + for name in sorted(set(p["col"] for p in to_process)): + col_ids[name] = get_or_create_collection(name) + print(f" Collection '{name}' : {col_ids[name][:8]}...") + + state = load_state() + # Reset state pour repartir propre (collections wipées) + state = {"done": [], "failed": []} + save_state(state) + + total_chunks = 0 + + for item in to_process: + slug = item["slug"] + if slug in state["done"]: + print(f"[résumé] {item['rel']}") + continue + + print(f"\n[->] {item['rel']}") + print(f" col={item['col']} cat={item['meta']['categorie']}" + f" lang={item['meta']['langue']} canon={item['meta']['canonique']}") + + try: + text = extract_text(item["path"]) + print(f" {len(text)} chars extraits") + + chunks = chunk_structural(text, item["cw"], item["ov"]) + print(f" {len(chunks)} chunks (~{item['cw']}w cible, overlap {item['ov']}w)") + + ids, embeddings, documents, metadatas = [], [], [], [] + for i, chunk in enumerate(chunks): + t0 = time.time() + emb = embed(chunk) + wc = len(chunk.split()) + ids.append(f"{slug}__{i:04d}") + embeddings.append(emb) + documents.append(chunk) + metadatas.append({**item["meta"], + "chunk": i, "total_chunks": len(chunks), + "word_count": wc}) + print(f" chunk {i+1}/{len(chunks)} ({wc}w) {time.time()-t0:.1f}s") + + if len(ids) >= BATCH_SIZE: + upsert_batch(col_ids[item["col"]], ids, embeddings, documents, metadatas) + ids, embeddings, documents, metadatas = [], [], [], [] + + if ids: + upsert_batch(col_ids[item["col"]], ids, embeddings, documents, metadatas) + + total_chunks += len(chunks) + state["done"].append(slug) + state["failed"] = [f for f in state["failed"] if f.get("slug") != slug] + print(f" OK — {len(chunks)} chunks") + + except Exception as e: + state["failed"].append({"slug": slug, "error": str(e)}) + print(f" ERREUR : {e}") + + save_state(state) + + print(f"\n{'='*60}") + print(f"TERMINÉ — {len(state['done'])}/{len(to_process)} PDFs") + print(f"Chunks totaux : {total_chunks}") + if state.get("failed"): + print("Échecs :") + for f in state["failed"]: + print(f" - {f['slug']}: {f['error']}") + + +if __name__ == "__main__": + main()