DrishtiSharma commited on
Commit
1b02f28
Β·
verified Β·
1 Parent(s): a609396

Create metadata_issue_debugging_statements.py

Browse files
lab/metadata_issue_debugging_statements.py ADDED
@@ -0,0 +1,308 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ Hugging Face's logo
2
+ Hugging Face
3
+ Models
4
+ Datasets
5
+ Spaces
6
+ Posts
7
+ Docs
8
+ Enterprise
9
+ Pricing
10
+
11
+
12
+
13
+ Spaces:
14
+
15
+ DrishtiSharma
16
+ /
17
+ test-deepseek
18
+
19
+
20
+ like
21
+ 0
22
+
23
+ Logs
24
+ App
25
+ Files
26
+ Community
27
+ Settings
28
+ test-deepseek
29
+ /
30
+ app.py
31
+
32
+ DrishtiSharma's picture
33
+ DrishtiSharma
34
+ Update app.py
35
+ a609396
36
+ verified
37
+ 5 minutes ago
38
+ raw
39
+
40
+ Copy download link
41
+ history
42
+ blame
43
+ edit
44
+ delete
45
+
46
+ 11.3 kB
47
+ import streamlit as st
48
+ import os
49
+ import json
50
+ import requests
51
+ import pdfplumber
52
+ import chromadb
53
+ import re
54
+ from langchain.document_loaders import PDFPlumberLoader
55
+ from langchain_huggingface import HuggingFaceEmbeddings
56
+ from langchain_experimental.text_splitter import SemanticChunker
57
+ from langchain_chroma import Chroma
58
+ from langchain.chains import LLMChain
59
+ from langchain.prompts import PromptTemplate
60
+ from langchain_groq import ChatGroq
61
+ from prompts import rag_prompt, relevancy_prompt, relevant_context_picker_prompt, response_synth
62
+
63
+ # ----------------- Streamlit UI Setup -----------------
64
+ st.set_page_config(page_title="Blah-1", layout="centered")
65
+
66
+ # ----------------- API Keys -----------------
67
+ os.environ["GROQ_API_KEY"] = st.secrets.get("GROQ_API_KEY", "")
68
+
69
+ # Load LLM models
70
+ llm_judge = ChatGroq(model="deepseek-r1-distill-llama-70b")
71
+ rag_llm = ChatGroq(model="mixtral-8x7b-32768")
72
+
73
+ llm_judge.verbose = True
74
+ rag_llm.verbose = True
75
+
76
+ # Clear ChromaDB cache to fix tenant issue
77
+ chromadb.api.client.SharedSystemClient.clear_system_cache()
78
+
79
+ # st.title("Blah")
80
+
81
+ # ----------------- ChromaDB Persistent Directory -----------------
82
+ CHROMA_DB_DIR = "/mnt/data/chroma_db"
83
+ os.makedirs(CHROMA_DB_DIR, exist_ok=True)
84
+
85
+ # ----------------- Initialize Session State -----------------
86
+ if "pdf_loaded" not in st.session_state:
87
+ st.session_state.pdf_loaded = False
88
+ if "chunked" not in st.session_state:
89
+ st.session_state.chunked = False
90
+ if "vector_created" not in st.session_state:
91
+ st.session_state.vector_created = False
92
+ if "processed_chunks" not in st.session_state:
93
+ st.session_state.processed_chunks = None
94
+ if "vector_store" not in st.session_state:
95
+ st.session_state.vector_store = None
96
+
97
+ # ----------------- Metadata Extraction -----------------
98
+ # ----------------- Metadata Extraction -----------------
99
+ def extract_metadata_llm(pdf_path):
100
+ """Extracts metadata using LLM instead of regex and logs progress in Streamlit UI."""
101
+
102
+ with pdfplumber.open(pdf_path) as pdf:
103
+ first_page_text = pdf.pages[0].extract_text() or "No text found." if pdf.pages else "No text found."
104
+
105
+ # Streamlit Debugging: Show extracted text
106
+ st.subheader("πŸ“„ Extracted First Page Text for Metadata")
107
+ st.text_area("First Page Text:", first_page_text, height=200)
108
+
109
+ # Define metadata prompt
110
+ metadata_prompt = PromptTemplate(
111
+ input_variables=["text"],
112
+ template="""
113
+ Given the following first page of a research paper, extract metadata **strictly in JSON format**.
114
+ - If no data is found for a field, return `"Unknown"` instead.
115
+ - Ensure the output is valid JSON (do not include markdown syntax).
116
+
117
+ Example output:
118
+ {
119
+ "Title": "Example Paper Title",
120
+ "Author": "John Doe, Jane Smith",
121
122
+ "Affiliations": "School of AI, University of Example"
123
+ }
124
+
125
+ Now, extract the metadata from this document:
126
+ {text}
127
+ """
128
+ )
129
+
130
+ # Run LLM Metadata Extraction
131
+ metadata_chain = LLMChain(llm=llm_judge, prompt=metadata_prompt, output_key="metadata")
132
+
133
+ # Debugging: Log the LLM input
134
+ st.subheader("πŸ” LLM Input for Metadata Extraction")
135
+ st.json({"text": first_page_text})
136
+
137
+ try:
138
+ metadata_response = metadata_chain.invoke({"text": first_page_text})
139
+
140
+ # Debugging: Log raw LLM response
141
+ st.subheader("πŸ” Raw LLM Response")
142
+ st.json(metadata_response)
143
+
144
+ # Handle JSON extraction from LLM response
145
+ try:
146
+ metadata_dict = json.loads(metadata_response["metadata"])
147
+ except json.JSONDecodeError:
148
+ try:
149
+ # Attempt to clean up JSON if needed
150
+ metadata_dict = json.loads(metadata_response["metadata"].strip("```json\n").strip("\n```"))
151
+ except json.JSONDecodeError:
152
+ metadata_dict = {
153
+ "Title": "Unknown",
154
+ "Author": "Unknown",
155
+ "Emails": "No emails found",
156
+ "Affiliations": "No affiliations found"
157
+ }
158
+
159
+ except Exception as e:
160
+ st.error(f"❌ LLM Metadata Extraction Failed: {e}")
161
+ metadata_dict = {
162
+ "Title": "Unknown",
163
+ "Author": "Unknown",
164
+ "Emails": "No emails found",
165
+ "Affiliations": "No affiliations found"
166
+ }
167
+
168
+ # Ensure all required fields exist
169
+ required_fields = ["Title", "Author", "Emails", "Affiliations"]
170
+ for field in required_fields:
171
+ metadata_dict.setdefault(field, "Unknown")
172
+
173
+ # Streamlit Debugging: Display Final Extracted Metadata
174
+ st.subheader("βœ… Extracted Metadata")
175
+ st.json(metadata_dict)
176
+
177
+ return metadata_dict
178
+
179
+
180
+ # ----------------- Step 1: Choose PDF Source -----------------
181
+ pdf_source = st.radio("Upload or provide a link to a PDF:", ["Upload a PDF file", "Enter a PDF URL"], index=0, horizontal=True)
182
+
183
+ if pdf_source == "Upload a PDF file":
184
+ uploaded_file = st.file_uploader("Upload your PDF file", type=["pdf"])
185
+ if uploaded_file:
186
+ st.session_state.pdf_path = "/mnt/data/temp.pdf"
187
+ with open(st.session_state.pdf_path, "wb") as f:
188
+ f.write(uploaded_file.getbuffer())
189
+ st.session_state.pdf_loaded = False
190
+ st.session_state.chunked = False
191
+ st.session_state.vector_created = False
192
+
193
+ elif pdf_source == "Enter a PDF URL":
194
+ pdf_url = st.text_input("Enter PDF URL:")
195
+ if pdf_url and not st.session_state.pdf_loaded:
196
+ with st.spinner("πŸ”„ Downloading PDF..."):
197
+ try:
198
+ response = requests.get(pdf_url)
199
+ if response.status_code == 200:
200
+ st.session_state.pdf_path = "/mnt/data/temp.pdf"
201
+ with open(st.session_state.pdf_path, "wb") as f:
202
+ f.write(response.content)
203
+ st.session_state.pdf_loaded = False
204
+ st.session_state.chunked = False
205
+ st.session_state.vector_created = False
206
+ st.success("βœ… PDF Downloaded Successfully!")
207
+ else:
208
+ st.error("❌ Failed to download PDF. Check the URL.")
209
+ except Exception as e:
210
+ st.error(f"Error downloading PDF: {e}")
211
+
212
+
213
+ # ----------------- Process PDF -----------------
214
+ if not st.session_state.pdf_loaded and "pdf_path" in st.session_state:
215
+ with st.spinner("πŸ”„ Processing document... Please wait."):
216
+ loader = PDFPlumberLoader(st.session_state.pdf_path)
217
+ docs = loader.load()
218
+ st.json(docs[0].metadata)
219
+
220
+ # Extract metadata
221
+ metadata = extract_metadata_llm(st.session_state.pdf_path)
222
+
223
+ # Display extracted-metadata
224
+ if isinstance(metadata, dict):
225
+ st.subheader("πŸ“„ Extracted Document Metadata")
226
+ st.write(f"**Title:** {metadata.get('Title', 'Unknown')}")
227
+ st.write(f"**Author:** {metadata.get('Author', 'Unknown')}")
228
+ st.write(f"**Emails:** {metadata.get('Emails', 'No emails found')}")
229
+ st.write(f"**Affiliations:** {metadata.get('Affiliations', 'No affiliations found')}")
230
+ else:
231
+ st.error("Metadata extraction failed. Check the LLM response format.")
232
+
233
+ # Embedding Model
234
+ model_name = "nomic-ai/modernbert-embed-base"
235
+ embedding_model = HuggingFaceEmbeddings(model_name=model_name, model_kwargs={"device": "cpu"}, encode_kwargs={'normalize_embeddings': False})
236
+
237
+ # Convert metadata into a retrievable chunk
238
+ metadata_doc = {"page_content": metadata, "metadata": {"source": "metadata"}}
239
+
240
+
241
+ # Prevent unnecessary re-chunking
242
+ if not st.session_state.chunked:
243
+ text_splitter = SemanticChunker(embedding_model)
244
+ document_chunks = text_splitter.split_documents(docs)
245
+ document_chunks.insert(0, metadata_doc) # Insert metadata as a retrievable document
246
+ st.session_state.processed_chunks = document_chunks
247
+ st.session_state.chunked = True
248
+
249
+ st.session_state.pdf_loaded = True
250
+ st.success("βœ… Document processed and chunked successfully!")
251
+
252
+ # ----------------- Setup Vector Store -----------------
253
+ if not st.session_state.vector_created and st.session_state.processed_chunks:
254
+ with st.spinner("πŸ”„ Initializing Vector Store..."):
255
+ st.session_state.vector_store = Chroma(
256
+ persist_directory=CHROMA_DB_DIR, # <-- Ensures persistence
257
+ collection_name="deepseek_collection",
258
+ collection_metadata={"hnsw:space": "cosine"},
259
+ embedding_function=embedding_model
260
+ )
261
+ st.session_state.vector_store.add_documents(st.session_state.processed_chunks)
262
+ st.session_state.vector_created = True
263
+ st.success("βœ… Vector store initialized successfully!")
264
+
265
+
266
+ # ----------------- Query Input -----------------
267
+ query = st.text_input("πŸ” Ask a question about the document:")
268
+
269
+ if query:
270
+ with st.spinner("πŸ”„ Retrieving relevant context..."):
271
+ retriever = st.session_state.vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 5})
272
+ retrieved_docs = retriever.invoke(query)
273
+ context = [d.page_content for d in retrieved_docs]
274
+ st.success("βœ… Context retrieved successfully!")
275
+
276
+ # ----------------- Run Individual Chains Explicitly -----------------
277
+ context_relevancy_chain = LLMChain(llm=llm_judge, prompt=PromptTemplate(input_variables=["retriever_query", "context"], template=relevancy_prompt), output_key="relevancy_response")
278
+ relevant_context_chain = LLMChain(llm=llm_judge, prompt=PromptTemplate(input_variables=["relevancy_response"], template=relevant_context_picker_prompt), output_key="context_number")
279
+ relevant_contexts_chain = LLMChain(llm=llm_judge, prompt=PromptTemplate(input_variables=["context_number", "context"], template=response_synth), output_key="relevant_contexts")
280
+ response_chain = LLMChain(llm=rag_llm, prompt=PromptTemplate(input_variables=["query", "context"], template=rag_prompt), output_key="final_response")
281
+
282
+ response_crisis = context_relevancy_chain.invoke({"context": context, "retriever_query": query})
283
+ relevant_response = relevant_context_chain.invoke({"relevancy_response": response_crisis["relevancy_response"]})
284
+ contexts = relevant_contexts_chain.invoke({"context_number": relevant_response["context_number"], "context": context})
285
+ final_response = response_chain.invoke({"query": query, "context": contexts["relevant_contexts"]})
286
+
287
+ # ----------------- Display All Outputs -----------------
288
+ st.markdown("### Context Relevancy Evaluation")
289
+ st.json(response_crisis["relevancy_response"])
290
+
291
+ st.markdown("### Picked Relevant Contexts")
292
+ st.json(relevant_response["context_number"])
293
+
294
+ st.markdown("### Extracted Relevant Contexts")
295
+ st.json(contexts["relevant_contexts"])
296
+
297
+ st.subheader("context_relevancy_evaluation_chain Statement")
298
+ st.json(final_response["relevancy_response"])
299
+
300
+ st.subheader("pick_relevant_context_chain Statement")
301
+ st.json(final_response["context_number"])
302
+
303
+ st.subheader("relevant_contexts_chain Statement")
304
+ st.json(final_response["relevant_contexts"])
305
+
306
+ st.subheader("RAG Response Statement")
307
+ st.json(final_response["final_response"])
308
+