Eric Furst
5a3294f74c
- Rename build_exp_claude.py → build_store.py - Rename query_hybrid_bm25_v4.py → query_hybrid.py - Rename retrieve_hybrid_raw.py → retrieve.py - Archive query_topk_prompt_engine_v3.py (superseded by hybrid) - Archive retrieve_raw.py (superseded by hybrid) - Move build_clippings.py, retrieve_clippings.py → clippings_search/ - Update run_query.sh, README.md, CLAUDE.md for new names
97 lines
2.9 KiB
Python
97 lines
2.9 KiB
Python
# retrieve_raw.py
|
|
# Verbatim chunk retrieval: vector search + cross-encoder re-ranking, no LLM.
|
|
#
|
|
# Returns the top re-ranked chunks with their full text, file metadata, and
|
|
# scores. Useful for browsing source material directly and verifying what
|
|
# the RAG pipeline retrieves before LLM synthesis.
|
|
#
|
|
# Uses the same vector store, embedding model, and re-ranker as
|
|
# query_topk_prompt_engine_v3.py, but skips the LLM step entirely.
|
|
#
|
|
# E.M.F. February 2026
|
|
|
|
# Environment vars must be set before importing huggingface/transformers
|
|
# libraries, because huggingface_hub.constants evaluates HF_HUB_OFFLINE
|
|
# at import time.
|
|
import os
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
os.environ["SENTENCE_TRANSFORMERS_HOME"] = "./models"
|
|
os.environ["HF_HUB_OFFLINE"] = "1"
|
|
|
|
from llama_index.core import (
|
|
StorageContext,
|
|
load_index_from_storage,
|
|
Settings,
|
|
)
|
|
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
|
|
from llama_index.core.postprocessor import SentenceTransformerRerank
|
|
import sys
|
|
import textwrap
|
|
|
|
#
|
|
# Globals
|
|
#
|
|
|
|
# Embedding model (must match build_exp_claude.py)
|
|
EMBED_MODEL = HuggingFaceEmbedding(cache_folder="./models", model_name="BAAI/bge-large-en-v1.5", local_files_only=True)
|
|
|
|
# Cross-encoder model for re-ranking (cached in ./models/)
|
|
RERANK_MODEL = "cross-encoder/ms-marco-MiniLM-L-12-v2"
|
|
RERANK_TOP_N = 15
|
|
RETRIEVE_TOP_K = 30
|
|
|
|
# Output formatting
|
|
WRAP_WIDTH = 80
|
|
|
|
|
|
def main():
|
|
# No LLM needed -- set embed model only
|
|
Settings.embed_model = EMBED_MODEL
|
|
|
|
# Load persisted vector store
|
|
storage_context = StorageContext.from_defaults(persist_dir="./storage_exp")
|
|
index = load_index_from_storage(storage_context)
|
|
|
|
# Build retriever (vector search only, no query engine / LLM)
|
|
retriever = index.as_retriever(similarity_top_k=RETRIEVE_TOP_K)
|
|
|
|
# Cross-encoder re-ranker
|
|
reranker = SentenceTransformerRerank(
|
|
model=RERANK_MODEL,
|
|
top_n=RERANK_TOP_N,
|
|
)
|
|
|
|
# Query
|
|
if len(sys.argv) < 2:
|
|
print("Usage: python retrieve_raw.py QUERY_TEXT")
|
|
sys.exit(1)
|
|
q = " ".join(sys.argv[1:])
|
|
|
|
# Retrieve and re-rank
|
|
nodes = retriever.retrieve(q)
|
|
reranked = reranker.postprocess_nodes(nodes, query_str=q)
|
|
|
|
# Output
|
|
print(f"\nQuery: {q}")
|
|
print(f"Retrieved {len(nodes)} chunks, re-ranked to top {len(reranked)}\n")
|
|
|
|
for i, node in enumerate(reranked, 1):
|
|
meta = getattr(node, "metadata", None) or node.node.metadata
|
|
score = getattr(node, "score", None)
|
|
file_name = meta.get("file_name", "unknown")
|
|
text = node.get_content()
|
|
|
|
print("="*WRAP_WIDTH)
|
|
print(f"=== [{i}] {file_name} (score: {score:.3f}) ")
|
|
print("="*WRAP_WIDTH)
|
|
# Wrap text for readability
|
|
for line in text.splitlines():
|
|
if line.strip():
|
|
print(textwrap.fill(line, width=WRAP_WIDTH))
|
|
else:
|
|
print()
|
|
print()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|