updata inference demo

dots_ocr/model/inference.py

@@ -14,11 +14,14 @@ def inference_with_vllm(
         top_p=0.9,
         max_completion_tokens=32768,
         model_name='rednote-hilab/dots.ocr',
+        system_prompt=None,
         ):
     
     addr = f"{protocol}://{ip}:{port}/v1"
     client = OpenAI(api_key="{}".format(os.environ.get("API_KEY", "0")), base_url=addr)
     messages = []
+    if system_prompt:
+        messages.append({"role": "system", "content": system_prompt})
     messages.append(
         {
             "role": "user",

dots_ocr/utils/prompts.py

@@ -22,12 +22,24 @@ dict_promptmode_to_prompt = {
     # prompt_layout_only_en: layout detection
     "prompt_layout_only_en": """Please output the layout information from this PDF image, including each layout's bbox and its category. The bbox should be in the format [x1, y1, x2, y2]. The layout categories for the PDF document include ['Caption', 'Footnote', 'Formula', 'List-item', 'Page-footer', 'Page-header', 'Picture', 'Section-header', 'Table', 'Text', 'Title']. Do not output the corresponding text. The layout result should be in JSON format.""",
 
-    # prompt_layout_only_en: parse ocr text except the Page-header and Page-footer
+    # prompt_ocr: parse ocr text except the Page-header and Page-footer
     "prompt_ocr": """Extract the text content from this image.""",
 
     # prompt_grounding_ocr: extract text content in the given bounding box
     "prompt_grounding_ocr": """Extract text from the given bounding box on the image (format: [x1, y1, x2, y2]).\nBounding Box:\n""",
 
+    # prompt_web_parsing: parse all webpage layout info in json format.
+    "prompt_web_parsing": """Parsing the layout info of this webpage image with format json:\n""",
+
+    # prompt_scene_spotting: scene spotting
+    "prompt_scene_spotting": """Detect and recognize the text in the image.""",
+    
+    # prompt_img2svg: generate the SVG code of the image
+    "prompt_image_to_svg": """Please generate the SVG code based on the image.viewBox="0 0 {width} {height}\"""",
+
+    # prompt_free_qa: general prompt 
+    "prompt_general": """ """,
+
     # "prompt_table_html": """Convert the table in this image to HTML.""",
     # "prompt_table_latex": """Convert the table in this image to LaTeX.""",
     # "prompt_formula_latex": """Convert the formula in this image to LaTeX.""",