Spaces:

baidu
/

Qianfan-OCR-Demo

Paused

App Files Files Community

Azure99 commited on Mar 20

Commit

2a0b35e

verified ·

1 Parent(s): b20a58d

Upload folder using huggingface_hub

Browse files

Files changed (1) hide show

app.py +48 -9

app.py CHANGED Viewed

@@ -1,3 +1,4 @@
 import os
 import re
 from pathlib import Path
@@ -12,9 +13,10 @@ from image_utils import load_image
 MODEL_PATH = os.getenv("MODEL_PATH", "baidu/Qianfan-OCR")
 MAX_TILES_PER_IMAGE = int(os.getenv("MAX_TILES_PER_IMAGE", "12"))
-ZEROGPU_DURATION = int(os.getenv("ZEROGPU_DURATION", "120"))
-DEFAULT_MAX_NEW_TOKENS = 4096
-MAX_NEW_TOKENS_LIMIT = 16384
 DEFAULT_PROMPT = "Please extract the text from the image."
 IMAGE_FILE_TYPES = [".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tif", ".tiff"]
 APP_DIR = Path(__file__).resolve().parent
@@ -29,6 +31,7 @@ LATEX_FENCE_RE = re.compile(
     r"(^|\n)```(?:latex|tex)[ \t]*\n(?P<body>.*?)(?:\n```)(?=\n|$)",
     re.IGNORECASE | re.DOTALL,
 )
 EXAMPLE_ITEMS = [
     {
@@ -106,6 +109,13 @@ def validate_prompt(prompt: str) -> str:
     return prompt
 def preview_uploaded_images(file_paths):
     if not file_paths:
         return []
@@ -113,7 +123,7 @@ def preview_uploaded_images(file_paths):
 def clear_form():
-    return None, [], DEFAULT_PROMPT, DEFAULT_MAX_NEW_TOKENS, ""
 def load_example(example_index: int):
@@ -126,6 +136,7 @@ def load_example(example_index: int):
         image_paths,
         image_paths,
         example["prompt"],
         DEFAULT_MAX_NEW_TOKENS,
         "",
     )
@@ -148,13 +159,35 @@ def normalize_markdown_math(text: str) -> str:
     return LATEX_FENCE_RE.sub(replace_latex_fence, text)
-@spaces.GPU(duration=ZEROGPU_DURATION)
-def run_inference(file_paths, prompt, max_new_tokens):
     if not file_paths:
         raise gr.Error("Please upload at least one image.")
     pixel_values = load_images(file_paths, _MODEL.device, _MODEL.dtype)
-    question = validate_prompt(prompt)
     generation_config = {"max_new_tokens": int(max_new_tokens)}
     with torch.no_grad():
@@ -164,7 +197,7 @@ def run_inference(file_paths, prompt, max_new_tokens):
             question=question,
             generation_config=generation_config,
         )
-    return normalize_markdown_math(response)
 def build_demo():
@@ -204,6 +237,10 @@ def build_demo():
                     value=DEFAULT_PROMPT,
                     placeholder="Describe the task for the uploaded image(s).",
                 )
                 max_new_tokens_input = gr.Slider(
                     label="max_new_tokens",
                     minimum=256,
@@ -249,6 +286,7 @@ def build_demo():
                     image_input,
                     image_preview,
                     prompt_input,
                     max_new_tokens_input,
                     output_box,
                 ],
@@ -261,7 +299,7 @@ def build_demo():
         )
         submit_button.click(
             fn=run_inference,
-            inputs=[image_input, prompt_input, max_new_tokens_input],
             outputs=output_box,
         )
         clear_button.click(
@@ -270,6 +308,7 @@ def build_demo():
                 image_input,
                 image_preview,
                 prompt_input,
                 max_new_tokens_input,
                 output_box,
             ],

