Add embed_fallout v3 — structural chunking + multi-collection routing
This commit is contained in:
@@ -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()
|
||||
Reference in New Issue
Block a user