from transformers import MllamaForConditionalGeneration, AutoProcessor, TextIteratorStreamer , AutoModel,Qwen2VLForConditionalGeneration, AutoModelForImageTextToText , Qwen2_5_VLForConditionalGeneration from qwen_vl_utils import process_vision_info 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 ="Qwen/Qwen2.5-VL-7B-Instruct" model = Qwen2_5_VLForConditionalGeneration.from_pretrained(ckpt, torch_dtype=torch.bfloat16,trust_remote_code=True).to("cuda") processor = AutoProcessor.from_pretrained(ckpt,trust_remote_code=True) 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 = 3 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(prompt_option, max_new_tokens=4096): try: # Define predetermined prompts prompts = { "NOC Timesheet": ( """Extract structured information from the provided timesheet. The extracted details should include: Name Position Title Work Location Contractor NOC ID Month and Year Regular Service Days (ONSHORE) Standby Days (ONSHORE in Doha) Offshore Days Standby & Extended Hitch Days (OFFSHORE) Extended Hitch Days (ONSHORE Rotational) Service during Weekends & Public Holidays ONSHORE Overtime Hours (Over 8 hours) OFFSHORE Overtime Hours (Over 12 hours) Per Diem Days (ONSHORE/OFFSHORE Rotational Personnel) Training Days Travel Days Noc representative appoval's name as approved_by Noc representative's date approval_date Noc representative status as approval_status The output should be formatted as a JSON instance that conforms to the JSON schema below.\n\nAs an example, for the schema {"properties": {"foo": {"title": "Foo", "description": "a list of strings", "type": "array", "items": {"type": "string"}}}, "required": ["foo"]}\nthe object {"foo": ["bar", "baz"]} is a well-formatted instance of the schema. The object {"properties": {"foo": ["bar", "baz"]}} is not well-formatted.""" ), "NOC Basic": ( "Based on the provided timesheet details, extract the following information:\n" " - Full name of the person\n" " - Position title of the person\n" " - Work location\n" " - Contractor's name\n" " - NOC ID\n" " - Month and year (in MM/YYYY format)" ), "NOC Structured test": ( "You are an advanced data extraction assistant. Your task is to parse structured input text and extract key data points into clearly defined categories. Focus only on the requested details, ensuring accuracy and proper grouping. Below is the format for extracting the data:\n\n" "---\n" "Project Information\n\n" "Project Name:\n\n" "Project and Package:\n\n" "RPO Number:\n\n" "PMC Name:\n\n" "Project Location:\n\n" "Year:\n\n" "Month:\n\n" "Timesheet Details\n\n" "Week X (Date)\n\n" "Holidays:\n\n" "Regular Hours:\n\n" "Overtime Hours:\n\n" "Total Hours:\n\n" "Comments:\n\n" "Additional Data\n\n" "Reviewed By:\n\n" "Date of Review:\n\n" "Position:\n\n" "Supervisor Business:\n\n" "Date of Approval:\n\n" "---\n\n" "Ensure the extracted data strictly follows the format above and is organized by category. Ignore unrelated text. Respond only with the formatted output." ), "Aramco Full structured": ( """You are a document parsing assistant designed to extract structured data from various document types, including invoices, timesheets, purchase orders, and travel bookings. Your goal is to return highly accurate, properly formatted JSON for each document type. General Rules: 1. Always return ONLY valid JSON—no explanations, comments, or additional text. 2. Use null for any fields that are not present or cannot be extracted. 3. Ensure all JSON keys are enclosed in double quotes and properly formatted. 4. Validate financial, time tracking, and contract details carefully before output. Extraction Instructions: 1. Invoice: - Parse and extract financial and invoice-specific details. - JSON structure: ```json { "invoice": { "date": null, "dueDate": null, "accountNumber": null, "invoiceNumber": null, "customerContact": null, "kintecContact": null, "accountsContact": null, "periodEnd": null, "contractNo": null, "specialistsName": null, "rpoNumber": null, "assignmentProject": null, "workLocation": null, "expenses": null, "regularHours": null, "overtime": null, "mobilisationAllowance": null, "dailyHousing": null, "opPipTechnical": null, "code": null, "vatBasis": null, "vatRate": null, "vatAmount": null, "totalExclVat": null, "totalInclVat": null } } ``` 2. Timesheet: - Extract time tracking, work details, and approvals. - JSON structure: ```json { "timesheet": { "Year": null, "RPO_Number": null, "PMC_Name": null, "Project_Location": null, "Project_and_Package": null, "Month": null, "Timesheet_Details": [ { "Week": null, "Regular_Hours": null, "Overtime_Hours": null, "Total_Hours": null, "Comments": null }, { "Week": null, "Regular_Hours": null, "Overtime_Hours": null, "Total_Hours": null, "Comments": null } ], "Monthly_Totals": { "Regular_Hours": null, "Overtime_Hours": null, "Total_Hours": null }, "reviewedBy": { "name": null, "position": null, "date": null }, "approvedBy": { "name": null, "position": null, "date": null } } } ``` 3. Purchase Order: - Extract contract and pricing details with precision. - JSON structure: ```json { "purchaseOrder": { "contractNo": null, "relPoNo": null, "version": null, "title": null, "startDate": null, "endDate": null, "costCenter": null, "purchasingGroup": null, "contractor": null, "location": null, "workDescription": null, "pricing": { "regularRate": null, "overtimeRate": null, "totalBudget": null } } } ``` 4. Travel Booking: - Parse travel-specific and employee information. - JSON structure: ```json { "travelBooking": { "requestId": null, "approvalStatus": null, "employee": { "name": null, "id": null, "email": null, "firstName": null, "lastName": null, "gradeCodeGroup": null }, "defaultManager": { "name": null, "email": null }, "sender": { "name": null, "email": null }, "travel": { "startDate": null, "endDate": null, "requestPolicy": null, "requestType": null, "employeeType": null, "travelActivity": null, "tripType": null }, "cost": { "companyCode": null, "costObject": null, "costObjectId": null }, "transport": { "type": null, "comments": null }, "changeRequired": null, "comments": null } } ``` Use these structures for parsing documents and ensure compliance with the rules and instructions provided for each type. """ ), "Aramco Timesheet only": ( """ Extract time tracking, work details, and approvals. - JSON structure: ```json { "timesheet": { "Year": null, "RPO_Number": null, "PMC_Name": null, "Project_Location": null, "Project_and_Package": null, "Month": null, "Timesheet_Details": [ { "Week": null, "Regular_Hours": null, "Overtime_Hours": null, "Total_Hours": null, "Comments": null }, { "Week": null, "Regular_Hours": null, "Overtime_Hours": null, "Total_Hours": null, "Comments": null } ], "Monthly_Totals": { "Regular_Hours": null, "Overtime_Hours": null, "Total_Hours": null }, "reviewedBy": { "name": null, "position": null, "date": null }, "approvedBy": { "name": null, "position": null, "date": null } } } ```""" ), "Aramco test": ( """You are a high-performance document parsing assistant, optimized for speed and accuracy. Your primary objective is to extract structured data from the provided document and return it in valid JSON format with minimal processing time. Guidelines for Speed Optimization: 1. Process the document with minimal computation and only extract the required fields. 2. Use null for any fields that are missing or not clearly identifiable. 3. Avoid redundant checks or deep parsing; rely on the most straightforward extraction methods. 4. Always return ONLY valid JSON—no additional text, explanations, or formatting errors. 5. Focus on precision for key-value pairs; skip over ambiguous or irrelevant information. Document-Specific JSON Structures: 1. **Invoice**: - Extract financial and customer details efficiently. - JSON format: ```json { "invoice": { "date": null, "dueDate": null, "accountNumber": null, "invoiceNumber": null, "customerContact": null, "kintecContact": null, "accountsContact": null, "periodEnd": null, "contractNo": null, "specialistsName": null, "rpoNumber": null, "assignmentProject": null, "workLocation": null, "expenses": null, "regularHours": null, "overtime": null, "mobilisationAllowance": null, "dailyHousing": null, "opPipTechnical": null, "code": null, "vatBasis": null, "vatRate": null, "vatAmount": null, "totalExclVat": null, "totalInclVat": null } } ``` 2. **Timesheet**: - Extract time tracking and approval data swiftly. - JSON format: ```json { "timesheet": { "Year": null, "RPO_Number": null, "PMC_Name": null, "Project_Location": null, "Project_and_Package": null, "Month": null, "Timesheet_Details": [ { "Week": null, "Regular_Hours": null, "Overtime_Hours": null, "Total_Hours": null, "Comments": null }, { "Week": null, "Regular_Hours": null, "Overtime_Hours": null, "Total_Hours": null, "Comments": null } ], "Monthly_Totals": { "Regular_Hours": null, "Overtime_Hours": null, "Total_Hours": null }, "reviewedBy": { "name": null, "position": null, "date": null }, "approvedBy": { "name": null, "position": null, "date": null } } } ``` 3. **Purchase Order**: - Extract contract and pricing details with minimal overhead. - JSON format: ```json { "purchaseOrder": { "contractNo": null, "relPoNo": null, "version": null, "title": null, "startDate": null, "endDate": null, "costCenter": null, "purchasingGroup": null, "contractor": null, "location": null, "workDescription": null, "pricing": { "regularRate": null, "overtimeRate": null, "totalBudget": null } } } ``` 4. **Travel Booking**: - Extract essential travel and employee data efficiently. - JSON format: ```json { "travelBooking": { "requestId": null, "approvalStatus": null, "employee": { "name": null, "id": null, "email": null, "firstName": null, "lastName": null, "gradeCodeGroup": null }, "defaultManager": { "name": null, "email": null }, "sender": { "name": null, "email": null }, "travel": { "startDate": null, "endDate": null, "requestPolicy": null, "requestType": null, "employeeType": null, "travelActivity": null, "tripType": null }, "cost": { "companyCode": null, "costObject": null, "costObjectId": null }, "transport": { "type": null, "comments": null }, "changeRequired": null, "comments": null } } ``` Ensure your parsing method balances accuracy and speed, prioritizing quick turnaround without compromising JSON validity or structural integrity. """ ) } # Get the selected prompt selected_prompt = prompts.get(prompt_option, "Invalid prompt selected.") messages = [] # 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"{selected_prompt}{context}" messages.append({"role": "user", "content": [{"type": "text", "text": current_msg}, {"type": "image"}]}) else: messages.append({"role": "user", "content": [{"type": "text", "text": selected_prompt}]}) # 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 Predetermined Prompts") gr.Markdown("Upload a PDF or image (PNG, JPG, JPEG, GIF, BMP, WEBP) and select a prompt to analyze 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 ) with gr.Row(): prompt_dropdown = gr.Dropdown( label="Select Prompt", choices=[ "NOC Timesheet", "NOC Basic", "NOC Structured test", "Aramco Full structured", "Aramco Timesheet only", "Aramco test" ], value="Options" ) generate_btn = gr.Button("Generate") clear_btn = gr.Button("Clear Document Context") output_text = gr.Textbox( label="Output", interactive=False ) file_upload.change( fn=process_uploaded_file, inputs=[file_upload], outputs=[upload_status] ) generate_btn.click( fn=bot_streaming, inputs=[prompt_dropdown], outputs=[output_text] ) clear_btn.click( fn=clear_context, outputs=[upload_status] ) # Launch the interface demo.launch(debug=True)