Update app.py

app.py

@@ -16,17 +16,42 @@ embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-b
 vector_store = None
 retriever = None
 
-def extract_text_from_pdf(filepath):
+def parse_page_input(page_input):
+    pages = set()
+    for part in page_input.split(","):
+        part = part.strip()
+        if '-' in part:  # Handle ranges
+            start, end = part.split('-')
+            try:
+                pages.update(range(int(start), int(end) + 1))
+            except ValueError:
+                continue  # Skip invalid ranges
+        else:  # Handle individual pages
+            try:
+                pages.add(int(part))
+            except ValueError:
+                continue  # Skip invalid page numbers
+    return sorted(pages)  # Return a sorted list of pages
+
+def extract_text_from_pdf(filepath, pages):
     chunk_size = 1000  # Example chunk size
     overlap = 100      # Example overlap
     loader = UnstructuredLoader([filepath], chunk_size=chunk_size, overlap=overlap)
-    pages = []
+    pages_to_load = parse_page_input(pages)  # Parse the input for page numbers
+    
+    # Filter pages according to user input
+    pages_data = []
     for doc in loader.lazy_load():
-        pages.append(doc)
-    return "\n".join([page.page_content for page in pages])
+        if doc.page_number in pages_to_load:  # Assuming doc.page_number exists
+            pages_data.append(doc.page_content)
+    
+    return "\n".join(pages_data)
 
-def extract_tables_from_pdf(filepath):
-    tables = camelot.read_pdf(filepath, pages='1-end')
+def extract_tables_from_pdf(filepath, pages):
+    if pages:
+        tables = camelot.read_pdf(filepath, pages=pages)
+    else:
+        tables = camelot.read_pdf(filepath, pages='1-end')
     return [table.df.to_string(index=False) for table in tables]
 
 def update_documents(text_input):
@@ -63,9 +88,9 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
     response = rag_chain({"question": message, "chat_history": history})
     return response['answer']
 
-def upload_file(filepath):
-    text = extract_text_from_pdf(filepath)
-    tables = extract_tables_from_pdf(filepath)
+def upload_file(filepath, pages):
+    text = extract_text_from_pdf(filepath, pages)
+    tables = extract_tables_from_pdf(filepath, pages)
     
     # Update documents in the vector store
     update_documents(text)
@@ -81,12 +106,13 @@ with demo:
     with gr.Row():
         u = gr.UploadButton("Upload a file", file_count="single")
         d = gr.DownloadButton("Download the file", visible=False)
+        page_input = gr.Textbox(label="Pages to Parse (e.g., 1, 2, 5-7)", placeholder="Enter page numbers or ranges")
 
     # Create a Textbox for the status message
     status_output = gr.Textbox(label="Status", visible=True)
 
     # Use the proper output components in the upload method
-    u.upload(upload_file, u, [u, d, status_output])
+    u.upload(upload_file, [u, page_input], [u, d, status_output])
     
     with gr.Row():
         chat = gr.ChatInterface(