Update app.py

app.py

@@ -80,7 +80,8 @@ def search_faiss_index(query_embedding, index, top_k=5):
     distances, indices = index.search(query_embedding, top_k)
     return indices[0], distances[0]
 
-# Streamlit App
+
+
 def main():
     st.title("Enhanced RAG Model with FAISS Indexing")
 
@@ -107,45 +108,66 @@ def main():
     # Input Prompt
     prompt = st.text_input("Enter your query:")
 
-    if pdf_file and prompt:
-        # Extract text from PDF
+    # State to hold intermediate results
+    if 'embeddings' not in st.session_state:
+        st.session_state.embeddings = None
+    if 'chunks' not in st.session_state:
+        st.session_state.chunks = []
+    if 'faiss_index' not in st.session_state:
+        st.session_state.faiss_index = None
+    if 'relevant_chunks' not in st.session_state:
+        st.session_state.relevant_chunks = []
+    if 'translated_queries' not in st.session_state:
+        st.session_state.translated_queries = []
+
+    # Button 1: Embed PDF
+    if st.button("1. Embed PDF") and pdf_file:
         text_lines = extract_text_from_pdf(pdf_file)
-
-        # Detect Language
-        lang = detect_language(" ".join(text_lines))
-        st.write(f"**Detected Language:** {lang}")
-
+        st.session_state.lang = detect_language(" ".join(text_lines))
+        st.write(f"**Detected Language:** {st.session_state.lang}")
+        
         # Chunk the text
-        chunks = split_text_into_chunks(text_lines)
+        st.session_state.chunks = split_text_into_chunks(text_lines)
 
         # Encode chunks
-        chunk_embeddings = embedder.encode(chunks, convert_to_tensor=False)
+        chunk_embeddings = embedder.encode(st.session_state.chunks, convert_to_tensor=False)
 
         # Build FAISS index
-        index = build_faiss_index(np.array(chunk_embeddings))
-
-        # Embed the query
-        query_embedding = embedder.encode([prompt], convert_to_tensor=False)
+        st.session_state.faiss_index = build_faiss_index(np.array(chunk_embeddings))
 
-        # Search for relevant chunks
-        top_k_indices, _ = search_faiss_index(np.array(query_embedding), index, top_k=5)
+        st.success("PDF Embedded Successfully")
 
-        # Retrieve relevant chunks
-        relevant_chunks = [chunks[i] for i in top_k_indices]
-
-        # Combine the context
-        context = "\n".join(relevant_chunks)
-
-        # Format the system prompt
+    # Button 2: Generate Translated Queries
+    if st.button("2. Query Translation") and prompt:
         formatted_prompt = DEFAULT_SYSTEM_PROMPTS[query_translation].format(question=prompt)
-
-        # Query LLM
-        llm_input = f"{formatted_prompt}\n\nContext: {context}\n\nAnswer this question: {prompt}"
-        response = query_huggingface_model(llm_input, max_new_tokens, temperature, top_k)
-
-        # Display the result
-        st.subheader("Response:")
-        st.write(response)
+        response = query_huggingface_model(formatted_prompt, max_new_tokens, temperature, top_k)
+        st.session_state.translated_queries = response.split("\n")
+        st.write("**Generated Queries:**")
+        st.write(st.session_state.translated_queries)
+
+    # Button 3: Retrieve Document Details
+    if st.button("3. Retrieve Documents") and st.session_state.translated_queries:
+        st.session_state.relevant_chunks = []
+        for query in st.session_state.translated_queries:
+            query_embedding = embedder.encode([query], convert_to_tensor=False)
+            top_k_indices, _ = search_faiss_index(np.array(query_embedding), st.session_state.faiss_index, top_k=5)
+            relevant_chunks = [st.session_state.chunks[i] for i in top_k_indices]
+            st.session_state.relevant_chunks.append(relevant_chunks)
+
+        st.write("**Retrieved Documents (for each query):**")
+        for i, relevant_chunks in enumerate(st.session_state.relevant_chunks):
+            st.write(f"**Query {i + 1}: {st.session_state.translated_queries[i]}**")
+            for chunk in relevant_chunks:
+                st.write(f"{chunk[:100]}...")
+
+    # Button 4: Generate Final Response
+    if st.button("4. Final Response") and st.session_state.relevant_chunks:
+        context = "\n".join([chunk for sublist in st.session_state.relevant_chunks for chunk in sublist])
+        llm_input = f"{DEFAULT_SYSTEM_PROMPTS[query_translation].format(question=prompt)}\n\nContext: {context}\n\nAnswer this question: {prompt}"
+        final_response = query_huggingface_model(llm_input, max_new_tokens, temperature, top_k)
+
+        st.subheader("Final Response:")
+        st.write(final_response)
 
 if __name__ == "__main__":
     main()