Update app.py

app.py

@@ -7,55 +7,75 @@ import gradio as gr
 from gradio import FileData
 import time
 import spaces
-ckpt ="Daemontatox/DocumentCogito"
+from pdf2image import convert_from_path
+import os
+from PyPDF2 import PdfReader
+import tempfile
+
+ckpt = "Daemontatox/DocumentCogito"
 model = MllamaForConditionalGeneration.from_pretrained(ckpt,
     torch_dtype=torch.bfloat16).to("cuda")
 processor = AutoProcessor.from_pretrained(ckpt)
 
+def process_pdf(pdf_path):
+    """Convert PDF pages to images and extract text."""
+    images = convert_from_path(pdf_path)
+    pdf_reader = PdfReader(pdf_path)
+    text = ""
+    for page in pdf_reader.pages:
+        text += page.extract_text() + "\n"
+    return images, text
+
+def is_pdf(file_path):
+    """Check if the file is a PDF."""
+    return file_path.lower().endswith('.pdf')
 
 @spaces.GPU()
 def bot_streaming(message, history, max_new_tokens=2048):
-    
     txt = message["text"]
     ext_buffer = f"{txt}"
     
-    messages= [] 
+    messages = []
     images = []
     
-
-    for i, msg in enumerate(history): 
+    # 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": "user", "content": [{"type": "text", "text": history[i+1][0]}, {"type": "text", "text": history[i+1][1]}]})
             messages.append({"role": "assistant", "content": [{"type": "text", "text": history[i+1][1]}]})
             images.append(Image.open(msg[0][0]).convert("RGB"))
         elif isinstance(history[i-1], tuple) and isinstance(msg[0], str):
-            # messages are already handled
             pass
-        elif isinstance(history[i-1][0], str) and isinstance(msg[0], str): # text only turn
+        elif isinstance(history[i-1][0], str) and isinstance(msg[0], str):
             messages.append({"role": "user", "content": [{"type": "text", "text": msg[0]}]})
             messages.append({"role": "assistant", "content": [{"type": "text", "text": msg[1]}]})
 
-    # add current message
+    # Process current message
     if len(message["files"]) == 1:
+        file_path = message["files"][0]["path"] if isinstance(message["files"][0], dict) else message["files"][0]
         
-        if isinstance(message["files"][0], str): # examples
-            image = Image.open(message["files"][0]).convert("RGB")
-        else: # regular input
-            image = Image.open(message["files"][0]["path"]).convert("RGB")
-        images.append(image)
+        if is_pdf(file_path):
+            # Handle PDF
+            pdf_images, pdf_text = process_pdf(file_path)
+            images.extend(pdf_images)
+            txt = f"{txt}\nExtracted text from PDF:\n{pdf_text}"
+        else:
+            # Handle regular image
+            image = Image.open(file_path).convert("RGB")
+            images.append(image)
+            
         messages.append({"role": "user", "content": [{"type": "text", "text": txt}, {"type": "image"}]})
     else:
         messages.append({"role": "user", "content": [{"type": "text", "text": txt}]})
 
-
     texts = processor.apply_chat_template(messages, add_generation_prompt=True)
 
-    if images == []:
+    if not images:
         inputs = processor(text=texts, return_tensors="pt").to("cuda")
     else:
         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)
     generated_text = ""
     
@@ -69,32 +89,34 @@ def bot_streaming(message, history, max_new_tokens=2048):
         time.sleep(0.01)
         yield buffer
 
-
-demo = gr.ChatInterface(fn=bot_streaming, title="Document Analyzer", examples=[
-    [{"text": "Which era does this piece belong to? Give details about the era.", "files":["./examples/rococo.jpg"]},
-    200],
-    [{"text": "Where do the droughts happen according to this diagram?", "files":["./examples/weather_events.png"]},
-    250],
-    [{"text": "What happens when you take out white cat from this chain?", "files":["./examples/ai2d_test.jpg"]},
-    250],
-    [{"text": "How long does it take from invoice date to due date? Be short and concise.", "files":["./examples/invoice.png"]},
-    250],
-    [{"text": "Where to find this monument? Can you give me other recommendations around the area?", "files":["./examples/wat_arun.jpg"]}, 
-    250],
+demo = gr.ChatInterface(
+    fn=bot_streaming,
+    title="Document Analyzer",
+    examples=[
+        [{"text": "Which era does this piece belong to? Give details about the era.", "files":["./examples/rococo.jpg"]}, 200],
+        [{"text": "Where do the droughts happen according to this diagram?", "files":["./examples/weather_events.png"]}, 250],
+        [{"text": "What happens when you take out white cat from this chain?", "files":["./examples/ai2d_test.jpg"]}, 250],
+        [{"text": "How long does it take from invoice date to due date? Be short and concise.", "files":["./examples/invoice.png"]}, 250],
+        [{"text": "Where to find this monument? Can you give me other recommendations around the area?", "files":["./examples/wat_arun.jpg"]}, 250],
     ],
-      textbox=gr.MultimodalTextbox(), 
-      additional_inputs = [gr.Slider(
-              minimum=10,
-              maximum=500,
-              value=2048,
-              step=10,
-              label="Maximum number of new tokens to generate",
-          )
-        ],
-      cache_examples=False,
-      description="MllM ",
-      stop_btn="Stop Generation", 
-      fill_height=True,
-    multimodal=True)
-    
+    textbox=gr.MultimodalTextbox(),
+    additional_inputs=[
+        gr.Slider(
+            minimum=10,
+            maximum=500,
+            value=2048,
+            step=10,
+            label="Maximum number of new tokens to generate",
+        )
+    ],
+    cache_examples=False,
+    description="MllM Document and PDF Analyzer",
+    stop_btn="Stop Generation",
+    fill_height=True,
+    multimodal=True
+)
+
+# Update file types to include PDFs
+demo.textbox.file_types = ["image", "pdf"]
+
 demo.launch(debug=True)