Upload app (6).py
Browse files- 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"])
|