Create app1.py
Browse files
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 |
+
------
|