Update app.py

app.py

@@ -1,24 +1,35 @@
 import streamlit as st
 from llama_index.core import Settings
-from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, StorageContext
+from llama_index.core import VectorStoreIndex, Document
 from llama_index.embeddings.gemini import GeminiEmbedding
 from llama_index.llms.gemini import Gemini
-from llama_index.core import Document
-import google.generativeai as genai
-#import streamlit_analytics2 as streamlit_analytics
-
-# Set up Google API key
 import os
+import PyPDF2
 
-# Configure Google Gemini
+# Function to chunk text into smaller pieces
+def chunk_text(text, chunk_size=1000):
+    """Split the text into chunks of specified size."""
+    return [text[i:i + chunk_size] for i in range(0, len(text), chunk_size)]
 
 # Load and index the legal document data
 def load_data(uploaded_files):
-    documents = [Document(text=t) for t in uploaded_files]
-    #documents = SimpleDirectoryReader(input_files=[uploaded_files]).load_data()
+    documents = []
+    for uploaded_file in uploaded_files:
+        document_text = ""
+        if uploaded_file.type == "application/pdf":
+            pdf_reader = PyPDF2.PdfReader(uploaded_file)
+            for page in pdf_reader.pages:
+                document_text += page.extract_text()
+        else:
+            document_text = uploaded_file.getvalue().decode("utf-8")
+        
+        # Chunk the document text
+        chunks = chunk_text(document_text)
+        for chunk in chunks:
+            documents.append(Document(text=chunk))
+    
     Settings.embed_model = GeminiEmbedding(api_key=os.getenv("GOOGLE_API_KEY"), model_name="models/embedding-001")
-    Settings.llm = Gemini(api_key=os.getenv("GOOGLE_API_KEY"), temperature=0.8, model_name="models/gemini-pro")
-    llm = Gemini(api_key=os.getenv("GOOGLE_API_KEY"), temperature=0.8, model_name="models/gemini-pro")
+    Settings.llm = Gemini(api_key=os.getenv("GOOGLE_API_KEY"), temperature=0.1, model_name="models/gemini-pro")
     index = VectorStoreIndex.from_documents(documents)
     return index
 
@@ -45,33 +56,37 @@ def generate_summary(index, document_text):
 # Streamlit app
 def main():
     st.title("Legal Document Summarizer")
-    st.write("Upload a legal document, and let our AI summarize it!")
+    st.write("Upload legal documents, and let our AI summarize them!")
 
     # File uploader
-    uploaded_file = st.file_uploader("Choose a legal document file", type=["txt", "pdf"])
+    uploaded_files = st.file_uploader("Choose legal document files", type=["txt", "pdf"], accept_multiple_files=True)
 
-    if uploaded_file is not None:
-        # Read file contents
-        if uploaded_file.type == "application/pdf":
-            # You'll need to install PyPDF2 for this
-            import PyPDF2
-            pdf_reader = PyPDF2.PdfReader(uploaded_file)
-            document_text = ""
-            l = []
-            for page in pdf_reader.pages:
-                document_text += page.extract_text()
-                l.append(page.extract_text())
-        else:
-            document_text = uploaded_file.getvalue().decode("utf-8")
+    if uploaded_files:
+        st.write("Analyzing legal documents...")
 
-        st.write("Analyzing legal document...")
-
-        # Load data and generate summary
-        index = load_data(l)
-        summary = generate_summary(index, document_text)
+        # Load data and generate summaries
+        index = load_data(uploaded_files)
+        summaries = []
+        
+        for uploaded_file in uploaded_files:
+            document_text = ""
+            if uploaded_file.type == "application/pdf":
+                pdf_reader = PyPDF2.PdfReader(uploaded_file)
+                for page in pdf_reader.pages:
+                    document_text += page.extract_text()
+            else:
+                document_text = uploaded_file.getvalue().decode("utf-8")
+                
+            # Chunk the document text for summarization
+            chunks = chunk_text(document_text)
+            for chunk in chunks:
+                summary = generate_summary(index, chunk)
+                summaries.append(summary)
 
-        st.write("## Legal Document Summary")
-        st.write(summary)
+        st.write("## Legal Document Summaries")
+        for i, summary in enumerate(summaries):
+            st.write(f"### Summary of Document {i + 1}")
+            st.write(summary)
 
 if __name__ == "__main__":
-    main()
+    main()