Spaces:
Build error
Build error
| import streamlit as st | |
| import os | |
| import requests | |
| import chromadb | |
| 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, SequentialChain | |
| 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="wide") | |
| st.image("https://huggingface.co/front/assets/huggingface_logo-noborder.svg", width=150) | |
| st.title("Blah-1") | |
| # ----------------- API Keys ----------------- | |
| os.environ["GROQ_API_KEY"] = st.secrets.get("GROQ_API_KEY", "") | |
| # ----------------- Ensure Vector Store Directory Exists ----------------- | |
| if not os.path.exists("./chroma_langchain_db"): | |
| os.makedirs("./chroma_langchain_db") | |
| # ----------------- 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 | |
| # ----------------- Load Models ------------------- | |
| llm_judge = ChatGroq(model="deepseek-r1-distill-llama-70b") | |
| rag_llm = ChatGroq(model="mixtral-8x7b-32768") | |
| # Enable verbose logging for debugging | |
| llm_judge.verbose = True | |
| rag_llm.verbose = True | |
| # ----------------- PDF Selection (Upload or URL) ----------------- | |
| st.sidebar.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.sidebar.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.sidebar.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() | |
| # Embedding Model (HF on CPU) | |
| model_name = "nomic-ai/modernbert-embed-base" | |
| embedding_model = HuggingFaceEmbeddings(model_name=model_name, model_kwargs={"device": "cpu"}) | |
| # Split into Chunks | |
| text_splitter = SemanticChunker(embedding_model) | |
| document_chunks = text_splitter.split_documents(docs) | |
| # Store chunks in session state | |
| st.session_state.processed_chunks = document_chunks | |
| 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..."): | |
| vector_store = Chroma( | |
| collection_name="deepseek_collection", | |
| collection_metadata={"hnsw:space": "cosine"}, | |
| embedding_function=embedding_model, | |
| persist_directory="./chroma_langchain_db" | |
| ) | |
| vector_store.add_documents(st.session_state.processed_chunks) | |
| st.session_state.vector_store = vector_store | |
| 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!") | |
| # ----------------- Full SequentialChain Execution ----------------- | |
| with st.spinner("🔄 Running full pipeline..."): | |
| context_relevancy_checker_prompt = PromptTemplate(input_variables=["retriever_query", "context"], template=relevancy_prompt) | |
| relevant_prompt = PromptTemplate(input_variables=["relevancy_response"], template=relevant_context_picker_prompt) | |
| context_prompt = PromptTemplate(input_variables=["context_number", "context"], template=response_synth) | |
| final_prompt = PromptTemplate(input_variables=["query", "context"], template=rag_prompt) | |
| context_relevancy_chain = LLMChain(llm=llm_judge, prompt=context_relevancy_checker_prompt, output_key="relevancy_response") | |
| relevant_context_chain = LLMChain(llm=llm_judge, prompt=relevant_prompt, output_key="context_number") | |
| relevant_contexts_chain = LLMChain(llm=llm_judge, prompt=context_prompt, output_key="relevant_contexts") | |
| response_chain = LLMChain(llm=rag_llm, prompt=final_prompt, output_key="final_response") | |
| context_management_chain = SequentialChain( | |
| chains=[context_relevancy_chain, relevant_context_chain, relevant_contexts_chain, response_chain], | |
| input_variables=["context", "retriever_query", "query"], | |
| output_variables=["relevancy_response", "context_number", "relevant_contexts", "final_response"] | |
| ) | |
| final_output = context_management_chain.invoke({"context": context, "retriever_query": query, "query": query}) | |
| st.success("✅ Full pipeline executed successfully!") | |
| # ----------------- Display All Outputs (Formatted) ----------------- | |
| st.markdown("### 🟥 Context Relevancy Evaluation") | |
| st.json(final_output["relevancy_response"]) | |
| st.markdown("### 🟦 Picked Relevant Contexts") | |
| st.json(final_output["context_number"]) | |
| st.markdown("### 🟥 Extracted Relevant Contexts") | |
| st.json(final_output["relevant_contexts"]) | |
| st.markdown("## 🟥 RAG Final Response") | |
| st.write(final_output["final_response"]) | |