Create metadata_fixed.py

lab/metadata_fixed.py

@@ -0,0 +1,148 @@
+import streamlit as st
+import os
+import requests
+import chromadb
+import pdfplumber
+from langchain.document_loaders import PDFPlumberLoader
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_experimental.text_splitter import SemanticChunker
+from langchain_chroma import Chroma
+from langchain.chains import LLMChain
+from langchain.prompts import PromptTemplate
+from langchain_groq import ChatGroq
+from prompts import rag_prompt, relevancy_prompt, relevant_context_picker_prompt, response_synth
+
+# ----------------- Streamlit UI Setup -----------------
+st.set_page_config(page_title="Blah", layout="centered")
+st.title("Blah-1")
+
+# ----------------- API Keys -----------------
+os.environ["GROQ_API_KEY"] = st.secrets.get("GROQ_API_KEY", "")
+os.environ["HF_TOKEN"] = st.secrets.get("HF_TOKEN", "")
+
+# ----------------- Clear ChromaDB Cache -----------------
+chromadb.api.client.SharedSystemClient.clear_system_cache()
+
+# ----------------- Initialize Session State -----------------
+if "pdf_loaded" not in st.session_state:
+    st.session_state.pdf_loaded = False
+if "chunked" not in st.session_state:
+    st.session_state.chunked = False
+if "vector_created" not in st.session_state:
+    st.session_state.vector_created = False
+if "processed_chunks" not in st.session_state:
+    st.session_state.processed_chunks = None
+if "vector_store" not in st.session_state:
+    st.session_state.vector_store = None
+
+# ----------------- Function to Extract PDF Title -----------------
+def extract_pdf_title(pdf_path):
+    """Extract title from PDF metadata or first page."""
+    try:
+        with pdfplumber.open(pdf_path) as pdf:
+            first_page = pdf.pages[0]
+            text = first_page.extract_text()
+            return text.split("\n")[0] if text else "Untitled Document"
+    except Exception as e:
+        return "Untitled Document"
+
+# ----------------- PDF Selection (Upload or URL) -----------------
+st.subheader("📂 PDF Selection")
+pdf_source = st.radio("Choose a PDF source:", ["Upload a PDF file", "Enter a PDF URL"], index=0, horizontal=True)
+
+if pdf_source == "Upload a PDF file":
+    uploaded_file = st.file_uploader("Upload your PDF file", type=["pdf"])
+    if uploaded_file:
+        st.session_state.pdf_path = "temp.pdf"
+        with open(st.session_state.pdf_path, "wb") as f:
+            f.write(uploaded_file.getbuffer())
+        st.session_state.pdf_loaded = False
+        st.session_state.chunked = False
+        st.session_state.vector_created = False
+
+elif pdf_source == "Enter a PDF URL":
+    pdf_url = st.text_input("Enter PDF URL:")
+    if pdf_url and not st.session_state.pdf_loaded:
+        with st.spinner("🔄 Downloading PDF..."):
+            try:
+                response = requests.get(pdf_url)
+                if response.status_code == 200:
+                    st.session_state.pdf_path = "temp.pdf"
+                    with open(st.session_state.pdf_path, "wb") as f:
+                        f.write(response.content)
+                    st.session_state.pdf_loaded = False
+                    st.session_state.chunked = False
+                    st.session_state.vector_created = False
+                    st.success("✅ PDF Downloaded Successfully!")
+                else:
+                    st.error("❌ Failed to download PDF. Check the URL.")
+            except Exception as e:
+                st.error(f"Error downloading PDF: {e}")
+
+# ----------------- Process PDF -----------------
+if not st.session_state.pdf_loaded and "pdf_path" in st.session_state:
+    with st.spinner("🔄 Processing document... Please wait."):
+        loader = PDFPlumberLoader(st.session_state.pdf_path)
+        docs = loader.load()
+        
+        # Extract metadata
+        metadata = docs[0].metadata
+
+        # Try to get title from metadata, fallback to first page
+        title = metadata.get("Title", "").strip() if metadata.get("Title") else extract_pdf_title(st.session_state.pdf_path)
+
+        # Display Title
+        st.subheader(f"📄 Document Title: {title}")
+
+        # Debugging: Show metadata
+        st.json(metadata)
+
+        # Embedding Model (HF on CPU)
+        model_name = "nomic-ai/modernbert-embed-base"
+        embedding_model = HuggingFaceEmbeddings(model_name=model_name, model_kwargs={"device": "cpu"})
+
+        # Prevent unnecessary re-chunking
+        if not st.session_state.chunked:
+            text_splitter = SemanticChunker(embedding_model)
+            document_chunks = text_splitter.split_documents(docs)
+            st.session_state.processed_chunks = document_chunks
+            st.session_state.chunked = True
+
+        st.session_state.pdf_loaded = True
+        st.success("✅ Document processed and chunked successfully!")
+
+# ----------------- Setup Vector Store -----------------
+if not st.session_state.vector_created and st.session_state.processed_chunks:
+    with st.spinner("🔄 Initializing Vector Store..."):
+        st.session_state.vector_store = Chroma(
+            collection_name="deepseek_collection",
+            collection_metadata={"hnsw:space": "cosine"},
+            embedding_function=embedding_model
+        )
+        st.session_state.vector_store.add_documents(st.session_state.processed_chunks)
+        st.session_state.vector_created = True
+        st.success("✅ Vector store initialized successfully!")
+
+# ----------------- Query Input -----------------
+query = st.text_input("🔍 Ask a question about the document:")
+
+if query:
+    with st.spinner("🔄 Retrieving relevant context..."):
+        retriever = st.session_state.vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 5})
+        retrieved_docs = retriever.invoke(query)
+        context = [d.page_content for d in retrieved_docs]
+        st.success("✅ Context retrieved successfully!")
+
+    # ----------------- Run Individual Chains Explicitly -----------------
+    context_relevancy_chain = LLMChain(llm=ChatGroq(model="deepseek-r1-distill-llama-70b"), prompt=PromptTemplate(input_variables=["retriever_query", "context"], template=relevancy_prompt), output_key="relevancy_response")
+    response_chain = LLMChain(llm=ChatGroq(model="mixtral-8x7b-32768"), prompt=PromptTemplate(input_variables=["query", "context"], template=rag_prompt), output_key="final_response")
+
+    response_crisis = context_relevancy_chain.invoke({"context": context, "retriever_query": query})
+    final_response = response_chain.invoke({"query": query, "context": context})
+
+    # ----------------- Display All Outputs -----------------
+    st.markdown("### 🟦 Picked Relevant Contexts")
+    st.json(response_crisis["relevancy_response"])
+
+    st.markdown("## 🟥 RAG Final Response")
+    st.write(final_response["final_response"])