Update app.py

app.py

@@ -1,128 +1,197 @@
 import gradio as gr
+import os
+import tempfile
+from pathlib import Path
+import secrets
 from transformers import Qwen2VLForConditionalGeneration, AutoProcessor, AutoModelForCausalLM, AutoTokenizer
+from PIL import Image
 import torch
 
-# Load the OCR model and processor
+# Set up models and processors
 ocr_model = Qwen2VLForConditionalGeneration.from_pretrained(
     "Qwen/Qwen2-VL-7B-Instruct",
     torch_dtype="auto",
     device_map="auto",
 )
-
 ocr_processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")
 
-# Load the Math model and tokenizer
 math_model = AutoModelForCausalLM.from_pretrained(
-    "Qwen/Qwen2.5-Math-72B-Instruct",
+    "Qwen/Qwen2.5-Math-7B-Instruct",
     torch_dtype="auto",
-    device_map="auto"
+    device_map="auto",
 )
+math_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Math-7B-Instruct")
 
-math_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Math-72B-Instruct")
+math_messages = []
 
-# OCR extraction function
-def ocr_and_query(image, question):
-    # Prepare image for the model
+def process_image(image, should_convert=False):
+    """
+    Processes the uploaded image and extracts math-related content using Qwen2-VL.
+    """
+    global math_messages
+    math_messages = []  # Reset when uploading a new image
+    uploaded_file_dir = os.environ.get("GRADIO_TEMP_DIR") or str(
+        Path(tempfile.gettempdir()) / "gradio"
+    )
+    os.makedirs(uploaded_file_dir, exist_ok=True)
+    
+    name = f"tmp{secrets.token_hex(20)}.jpg"
+    filename = os.path.join(uploaded_file_dir, name)
+    
+    if should_convert:
+        # Convert image to RGB if required
+        new_img = Image.new('RGB', size=(image.width, image.height), color=(255, 255, 255))
+        new_img.paste(image, (0, 0), mask=image)
+        image = new_img
+    image.save(filename)
+    
+    # Prepare OCR input
     messages = [
         {
-            "role": "user",
-            "content": [
-                {"type": "image"},
-                {
-                    "type": "text",
-                    "text": question
-                },
-            ],
+            'role': 'system',
+            'content': [{'text': 'You are a helpful assistant.'}]
+        },
+        {
+            'role': 'user',
+            'content': [
+                {'image': f'file://{filename}'},
+                {'text': 'Please describe the math-related content in this image, ensuring that any LaTeX formulas are correctly transcribed. Non-mathematical details do not need to be described.'}
+            ]
         }
     ]
     
-    # Process image and text prompt
+    # Generate OCR output
     text_prompt = ocr_processor.apply_chat_template(messages, add_generation_prompt=True)
     inputs = ocr_processor(text=[text_prompt], images=[image], padding=True, return_tensors="pt")
-    
-    # Run the model to generate OCR results
-    inputs = inputs.to("cuda")
+    inputs = inputs.to("cuda")  # Use CPU if GPU is unavailable
     output_ids = ocr_model.generate(**inputs, max_new_tokens=1024)
+    output_text = ocr_processor.batch_decode(output_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)[0]
     
-    # Decode the generated text
-    generated_ids = [
-        output_ids[len(input_ids):]
-        for input_ids, output_ids in zip(inputs.input_ids, output_ids)
-    ]
-    output_text = ocr_processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)[0]
-    
+    os.remove(filename)
     return output_text
 
-# Math problem solving function
-def solve_math_problem(prompt):
-    # CoT (Chain of Thought)
-    messages = [
-        {"role": "system", "content": "Please reason step by step, and put your final answer within \\boxed{}."},
-        {"role": "user", "content": prompt}
-    ]
+def get_math_response(image_description, user_question):
+    """
+    Sends the OCR output and user question to Qwen2-Math and retrieves the solution.
+    """
+    global math_messages
 
-    text = math_tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True
-    )
-    model_inputs = math_tokenizer([text], return_tensors="pt").to("cuda")
+    # Initialize the math assistant role
+    if not math_messages:
+        math_messages.append({'role': 'system', 'content': 'You are a helpful math assistant.'})
+    math_messages = math_messages[:1]  # Retain only the system prompt
+
+    # Format the input question
+    if image_description is not None:
+        content = f'Image description: {image_description}\n\n'
+    else:
+        content = ''
+    query = f"{content}User question: {user_question}"
+    math_messages.append({'role': 'user', 'content': query})
+
+    # Prepare math model input
+    inputs = math_tokenizer(
+        text=query,
+        padding=True,
+        return_tensors="pt"
+    ).to("cuda")  # Use CPU if GPU is unavailable
 
