pmock/localRAG
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add Docling corpus pipeline, retrieval evaluation and hybrid search with reranking

Browse Source 
- prepare_corpus.py: split double-page scans (split_spreads.py), re-OCR with
  OCRmyPDF/Tesseract only where the existing text layer is worse, Docling
  layout extraction with heading hierarchy, QA report per PDF
- docling_loader.py: cleanup (margin noise, split paragraphs, footnote marks,
  OCR comma/period repair), page mapping back to the original PDF, chunk
  headers with author, title and section
- local_rag.py: bge-m3 embeddings, configurable chunking and loader, index
  settings check (index_info.json), hybrid/union BM25 search with Reciprocal
  Rank Fusion and bge-reranker-v2-m3 cross-encoder; query instruction text
  removed after it lowered recall
- eval_retrieval.py: CLI options, same retrieval path as the chat, tolerant
  file name matching; goldsets and evaluation reports for the Kafka corpus
- grobid_test.py, grobid_qa.py, extract_docling.py, compare_extractors.py:
  extractor tests and quality checks
- .gitignore: corpus PDFs, vector stores and extractor output stay local
- server.py: uses ./vectorstore_docling_hdr

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Philipp Mock
2026-10-07 16:48:12 +02:00
co-authored by Claude Opus 5.5
parent 7d70169ee3
commit d54b3e5b8b
49 changed files with 5007 additions and 75766 deletions
Download Patch File Download Diff File Expand all files Collapse all files
local_rag.py
+179 -14
View File
@@ -2,9 +2,13 @@
Local RAG setup with LangChain, Ollama/OpenAI, and FAISS
Minimal dependencies, simple code
"""
import json
import os
import re
from pathlib import Path
import numpy as np
from langchain_community.document_loaders import PyPDFLoader, TextLoader
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
@@ -12,17 +16,52 @@ from langchain_ollama import ChatOllama
from langchain_openai import ChatOpenAI
from langchain_text_splitters import RecursiveCharacterTextSplitter
import docling_loader
# Pipeline tuning: used when adding documents (rebuild index after changing).
CHUNK_SIZE = 2000
CHUNK_OVERLAP = 400
# Characters, not tokens: 1500 characters of German prose are roughly 400 tokens.
CHUNK_SIZE = 1500
CHUNK_OVERLAP = 300
# Multilingual embedding model with a long input window (8192 tokens), so chunks
# are never truncated. Changing it requires rebuilding the index.
EMBEDDING_MODEL = "BAAI/bge-m3"
# How PDFs are read: "pypdf" (PyPDFLoader, one unit per page) or "docling"
# (structured output from prepare_corpus.py, cleaned by docling_loader.py).
PDF_LOADER = "pypdf"
# Where prepare_corpus.py put the Docling output for the PDFs being indexed.
PROCESSED_DIR = "corpus_processed/data_split"
# Prefix every chunk with author and title from the file name (Docling: plus section),
# so questions like "according to Lehmann ..." can match. See docling_loader.chunk_header.
CHUNK_HEADERS = True
# Retrieval / query tuning
RETRIEVAL_K = 5
# Hybrid search: besides the vector search, a keyword search (BM25 on German word stems)
# finds names, titles and quotations that embeddings blur.
# "dense" vector search alone
# "hybrid" both lists merged with Reciprocal Rank Fusion, which also sets the order
# "union" BM25 only adds candidates (BM25_CANDIDATES) to the vector candidates;
# the order comes from the vector search or, if enabled, from the reranker
RETRIEVAL_MODE = "hybrid"
CANDIDATES_K = 50 # candidates per search before merging and reranking
BM25_CANDIDATES = 20 # extra keyword candidates in "union" mode
RRF_K = 60 # standard constant of Reciprocal Rank Fusion
# Cross-encoder that rereads question and candidate together and re-sorts the candidates.
# None switches reranking off.
RERANKER_MODEL = "BAAI/bge-reranker-v2-m3"
# Appended to the embedding query for FAISS (does not affect the answer-generation prompt).
QUERY_ADDITIONAL_INSTRUCTIONS = (
"Do not return a list of references but prioritize meaningful text from abstracts, results and discussion sections."
)
# Left empty on purpose: embedding models do not follow instructions, and on the Kafka
# corpus the former text ("Do not return a list of references ...") lowered Recall@5
# from 0.67 to 0.53 (eval_retrieval.py --query-instructions). Measure before re-enabling.
QUERY_ADDITIONAL_INSTRUCTIONS = ""
# Written next to every index; an index built with another embedding model is not loaded.
INDEX_INFO_FILE = "index_info.json"
# Opening instructions for the final answer LLM call only. The rest (history, context, question)
# is assembled in query_with_history.
@@ -38,23 +77,41 @@ class LocalRAG:
ollama_model="gpt-oss:20b",
openai_model="gpt-5-mini",
ollama_base_url="http://localhost:11434",
pdf_loader=PDF_LOADER,
processed_dir=PROCESSED_DIR,
retrieval_mode=RETRIEVAL_MODE,
reranker_model=RERANKER_MODEL,
chunk_size=CHUNK_SIZE,
chunk_overlap=CHUNK_OVERLAP,
chunk_headers=CHUNK_HEADERS,
):
"""Initialize local RAG system. llm_provider: 'ollama' or 'openai'."""
"""Initialize local RAG system. llm_provider: 'ollama' or 'openai'.
pdf_loader: 'pypdf' or 'docling' (see PDF_LOADER)."""
self.vectorstore_path = vectorstore_path
self.llm_provider = llm_provider
self.pdf_loader = pdf_loader
self.processed_dir = Path(processed_dir)
self.chunk_size, self.chunk_overlap, self.chunk_headers = chunk_size, chunk_overlap, chunk_headers
self.retrieval_mode = retrieval_mode
self.reranker_model = reranker_model
self._bm25 = None # built on first use from the documents in the index
self._reranker = None # loaded on first use
# Embeddings
print("Loading embeddings model...")
# Normalized vectors: FAISS's L2 distance then ranks exactly like cosine similarity.
self.embeddings = HuggingFaceEmbeddings(
model_name="sentence-transformers/all-MiniLM-L6-v2"
model_name=EMBEDDING_MODEL,
encode_kwargs={"normalize_embeddings": True},
query_encode_kwargs={"normalize_embeddings": True},
)
# Text splitter (used when adding documents; tune CHUNK_SIZE / CHUNK_OVERLAP above)
self.text_splitter = RecursiveCharacterTextSplitter(
chunk_size=CHUNK_SIZE,
chunk_overlap=CHUNK_OVERLAP,
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
)
print(f"Text splitter: chunk_size={CHUNK_SIZE}, chunk_overlap={CHUNK_OVERLAP}")
print(f"Text splitter: chunk_size={chunk_size}, chunk_overlap={chunk_overlap}")
# LLM (Ollama or OpenAI)
if llm_provider == "openai":
@@ -76,10 +133,41 @@ class LocalRAG:
self.vectorstore = None
self._load_vectorstore()
def _index_info(self):
"""Settings that must match between the index on disk and this code."""
return {
"embedding_model": EMBEDDING_MODEL,
"chunk_size": self.chunk_size,
"chunk_overlap": self.chunk_overlap,
"pdf_loader": self.pdf_loader,
"chunk_headers": self.chunk_headers,
}
def _index_matches(self, index_file):
"""False if the index was built with a different embedding model than EMBEDDING_MODEL."""
info_file = os.path.join(self.vectorstore_path, INDEX_INFO_FILE)
if os.path.exists(info_file):
with open(info_file, encoding="utf-8") as f:
built_with = json.load(f).get("embedding_model")
if built_with != EMBEDDING_MODEL:
print(f"Index in {self.vectorstore_path} was built with {built_with}, "
f"but EMBEDDING_MODEL is {EMBEDDING_MODEL}. Rebuild it (add_pdfs.py).")
return False
return True
# Older index without info file: at least the vector dimensions must agree
import faiss
index_dim = faiss.read_index(index_file).d
model_dim = len(self.embeddings.embed_query("dimension check"))
if index_dim != model_dim:
print(f"Index in {self.vectorstore_path} has {index_dim} dimensions, {EMBEDDING_MODEL} "
f"produces {model_dim}. It was built with another model; rebuild it (add_pdfs.py).")
return False
return True
def _load_vectorstore(self):
"""Load existing vector store if available"""
index_file = os.path.join(self.vectorstore_path, "index.faiss")
if os.path.exists(index_file):
if os.path.exists(index_file) and self._index_matches(index_file):
try:
self.vectorstore = FAISS.load_local(
self.vectorstore_path,
@@ -102,6 +190,17 @@ class LocalRAG:
print(f"Warning: {file_path} not found, skipping")
continue
# Docling: already structured and cleaned, the loader returns finished chunks
if path.suffix.lower() == '.pdf' and self.pdf_loader == "docling":
processed = self.processed_dir / f"{path.stem}.json"
if processed.exists():
chunks = docling_loader.load(processed, str(path), self.text_splitter,
header=self.chunk_headers)
all_docs.extend(chunks)
print(f" - {path.name}: {len(chunks)} chunks (docling)")
continue
print(f"Warning: no Docling output for {path.name} in {self.processed_dir}, using PyPDFLoader")
# Load document
if path.suffix.lower() == '.pdf':
loader = PyPDFLoader(str(path))
@@ -113,6 +212,10 @@ class LocalRAG:
docs = loader.load()
chunks = self.text_splitter.split_documents(docs)
if self.chunk_headers:
header = docling_loader.chunk_header(str(path))
for chunk in chunks:
chunk.page_content = header + chunk.page_content
all_docs.extend(chunks)
print(f" - {path.name}: {len(chunks)} chunks")
@@ -131,6 +234,8 @@ class LocalRAG:
# Save
os.makedirs(self.vectorstore_path, exist_ok=True)
self.vectorstore.save_local(self.vectorstore_path)
with open(os.path.join(self.vectorstore_path, INDEX_INFO_FILE), "w", encoding="utf-8") as f:
json.dump(self._index_info(), f, indent=2)
print(f"Vector store saved to {self.vectorstore_path}")
def list_documents(self):
@@ -204,6 +309,64 @@ class LocalRAG:
for doc, score in docs_with_scores
]
# ---------- retrieval: vector search, keyword search, reranking ----------
_stemmer = None
@classmethod
def _bm25_tokens(cls, text):
"""Lowercase German word stems, so "ohnmächtigen" matches "Ohnmacht"."""
if cls._stemmer is None:
import snowballstemmer
cls._stemmer = snowballstemmer.stemmer("german")
return cls._stemmer.stemWords(re.findall(r"[a-zäöüß0-9]{2,}", text.lower()))
def _bm25_index(self):
if self._bm25 is None:
from rank_bm25 import BM25Okapi
docs = list(self.vectorstore.docstore._dict.values())
self._bm25 = (BM25Okapi([self._bm25_tokens(d.page_content) for d in docs]), docs)
return self._bm25
def retrieve(self, query, k=None):
"""The k best chunks for a query as [(Document, score)].
dense: vector search only (score = L2 distance, lower is better)
hybrid: vector search + BM25, merged by Reciprocal Rank Fusion (higher is better)
union: vector candidates plus BM25_CANDIDATES keyword candidates, vector order first
With a reranker, the merged candidates are re-sorted by the cross-encoder
(score = relevance, higher is better)."""
k = k or RETRIEVAL_K
pool = max(k, CANDIDATES_K) if (self.retrieval_mode != "dense" or self.reranker_model) else k
dense = self.vectorstore.similarity_search_with_score(query, k=pool)
key = lambda doc: (doc.metadata.get("source"), doc.metadata.get("page"), doc.page_content[:200])
if self.retrieval_mode in ("hybrid", "union"):
bm25, docs = self._bm25_index()
scores = bm25.get_scores(self._bm25_tokens(query))
keyword = [docs[i] for i in np.argsort(scores)[::-1][:pool]]
if self.retrieval_mode == "union":
seen = {key(d) for d, _ in dense}
extra = [d for d in keyword if key(d) not in seen][:BM25_CANDIDATES]
candidates = [[d, s] for d, s in dense] + [[d, float("nan")] for d in extra]
elif self.retrieval_mode == "hybrid":
fused = {}
for ranking in ([d for d, _ in dense], keyword):
for rank, doc in enumerate(ranking):
entry = fused.setdefault(key(doc), [doc, 0.0])
entry[1] += 1 / (RRF_K + rank + 1)
candidates = sorted(fused.values(), key=lambda e: e[1], reverse=True)[:pool]
else:
candidates = [[d, s] for d, s in dense]
if self.reranker_model and candidates:
if self._reranker is None:
from sentence_transformers import CrossEncoder
print(f"Loading reranker {self.reranker_model}...")
self._reranker = CrossEncoder(self.reranker_model, max_length=1024)
relevance = self._reranker.predict([(query, d.page_content) for d, _ in candidates])
candidates = [[d, float(r)] for (d, _), r in zip(candidates, relevance)]
candidates.sort(key=lambda e: e[1], reverse=True)
return [(d, s) for d, s in candidates[:k]]
def query(self, question, k=None):
"""Query the RAG system (no conversation history). Returns dict with 'answer' and 'retrieved'."""
if k is None:
@@ -243,12 +406,14 @@ Standalone question:"""
search_query = (rephrase_response.content if hasattr(rephrase_response, "content") else str(rephrase_response)).strip() or question
print(f"[RAG] Standalone search query (rephrased): {search_query!r}")
retrieval_query = f"{search_query}\n\n{QUERY_ADDITIONAL_INSTRUCTIONS}"
retrieval_query = search_query
if QUERY_ADDITIONAL_INSTRUCTIONS:
retrieval_query += f"\n\n{QUERY_ADDITIONAL_INSTRUCTIONS}"
print(f"[RAG] Search query: {search_query!r}")
print(f"[RAG] Retrieval query sent to vector store: {retrieval_query!r}")
# 2) Retrieve documents with scores (FAISS: L2 distance, lower = more similar)
docs_with_scores = self.vectorstore.similarity_search_with_score(retrieval_query, k=k)
# 2) Retrieve documents with scores (see retrieve() for what the score means)
docs_with_scores = self.retrieve(retrieval_query, k=k)
docs = [doc for doc, _ in docs_with_scores]
retrieved = self._docs_scores_to_retrieved(docs_with_scores)
if docs_with_scores: