Update app.py

app.py

@@ -1,7 +1,6 @@
 import os
 import time
 
-import openai
 import streamlit as st
 from dotenv import load_dotenv
 
@@ -10,12 +9,9 @@ from generate import generate_response
 from preprocess import preprocess_text
 from retrieve import create_vectorizer, retrieve
 
-# Load environment variables from .env file
+# Load environment variables from .env file (if needed)
 load_dotenv()
 
-# Set OpenAI API key
-openai.api_key = os.getenv('api_key')
-
 # Initialize session state
 if "messages" not in st.session_state:
     st.session_state.messages = []
@@ -32,8 +28,80 @@ st.title("RAG-based PDF Query System")
 uploaded_files = st.file_uploader("Upload PDFs", type=["pdf"], accept_multiple_files=True)
 
 if uploaded_files:
+    # Check if new files were uploaded (clear old data if new ones are uploaded)
     if "uploaded_files" not in st.session_state or uploaded_files != st.session_state.uploaded_files:
         st.session_state.uploaded_files = uploaded_files
+        st.session_state.messages = []
+        st.session_state.pdf_files = []
+        st.session_state.processed_texts = []
+
+        # Initialize status container
+        with st.status("Processing the uploaded PDFs...", state="running") as status:
+            # Save uploaded files to disk
+            for uploaded_file in uploaded_files:
+                with open(uploaded_file.name, "wb") as f:
+                    f.write(uploaded_file.getbuffer())
+                st.session_state.pdf_files.append(uploaded_file.name)
+
+            # Extract text from PDFs
+            num_files = len(st.session_state.pdf_files)
+            texts = []
+            for i, pdf_file in enumerate(st.session_state.pdf_files):
+                st.write(f"Extracting text from file {i + 1} of {num_files}...")
+                text = extract_text_from_pdfs([pdf_file])
+                texts.extend(text)
+                time.sleep(0.1)
+
+            # Preprocess text
+            st.write("Preprocessing text...")
+            st.session_state.processed_texts = preprocess_text(texts)
+            time.sleep(0.1)
+
+            # Create vectorizer and transform texts
+            st.write("Creating vectorizer and transforming texts...")
+            st.session_state.vectorizer, st.session_state.X = create_vectorizer(st.session_state.processed_texts)
+            time.sleep(0.1)
+
+            # Update status to complete
+            status.update(label="Processing complete!", state="complete")
+
+else:
+    st.stop()
+
+# Chat interface
+st.write("### Ask a question about the uploaded PDFs")
+
+# Display chat messages
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.write(message["content"])
+
+# Chat input
+prompt = st.chat_input("Ask something about the uploaded PDFs")
+if prompt:
+    # Add user message to session state
+    st.session_state.messages.append({"role": "user", "content": prompt})
+
+    # Retrieve relevant texts
+    top_indices = retrieve(prompt, st.session_state.X, st.session_state.vectorizer)
+    retrieved_texts = [" ".join(st.session_state.processed_texts[i]) for i in top_indices]
+
+    # Generate response using Qwen2.5-7B-Instruct-1M
+    response = generate_response(retrieved_texts, prompt)
+    st.session_state.messages.append({"role": "assistant", "content": response})
+
+    # Display user message
+    with st.chat_message("user"):
+        st.write(prompt)
+
+    # Display assistant message
+    with st.chat_message("assistant"):
+        st.write(response)
+
+# Clean up uploaded files
+for pdf_file in st.session_state.pdf_files:
+    if os.path.exists(pdf_file):
+        os.remove(pdf_file)
         st.session_state.messages = []  # Clear previous messages
         st.session_state.pdf_files = []
         st.session_state.processed_texts = []