Upload 12 files (#1)

- Upload 12 files (e998abec7346526da097cc568cf90badf620d1bd)

backend/embed_models.py

@@ -0,0 +1,12 @@
+from fastembed import TextEmbedding, LateInteractionTextEmbedding
+from fastembed.rerank.cross_encoder import TextCrossEncoder
+from dotenv import load_dotenv
+import os
+load_dotenv()
+from huggingface_hub import login
+
+login(token=os.getenv("HUGGINGFACE_HUB_TOKEN"))
+
+dense_model = TextEmbedding("sentence-transformers/all-MiniLM-L6-v2")
+colbert_model = LateInteractionTextEmbedding("colbert-ir/colbertv2.0")
+cross_encoder = TextCrossEncoder("jinaai/jina-reranker-v2-base-multilingual")

backend/indexer.py

@@ -0,0 +1,40 @@
+from backend.qdrant import client
+from backend.embed_models import dense_model, colbert_model
+from qdrant_client import models
+
+def setup_collections():
+    dense_dim = client.get_embedding_size("sentence-transformers/all-MiniLM-L6-v2")
+    print("creating collection")
+    client.recreate_collection(
+        collection_name="pdf_dense",
+        vectors_config={"embedding": models.VectorParams(size=dense_dim, distance=models.Distance.COSINE)}
+    )
+    print("creating collection")
+
+    client.recreate_collection(
+        collection_name="pdf_colbert",
+        vectors_config=models.VectorParams(
+            size=128,
+            distance=models.Distance.COSINE,
+            multivector_config=models.MultiVectorConfig(
+                comparator=models.MultiVectorComparator.MAX_SIM
+            )
+        )
+    )
+
+def index_documents(chunks):
+    dense_embs = list(dense_model.embed(chunks))
+    colbert_embs = list(colbert_model.embed(chunks))
+
+    dense_points = [
+        models.PointStruct(id=idx, payload={"chunk": chunk}, vector={"embedding": vec})
+        for idx, (chunk, vec) in enumerate(zip(chunks, dense_embs))
+    ]
+    colbert_points = [
+        models.PointStruct(id=idx, payload={"chunk": chunk}, vector=vec)
+        for idx, (chunk, vec) in enumerate(zip(chunks, colbert_embs))
+    ]
+
+
+    client.upload_points("pdf_dense", dense_points)
+    client.upload_points("pdf_colbert", colbert_points)

backend/pdf_utils.py

@@ -0,0 +1,31 @@
+# import pdfplumber
+#
+# def extract_chunks_from_pdf(pdf_path, chunk_size=100):
+#     text = ""
+#     with pdfplumber.open(pdf_path) as pdf:
+#         for page in pdf.pages:
+#             content = page.extract_text()
+#             if content:
+#                 text += content + "\n"
+#     words = text.split()
+#     return [" ".join(words[i:i+chunk_size]) for i in range(0, len(words), chunk_size)]
+from langchain.document_loaders import PyPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+
+
+def extract_chunks_with_langchain(pdf_path, chunk_size=500, chunk_overlap=100):
+    # Step 1: Load the PDF
+    loader = PyPDFLoader(pdf_path)
+    documents = loader.load()  # Returns a list of Document objects
+
+    # Step 2: Split the text intelligently
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=chunk_size,
+        chunk_overlap=chunk_overlap,
+        separators=["\n\n", "\n", ".", "!", "?", ",", " ", ""]
+    )
+
+    chunks = splitter.split_documents(documents)
+
+    # Optional: Return just the text content
+    return [chunk.page_content for chunk in chunks]

backend/qdrant.py

@@ -0,0 +1,14 @@
+import os
+from dotenv import load_dotenv
+from qdrant_client import QdrantClient
+# Load .env file
+load_dotenv()
+
+# Fetch credentials from environment
+QDRANT_URL = os.getenv("QDRANT_URL")
+QDRANT_API_KEY = os.getenv("QDRANT_API_KEY")
+
+client = QdrantClient(
+    url=QDRANT_URL,
+    api_key=QDRANT_API_KEY,
+)

backend/search.py

@@ -0,0 +1,30 @@
+from backend.qdrant import client
+from backend.embed_models import dense_model, colbert_model, cross_encoder
+
+def search_and_rerank(query, top_k=5):
+    query_vec = list(dense_model.query_embed(query))[0]
+    raw_results = client.query_points(
+        collection_name="pdf_dense",
+        using="embedding",
+        query=query_vec,
+        limit=top_k,
+        with_payload=True
+    )
+    chunks = [pt.payload["chunk"] for pt in raw_results.points]
+    cross_scores = list(cross_encoder.rerank(query, chunks))
+    cross_ranks = sorted(zip(chunks, cross_scores), key=lambda x: x[1], reverse=True)
+
+    query_colbert = list(colbert_model.query_embed(query))[0]
+    colbert_results = client.query_points(
+        collection_name="pdf_colbert",
+        query=query_colbert,
+        limit=top_k,
+        with_payload=True
+    )
+    colbert_ranks = [(pt.payload["chunk"], pt.score) for pt in colbert_results.points]
+
+    return {
+        "raw": [(pt.payload["chunk"], pt.score) for pt in raw_results.points],
+        "cross": cross_ranks,
+        "colbert": colbert_ranks
+    }

main.py

@@ -0,0 +1,48 @@
+import streamlit as st
+import tempfile
+from backend.pdf_utils import extract_chunks_with_langchain
+from backend.indexer import setup_collections, index_documents
+from backend.search import search_and_rerank
+
+st.title("Qdrant PDF Search")
+
+if "indexed" not in st.session_state:
+    st.session_state.indexed = False
+
+uploaded = st.file_uploader("Upload a PDF", type=["pdf"])
+
+if uploaded:
+    with tempfile.NamedTemporaryFile(delete=False) as tmp:
+        tmp.write(uploaded.read())
+        pdf_path = tmp.name
+
+    st.success("PDF uploaded!")
+    chunks = extract_chunks_with_langchain(pdf_path)
+
+    if st.button("Index PDF in Qdrant Cloud"):
+        with st.spinner("Indexing..."):
+            setup_collections()
+            index_documents(chunks)
+            st.session_state.indexed = True  # Mark as indexed
+        st.success("Indexed successfully!")
+
+# Only show query input *after* indexing is done
+if st.session_state.indexed:
+    query = st.text_input("Enter your search query:")
+
+    if query:
+        results = search_and_rerank(query)
+
+        st.subheader("Raw Dense Results")
+        for chunk, score in results["raw"]:
+            st.markdown(f"**{score:.3f}** - {chunk[:200]}...")
+
+        st.subheader("Cross-Encoder Reranked")
+        for chunk, score in results["cross"]:
+            st.markdown(f"**{score:.3f}** - {chunk[:200]}...")
+
+        st.subheader("ColBERT Reranked")
+        for chunk, score in results["colbert"]:
+            st.markdown(f"**{score:.3f}** - {chunk[:200]}...")
+else:
+    st.info("Please upload and index a PDF before searching.")

requirements.txt

@@ -1,3 +1,9 @@
-altair
-pandas
-streamlit
+qdrant-client[fastembed]>=1.14.2
+pdfplumber
+sentence-transformers
+streamlit
+python-dotenv
+langchain
+pypdf
+tiktoken
+langchain-community