from transformers import MllamaForConditionalGeneration, AutoProcessor, TextIteratorStreamer from PIL import Image import requests import torch from threading import Thread import gradio as gr from gradio import FileData import time import spaces import fitz # PyMuPDF import io import numpy as np import logging # Set up logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Load model and processor ckpt = "Daemontatox/DocumentCogito" model = MllamaForConditionalGeneration.from_pretrained(ckpt, torch_dtype=torch.bfloat16).to("cuda") processor = AutoProcessor.from_pretrained(ckpt) class DocumentState: def __init__(self): self.current_doc_images = [] self.current_doc_text = "" self.doc_type = None def clear(self): self.current_doc_images = [] self.current_doc_text = "" self.doc_type = None doc_state = DocumentState() def process_pdf_file(file_path): """Convert PDF to images and extract text using PyMuPDF.""" try: doc = fitz.open(file_path) images = [] text = "" for page_num in range(doc.page_count): try: page = doc[page_num] page_text = page.get_text("text") if page_text.strip(): text += f"Page {page_num + 1}:\n{page_text}\n\n" zoom = 2 mat = fitz.Matrix(zoom, zoom) pix = page.get_pixmap(matrix=mat, alpha=False) img_data = pix.tobytes("png") img = Image.open(io.BytesIO(img_data)) img = img.convert("RGB") max_size = 1600 if max(img.size) > max_size: ratio = max_size / max(img.size) new_size = tuple(int(dim * ratio) for dim in img.size) img = img.resize(new_size, Image.Resampling.LANCZOS) images.append(img) except Exception as e: logger.error(f"Error processing page {page_num}: {str(e)}") continue doc.close() if not images: raise ValueError("No valid images could be extracted from the PDF") return images, text except Exception as e: logger.error(f"Error processing PDF file: {str(e)}") raise def process_uploaded_file(file): """Process uploaded file and update document state.""" try: doc_state.clear() if file is None: return "No file uploaded. Please upload a file." # Get the file path and extension if isinstance(file, dict): file_path = file["name"] else: file_path = file.name # Get file extension file_ext = file_path.lower().split('.')[-1] # Define allowed extensions image_extensions = {'png', 'jpg', 'jpeg', 'gif', 'bmp', 'webp'} if file_ext == 'pdf': doc_state.doc_type = 'pdf' try: doc_state.current_doc_images, doc_state.current_doc_text = process_pdf_file(file_path) return f"PDF processed successfully. Total pages: {len(doc_state.current_doc_images)}. You can now ask questions about the content." except Exception as e: return f"Error processing PDF: {str(e)}. Please try a different PDF file." elif file_ext in image_extensions: doc_state.doc_type = 'image' try: img = Image.open(file_path).convert("RGB") max_size = 1600 if max(img.size) > max_size: ratio = max_size / max(img.size) new_size = tuple(int(dim * ratio) for dim in img.size) img = img.resize(new_size, Image.Resampling.LANCZOS) doc_state.current_doc_images = [img] return "Image loaded successfully. You can now ask questions about the content." except Exception as e: return f"Error processing image: {str(e)}. Please try a different image file." else: return f"Unsupported file type: {file_ext}. Please upload a PDF or image file (PNG, JPG, JPEG, GIF, BMP, WEBP)." except Exception as e: logger.error(f"Error in process_file: {str(e)}") return "An error occurred while processing the file. Please try again." @spaces.GPU() def bot_streaming(message, history, max_new_tokens=8192): try: messages = [] # Process history for i, msg in enumerate(history): try: messages.append({"role": "user", "content": [{"type": "text", "text": msg[0]}]}) messages.append({"role": "assistant", "content": [{"type": "text", "text": msg[1]}]}) except Exception as e: logger.error(f"Error processing history message {i}: {str(e)}") continue # Include document context if doc_state.current_doc_images: context = f"\nDocument context:\n{doc_state.current_doc_text}" if doc_state.current_doc_text else "" current_msg = f"{message}{context}" messages.append({"role": "user", "content": [{"type": "text", "text": current_msg}, {"type": "image"}]}) else: messages.append({"role": "user", "content": [{"type": "text", "text": message}]}) # Process inputs texts = processor.apply_chat_template(messages, add_generation_prompt=True) try: if doc_state.current_doc_images: inputs = processor( text=texts, images=doc_state.current_doc_images[0:1], return_tensors="pt" ).to("cuda") else: inputs = processor(text=texts, 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) thread = Thread(target=model.generate, kwargs=generation_kwargs) thread.start() buffer = "" for new_text in streamer: buffer += new_text time.sleep(0.01) yield buffer except Exception as e: logger.error(f"Error in model processing: {str(e)}") yield "An error occurred while processing your request. Please try again." except Exception as e: logger.error(f"Error in bot_streaming: {str(e)}") yield "An error occurred. Please try again." def clear_context(): """Clear the current document context.""" doc_state.clear() return "Document context cleared. You can upload a new document." # Create the Gradio interface with gr.Blocks() as demo: gr.Markdown("# Document Analyzer with Chat Support") gr.Markdown("Upload a PDF or image (PNG, JPG, JPEG, GIF, BMP, WEBP) and chat about its contents.") with gr.Row(): file_upload = gr.File( label="Upload Document", file_types=[".pdf", ".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp"] ) upload_status = gr.Textbox( label="Upload Status", interactive=False ) clear_btn = gr.Button("Clear Document Context") chatbot = gr.ChatInterface( fn=bot_streaming, title="Document Chat", additional_inputs=[ gr.Slider( minimum=10, maximum=2048, value=8192, step=10, label="Maximum number of new tokens to generate", ) ], stop_btn="Stop Generation", fill_height=True ) file_upload.change( fn=process_uploaded_file, inputs=[file_upload], outputs=[upload_status] ) clear_btn.click( fn=clear_context, outputs=[upload_status] ) # Launch the interface demo.launch(debug=True)