Update app.py

app.py

@@ -125,26 +125,80 @@ def process_uploaded_file(file):
         return "An error occurred while processing the file. Please try again."
 
 @spaces.GPU()
-def bot_streaming(message, history, max_new_tokens=8192):
+def bot_streaming(prompt_option, max_new_tokens=8192):
     try:
+        # Define predetermined prompts
+        prompts = {
+            "Timesheet Details (Full Extraction)": (
+                "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)\n"
+                "And from the bottom table:\n"
+                "   - Number of service days onshore\n"
+                "   - Number of standby days onshore in Doha\n"
+                "   - Number of service days offshore\n"
+                "   - Number of service days during weekends or public holidays\n"
+                "   - Number of standby and extended hitch days offshore\n"
+                "   - Number of extended hitch days onshore for rotational personnel\n"
+                "   - Number of Per Diem days for onshore/offshore rotational personnel\n"
+                "   - Number of training days\n"
+                "   - Number of travel days\n"
+                "   - Number of offshore days"
+            ),
+            "Timesheet Details (Basic Extraction)": (
+                "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)"
+            ),
+            "Structured Data Extraction": (
+                "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."
+            )
+        }
+        
+        # Get the selected prompt
+        selected_prompt = prompts.get(prompt_option, "Invalid prompt selected.")
+        
         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}"
+            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": message}]})
+            messages.append({"role": "user", "content": [{"type": "text", "text": selected_prompt}]})
 
         # Process inputs
         texts = processor.apply_chat_template(messages, add_generation_prompt=True)
@@ -186,8 +240,8 @@ def clear_context():
 
 # 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.")
+    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(
@@ -199,22 +253,23 @@ with gr.Blocks() as demo:
             interactive=False
         )
     
+    with gr.Row():
+        prompt_dropdown = gr.Dropdown(
+            label="Select Prompt",
+            choices=[
+                "Timesheet Details (Full Extraction)",
+                "Timesheet Details (Basic Extraction)",
+                "Structured Data Extraction"
+            ],
+            value="Timesheet Details (Full Extraction)"
+        )
+        generate_btn = gr.Button("Generate")
+    
     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
+    output_text = gr.Textbox(
+        label="Output",
+        interactive=False
     )
     
     file_upload.change(
@@ -223,6 +278,12 @@ with gr.Blocks() as demo:
         outputs=[upload_status]
     )
     
+    generate_btn.click(
+        fn=bot_streaming,
+        inputs=[prompt_dropdown],
+        outputs=[output_text]
+    )
+    
     clear_btn.click(
         fn=clear_context,
         outputs=[upload_status]