DrishtiSharma commited on
Commit
8a82b65
Β·
verified Β·
1 Parent(s): feb6b77

Create app1.py

Browse files
Files changed (1) hide show
  1. app1.py +162 -0
app1.py ADDED
@@ -0,0 +1,162 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import requests
3
+ import streamlit as st
4
+ import pickle
5
+ from langchain.chains import LLMChain
6
+ from langchain.prompts import PromptTemplate
7
+ from langchain_groq import ChatGroq
8
+ from langchain.document_loaders import PDFPlumberLoader
9
+ from langchain_experimental.text_splitter import SemanticChunker
10
+ from langchain_huggingface import HuggingFaceEmbeddings
11
+ from langchain_chroma import Chroma
12
+ from prompts import rag_prompt, relevancy_prompt, relevant_context_picker_prompt, response_synth
13
+
14
+ # Set API Keys
15
+ os.environ["GROQ_API_KEY"] = st.secrets.get("GROQ_API_KEY", "")
16
+
17
+ # Load LLM models
18
+ llm_judge = ChatGroq(model="deepseek-r1-distill-llama-70b")
19
+ rag_llm = ChatGroq(model="mixtral-8x7b-32768")
20
+
21
+ llm_judge.verbose = True
22
+ rag_llm.verbose = True
23
+
24
+ VECTOR_DB_PATH = "/tmp/chroma_db"
25
+ CHUNKS_FILE = "/tmp/chunks.pkl"
26
+
27
+ # Session State Initialization
28
+ if "vector_store" not in st.session_state:
29
+ st.session_state.vector_store = None
30
+ if "documents" not in st.session_state:
31
+ st.session_state.documents = None
32
+ if "pdf_path" not in st.session_state:
33
+ st.session_state.pdf_path = None
34
+ if "pdf_loaded" not in st.session_state:
35
+ st.session_state.pdf_loaded = False
36
+ if "chunked" not in st.session_state:
37
+ st.session_state.chunked = False
38
+ if "vector_created" not in st.session_state:
39
+ st.session_state.vector_created = False
40
+
41
+ st.title("Blah-2")
42
+
43
+ # Step 1: Choose PDF Source
44
+ pdf_source = st.radio("Upload or provide a link to a PDF:", ["Upload a PDF file", "Enter a PDF URL"], index=0, horizontal=True)
45
+
46
+ if pdf_source == "Upload a PDF file":
47
+ uploaded_file = st.file_uploader("Upload your PDF file", type="pdf")
48
+ if uploaded_file:
49
+ st.session_state.pdf_path = "temp.pdf"
50
+ with open(st.session_state.pdf_path, "wb") as f:
51
+ f.write(uploaded_file.getbuffer())
52
+ st.session_state.pdf_loaded = False
53
+ st.session_state.chunked = False
54
+ st.session_state.vector_created = False
55
+
56
+ elif pdf_source == "Enter a PDF URL":
57
+ pdf_url = st.text_input("Enter PDF URL:")
58
+ if pdf_url and not st.session_state.pdf_path:
59
+ with st.spinner("Downloading PDF..."):
60
+ try:
61
+ response = requests.get(pdf_url)
62
+ if response.status_code == 200:
63
+ st.session_state.pdf_path = "temp.pdf"
64
+ with open(st.session_state.pdf_path, "wb") as f:
65
+ f.write(response.content)
66
+ st.session_state.pdf_loaded = False
67
+ st.session_state.chunked = False
68
+ st.session_state.vector_created = False
69
+ st.success("βœ… PDF Downloaded Successfully!")
70
+ else:
71
+ st.error("❌ Failed to download PDF. Check the URL.")
72
+ except Exception as e:
73
+ st.error(f"❌ Error downloading PDF: {e}")
74
+
75
+ # Step 2: Load & Process PDF (Only Once)
76
+ if st.session_state.pdf_path and not st.session_state.pdf_loaded:
77
+ with st.spinner("Loading PDF..."):
78
+ try:
79
+ loader = PDFPlumberLoader(st.session_state.pdf_path)
80
+ docs = loader.load()
81
+ st.session_state.documents = docs
82
+ st.session_state.pdf_loaded = True
83
+ st.success(f"βœ… **PDF Loaded!** Total Pages: {len(docs)}")
84
+ except Exception as e:
85
+ st.error(f"❌ Error processing PDF: {e}")
86
+
87
+ # Load Cached Chunks if Available
88
+ def load_chunks():
89
+ if os.path.exists(CHUNKS_FILE):
90
+ with open(CHUNKS_FILE, "rb") as f:
91
+ return pickle.load(f)
92
+ return None
93
+
94
+ if not st.session_state.chunked: # Ensure chunking only happens once
95
+ cached_chunks = load_chunks()
96
+ if cached_chunks:
97
+ st.session_state.documents = cached_chunks
98
+ st.session_state.chunked = True
99
+
100
+ # Step 3: Chunking (Only Happens Once)
101
+ if st.session_state.pdf_loaded and not st.session_state.chunked:
102
+ with st.spinner("Chunking the document..."):
103
+ try:
104
+ model_name = "nomic-ai/modernbert-embed-base"
105
+ embedding_model = HuggingFaceEmbeddings(model_name=model_name, model_kwargs={'device': 'cpu'})
106
+ text_splitter = SemanticChunker(embedding_model)
107
+
108
+ if st.session_state.documents:
109
+ documents = text_splitter.split_documents(st.session_state.documents)
110
+ st.session_state.documents = documents
111
+ st.session_state.chunked = True
112
+
113
+ # Save chunks for persistence
114
+ with open(CHUNKS_FILE, "wb") as f:
115
+ pickle.dump(documents, f)
116
+
117
+ st.success(f"βœ… **Document Chunked!** Total Chunks: {len(documents)}")
118
+ except Exception as e:
119
+ st.error(f"❌ Error chunking document: {e}")
120
+
121
+ # Step 4: Setup Vectorstore
122
+ def load_vector_store():
123
+ return Chroma(persist_directory=VECTOR_DB_PATH, collection_name="deepseek_collection", embedding_function=HuggingFaceEmbeddings(model_name="nomic-ai/modernbert-embed-base"))
124
+
125
+ if st.session_state.chunked and not st.session_state.vector_created:
126
+ with st.spinner("Creating vector store..."):
127
+ try:
128
+ if st.session_state.vector_store is None: # Prevent unnecessary reloading
129
+ st.session_state.vector_store = load_vector_store()
130
+
131
+ if len(st.session_state.vector_store.get()["documents"]) == 0: # Prevent duplicate insertions
132
+ st.session_state.vector_store.add_documents(st.session_state.documents)
133
+
134
+ num_documents = len(st.session_state.vector_store.get()["documents"])
135
+ st.session_state.vector_created = True
136
+ st.success(f"βœ… **Vector Store Created!** Total documents stored: {num_documents}")
137
+ except Exception as e:
138
+ st.error(f"❌ Error creating vector store: {e}")
139
+
140
+ # Debugging Logs
141
+ st.write("πŸ“„ **PDF Loaded:**", st.session_state.pdf_loaded)
142
+ st.write("πŸ”Ή **Chunked:**", st.session_state.chunked)
143
+ st.write("πŸ“‚ **Vector Store Created:**", st.session_state.vector_created)
144
+
145
+
146
+ # ----------------- Query Input -----------------
147
+ query = st.text_input("πŸ” Ask a question about the document:")
148
+ if query:
149
+ with st.spinner("πŸ”„ Retrieving relevant context..."):
150
+ retriever = st.session_state.vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 5})
151
+ contexts = retriever.invoke(query)
152
+ # Debugging: Check what was retrieved
153
+ st.write("Retrieved Contexts:", contexts)
154
+ st.write("Number of Contexts:", len(contexts))
155
+
156
+ context = [d.page_content for d in contexts]
157
+ # Debugging: Check extracted context
158
+ st.write("Extracted Context (page_content):", context)
159
+ st.write("Number of Extracted Contexts:", len(context))
160
+
161
+
162
+ ------