Create combined_pipeline.py

combined_pipeline.py

@@ -0,0 +1,70 @@
+# combined_pipeline.py
+import os, subprocess
+from sentence_transformers import CrossEncoder
+import fitz
+import pandas as pd
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import Chroma
+from langchain.docstore.document import Document
+
+def extract_styled_blocks(pdf_path):
+    doc = fitz.open(pdf_path)
+    pages = []
+    for page_num, page in enumerate(doc):
+        blocks = page.get_text("dict")["blocks"]
+        for block in blocks:
+            if "lines" in block:
+                for line in block["lines"]:
+                    for span in line["spans"]:
+                        pages.append({
+                            "text": span["text"].strip(),
+                            "is_bold": "Bold" in span["font"],
+                            "is_italic": "Italic" in span["font"],
+                            "page": page_num + 1
+                        })
+    return pages
+
+def parse_pdf_with_headings(pdf_path):
+    spans = extract_styled_blocks(pdf_path)
+    docs, cur_h, cur_s = [], "No heading", None
+    for span in spans:
+        t = span["text"]
+        if not t: 
+            continue
+        if span["is_bold"] and len(t) > 3:
+            cur_h, cur_s = t, None
+        elif span["is_italic"] and len(t) > 3:
+            cur_s = t
+        docs.append(Document(page_content=t, metadata={"page": span["page"], "heading": cur_h, "subheading": cur_s}))
+    return docs
+
+def run_rag(pdf_path, df_items):
+    raw = parse_pdf_with_headings(pdf_path)
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
+    docs = splitter.split_documents(raw)
+
+    embedder = HuggingFaceEmbeddings(model_name="nomic-ai/nomic-embed-text-v1", model_kwargs={"trust_remote_code": True})
+    vectordb = Chroma.from_documents(documents=docs, embedding=embedder, persist_directory="/tmp/chroma_db")
+
+    reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
+    results = []
+    for idx, row in df_items.iterrows():
+        query = row["item"]
+        tops = vectordb.similarity_search_with_score(query, k=10)
+        inps = [[query, d.page_content] for d, _ in tops]
+        scores = reranker.predict(inps)
+        ranked = sorted(zip(scores, tops), key=lambda x: x[0], reverse=True)[:3]
+        pages = [d.metadata["page"] for _, (d, _) in ranked]
+        headings = [d.metadata["heading"] for _, (d, _) in ranked]
+        subs = [d.metadata.get("subheading") or "" for _, (d, _) in ranked]
+        results.append({
+            "item_title": row["item_title"],
+            **{f"sim_{i+1}": pages[i] for i in range(3)},
+            **{f"heading_{i+1}": headings[i] for i in range(3)},
+            **{f"subheading_{i+1}": subs[i] for i in range(3)},
+            **{f"rerank_score_{i+1}": round(ranked[i][0], 2) for i in range(3)}
+        })
+
+    vectordb.delete_collection()
+    return pd.DataFrame(results)