Update app.py

app.py

@@ -1,7 +1,6 @@
 import os
 import streamlit as st
 import pdfplumber
-from concurrent.futures import ThreadPoolExecutor
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
@@ -13,8 +12,7 @@ st.set_page_config(page_title="RAG-based PDF Chat", layout="centered", page_icon
 # Load the summarization pipeline model
 @st.cache_resource
 def load_summarization_pipeline():
-    summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
-    return summarizer
+    return pipeline("summarization", model="facebook/bart-large-cnn")
 
 summarizer = load_summarization_pipeline()
 
@@ -41,8 +39,7 @@ LANGUAGES = {
 @st.cache_data
 def get_text_chunks(text):
     text_splitter = RecursiveCharacterTextSplitter(chunk_size=10000, chunk_overlap=1000)
-    chunks = text_splitter.split_text(text)
-    return chunks
+    return text_splitter.split_text(text)
 
 # Initialize embedding function
 embedding_function = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
@@ -50,11 +47,7 @@ embedding_function = HuggingFaceEmbeddings(model_name="sentence-transformers/all
 # Create a FAISS vector store with embeddings
 @st.cache_resource
 def load_or_create_vector_store(text_chunks):
-    if not text_chunks:
-        st.error("No valid text chunks found to create a vector store. Please check your PDF files.")
-        return None
-    vector_store = FAISS.from_texts(text_chunks, embedding=embedding_function)
-    return vector_store
+    return FAISS.from_texts(text_chunks, embedding=embedding_function) if text_chunks else None
 
 # Helper function to process a single PDF
 def process_single_pdf(file_path):
@@ -73,57 +66,67 @@ def process_single_pdf(file_path):
 def load_pdfs_with_progress(folder_path):
     if not os.path.exists(folder_path):
         st.error(f"The folder '{folder_path}' does not exist. Please create it and add PDF files.")
-        st.session_state['vector_store'] = None
-        st.session_state['loading'] = False
-        return
+        return None
 
     all_text = ""
     pdf_files = [os.path.join(folder_path, filename) for filename in os.listdir(folder_path) if filename.endswith('.pdf')]
-    num_files = len(pdf_files)
-
-    if num_files == 0:
+    if not pdf_files:
         st.error("No PDF files found in the specified folder.")
-        st.session_state['vector_store'] = None
-        st.session_state['loading'] = False
-        return
+        return None
 
     st.markdown("### Loading data...")
     progress_bar = st.progress(0)
-    status_text = st.empty()
 
-    processed_count = 0
-
-    for file_path in pdf_files:
-        result = process_single_pdf(file_path)
-        all_text += result
-        processed_count += 1
-        progress_percentage = int((processed_count / num_files) * 100)
-        progress_bar.progress(processed_count / num_files)
-        status_text.text(f"Loading documents: {progress_percentage}% completed")
+    for i, file_path in enumerate(pdf_files):
+        all_text += process_single_pdf(file_path)
+        progress_bar.progress((i + 1) / len(pdf_files))
 
     progress_bar.empty()
-    status_text.text("Document loading completed!")
-
-    if all_text:
-        text_chunks = get_text_chunks(all_text)
-        vector_store = load_or_create_vector_store(text_chunks)
-        st.session_state['vector_store'] = vector_store
-    else:
-        st.session_state['vector_store'] = None
-
-    st.session_state['loading'] = False
+    return load_or_create_vector_store(get_text_chunks(all_text)) if all_text else None
 
 # Generate summary based on retrieved text
-def generate_summary_with_huggingface(query, retrieved_text):
-    summarization_input = f"{query} Related information:{retrieved_text}"
-    max_input_length = 1024
-    summarization_input = summarization_input[:max_input_length]
+def generate_summary(query, retrieved_text):
+    summarization_input = f"{query} Related information:{retrieved_text}"[:1024]
     summary = summarizer(summarization_input, max_length=500, min_length=50, do_sample=False)
     return summary[0]["summary_text"]
 
-# Generate response for user query
-def user_input(user_question):
-    vector_store = st.session_state.get('vector_store')
-    if vector_store is None:
-        return "The app is still loading documents or no documents were successfully loaded."
-    docs = vector_store.simi
+# Translate text to selected language
+def translate_text(text, target_lang):
+    translation_tokenizer.tgt_lang = target_lang
+    encoded_text = translation_tokenizer(text, return_tensors="pt")
+    generated_tokens = translation_model.generate(**encoded_text)
+    return translation_tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+
+# Main function to run the Streamlit app
+def main():
+    st.markdown(
+        """
+        <h1 style="font-size:30px; text-align: center;">
+        📄 JusticeCompass: Your AI-Powered Legal Navigator for Swift, Accurate Guidance.
+        </h1>
+        """,
+        unsafe_allow_html=True
+    )
+
+    if "vector_store" not in st.session_state:
+        st.session_state["vector_store"] = load_pdfs_with_progress('documents1')
+        if st.session_state["vector_store"] is None:
+            return
+
+    # Prompt input
+    user_question = st.text_input("Ask a Question:", placeholder="Type your question here...")
+
+    # Language selection dropdown
+    selected_language = st.selectbox("Select output language:", list(LANGUAGES.keys()))
+
+    if user_question and st.button("Get Response"):
+        with st.spinner("Generating response..."):
+            docs = st.session_state["vector_store"].similarity_search(user_question)
+            context_text = " ".join([doc.page_content for doc in docs])
+            answer = generate_summary(user_question, context_text)
+            translated_answer = translate_text(answer, LANGUAGES[selected_language])
+            st.markdown(f"**🤖 AI ({selected_language}):** {translated_answer}")
+
+if __name__ == "__main__":
+    main()
+