Notes

import os
import uuid
from typing import List, Dict
from sentence_transformers import SentenceTransformer, CrossEncoder
from qdrant_client import QdrantClient
from qdrant_client.http.models import VectorParams, Distance, PointStruct
from langchain.text_splitter import RecursiveCharacterTextSplitter
from unstructured.partition.pdf import partition_pdf
from unstructured.partition.docx import partition_docx
from unstructured.cleaners.core import clean
from rank_bm25 import BM25Okapi

CHUNK_SIZE = 1000
CHUNK_OVERLAP = 150
EMBED_MODEL = "sentence-transformers/all-mpnet-base-v2" # 768-dim
CROSS_ENCODER_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"
QDRANT_PATH = "/home/parveenk/Downloads/rfp_proposal/qdrant"
COLLECTION_NAME = "rfp"
CHUNKS_FILE = "chunks_debug.txt"

def load_documents(file_paths: List[str]) -> List[Dict]:
docs = []
for path in file_paths:
ext = os.path.splitext(path)[-1].lower()
if ext == ".pdf":
elements = partition_pdf(path)
elif ext in [".docx", ".doc"]:
elements = partition_docx(path)
else:
raise ValueError(f"Unsupported file type: {ext}")

raw_text = "n".join([el.text for el in elements if el.text])
cleaned_text = clean(raw_text, bullets=True, dashes=True)
docs.append({"filename": os.path.basename(path), "text": cleaned_text})
return docs

def chunk_documents(docs: List[Dict], chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP) -> List[Dict]:
splitter = RecursiveCharacterTextSplitter(
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
separators=["nn", "n", ".", "?", "!", " ", ""]
)
chunks = []
for doc in docs:
filename = doc["filename"]
for chunk in splitter.split_text(doc["text"]):
chunks.append({
"id": str(uuid.uuid4()),
"text": chunk,
"source": filename
})
# Save chunks for debugging
with open(CHUNKS_FILE, "w", encoding="utf-8") as f:
for c in chunks:
f.write(f"ID: {c['id']} | FILE: {c['source']}nTEXT: {c['text']}nn{'='*80}nn")
return chunks

def embed_chunks(chunks: List[Dict], embedder) -> List[List[float]]:
texts = [c["text"] for c in chunks]
return embedder.encode(texts, convert_to_tensor=False)

def setup_qdrant():
client = QdrantClient(path=QDRANT_PATH)
try:
client.create_collection(
collection_name=COLLECTION_NAME,
vectors_config=VectorParams(size=768, distance=Distance.COSINE),
)
print(f"✅ Collection '{COLLECTION_NAME}' created.")
except:
print(f"⚠️ Collection '{COLLECTION_NAME}' already exists.")
return client

def store_in_qdrant(client, chunks, embeddings):
points = [
PointStruct(
id=c["id"],
vector=emb,
payload={"source": c["source"], "id": c["id"], "text": c["text"]}
)
for c, emb in zip(chunks, embeddings)
]
client.upsert(collection_name=COLLECTION_NAME, points=points)
print(f"✅ Stored {len(points)} chunks in Qdrant.")

def setup_bm25(chunks: List[Dict]):
tokenized_corpus = [c["text"].split() for c in chunks]
return BM25Okapi(tokenized_corpus)

def rerank_cross_encoder(query, candidates, cross_encoder):
pairs = [[query, c[0]] for c in candidates] # c[0] = text
scores = cross_encoder.predict(pairs)
ranked = sorted(zip(candidates, scores), key=lambda x: x[1], reverse=True)
return ranked

def hybrid_search(query, chunks, client, bm25, cross_encoder, top_k=1):
# Vector search
search_result = client.search(
collection_name=COLLECTION_NAME,
query_vector=embedder.encode(query).tolist(),
limit=top_k * 3 # retrieve more for reranking
)
vector_candidates = [(hit.payload["text"], hit.payload["source"]) for hit in search_result]

# BM25 search
bm25_texts = bm25.get_top_n(query.split(), [c["text"] for c in chunks], n=top_k * 3)
bm25_candidates = [(c["text"], c["source"]) for c in chunks if c["text"] in bm25_texts]

# Combine and deduplicate
seen = set()
combined = []
for text, source in vector_candidates + bm25_candidates:
if text not in seen:
combined.append((text, source))
seen.add(text)

# Cross-Encoder rerank
reranked = rerank_cross_encoder(query, combined, cross_encoder)
return reranked[:top_k]

def attach_next_chunk(results, all_chunks):
"""
For each result (text, source), attach its immediate next chunk if from the same file.
"""
expanded = []
seen = set()

for (text, source), score in results:
# Add the main result
expanded.append(((text, source), score))
seen.add(text)

# Find index in all_chunks and attach next chunk
for i, c in enumerate(all_chunks):
if c["text"] == text and c["source"] == source:
if i + 1 < len(all_chunks) and all_chunks[i + 1]["source"] == source:
next_text = all_chunks[i + 1]["text"]
if next_text not in seen:
expanded.append(((next_text, source), score - 0.0001)) # Slightly lower score
seen.add(next_text)
break
return expanded

if __name__ == "__main__":
file_paths = [
"/home/parveenk/Downloads/rfp_proposal/docs/HOLLYWOOD - Response_061024_removed.pdf"
]

# Ingest pipeline
docs = load_documents(file_paths)
chunks = chunk_documents(docs)
embedder = SentenceTransformer(EMBED_MODEL)
embeddings = embed_chunks(chunks, embedder)
client = setup_qdrant()
store_in_qdrant(client, chunks, embeddings)
bm25 = setup_bm25(chunks)

# Search
cross_encoder = CrossEncoder(CROSS_ENCODER_MODEL)
query = "PARTNERSHIPS & MAJOR CLIENTS"
results = hybrid_search(query, chunks, client, bm25, cross_encoder, top_k=5)

print("n🔍 Top Results (Before Adding Next Chunks):")
for ((text, source), score) in results:
print(f"[Score: {score:.4f}] (File: {source}) {text[:250]}...")

# Attach next chunk for each result
expanded_results = attach_next_chunk(results, chunks)

print("n✅ Final Results (With Next Chunks):")
for ((text, source), score) in expanded_results:
print(f"(File: {source}) {text[:250]}...")