pmock/localRAG
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
d54b3e5b8b32ed258266fdadf0cd5ff45c7241a5
localRAG/local_rag.py
T
Philipp MockandClaude Opus 5.5 d54b3e5b8b Add Docling corpus pipeline, retrieval evaluation and hybrid search with reranking 
- 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>
2026-10-07 16:48:12 +02:00

477 lines
20 KiB
Python
Raw Blame History

"""
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
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).
# 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).
# 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.
ANSWER_PROMPT = """You are an assistant for question-answering. Use the chat history (if any) and the retrieved context below to answer the current question.
If you don't know the answer, say so. Keep the conversation coherent."""
class LocalRAG:
def __init__(
self,
vectorstore_path="./vectorstore",
llm_provider="ollama",
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'.
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=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,
)
print(f"Text splitter: chunk_size={chunk_size}, chunk_overlap={chunk_overlap}")
# LLM (Ollama or OpenAI)
if llm_provider == "openai":
api_key = os.environ.get("OPENAI_API_KEY")
if not api_key:
raise ValueError(
"OPENAI_API_KEY environment variable is required when llm_provider='openai'"
)
print(f"Using OpenAI (model: {openai_model})...")
self.llm = ChatOpenAI(model=openai_model, api_key=api_key)
else:
print(f"Using Ollama (model: {ollama_model})...")
self.llm = ChatOllama(
model=ollama_model,
base_url=ollama_base_url
)
# Vector store (load if exists, otherwise None)
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) and self._index_matches(index_file):
try:
self.vectorstore = FAISS.load_local(
self.vectorstore_path,
self.embeddings,
allow_dangerous_deserialization=True
)
print(f"Loaded existing vector store from {self.vectorstore_path}")
except Exception as e:
print(f"Could not load vector store: {e}")
self.vectorstore = None
def add_documents(self, file_paths):
"""Add documents to the vector store"""
print(f"\nLoading {len(file_paths)} document(s)...")
all_docs = []
for file_path in file_paths:
path = Path(file_path)
if not path.exists():
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))
elif path.suffix.lower() in ['.txt', '.md']:
loader = TextLoader(str(path))
else:
print(f"Warning: Unsupported file type {path.suffix}, skipping")
continue
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")
if not all_docs:
print("No documents loaded!")
return
# Create or update vector store
print(f"\nCreating embeddings for {len(all_docs)} chunks...")
if self.vectorstore is None:
self.vectorstore = FAISS.from_documents(all_docs, self.embeddings)
else:
new_store = FAISS.from_documents(all_docs, self.embeddings)
self.vectorstore.merge_from(new_store)
# 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):
"""List all documents in the vector store"""
if self.vectorstore is None:
print("No documents in vector store.")
return []
# Get all documents from the vector store
# We'll retrieve a large number to get all documents
all_docs = self.vectorstore.similarity_search("", k=10000) # Large k to get all
# Extract unique document sources from metadata
documents = {}
for doc in all_docs:
source = doc.metadata.get('source', 'Unknown')
if source not in documents:
documents[source] = {
'source': source,
'chunks': 0,
'page': doc.metadata.get('page', None)
}
documents[source]['chunks'] += 1
# Convert to list and sort
doc_list = list(documents.values())
doc_list.sort(key=lambda x: x['source'])
print(f"\nDocuments in vector store ({len(doc_list)} unique documents):")
print("-" * 60)
for doc_info in doc_list:
print(f" - {doc_info['source']}")
print(f" Chunks: {doc_info['chunks']}")
if doc_info['page'] is not None:
print(f" Page: {doc_info['page']}")
return doc_list
def _format_history(self, chat_history):
"""Format chat history as a string for prompts."""
lines = []
for turn in chat_history or []:
role = (turn.get("role") or "").lower()
content = (turn.get("content") or "").strip()
if role == "user":
lines.append(f"User: {content}")
elif role == "assistant":
lines.append(f"Assistant: {content}")
return "\n".join(lines) if lines else ""
def _docs_to_retrieved(self, docs):
"""Convert document list to retrieved chunks format for API."""
return [
{
"content": doc.page_content,
"source": doc.metadata.get("source", ""),
"page": doc.metadata.get("page"),
}
for doc in docs
]
def _docs_scores_to_retrieved(self, docs_with_scores):
"""Convert (Document, score) list to retrieved chunks format with score. FAISS returns L2 distance (lower = more similar)."""
return [
{
"content": doc.page_content,
"source": doc.metadata.get("source", ""),
"page": doc.metadata.get("page"),
"score": float(score),
}
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:
k = RETRIEVAL_K
return self.query_with_history(question, chat_history=[], k=k)
def query_with_history(self, question, chat_history=None, k=None):
"""Query the RAG with conversation history: rephrase question using history for retrieval,
then answer with full conversation + retrieved context in the prompt.
Returns dict with 'answer' and 'retrieved' (list of chunks with content, source, page).
"""
if k is None:
k = RETRIEVAL_K
if self.vectorstore is None:
return {
"answer": "Error: No documents loaded. Please add documents first.",
"retrieved": [],
}
history_str = self._format_history(chat_history)
search_query = question
print(f"[RAG] User question: {question!r}")
# 1) If we have history, rephrase the question into a standalone query for better retrieval
if history_str.strip():
rephrase_prompt = f"""Given this chat history and the latest user question, write a single standalone question that captures what the user is asking.
Do not answer it; only output the standalone question. If the latest question is already clear on its own, output it unchanged.
Chat history:
{history_str}
Latest user question: {question}
Standalone question:"""
rephrase_response = self.llm.invoke(rephrase_prompt)
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 = 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 (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:
scores = [f"{s:.3f}" for _, s in docs_with_scores]
print(f"[RAG] Retrieved {len(docs)} chunk(s), scores (L2 dist): [{', '.join(scores)}]")
context = "\n\n".join([doc.page_content for doc in docs])
# 3) Answer using conversation history + retrieved context
history_block = f"Chat history:\n{history_str}\n\n" if history_str else ""
answer_prompt = (
ANSWER_PROMPT
+ f"""
{history_block}Relevant context from documents:
{context}
Current question: {question}
Answer:"""
)
print(f"[RAG] Composed answer prompt:\n{answer_prompt}")
response = self.llm.invoke(answer_prompt)
answer = response.content if hasattr(response, "content") else str(response)
return {"answer": answer, "retrieved": retrieved}
def main():
"""Example usage"""
print("=" * 60)
print("Local RAG with LangChain, Ollama/OpenAI, and FAISS")
print("=" * 60)
# Initialize
rag = LocalRAG(ollama_model="gpt-oss:20b")
# Add documents (uncomment and add your file paths)
# rag.add_documents([
# "data/dok1.pdf",
# "data/dok2.pdf",
# "data/dok3.pdf"
# ])
# List documents
rag.list_documents()
# Query
question = "What do you knowabout modality for perceived message perception?"
result = rag.query(question)
print(f"\nQuestion: {question}")
print(f"Answer: {result['answer']}")
if result.get("retrieved"):
print(f"Retrieved {len(result['retrieved'])} chunks")
# print("\nSetup complete! Uncomment the code above to add documents and query.")
if __name__ == "__main__":
main()
Reference in New Issue View Git Blame Copy Permalink