-    generated_ids = math_model.generate(
-        **model_inputs,
-        max_new_tokens=512
+    # Generate the math reasoning response
+    output_ids = math_model.generate(
+        **inputs,
+        max_new_tokens=1024,
+        pad_token_id=math_tokenizer.pad_token_id
     )
-    generated_ids = [
-        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
-    ]
+    response = math_tokenizer.batch_decode(output_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)[0]
+    math_messages.append({'role': 'assistant', 'content': response})  # Append assistant response
 
-    response = math_tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
-    
     return response
 
-# Function to clear inputs and output
-def clear_inputs():
-    return None, "", ""
+def math_chat_bot(image, sketchpad, question, state):
+    """
+    Orchestrates the OCR (image processing) and math reasoning based on user input.
+    """
+    current_tab_index = state["tab_index"]
+    image_description = None
+    # Upload tab
+    if current_tab_index == 0:
+        if image is not None:
+            image_description = process_image(image)
+    # Sketch tab
+    elif current_tab_index == 1:
+        if sketchpad and sketchpad["composite"]:
+            image_description = process_image(sketchpad["composite"], True)
+    
+    response = get_math_response(image_description, question)
+    yield response
 
-# Gradio interface setup
-def gradio_app(image, question, task):
-    if task == "OCR and Query":
-        return image, question, ocr_and_query(image, question)
-    elif task == "Solve Math Problem from Image":
-        if image is None:
-            return image, question, "Please upload an image."
-        extracted_text = ocr_and_query(image, "")
-        math_solution = solve_math_problem(extracted_text)
-        return image, extracted_text, math_solution
-    elif task == "Solve Math Problem from Text":
-        if question.strip() == "":
-            return image, question, "Please enter a math problem."
-        math_solution = solve_math_problem(question)
-        return image, question, math_solution
-    else:
-        return image, question, "Please select a task."
+css = """
+#qwen-md .katex-display { display: inline; }
+#qwen-md .katex-display>.katex { display: inline; }
+#qwen-md .katex-display>.katex>.katex-html { display: inline; }
+"""
 
-# Gradio interface
-with gr.Blocks() as app:
-    gr.Markdown("# Image OCR and Math Solver")
-    gr.Markdown("Upload an image, enter your question or math problem, and select the appropriate task.")
-    
-    with gr.Row():
-        image_input = gr.Image(type="pil", label="Upload Image")
-        text_input = gr.Textbox(lines=2, placeholder="Enter your question or math problem here...", label="Input")
-    
-    with gr.Row():
-        task_radio = gr.Radio(["OCR and Query", "Solve Math Problem from Image", "Solve Math Problem from Text"], label="Task")
-    
+def tabs_select(e: gr.SelectData, _state):
+    _state["tab_index"] = e.index
+
+# Create Gradio interface
+with gr.Blocks(css=css) as demo:
+    gr.HTML("""\
+<p align="center"><img src="https://modelscope.oss-cn-beijing.aliyuncs.com/resource/qwen.png" style="height: 60px"/><p>"""
+            """<center><font size=8>📖 Qwen2-Math Demo</center>"""
+            """\
+<center><font size=3>This WebUI is based on Qwen2-VL for OCR and Qwen2-Math for mathematical reasoning. You can input either images or texts of mathematical or arithmetic problems.</center>"""
+            )
+    state = gr.State({"tab_index": 0})
     with gr.Row():
-        complete_button = gr.Button("Complete")
-        clear_button = gr.Button("Clear")
-    
-    output = gr.Markdown(label="Output")
-    
-    # Event listeners
-    complete_button.click(fn=gradio_app, inputs=[image_input, text_input, task_radio], outputs=[image_input, text_input, output])
-    clear_button.click(fn=clear_inputs, outputs=[image_input, text_input, output])
+        with gr.Column():
+            with gr.Tabs() as input_tabs:
+                with gr.Tab("Upload"):
+                    input_image = gr.Image(type="pil", label="Upload")
+                with gr.Tab("Sketch"):
+                    input_sketchpad = gr.Sketchpad(type="pil", label="Sketch", layers=False)
+            input_tabs.select(fn=tabs_select, inputs=[state])
+            input_text = gr.Textbox(label="Input your question")
+            with gr.Row():
+                with gr.Column():
+                    clear_btn = gr.ClearButton([input_image, input_sketchpad, input_text])
+                with gr.Column():
+                    submit_btn = gr.Button("Submit", variant="primary")
+        with gr.Column():
+            output_md = gr.Markdown(label="answer",
+                                    latex_delimiters=[{
+                                        "left": "\\(",
+                                        "right": "\\)",
+                                        "display": True
+                                    }, {
+                                        "left": "\\begin\{equation\}",
+                                        "right": "\\end\{equation\}",
+                                        "display": True
+                                    }, {
+                                        "left": "\\begin\{align\}",
+                                        "right": "\\end\{align\}",
+                                        "display": True
+                                    }, {
+                                        "left": "\\begin\{alignat\}",
+                                        "right": "\\end\{alignat\}",
+                                        "display": True
+                                    }, {
+                                        "left": "\\begin\{gather\}",
+                                        "right": "\\end\{gather\}",
+                                        "display": True
+                                    }, {
+                                        "left": "\\begin\{CD\}",
+                                        "right": "\\end\{CD\}",
+                                        "display": True
+                                    }, {
+                                        "left": "\\[",
+                                        "right": "\\]",
+                                        "display": True
+                                    }],
+                                    elem_id="qwen-md")
+        submit_btn.click(
+            fn=math_chat_bot,
+            inputs=[input_image, input_sketchpad, input_text, state],
+            outputs=output_md)
 
-# Launch the app
-app.launch(share=True)
+demo.launch(share=True)