DrishtiSharma commited on
Commit
2735883
Β·
verified Β·
1 Parent(s): 3f05b9b

Upload app (6).py

Browse files
Files changed (1) hide show
  1. app (6).py +157 -0
app (6).py ADDED
@@ -0,0 +1,157 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ import os
3
+ import requests
4
+ import chromadb
5
+ from langchain.document_loaders import PDFPlumberLoader
6
+ from langchain_huggingface import HuggingFaceEmbeddings
7
+ from langchain_experimental.text_splitter import SemanticChunker
8
+ from langchain_chroma import Chroma
9
+ from langchain.chains import LLMChain, SequentialChain
10
+ from langchain.prompts import PromptTemplate
11
+ from langchain_groq import ChatGroq
12
+ from prompts import rag_prompt, relevancy_prompt, relevant_context_picker_prompt, response_synth
13
+
14
+ # ----------------- Streamlit UI Setup -----------------
15
+ st.set_page_config(page_title="Blah", layout="wide")
16
+ st.image("https://huggingface.co/front/assets/huggingface_logo-noborder.svg", width=150)
17
+ st.title("Blah-1")
18
+
19
+ # ----------------- API Keys -----------------
20
+ os.environ["GROQ_API_KEY"] = st.secrets.get("GROQ_API_KEY", "")
21
+
22
+ # ----------------- Ensure Vector Store Directory Exists -----------------
23
+ if not os.path.exists("./chroma_langchain_db"):
24
+ os.makedirs("./chroma_langchain_db")
25
+
26
+ # ----------------- Clear ChromaDB Cache -----------------
27
+ chromadb.api.client.SharedSystemClient.clear_system_cache()
28
+
29
+ # ----------------- Initialize Session State -----------------
30
+ if "pdf_loaded" not in st.session_state:
31
+ st.session_state.pdf_loaded = False
32
+ if "chunked" not in st.session_state:
33
+ st.session_state.chunked = False
34
+ if "vector_created" not in st.session_state:
35
+ st.session_state.vector_created = False
36
+ if "processed_chunks" not in st.session_state:
37
+ st.session_state.processed_chunks = None
38
+ if "vector_store" not in st.session_state:
39
+ st.session_state.vector_store = None
40
+
41
+ # ----------------- Load Models -------------------
42
+ llm_judge = ChatGroq(model="deepseek-r1-distill-llama-70b")
43
+ rag_llm = ChatGroq(model="mixtral-8x7b-32768")
44
+
45
+ # Enable verbose logging for debugging
46
+ llm_judge.verbose = True
47
+ rag_llm.verbose = True
48
+
49
+ # ----------------- PDF Selection (Upload or URL) -----------------
50
+ st.sidebar.subheader("πŸ“‚ PDF Selection")
51
+ pdf_source = st.radio("Choose a PDF source:", ["Upload a PDF file", "Enter a PDF URL"], index=0, horizontal=True)
52
+
53
+ if pdf_source == "Upload a PDF file":
54
+ uploaded_file = st.sidebar.file_uploader("Upload your PDF file", type=["pdf"])
55
+ if uploaded_file:
56
+ st.session_state.pdf_path = "temp.pdf"
57
+ with open(st.session_state.pdf_path, "wb") as f:
58
+ f.write(uploaded_file.getbuffer())
59
+ st.session_state.pdf_loaded = False
60
+ st.session_state.chunked = False
61
+ st.session_state.vector_created = False
62
+
63
+ elif pdf_source == "Enter a PDF URL":
64
+ pdf_url = st.sidebar.text_input("Enter PDF URL:")
65
+ if pdf_url and not st.session_state.pdf_loaded:
66
+ with st.spinner("πŸ”„ Downloading PDF..."):
67
+ try:
68
+ response = requests.get(pdf_url)
69
+ if response.status_code == 200:
70
+ st.session_state.pdf_path = "temp.pdf"
71
+ with open(st.session_state.pdf_path, "wb") as f:
72
+ f.write(response.content)
73
+ st.session_state.pdf_loaded = False
74
+ st.session_state.chunked = False
75
+ st.session_state.vector_created = False
76
+ st.success("βœ… PDF Downloaded Successfully!")
77
+ else:
78
+ st.error("❌ Failed to download PDF. Check the URL.")
79
+ except Exception as e:
80
+ st.error(f"Error downloading PDF: {e}")
81
+
82
+ # ----------------- Process PDF -----------------
83
+ if not st.session_state.pdf_loaded and "pdf_path" in st.session_state:
84
+ with st.spinner("πŸ”„ Processing document... Please wait."):
85
+ loader = PDFPlumberLoader(st.session_state.pdf_path)
86
+ docs = loader.load()
87
+
88
+ # Embedding Model (HF on CPU)
89
+ model_name = "nomic-ai/modernbert-embed-base"
90
+ embedding_model = HuggingFaceEmbeddings(model_name=model_name, model_kwargs={"device": "cpu"})
91
+
92
+ # Split into Chunks
93
+ text_splitter = SemanticChunker(embedding_model)
94
+ document_chunks = text_splitter.split_documents(docs)
95
+
96
+ # Store chunks in session state
97
+ st.session_state.processed_chunks = document_chunks
98
+ st.session_state.pdf_loaded = True
99
+ st.success("βœ… Document processed and chunked successfully!")
100
+
101
+ # ----------------- Setup Vector Store -----------------
102
+ if not st.session_state.vector_created and st.session_state.processed_chunks:
103
+ with st.spinner("πŸ”„ Initializing Vector Store..."):
104
+ vector_store = Chroma(
105
+ collection_name="deepseek_collection",
106
+ collection_metadata={"hnsw:space": "cosine"},
107
+ embedding_function=embedding_model,
108
+ persist_directory="./chroma_langchain_db"
109
+ )
110
+ vector_store.add_documents(st.session_state.processed_chunks)
111
+ st.session_state.vector_store = vector_store
112
+ st.session_state.vector_created = True
113
+ st.success("βœ… Vector store initialized successfully!")
114
+
115
+ # ----------------- Query Input -----------------
116
+ query = st.text_input("πŸ” Ask a question about the document:")
117
+
118
+ if query:
119
+ with st.spinner("πŸ”„ Retrieving relevant context..."):
120
+ retriever = st.session_state.vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 5})
121
+ retrieved_docs = retriever.invoke(query)
122
+ context = [d.page_content for d in retrieved_docs]
123
+ st.success("βœ… Context retrieved successfully!")
124
+
125
+ # ----------------- Full SequentialChain Execution -----------------
126
+ with st.spinner("πŸ”„ Running full pipeline..."):
127
+ context_relevancy_checker_prompt = PromptTemplate(input_variables=["retriever_query", "context"], template=relevancy_prompt)
128
+ relevant_prompt = PromptTemplate(input_variables=["relevancy_response"], template=relevant_context_picker_prompt)
129
+ context_prompt = PromptTemplate(input_variables=["context_number", "context"], template=response_synth)
130
+ final_prompt = PromptTemplate(input_variables=["query", "context"], template=rag_prompt)
131
+
132
+ context_relevancy_chain = LLMChain(llm=llm_judge, prompt=context_relevancy_checker_prompt, output_key="relevancy_response")
133
+ relevant_context_chain = LLMChain(llm=llm_judge, prompt=relevant_prompt, output_key="context_number")
134
+ relevant_contexts_chain = LLMChain(llm=llm_judge, prompt=context_prompt, output_key="relevant_contexts")
135
+ response_chain = LLMChain(llm=rag_llm, prompt=final_prompt, output_key="final_response")
136
+
137
+ context_management_chain = SequentialChain(
138
+ chains=[context_relevancy_chain, relevant_context_chain, relevant_contexts_chain, response_chain],
139
+ input_variables=["context", "retriever_query", "query"],
140
+ output_variables=["relevancy_response", "context_number", "relevant_contexts", "final_response"]
141
+ )
142
+
143
+ final_output = context_management_chain.invoke({"context": context, "retriever_query": query, "query": query})
144
+ st.success("βœ… Full pipeline executed successfully!")
145
+
146
+ # ----------------- Display All Outputs (Formatted) -----------------
147
+ st.markdown("### πŸŸ₯ Context Relevancy Evaluation")
148
+ st.json(final_output["relevancy_response"])
149
+
150
+ st.markdown("### 🟦 Picked Relevant Contexts")
151
+ st.json(final_output["context_number"])
152
+
153
+ st.markdown("### πŸŸ₯ Extracted Relevant Contexts")
154
+ st.json(final_output["relevant_contexts"])
155
+
156
+ st.markdown("## πŸŸ₯ RAG Final Response")
157
+ st.write(final_output["final_response"])