+import math
 import os
 import re
 from pathlib import Path
 MODEL_PATH = os.getenv("MODEL_PATH", "baidu/Qianfan-OCR")
 MAX_TILES_PER_IMAGE = int(os.getenv("MAX_TILES_PER_IMAGE", "12"))
+DEFAULT_MAX_NEW_TOKENS = 2048
+MAX_NEW_TOKENS_LIMIT = 8192
+ZEROGPU_DURATION_MIN = 60
+ZEROGPU_DURATION_MAX = 300
 DEFAULT_PROMPT = "Please extract the text from the image."
 IMAGE_FILE_TYPES = [".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tif", ".tiff"]
 APP_DIR = Path(__file__).resolve().parent
     r"(^|\n)```(?:latex|tex)[ \t]*\n(?P<body>.*?)(?:\n```)(?=\n|$)",
     re.IGNORECASE | re.DOTALL,
 )
+THINK_BLOCK_RE = re.compile(r"<think>.*?</think>", re.DOTALL | re.IGNORECASE)
 EXAMPLE_ITEMS = [
     {
     return prompt
+def build_question(prompt: str, layout_as_thought: bool) -> str:
+    question = validate_prompt(prompt)
+    if layout_as_thought and not question.endswith("<think>"):
+        question = f"{question}<think>"
+    return question
 def preview_uploaded_images(file_paths):
     if not file_paths:
         return []
 def clear_form():
+    return None, [], DEFAULT_PROMPT, False, DEFAULT_MAX_NEW_TOKENS, ""
 def load_example(example_index: int):
         image_paths,
         image_paths,
         example["prompt"],
+        False,
         DEFAULT_MAX_NEW_TOKENS,
         "",
     )
     return LATEX_FENCE_RE.sub(replace_latex_fence, text)
+def wrap_think_blocks(text: str) -> str:
+    if not text:
+        return text
+    def replace_think_block(match: re.Match[str]) -> str:
+        block = match.group(0).strip()
+        return f"\n```text\n{block}\n```\n"
+    return THINK_BLOCK_RE.sub(replace_think_block, text)
+def estimate_zerogpu_duration(file_paths, prompt, layout_as_thought, max_new_tokens):
+    del file_paths, prompt, layout_as_thought
+    estimated_duration = math.ceil(int(max_new_tokens) / 25 + 15)
+    final_duration = max(
+        ZEROGPU_DURATION_MIN,
+        min(ZEROGPU_DURATION_MAX, estimated_duration),
+    )
+    return final_duration
+@spaces.GPU(duration=estimate_zerogpu_duration)
+def run_inference(file_paths, prompt, layout_as_thought, max_new_tokens):
     if not file_paths:
         raise gr.Error("Please upload at least one image.")
     pixel_values = load_images(file_paths, _MODEL.device, _MODEL.dtype)
+    question = build_question(prompt, layout_as_thought)
     generation_config = {"max_new_tokens": int(max_new_tokens)}
     with torch.no_grad():
             question=question,
             generation_config=generation_config,
         )
+    return normalize_markdown_math(wrap_think_blocks(response))
 def build_demo():
                     value=DEFAULT_PROMPT,
                     placeholder="Describe the task for the uploaded image(s).",
                 )
+                layout_as_thought_input = gr.Checkbox(
+                    label="Layout-as-Thought",
+                    value=False,
+                )
                 max_new_tokens_input = gr.Slider(
                     label="max_new_tokens",
                     minimum=256,
                     image_input,
                     image_preview,
                     prompt_input,
+                    layout_as_thought_input,
                     max_new_tokens_input,
                     output_box,
                 ],
         )
         submit_button.click(
             fn=run_inference,
+            inputs=[image_input, prompt_input, layout_as_thought_input, max_new_tokens_input],
             outputs=output_box,
         )
         clear_button.click(
                 image_input,
                 image_preview,
                 prompt_input,
+                layout_as_thought_input,
                 max_new_tokens_input,
                 output_box,
             ],