Update app.py

app.py

@@ -20,7 +20,7 @@ class DocumentState:
     def __init__(self):
         self.current_doc_images = []
         self.current_doc_text = ""
-        self.doc_type = None  # 'pdf' or 'image'
+        self.doc_type = None
         
     def clear(self):
         self.current_doc_images = []
@@ -35,17 +35,18 @@ def process_pdf_file(file_path):
     images = []
     text = ""
     
-    for page_num, page in enumerate(doc):
-        # Extract text
-        text += f"\n=== Page {page_num + 1} ===\n"
-        text += page.get_text() + "\n"
-        
-        # Convert page to image
-        pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72))  # 300 DPI
+    # Take first page only for initial processing
+    if doc.page_count > 0:
+        page = doc[0]
+        text = f"First page content:\n{page.get_text()}\n"
+        pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72))
         img_data = pix.tobytes("png")
         img = Image.open(io.BytesIO(img_data))
         images.append(img.convert("RGB"))
         
+        if doc.page_count > 1:
+            text += f"\nTotal pages in document: {doc.page_count}\n"
+    
     doc.close()
     return images, text
 
@@ -61,7 +62,7 @@ def process_file(file):
     if file_path.lower().endswith('.pdf'):
         doc_state.doc_type = 'pdf'
         doc_state.current_doc_images, doc_state.current_doc_text = process_pdf_file(file_path)
-        return f"PDF processed successfully. {len(doc_state.current_doc_images)} pages loaded. You can now ask questions about the content."
+        return f"PDF first page processed. You can now ask questions about the content."
     else:
         doc_state.doc_type = 'image'
         doc_state.current_doc_images = [Image.open(file_path).convert("RGB")]
@@ -71,48 +72,35 @@ def process_file(file):
 def bot_streaming(message, history, max_new_tokens=2048):
     txt = message["text"]
     messages = []
-    images = []
     
     # Process new file if provided
     if message.get("files") and len(message["files"]) > 0:
         process_file(message["files"][0])
     
-    # Process history and maintain context
+    # Process history
     for i, msg in enumerate(history):
         if isinstance(msg[0], tuple):
-            messages.append({"role": "user", "content": [{"type": "text", "text": history[i+1][0]}, {"type": "image"}]})
-            messages.append({"role": "assistant", "content": [{"type": "text", "text": history[i+1][1]}]})
-        elif isinstance(history[i-1], tuple) and isinstance(msg[0], str):
-            pass
-        elif isinstance(history[i-1][0], str) and isinstance(msg[0], str):
+            messages.append({"role": "user", "content": [{"type": "text", "text": msg[0][1]}, {"type": "image"}]})
+            messages.append({"role": "assistant", "content": [{"type": "text", "text": msg[1]}]})
+        elif isinstance(msg[0], str):
             messages.append({"role": "user", "content": [{"type": "text", "text": msg[0]}]})
             messages.append({"role": "assistant", "content": [{"type": "text", "text": msg[1]}]})
 
     # Include document context in the current message
     if doc_state.current_doc_images:
-        images.extend(doc_state.current_doc_images)
-        context = ""
-        if doc_state.doc_type == 'pdf':
-            context = f"\nContext from PDF:\n{doc_state.current_doc_text}"
+        context = f"\nDocument context:\n{doc_state.current_doc_text}" if doc_state.current_doc_text else ""
         current_msg = f"{txt}{context}"
         messages.append({"role": "user", "content": [{"type": "text", "text": current_msg}, {"type": "image"}]})
+        
+        # Process with single image
+        inputs = processor(
+            text=texts,
+            images=doc_state.current_doc_images[0:1],  # Only use first image
+            return_tensors="pt"
+        ).to("cuda")
     else:
         messages.append({"role": "user", "content": [{"type": "text", "text": txt}]})
-
-    texts = processor.apply_chat_template(messages, add_generation_prompt=True)
-
-    if not images:
         inputs = processor(text=texts, return_tensors="pt").to("cuda")
-    else:
-        # Process images in batches if needed
-        max_images = 12  # Increased maximum number of images/pages
-        if len(images) > max_images:
-            # Take evenly spaced samples if we have too many pages
-            indices = np.linspace(0, len(images) - 1, max_images, dtype=int)
-            images = [images[i] for i in indices]
-            txt += f"\n(Note: Analyzing {max_images} evenly distributed pages from the document)"
-        
-        inputs = processor(text=texts, images=images, return_tensors="pt").to("cuda")
     
     streamer = TextIteratorStreamer(processor, skip_special_tokens=True, skip_prompt=True)
     generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=max_new_tokens)
@@ -131,10 +119,10 @@ def clear_context():
     doc_state.clear()
     return "Document context cleared. You can upload a new document."
 
-# Create the Gradio interface with enhanced features
+# Create the Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Document Analyzer with Chat Support")
-    gr.Markdown("Upload a PDF or image and chat about its contents. The context is maintained throughout the conversation.")
+    gr.Markdown("Upload a PDF or image and chat about its contents. For PDFs, the first page will be processed for visual analysis.")
     
     chatbot = gr.ChatInterface(
         fn=bot_streaming,