Add all project files with proper LFS tracking

.gitattributes

@@ -33,3 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.wav filter=lfs diff=lfs merge=lfs -text
+*.gif filter=lfs diff=lfs merge=lfs -text

chat.py

@@ -0,0 +1,198 @@
+""" 
+Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+SPDX-License-Identifier: MIT
+"""
+
+import os
+import warnings
+from collections import OrderedDict
+
+from omegaconf import ListConfig
+
+warnings.filterwarnings("ignore", category=UserWarning)
+warnings.filterwarnings("ignore", category=FutureWarning)
+os.environ.setdefault("PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION", "python")
+
+import torch
+from PIL import Image
+from transformers import PreTrainedTokenizerFast
+
+from utils.model import DonutConfig, DonutModel, SwinEncoder
+from utils.processor import DolphinProcessor
+
+
+def try_rename_lagacy_weights(ckpt, output_path=""):
+    if "state_dict" in ckpt.keys():
+        ckpt = ckpt["state_dict"]
+    if "module" in ckpt.keys():
+        ckpt = ckpt["module"]
+    new_ckpt = OrderedDict()
+    for k, v in ckpt.items():
+        if k.startswith("model."):
+            k = k[len("model.") :]
+        if k.startswith("encoder"):
+            new_ckpt["vpm" + k[len("encoder") :]] = v
+        elif k.startswith("decoder"):
+            new_ckpt["llm" + k[len("encoder") :]] = v
+        else:
+            new_ckpt[k] = v
+    if output_path:
+        torch.save(new_ckpt, output_path)
+    return new_ckpt
+
+
+def convert_listconfig_to_list(config):
+    new_config = {}
+    for k, v in config.items():
+        if isinstance(v, ListConfig):
+            new_config[k] = list(v)
+        else:
+            new_config[k] = v
+    return new_config
+
+
+class DOLPHIN:
+    def __init__(self, config, ckpt_path="") -> None:
+        self.model_args = config.model
+        self.swin_args = config.model.pop("swin_args")
+        self.swin_args = convert_listconfig_to_list(self.swin_args)
+
+        vision_tower = SwinEncoder(
+            input_size=self.swin_args["img_size"],
+            patch_size=self.swin_args["patch_size"],
+            embed_dim=self.swin_args["embed_dim"],
+            window_size=self.swin_args["window_size"],
+            encoder_layer=self.swin_args["encoder_layer"],
+            num_heads=self.swin_args["num_heads"],
+            align_long_axis=self.swin_args["align_long_axis"],
+        )
+
+        self.tokenizer = PreTrainedTokenizerFast(tokenizer_file=self.model_args.tokenizer_path)
+        self.tokenizer.pad_token = "<pad>"
+        self.tokenizer.bos_token = "<s>"
+        self.tokenizer.eos_token = "</s>"
+        self.tokenizer.unk_token = "<unk>"
+
+        if self.model_args.get("extra_answer_tokens", False):
+            # print("Allowing multitask training: adding <Answer/> to the tokenizer.")
+            prompt_end_token = " <Answer/>"
+            self.tokenizer.add_special_tokens({"additional_special_tokens": sorted(set([prompt_end_token]))})
+            self.tokenizer._prompt_end_token = prompt_end_token
+            self.tokenizer._prompt_end_token_id = self.tokenizer.convert_tokens_to_ids(prompt_end_token)
+
+        donut_config = DonutConfig(
+            decoder_layer=self.model_args.decoder_layer,
+            max_length=self.model_args.max_length,
+            max_position_embeddings=self.model_args.max_position_embeddings,
+            hidden_dimension=self.model_args.hidden_dimension,
+        )
+
+        self.model = DonutModel(config=donut_config, vision_tower=vision_tower, tokenizer=self.tokenizer)
+        if self.model_args.model_name_or_path:
+            ckpt = torch.load(self.model_args.model_name_or_path)
+            ckpt = try_rename_lagacy_weights(ckpt)
+            self.model.load_state_dict(ckpt, strict=True)
+
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model.to(device)
+        self.model.eval()
+        transform_args = {
+            "input_size": self.swin_args["img_size"],
+            "max_length": self.model_args.max_length,
+        }
+        self.processor = DolphinProcessor({}, self.tokenizer, transform_args=transform_args)
+
+    def chat(
+        self,
+        question,
+        image,
+        return_raw=False,
+        return_score=False,
+        return_img_size=False,
+        only_return_img_size=False,
+        max_batch_size=16,
+    ):
+
+        def _preprocess_image(image):
+            if isinstance(image, str):
+                image = Image.open(image).convert("RGB")
+            if return_img_size or only_return_img_size:
+                image_tensor, ori_size = self.processor.process_image_for_inference(image, return_img_size=True)
+            else:
+                image_tensor = self.processor.process_image_for_inference(image, return_img_size=False)
+                ori_size = None
+            return image_tensor, ori_size
+
+        def _preprocess_prompt(question):
+            if self.model_args.get("extra_answer_tokens", False):
+                if self.tokenizer._prompt_end_token not in question:
+                    question = question + self.tokenizer._prompt_end_token
+            prompt_ids = self.processor.process_prompt_for_inference(question)
+            return prompt_ids
+
+        def _preprocess_prompt_batch(question):
+            if self.model_args.get("extra_answer_tokens", False):
+                for i in range(len(question)):
+                    if self.tokenizer._prompt_end_token not in question[i]:
+                        question[i] = question[i] + self.tokenizer._prompt_end_token
+                    if not question[i].startswith("<s>"):
+                        question[i] = "<s>" + question[i]
+            return question
+
+        def _postprocess(output, question):
+            output = output.replace("<s>", "").replace(question, "").replace("</s>", "").replace("<pad>", "")
+            if self.model_args.get("extra_answer_tokens", False):
+                output = output.split(self.tokenizer._prompt_end_token)[-1]
+            return output
+
+        if isinstance(question, list):
+            image_tensor_list = []
+            for i in image:
+                image_tensor, ori_size = _preprocess_image(i)
+                image_tensor_list.append(image_tensor)
+            image_tensor = torch.cat(image_tensor_list, dim=0)
+
+            question = _preprocess_prompt_batch(question)
+            self.processor.tokenizer.padding_side = "left"
+            prompt_ids = self.processor.tokenizer(
+                question, add_special_tokens=False, return_tensors="pt", padding=True
+            ).input_ids
+        else:
+            image_tensor, ori_size = _preprocess_image(image)
+            prompt_ids = _preprocess_prompt(question)
+
+        if only_return_img_size:
+            return ori_size
+
+        model_output_batch = []
+        for i in range(0, image_tensor.shape[0], max_batch_size):
+            image_tensor_batch = image_tensor[i : i + max_batch_size]
+            prompt_ids_batch = prompt_ids[i : i + max_batch_size]
+            model_output = self.model.inference(image_tensors=image_tensor_batch, prompt_ids=prompt_ids_batch)
+            model_output_batch.append(model_output)
+        model_output = {}
+        for k, v in model_output_batch[0].items():
+            if isinstance(v, torch.Tensor):
+                model_output[k] = sum(
+                    [v_batch[k].cpu().numpy().tolist() for v_batch in model_output_batch],
+                    [],
+                )
+            else:
+                model_output[k] = sum([v_batch[k] for v_batch in model_output_batch], [])
+
+        if return_raw:
+            if return_img_size:
+                return model_output, ori_size
+            return model_output
+        else:
+            if isinstance(question, list):
+                output = [_postprocess(model_output["repetitions"][i], question[i]) for i in range(len(question))]
+                score = model_output["scores"]
+            else:
+                output = _postprocess(model_output["repetitions"][0], question)
+                score = model_output["scores"][0]
+            if return_score:
+                return output, score
+            if return_img_size:
+                return output, ori_size
+            return output

config/Dolphin.yaml

@@ -0,0 +1,17 @@
+model:
+  model_name_or_path: "./checkpoints/dolphin_model.bin"
+  tokenizer_path: "./checkpoints/dolphin_tokenizer.json"
+  extra_answer_tokens: True   # add <Answer/> token
+  max_length: 4096
+  decoder_layer: 10
+  max_position_embeddings: 4096
+  hidden_dimension: 1024
+  swin_args:
+    name: 'swin'
+    img_size: [896, 896]
+    patch_size: 4
+    embed_dim: 128
+    align_long_axis: False
+    window_size: 7
+    encoder_layer: [2, 2, 14, 2]
+    num_heads: [4, 8, 16, 32]

demo/element_imgs/markdown/table_1.md

@@ -0,0 +1,2 @@
+<table><tr><td></td><td></td><td>100-class (top-1 acc.)</td><td>1000-class (top-1 acc.)</td></tr><tr><td colspan="2">4096-d (float)</td><td>77.1 ± 1.5</td><td>65.0</td></tr><tr><td rowspan="3">1024 bits</td><td>BP</td><td>72.9 ± 1.3</td><td>58.1</td></tr><tr><td>CBE</td><td>73.0 ± 1.3</td><td>59.2</td></tr><tr><td>SP</td><td>73.8 ± 1.3</td><td>60.1</td></tr><tr><td rowspan="4">4096 bits</td><td>threshold [1]</td><td>73.5 ± 1.4</td><td>59.1</td></tr><tr><td>BP</td><td>76.0 ± 1.5</td><td>63.2</td></tr><tr><td>CBE</td><td>75.9 ± 1.4</td><td>63.0</td></tr><tr><td>SP</td><td>76.3 ± 1.5</td><td>63.3</td></tr><tr><td>8192 bits</td><td>SP</td><td>76.8 ± 1.4</td><td>64.2</td></tr><tr><td>16384 bits</td><td>SP</td><td>77.1 ± 1.6</td><td>64.5</td></tr></table>
+

demo/element_imgs/recognition_json/table_1.json

@@ -0,0 +1,6 @@
+[
+  {
+    "label": "tab",
+    "text": "<table><tr><td></td><td></td><td>100-class (top-1 acc.)</td><td>1000-class (top-1 acc.)</td></tr><tr><td colspan=\"2\">4096-d (float)</td><td>77.1 ± 1.5</td><td>65.0</td></tr><tr><td rowspan=\"3\">1024 bits</td><td>BP</td><td>72.9 ± 1.3</td><td>58.1</td></tr><tr><td>CBE</td><td>73.0 ± 1.3</td><td>59.2</td></tr><tr><td>SP</td><td>73.8 ± 1.3</td><td>60.1</td></tr><tr><td rowspan=\"4\">4096 bits</td><td>threshold [1]</td><td>73.5 ± 1.4</td><td>59.1</td></tr><tr><td>BP</td><td>76.0 ± 1.5</td><td>63.2</td></tr><tr><td>CBE</td><td>75.9 ± 1.4</td><td>63.0</td></tr><tr><td>SP</td><td>76.3 ± 1.5</td><td>63.3</td></tr><tr><td>8192 bits</td><td>SP</td><td>76.8 ± 1.4</td><td>64.2</td></tr><tr><td>16384 bits</td><td>SP</td><td>77.1 ± 1.6</td><td>64.5</td></tr></table>"
+  }
+]

demo/page_imgs/markdown/test_page3.md

@@ -0,0 +1,22 @@
+![Figure](figures/test_page3_figure_000.png)
+
+Figure 2: (left) Scaled Dot-Product Attention. (right) Multi-Head Attention consists of several attention layers running in parallel.
+
+query with all keys, divide each by $\sqrt{d_k}$ , and apply a softmax function to obtain the weights on the values.
+
+In practice, we compute the attention function on a set of queries simultaneously, packed together into a matrix $Q$ . The keys and values are also packed together into matrices $K$ and $V$ . We compute the matrix of outputs as: $$ \\     \text{Attention}(Q, K, V) = \mathrm{softmax}(\frac{QK^T}{\sqrt{d_k}})V \\ $$
+
+The two most commonly used attention functions are additive attention [2] , and dot-product (multiplicative) attention. Dot-product attention is identical to our algorithm, except for the scaling factor of $\frac{1}{\sqrt{d_k}}$ . Additive attention computes the compatibility function using a feed-forward network with a single hidden layer. While the two are similar in theoretical complexity, dot-product attention is much faster and more space-efficient in practice, since it can be implemented using highly optimized matrix multiplication code.
+
+While for small values of $d_k$ the two mechanisms perform similarly, additive attention outperforms dot product attention without scaling for larger values of $d_k$ [ 3 ] . We suspect that for large values of $d_k$ , the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients 4 To counteract this effect, we scale the dot products by $\frac{1}{\sqrt{d_k}}$ .
+
+3.2.2 Multi-Head Attention
+
+Instead of performing a single attention function with $d_{\text{model}}$ -dimensional keys, values and queries, we found it beneficial to linearly project the queries, keys and values $h$ times with different, learned linear projections to $d_k$ , $d_k$ and $d_v$ dimensions, respectively. On each of these projected versions of queries, keys and values we then perform the attention function in parallel, yielding $d_v$ -dimensional output values. These are concatenated and once again projected, resulting in the final values, as depicted in Figure 2 .
+
+Multi­head attention allows the model to jointly attend to information from different representation subspaces at different positions. With a single attention head, averaging inhibits this.
+
+${ }^{4}$ To illustrate why the dot products get large, assume that the components of $q$ and $k$ are independent random variables with mean 0 and variance 1 . Then their dot product, $q \cdot k=\sum_{i=1}^{d_{k}} q_{i} k_{i}$, has mean 0 and variance $d_{k}$.
+
+4
+

demo/page_imgs/recognition_json/page_1.json

@@ -0,0 +1,178 @@
+[
+  {
+    "label": "title",
+    "bbox": [
+      271,
+      188,
+      1194,
+      221
+    ],
+    "text": "LLaMA: Open and Efficient Foundation Language Models",
+    "reading_order": 0
+  },
+  {
+    "label": "author",
+    "bbox": [
+      313,
+      272,
+      1154,
+      317
+    ],
+    "text": "Hugo Touvron; Thibaut Lavril*, Gautier Izacard*, Xavier Martinet",
+    "reading_order": 1
+  },
+  {
+    "label": "para",
+    "bbox": [
+      269,
+      317,
+      1201,
+      425
+    ],
+    "text": "Marie-Anne Lachaux, Timothee Lacroix, Baptiste Rozière, Naman Goyal\nEric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin\nEdouard Grave*Guillaume Lample*",
+    "reading_order": 2
+  },
+  {
+    "label": "para",
+    "bbox": [
+      685,
+      440,
+      795,
+      482
+    ],
+    "text": "Meta AI",
+    "reading_order": 3
+  },
+  {
+    "label": "sec",
+    "bbox": [
+      376,
+      524,
+      502,
+      565
+    ],
+    "text": "\\begin{abstract}",
+    "reading_order": 4
+  },
+  {
+    "label": "para",
+    "bbox": [
+      209,
+      586,
+      675,
+      946
+    ],
+    "text": "We introduce LLaMA, a collection of founda-\ntion language models ranging from 7B to 65B\nparameters. We train our models on trillions\nof tokens, and show that it is possible to train\nstate-of-the-art models using publicly avail-\nable datasets exclusively, without resorting\nto proprietary and inaccessible datasets. In\nparticular, LLaMA-13B outperforms GPT-3\n(175B) on most benchmarks, and LLaMA-\n65B is competitive with the best models,\nChinchilla-70B and PaLM-540B. We release\nall our models to the research community $^1$ .",
+    "reading_order": 5
+  },
+  {
+    "label": "sec",
+    "bbox": [
+      167,
+      964,
+      376,
+      1006
+    ],
+    "text": "1 Introduction",
+    "reading_order": 6
+  },
+  {
+    "label": "para",
+    "bbox": [
+      167,
+      1027,
+      718,
+      1498
+    ],
+    "text": "Large Languages Models (LLMs) trained on mas-\nsive corpora of texts have shown their ability to per-\nform new tasks from textual instructions or from a\nfew examples ( Brown et al. , 2020 ) . These few-shot\nproperties first appeared when scaling models to a\nsufficient size ( Kaplan et al. , 2020 ) , resulting in a\nline of work that focuses on further scaling these\nmodels ( Chowdhery et al. , 2022 ; Rae et al. , 2021 ) .\nThese efforts are based on the assumption that\nmore parameters will lead to better performance.\nHowever, recent work from Hoffmann et al. ( 2022 )\nshows that, for a given compute budget, the best\nperformances are not achieved by the largest mod-\nels, but by smaller models trained on more data.",
+    "reading_order": 7
+  },
+  {
+    "label": "para",
+    "bbox": [
+      167,
+      1506,
+      717,
+      1844
+    ],
+    "text": "The objective of the scaling laws from Hoff-\nmann et al. ( 2022 ) is to determine how to best\nscale the dataset and model sizes for a particular\ntraining compute budget. However, this objective\ndisregards the inference budget, which becomes\ncritical when serving a language model at scale.\nIn this context, given a target level of performance,\nthe preferred model is not the fastest to train but the\nfastest at inference, and although it may be cheaper\nto train a large model to reach a certain level of",
+    "reading_order": 8
+  },
+  {
+    "label": "para",
+    "bbox": [
+      753,
+      539,
+      1304,
+      734
+    ],
+    "text": "performance, a smaller one trained longer will\nultimately be cheaper at inference. For instance,\nalthough Hoffmann et al. ( 2022 ) recommends\ntraining a 10B model on 200B tokens, we find\nthat the performance of a 7B model continues to\nimprove even after 1T tokens.",
+    "reading_order": 9
+  },
+  {
+    "label": "para",
+    "bbox": [
+      753,
+      769,
+      1305,
+      1236
+    ],
+    "text": "The focus of this work is to train a series of\nlanguage models that achieve the best possible per-\nformance at various inference budgets, by training\non more tokens than what is typically used. The\nresulting models, called LLaMA , ranges from 7B\nto 65B parameters with competitive performance\ncompared to the best existing LLMs. For instance,\nLLaMA-13B outperforms GPT-3 on most bench-\nmarks, despite being 10 $\\times$ smaller. We believe that\nthis model will help democratize the access and\nstudy of LLMs, since it can be run on a single GPU.\nAt the higher-end of the scale, our 65B-parameter\nmodel is also competitive with the best large lan-\nguage models such as Chinchilla or PaLM-540B.",
+    "reading_order": 10
+  },
+  {
+    "label": "para",
+    "bbox": [
+      753,
+      1257,
+      1305,
+      1601
+    ],
+    "text": "Unlike Chinchilla, PaLM, or GPT-3, we only\nuse publicly available data, making our work com-\npatible with open-sourcing, while most existing\nmodels rely on data which is either not publicly\navailable or undocumented (e.g. “ Books – 2TB ” or\n“ Social media conversations ” ). There exist some\nexceptions, notably OPT ( Zhang et al. , 2022 ) ,\nGPT-NeoX ( Black et al. , 2022 ) , BLOOM ( Scao\net al. , 2022 ) and GLM ( Zeng et al. , 2022 ) , but none\nthat are competitive with PaLM-62B or Chinchilla.",
+    "reading_order": 11
+  },
+  {
+    "label": "para",
+    "bbox": [
+      753,
+      1634,
+      1304,
+      1933
+    ],
+    "text": "In the rest of this paper, we present an overview\nof the modifications we made to the transformer\narchitecture ( Vaswani et al. , 2017 ) , as well as our\ntraining method. We then report the performance of\nour models and compare with others LLMs on a set\nof standard benchmarks. Finally, we expose some\nof the biases and toxicity encoded in our models,\nusing some of the most recent benchmarks from\nthe responsible AI community.",
+    "reading_order": 12
+  },
+  {
+    "label": "fnote",
+    "bbox": [
+      167,
+      1844,
+      712,
+      1907
+    ],
+    "text": "* Equal contribution.\nCorrespondence:\n{htouvron\nthibautlav,gizacard,egrave,glample}@meta.com",
+    "reading_order": 13
+  },
+  {
+    "label": "fnote",
+    "bbox": [
+      209,
+      1907,
+      632,
+      1931
+    ],
+    "text": "https://github.com/facebookresearch/llama",
+    "reading_order": 14
+  },
+  {
+    "label": "watermark",
+    "bbox": [
+      20,
+      649,
+      83,
+      1530
+    ],
+    "text": "arXiv:2302.1397lvl [cs.CL] 27 Feb 2023",
+    "reading_order": 15
+  }
+]

demo/page_imgs/recognition_json/test_page.json

@@ -0,0 +1,47 @@
+[
+  {
+    "label": "header",
+    "bbox": [
+      291,
+      90,
+      675,
+      120
+    ],
+    "text": "Scaled Dot-Product Attention",
+    "reading_order": 0
+  },
+  {
+    "label": "fig",
+    "text": "![Figure](figures/test_page_figure_001.png)",
+    "figure_path": "figures/test_page_figure_001.png",
+    "bbox": [
+      1274,
+      105,
+      1536,
+      627
+    ],
+    "reading_order": 1
+  },
+  {
+    "label": "cap",
+    "bbox": [
+      168,
+      719,
+      1413,
+      789
+    ],
+    "text": "Figure 2: (left) Scaled Dot-Product Attention. (right) Multi-Head Attention consists of several\nattention layers running in parallel.",
+    "reading_order": 2
+  },
+  {
+    "label": "para",
+    "bbox": [
+      168,
+      858,
+      1413,
+      934
+    ],
+    "text": "query with all keys, divide each by $\\sqrt{d_{k}}$, and apply a softmax function to obtain the weights on the\nvalues.",
+    "reading_order": 3
+  }
+]

demo/page_imgs/recognition_json/test_page2.json

@@ -0,0 +1,102 @@
+[
+  {
+    "label": "fig",
+    "text": "![Figure](figures/test_page2_figure_000.png)",
+    "figure_path": "figures/test_page2_figure_000.png",
+    "bbox": [
+      394,
+      117,
+      897,
+      837
+    ],
+    "reading_order": 0
+  },
+  {
+    "label": "cap",
+    "bbox": [
+      445,
+      852,
+      856,
+      873
+    ],
+    "text": "Figure 1: The Transformer - model architecture",
+    "reading_order": 1
+  },
+  {
+    "label": "para",
+    "bbox": [
+      218,
+      920,
+      1086,
+      1044
+    ],
+    "text": "wise fully connected feed-forward network. We employ a residual connection [ 10 ] around each of\nthe two sub-layers, followed by layer normalization [ 1 ] . That is, the output of each sub-layer is\n$\\mathrm{LayerNorm}(x+\\mathrm{Sublayer}(x))$ , where $\\mathrm{Sublayer}(x)$ is the function implemented by the sub-layer\nitself. To facilitate these residual connections, all sub-layers in the model, as well as the embedding\nlayers, produce outputs of dimension $d_{\\text{model}}=512$ .",
+    "reading_order": 2
+  },
+  {
+    "label": "para",
+    "bbox": [
+      218,
+      1071,
+      1085,
+      1244
+    ],
+    "text": "The The decoder is also composed of a stack of $N=6$ identical layers. In addition to the two\nsub-layers in each encoder layer, the decoder inserts a third sub-layer, which performs multi-head\nattention over the output of the encoder stack. Similar to the encoder, we employ residual connections\naround each of the sub-layers, followed by layer normalization. We also modify the self-attention\nsub-layer in the decoder stack to prevent positions from attending to subsequent positions. This\nmasking, combined with fact that the output embeddings are offset by one position, ensures that the\npredictions for position $i$ can depend only on the known outputs at positions less than $i$ .",
+    "reading_order": 3
+  },
+  {
+    "label": "sub_sec",
+    "bbox": [
+      226,
+      1283,
+      344,
+      1305
+    ],
+    "text": "3.2 Attention",
+    "reading_order": 4
+  },
+  {
+    "label": "para",
+    "bbox": [
+      218,
+      1322,
+      1087,
+      1422
+    ],
+    "text": "An attention function can be described as mapping a query and a set of key-value pairs to an output,\nwhere the query, keys, values, and output are all vectors. The output is computed as a weighted sum\nof the values, where the weight assigned to each value is computed by a compatibility function of the\nquery with the corresponding key.",
+    "reading_order": 5
+  },
+  {
+    "label": "sub_sub_sec",
+    "bbox": [
+      218,
+      1456,
+      562,
+      1474
+    ],
+    "text": "3.2.1 Scaled Dot-Product Attention",
+    "reading_order": 6
+  },
+  {
+    "label": "para",
+    "bbox": [
+      218,
+      1498,
+      1085,
+      1546
+    ],
+    "text": "We call our particular attention \"Scaled Dot-Product Attention\" (Figure 2 ). The input consists of\nqueries and keys of dimension $d_k$ , and values of dimension $d_v$ . We compute the dot products of the",
+    "reading_order": 7
+  },
+  {
+    "label": "foot",
+    "bbox": [
+      646,
+      1590,
+      662,
+      1607
+    ],
+    "text": "3",
+    "reading_order": 8
+  }
+]

demo/page_imgs/recognition_json/test_page3.json

@@ -0,0 +1,124 @@
+[
+  {
+    "label": "fig",
+    "text": "![Figure](figures/test_page3_figure_000.png)",
+    "figure_path": "figures/test_page3_figure_000.png",
+    "bbox": [
+      331,
+      134,
+      984,
+      489
+    ],
+    "reading_order": 0
+  },
+  {
+    "label": "cap",
+    "bbox": [
+      198,
+      554,
+      1065,
+      603
+    ],
+    "text": "Figure 2: (left) Scaled Dot-Product Attention. (right) Multi-Head Attention consists of several\nattention layers running in parallel.",
+    "reading_order": 1
+  },
+  {
+    "label": "para",
+    "bbox": [
+      198,
+      652,
+      1065,
+      701
+    ],
+    "text": "query with all keys, divide each by $\\sqrt{d_k}$ , and apply a softmax function to obtain the weights on the\nvalues.",
+    "reading_order": 2
+  },
+  {
+    "label": "para",
+    "bbox": [
+      198,
+      715,
+      1065,
+      881
+    ],
+    "text": "In practice, we compute the attention function on a set of queries simultaneously, packed together\ninto a matrix $Q$ . The keys and values are also packed together into matrices $K$ and $V$ . We compute\nthe matrix of outputs as:\n\\[\n    \\text{Attention}(Q, K, V) = \\mathrm{softmax}(\\frac{QK^T}{\\sqrt{d_k}})V\n\\]",
+    "reading_order": 3
+  },
+  {
+    "label": "para",
+    "bbox": [
+      198,
+      913,
+      1068,
+      1060
+    ],
+    "text": "The two most commonly used attention functions are additive attention [2] , and dot-product (multi-\nplicative) attention. Dot-product attention is identical to our algorithm, except for the scaling factor\nof $\\frac{1}{\\sqrt{d_k}}$ . Additive attention computes the compatibility function using a feed-forward network with\na single hidden layer. While the two are similar in theoretical complexity, dot-product attention is\nmuch faster and more space-efficient in practice, since it can be implemented using highly optimized\nmatrix multiplication code.",
+    "reading_order": 4
+  },
+  {
+    "label": "para",
+    "bbox": [
+      198,
+      1074,
+      1066,
+      1175
+    ],
+    "text": "While for small values of $d_k$ the two mechanisms perform similarly, additive attention outperforms\ndot product attention without scaling for larger values of $d_k$ [ 3 ] . We suspect that for large values of\n$d_k$ , the dot products grow large in magnitude, pushing the softmax function into regions where it has\nextremely small gradients 4 To counteract this effect, we scale the dot products by $\\frac{1}{\\sqrt{d_k}}$ .",
+    "reading_order": 5
+  },
+  {
+    "label": "sub_sub_sec",
+    "bbox": [
+      198,
+      1207,
+      467,
+      1225
+    ],
+    "text": "3.2.2 Multi-Head Attention",
+    "reading_order": 6
+  },
+  {
+    "label": "para",
+    "bbox": [
+      198,
+      1253,
+      1067,
+      1395
+    ],
+    "text": "Instead of performing a single attention function with $d_{\\text{model}}$ -dimensional keys, values and queries,\nwe found it beneficial to linearly project the queries, keys and values $h$ times with different, learned\nlinear projections to $d_k$ , $d_k$ and $d_v$ dimensions, respectively. On each of these projected versions of\nqueries, keys and values we then perform the attention function in parallel, yielding $d_v$ -dimensional\noutput values. These are concatenated and once again projected, resulting in the final values, as\ndepicted in Figure 2 .",
+    "reading_order": 7
+  },
+  {
+    "label": "para",
+    "bbox": [
+      198,
+      1403,
+      1065,
+      1453
+    ],
+    "text": "Multi­head attention allows the model to jointly attend to information from different representation\nsubspaces at different positions. With a single attention head, averaging inhibits this.",
+    "reading_order": 8
+  },
+  {
+    "label": "fnote",
+    "bbox": [
+      198,
+      1485,
+      1065,
+      1535
+    ],
+    "text": "${ }^{4}$ To illustrate why the dot products get large, assume that the components of $q$ and $k$ are independent random\nvariables with mean 0 and variance 1 . Then their dot product, $q \\cdot k=\\sum_{i=1}^{d_{k}} q_{i} k_{i}$, has mean 0 and variance $d_{k}$.",
+    "reading_order": 9
+  },
+  {
+    "label": "foot",
+    "bbox": [
+      625,
+      1578,
+      641,
+      1599
+    ],
+    "text": "4",
+    "reading_order": 10
+  }
+]

demo_element.py

@@ -0,0 +1,129 @@
+""" 
+Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+SPDX-License-Identifier: MIT
+"""
+
+import argparse
+import glob
+import os
+
+from omegaconf import OmegaConf
+from PIL import Image
+
+from chat import DOLPHIN
+from utils.utils import *
+
+
+def process_element(image_path, model, element_type, save_dir=None):
+    """Process a single element image (text, table, formula)
+
+    Args:
+        image_path: Path to the element image
+        model: DOLPHIN model instance
+        element_type: Type of element ('text', 'table', 'formula')
+        save_dir: Directory to save results (default: same as input directory)
+
+    Returns:
+        Parsed content of the element and recognition results
+    """
+    # Load and prepare image
+    pil_image = Image.open(image_path).convert("RGB")
+    pil_image = crop_margin(pil_image)
+
+    # Select appropriate prompt based on element type
+    if element_type == "table":
+        prompt = "Parse the table in the image."
+        label = "tab"
+    elif element_type == "formula":
+        prompt = "Read text in the image."
+        label = "formula"
+    else:  # Default to text
+        prompt = "Read text in the image."
+        label = "text"
+
+    # Process the element
+    result = model.chat(prompt, pil_image)
+
+    # Create recognition result in the same format as the document parser
+    recognition_result = [
+        {
+            "label": label,
+            "text": result.strip(),
+        }
+    ]
+
+    # Save results if save_dir is provided
+    if save_dir:
+        save_outputs(recognition_result, image_path, save_dir)
+        print(f"Results saved to {save_dir}")
+
+    return result, recognition_result
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Element-level processing using DOLPHIN model")
+    parser.add_argument("--config", default="./config/Dolphin.yaml", help="Path to configuration file")
+    parser.add_argument("--input_path", type=str, required=True, help="Path to input image or directory of images")
+    parser.add_argument(
+        "--element_type",
+        type=str,
+        choices=["text", "table", "formula"],
+        default="text",
+        help="Type of element to process (text, table, formula)",
+    )
+    parser.add_argument(
+        "--save_dir",
+        type=str,
+        default=None,
+        help="Directory to save parsing results (default: same as input directory)",
+    )
+    parser.add_argument("--print_results", action="store_true", help="Print recognition results to console")
+    args = parser.parse_args()
+
+    # Load Model
+    config = OmegaConf.load(args.config)
+    model = DOLPHIN(config)
+
+    # Set save directory
+    save_dir = args.save_dir or (
+        args.input_path if os.path.isdir(args.input_path) else os.path.dirname(args.input_path)
+    )
+    setup_output_dirs(save_dir)
+
+    # Collect Images
+    if os.path.isdir(args.input_path):
+        image_files = []
+        for ext in [".jpg", ".jpeg", ".png", ".JPG", ".JPEG", ".PNG"]:
+            image_files.extend(glob.glob(os.path.join(args.input_path, f"*{ext}")))
+        image_files = sorted(image_files)
+    else:
+        if not os.path.exists(args.input_path):
+            raise FileNotFoundError(f"Input path {args.input_path} does not exist")
+        image_files = [args.input_path]
+
+    total_samples = len(image_files)
+    print(f"\nTotal samples to process: {total_samples}")
+
+    # Process images one by one
+    for image_path in image_files:
+        print(f"\nProcessing {image_path}")
+        try:
+            result, recognition_result = process_element(
+                image_path=image_path,
+                model=model,
+                element_type=args.element_type,
+                save_dir=save_dir,
+            )
+
+            if args.print_results:
+                print("\nRecognition result:")
+                print(result)
+                print("-" * 40)
+
+        except Exception as e:
+            print(f"Error processing {image_path}: {str(e)}")
+            continue
+
+
+if __name__ == "__main__":
+    main()

demo_element_hf.py

@@ -0,0 +1,195 @@
+"""
+Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+SPDX-License-Identifier: MIT
+"""
+
+import argparse
+import glob
+import os
+
+import torch
+from PIL import Image
+from transformers import AutoProcessor, VisionEncoderDecoderModel
+
+from utils.utils import *
+
+
+class DOLPHIN:
+    def __init__(self, model_id_or_path):
+        """Initialize the Hugging Face model
+        
+        Args:
+            model_id_or_path: Path to local model or Hugging Face model ID
+        """
+        # Load model from local path or Hugging Face hub
+        self.processor = AutoProcessor.from_pretrained(model_id_or_path)
+        self.model = VisionEncoderDecoderModel.from_pretrained(model_id_or_path)
+        self.model.eval()
+        
+        # Set device and precision
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model.to(self.device)
+        self.model = self.model.half()  # Always use half precision by default
+        
+        # set tokenizer
+        self.tokenizer = self.processor.tokenizer
+        
+    def chat(self, prompt, image):
+        """Process an image with the given prompt
+        
+        Args:
+            prompt: Text prompt to guide the model
+            image: PIL Image to process
+            
+        Returns:
+            Generated text from the model
+        """
+        # Prepare image
+        pixel_values = self.processor(image, return_tensors="pt").pixel_values
+        pixel_values = pixel_values.half()
+            
+        # Prepare prompt
+        prompt = f"<s>{prompt} <Answer/>"
+        prompt_ids = self.tokenizer(
+            prompt, 
+            add_special_tokens=False, 
+            return_tensors="pt"
+        ).input_ids.to(self.device)
+        
+        decoder_attention_mask = torch.ones_like(prompt_ids)
+        
+        # Generate text
+        outputs = self.model.generate(
+            pixel_values=pixel_values.to(self.device),
+            decoder_input_ids=prompt_ids,
+            decoder_attention_mask=decoder_attention_mask,
+            min_length=1,
+            max_length=4096,
+            pad_token_id=self.tokenizer.pad_token_id,
+            eos_token_id=self.tokenizer.eos_token_id,
+            use_cache=True,
+            bad_words_ids=[[self.tokenizer.unk_token_id]],
+            return_dict_in_generate=True,
+            do_sample=False,
+            num_beams=1,
+            repetition_penalty=1.1,
+            temperature=1.0
+        )
+        
+        # Process the output
+        sequence = self.tokenizer.batch_decode(outputs.sequences, skip_special_tokens=False)[0]
+        sequence = sequence.replace(prompt, "").replace("<pad>", "").replace("</s>", "").strip()
+        
+        return sequence
+
+def process_element(image_path, model, element_type, save_dir=None):
+    """Process a single element image (text, table, formula)
+    
+    Args:
+        image_path: Path to the element image
+        model: HFModel model instance
+        element_type: Type of element ('text', 'table', 'formula')
+        save_dir: Directory to save results (default: same as input directory)
+        
+    Returns:
+        Parsed content of the element and recognition results
+    """
+    # Load and prepare image
+    pil_image = Image.open(image_path).convert("RGB")
+    pil_image = crop_margin(pil_image)
+    
+    # Select appropriate prompt based on element type
+    if element_type == "table":
+        prompt = "Parse the table in the image."
+        label = "tab"
+    elif element_type == "formula":
+        prompt = "Read text in the image."
+        label = "formula"
+    else:  # Default to text
+        prompt = "Read text in the image."
+        label = "text"
+    
+    # Process the element
+    result = model.chat(prompt, pil_image)
+    
+    # Create recognition result in the same format as the document parser
+    recognition_result = [
+        {
+            "label": label,
+            "text": result.strip(),
+        }
+    ]
+    
+    # Save results if save_dir is provided
+    if save_dir:
+        save_outputs(recognition_result, image_path, save_dir)
+        print(f"Results saved to {save_dir}")
+    
+    return result, recognition_result
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Element-level processing using DOLPHIN model")
+    parser.add_argument("--model_path", default="./hf_model", help="Path to Hugging Face model")
+    parser.add_argument("--input_path", type=str, required=True, help="Path to input image or directory of images")
+    parser.add_argument(
+        "--element_type",
+        type=str,
+        choices=["text", "table", "formula"],
+        default="text",
+        help="Type of element to process (text, table, formula)",
+    )
+    parser.add_argument(
+        "--save_dir",
+        type=str,
+        default=None,
+        help="Directory to save parsing results (default: same as input directory)",
+    )
+    parser.add_argument("--print_results", action="store_true", help="Print recognition results to console")
+    args = parser.parse_args()
+    
+    # Load Model
+    model = DOLPHIN(args.model_path)
+    
+    # Set save directory
+    save_dir = args.save_dir or (
+        args.input_path if os.path.isdir(args.input_path) else os.path.dirname(args.input_path)
+    )
+    setup_output_dirs(save_dir)
+    
+    # Collect Images
+    if os.path.isdir(args.input_path):
+        image_files = []
+        for ext in [".jpg", ".jpeg", ".png", ".JPG", ".JPEG", ".PNG"]:
+            image_files.extend(glob.glob(os.path.join(args.input_path, f"*{ext}")))
+        image_files = sorted(image_files)
+    else:
+        if not os.path.exists(args.input_path):
+            raise FileNotFoundError(f"Input path {args.input_path} does not exist")
+        image_files = [args.input_path]
+    
+    total_samples = len(image_files)
+    print(f"\nTotal samples to process: {total_samples}")
+    
+    # Process images one by one
+    for image_path in image_files:
+        print(f"\nProcessing {image_path}")
+        try:
+            result, recognition_result = process_element(
+                image_path=image_path,
+                model=model,
+                element_type=args.element_type,
+                save_dir=save_dir,
+            )
+
+            if args.print_results:
+                print("\nRecognition result:")
+                print(result)
+                print("-" * 40)
+        except Exception as e:
+            print(f"Error processing {image_path}: {str(e)}")
+            continue
+
+
+if __name__ == "__main__":
+    main()

demo_page.py

@@ -0,0 +1,247 @@
+""" 
+Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+SPDX-License-Identifier: MIT
+"""
+
+import argparse
+import glob
+import os
+
+import cv2
+from omegaconf import OmegaConf
+from PIL import Image
+
+from chat import DOLPHIN
+from utils.utils import *
+
+
+def process_document(document_path, model, save_dir, max_batch_size):
+    """Parse documents - Handles both images and PDFs"""
+    file_ext = os.path.splitext(document_path)[1].lower()
+    
+    if file_ext == '.pdf':
+        # Process PDF file
+        # Convert PDF to images
+        images = convert_pdf_to_images(document_path)
+        if not images:
+            raise Exception(f"Failed to convert PDF {document_path} to images")
+        
+        all_results = []
+        
+        # Process each page
+        for page_idx, pil_image in enumerate(images):
+            print(f"Processing page {page_idx + 1}/{len(images)}")
+            
+            # Generate output name for this page
+            base_name = os.path.splitext(os.path.basename(document_path))[0]
+            page_name = f"{base_name}_page_{page_idx + 1:03d}"
+            
+            # Process this page (don't save individual page results)
+            json_path, recognition_results = process_single_image(
+                pil_image, model, save_dir, page_name, max_batch_size, save_individual=False
+            )
+            
+            # Add page information to results
+            page_results = {
+                "page_number": page_idx + 1,
+                "elements": recognition_results
+            }
+            all_results.append(page_results)
+        
+        # Save combined results for multi-page PDF
+        combined_json_path = save_combined_pdf_results(all_results, document_path, save_dir)
+        
+        return combined_json_path, all_results
+
+    else:
+        # Process regular image file
+        pil_image = Image.open(document_path).convert("RGB")
+        base_name = os.path.splitext(os.path.basename(document_path))[0]
+        return process_single_image(pil_image, model, save_dir, base_name, max_batch_size)
+
+
+def process_single_image(image, model, save_dir, image_name, max_batch_size, save_individual=True):
+    """Process a single image (either from file or converted from PDF page)
+    
+    Args:
+        image: PIL Image object
+        model: DOLPHIN model instance
+        save_dir: Directory to save results
+        image_name: Name for the output file
+        max_batch_size: Maximum batch size for processing
+        save_individual: Whether to save individual results (False for PDF pages)
+        
+    Returns:
+        Tuple of (json_path, recognition_results)
+    """
+    # Stage 1: Page-level layout and reading order parsing
+    layout_output = model.chat("Parse the reading order of this document.", image)
+
+    # Stage 2: Element-level content parsing
+    padded_image, dims = prepare_image(image)
+    recognition_results = process_elements(layout_output, padded_image, dims, model, max_batch_size, save_dir, image_name)
+
+    # Save outputs only if requested (skip for PDF pages)
+    json_path = None
+    if save_individual:
+        # Create a dummy image path for save_outputs function
+        dummy_image_path = f"{image_name}.jpg"  # Extension doesn't matter, only basename is used
+        json_path = save_outputs(recognition_results, dummy_image_path, save_dir)
+
+    return json_path, recognition_results
+
+
+def process_elements(layout_results, padded_image, dims, model, max_batch_size, save_dir=None, image_name=None):
+    """Parse all document elements with parallel decoding"""
+    layout_results = parse_layout_string(layout_results)
+
+    text_table_elements = []  # Elements that need processing
+    figure_results = []  # Figure elements (no processing needed)
+    previous_box = None
+    reading_order = 0
+
+    # Collect elements for processing
+    for bbox, label in layout_results:
+        try:
+            # Adjust coordinates
+            x1, y1, x2, y2, orig_x1, orig_y1, orig_x2, orig_y2, previous_box = process_coordinates(
+                bbox, padded_image, dims, previous_box
+            )
+
+            # Crop and parse element
+            cropped = padded_image[y1:y2, x1:x2]
+            if cropped.size > 0 and cropped.shape[0] > 3 and cropped.shape[1] > 3:
+                if label == "fig":
+                    pil_crop = Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))
+                    
+                    figure_filename = save_figure_to_local(pil_crop, save_dir, image_name, reading_order)
+                    
+                    # For figure regions, store relative path instead of base64
+                    figure_results.append(
+                        {
+                            "label": label,
+                            "text": f"![Figure](figures/{figure_filename})",
+                            "figure_path": f"figures/{figure_filename}",
+                            "bbox": [orig_x1, orig_y1, orig_x2, orig_y2],
+                            "reading_order": reading_order,
+                        }
+                    )
+                else:
+                    # For text or table regions, prepare for parsing
+                    pil_crop = Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))
+                    prompt = "Parse the table in the image." if label == "tab" else "Read text in the image."
+                    text_table_elements.append(
+                        {
+                            "crop": pil_crop,
+                            "prompt": prompt,
+                            "label": label,
+                            "bbox": [orig_x1, orig_y1, orig_x2, orig_y2],
+                            "reading_order": reading_order,
+                        }
+                    )
+
+            reading_order += 1
+
+        except Exception as e:
+            print(f"Error processing bbox with label {label}: {str(e)}")
+            continue
+
+    # Parse text/table elements in parallel
+    recognition_results = figure_results
+    if text_table_elements:
+        crops_list = [elem["crop"] for elem in text_table_elements]
+        prompts_list = [elem["prompt"] for elem in text_table_elements]
+
+        # Inference in batch
+        batch_results = model.chat(prompts_list, crops_list, max_batch_size=max_batch_size)
+
+        # Add batch results to recognition_results
+        for i, result in enumerate(batch_results):
+            elem = text_table_elements[i]
+            recognition_results.append(
+                {
+                    "label": elem["label"],
+                    "bbox": elem["bbox"],
+                    "text": result.strip(),
+                    "reading_order": elem["reading_order"],
+                }
+            )
+
+    # Sort elements by reading order
+    recognition_results.sort(key=lambda x: x.get("reading_order", 0))
+
+    return recognition_results
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Document parsing based on DOLPHIN")
+    parser.add_argument("--config", default="./config/Dolphin.yaml", help="Path to configuration file")
+    parser.add_argument("--input_path", type=str, default="./demo", help="Path to input image/PDF or directory of files")
+    parser.add_argument(
+        "--save_dir",
+        type=str,
+        default=None,
+        help="Directory to save parsing results (default: same as input directory)",
+    )
+    parser.add_argument(
+        "--max_batch_size",
+        type=int,
+        default=4,
+        help="Maximum number of document elements to parse in a single batch (default: 4)",
+    )
+    args = parser.parse_args()
+
+    # Load Model
+    config = OmegaConf.load(args.config)
+    model = DOLPHIN(config)
+
+    # Collect Document Files (images and PDFs)
+    if os.path.isdir(args.input_path):
+        # Support both image and PDF files
+        file_extensions = [".jpg", ".jpeg", ".png", ".JPG", ".JPEG", ".PNG", ".pdf", ".PDF"]
+        
+        document_files = []
+        for ext in file_extensions:
+            document_files.extend(glob.glob(os.path.join(args.input_path, f"*{ext}")))
+        document_files = sorted(document_files)
+    else:
+        if not os.path.exists(args.input_path):
+            raise FileNotFoundError(f"Input path {args.input_path} does not exist")
+        
+        # Check if it's a supported file type
+        file_ext = os.path.splitext(args.input_path)[1].lower()
+        supported_exts = ['.jpg', '.jpeg', '.png', '.pdf']
+        
+        if file_ext not in supported_exts:
+            raise ValueError(f"Unsupported file type: {file_ext}. Supported types: {supported_exts}")
+        
+        document_files = [args.input_path]
+
+    save_dir = args.save_dir or (
+        args.input_path if os.path.isdir(args.input_path) else os.path.dirname(args.input_path)
+    )
+    setup_output_dirs(save_dir)
+
+    total_samples = len(document_files)
+    print(f"\nTotal files to process: {total_samples}")
+
+    # Process All Document Files
+    for file_path in document_files:
+        print(f"\nProcessing {file_path}")
+        try:
+            json_path, recognition_results = process_document(
+                document_path=file_path,
+                model=model,
+                save_dir=save_dir,
+                max_batch_size=args.max_batch_size,
+            )
+
+            print(f"Processing completed. Results saved to {save_dir}")
+
+        except Exception as e:
+            print(f"Error processing {file_path}: {str(e)}")
+            continue
+
+
+if __name__ == "__main__":
+    main()

demo_page_hf.py

@@ -0,0 +1,365 @@
+""" 
+Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+SPDX-License-Identifier: MIT
+"""
+
+import argparse
+import glob
+import os
+
+import cv2
+import torch
+from PIL import Image
+from transformers import AutoProcessor, VisionEncoderDecoderModel
+
+from utils.utils import *
+
+
+class DOLPHIN:
+    def __init__(self, model_id_or_path):
+        """Initialize the Hugging Face model
+        
+        Args:
+            model_id_or_path: Path to local model or Hugging Face model ID
+        """
+        # Load model from local path or Hugging Face hub
+        self.processor = AutoProcessor.from_pretrained(model_id_or_path)
+        self.model = VisionEncoderDecoderModel.from_pretrained(model_id_or_path)
+        self.model.eval()
+        
+        # Set device and precision
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model.to(self.device)
+        self.model = self.model.half()  # Always use half precision by default
+        
+        # set tokenizer
+        self.tokenizer = self.processor.tokenizer
+        
+    def chat(self, prompt, image):
+        """Process an image or batch of images with the given prompt(s)
+        
+        Args:
+            prompt: Text prompt or list of prompts to guide the model
+            image: PIL Image or list of PIL Images to process
+            
+        Returns:
+            Generated text or list of texts from the model
+        """
+        # Check if we're dealing with a batch
+        is_batch = isinstance(image, list)
+        
+        if not is_batch:
+            # Single image, wrap it in a list for consistent processing
+            images = [image]
+            prompts = [prompt]
+        else:
+            # Batch of images
+            images = image
+            prompts = prompt if isinstance(prompt, list) else [prompt] * len(images)
+        
+        # Prepare image
+        batch_inputs = self.processor(images, return_tensors="pt", padding=True)
+        batch_pixel_values = batch_inputs.pixel_values.half().to(self.device)
+        
+        # Prepare prompt
+        prompts = [f"<s>{p} <Answer/>" for p in prompts]
+        batch_prompt_inputs = self.tokenizer(
+            prompts,
+            add_special_tokens=False,
+            return_tensors="pt"
+        )
+
+        batch_prompt_ids = batch_prompt_inputs.input_ids.to(self.device)
+        batch_attention_mask = batch_prompt_inputs.attention_mask.to(self.device)
+        
+        # Generate text
+        outputs = self.model.generate(
+            pixel_values=batch_pixel_values,
+            decoder_input_ids=batch_prompt_ids,
+            decoder_attention_mask=batch_attention_mask,
+            min_length=1,
+            max_length=4096,
+            pad_token_id=self.tokenizer.pad_token_id,
+            eos_token_id=self.tokenizer.eos_token_id,
+            use_cache=True,
+            bad_words_ids=[[self.tokenizer.unk_token_id]],
+            return_dict_in_generate=True,
+            do_sample=False,
+            num_beams=1,
+            repetition_penalty=1.1,
+            temperature=1.0
+        )
+        
+        # Process output
+        sequences = self.tokenizer.batch_decode(outputs.sequences, skip_special_tokens=False)
+        
+        # Clean prompt text from output
+        results = []
+        for i, sequence in enumerate(sequences):
+            cleaned = sequence.replace(prompts[i], "").replace("<pad>", "").replace("</s>", "").strip()
+            results.append(cleaned)
+            
+        # Return a single result for single image input
+        if not is_batch:
+            return results[0]
+        return results
+
+
+def process_document(document_path, model, save_dir, max_batch_size=None):
+    """Parse documents with two stages - Handles both images and PDFs"""
+    file_ext = os.path.splitext(document_path)[1].lower()
+    
+    if file_ext == '.pdf':
+        # Process PDF file
+        # Convert PDF to images
+        images = convert_pdf_to_images(document_path)
+        if not images:
+            raise Exception(f"Failed to convert PDF {document_path} to images")
+        
+        all_results = []
+        
+        # Process each page
+        for page_idx, pil_image in enumerate(images):
+            print(f"Processing page {page_idx + 1}/{len(images)}")
+            
+            # Generate output name for this page
+            base_name = os.path.splitext(os.path.basename(document_path))[0]
+            page_name = f"{base_name}_page_{page_idx + 1:03d}"
+            
+            # Process this page (don't save individual page results)
+            json_path, recognition_results = process_single_image(
+                pil_image, model, save_dir, page_name, max_batch_size, save_individual=False
+            )
+            
+            # Add page information to results
+            page_results = {
+                "page_number": page_idx + 1,
+                "elements": recognition_results
+            }
+            all_results.append(page_results)
+        
+        # Save combined results for multi-page PDF
+        combined_json_path = save_combined_pdf_results(all_results, document_path, save_dir)
+        
+        return combined_json_path, all_results
+    
+    else:
+        # Process regular image file
+        pil_image = Image.open(document_path).convert("RGB")
+        base_name = os.path.splitext(os.path.basename(document_path))[0]
+        return process_single_image(pil_image, model, save_dir, base_name, max_batch_size)
+
+
+def process_single_image(image, model, save_dir, image_name, max_batch_size=None, save_individual=True):
+    """Process a single image (either from file or converted from PDF page)
+    
+    Args:
+        image: PIL Image object
+        model: DOLPHIN model instance
+        save_dir: Directory to save results
+        image_name: Name for the output file
+        max_batch_size: Maximum batch size for processing
+        save_individual: Whether to save individual results (False for PDF pages)
+        
+    Returns:
+        Tuple of (json_path, recognition_results)
+    """
+    # Stage 1: Page-level layout and reading order parsing
+    layout_output = model.chat("Parse the reading order of this document.", image)
+
+    # Stage 2: Element-level content parsing
+    padded_image, dims = prepare_image(image)
+    recognition_results = process_elements(layout_output, padded_image, dims, model, max_batch_size, save_dir, image_name)
+
+    # Save outputs only if requested (skip for PDF pages)
+    json_path = None
+    if save_individual:
+        # Create a dummy image path for save_outputs function
+        dummy_image_path = f"{image_name}.jpg"  # Extension doesn't matter, only basename is used
+        json_path = save_outputs(recognition_results, dummy_image_path, save_dir)
+
+    return json_path, recognition_results
+
+
+def process_elements(layout_results, padded_image, dims, model, max_batch_size, save_dir=None, image_name=None):
+    """Parse all document elements with parallel decoding"""
+    layout_results = parse_layout_string(layout_results)
+
+    # Store text and table elements separately
+    text_elements = []  # Text elements
+    table_elements = []  # Table elements
+    figure_results = []  # Image elements (no processing needed)
+    previous_box = None
+    reading_order = 0
+
+    # Collect elements to process and group by type
+    for bbox, label in layout_results:
+        try:
+            # Adjust coordinates
+            x1, y1, x2, y2, orig_x1, orig_y1, orig_x2, orig_y2, previous_box = process_coordinates(
+                bbox, padded_image, dims, previous_box
+            )
+
+            # Crop and parse element
+            cropped = padded_image[y1:y2, x1:x2]
+            if cropped.size > 0 and cropped.shape[0] > 3 and cropped.shape[1] > 3:
+                if label == "fig":
+                    pil_crop = Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))
+                    
+                    figure_filename = save_figure_to_local(pil_crop, save_dir, image_name, reading_order)
+                    
+                    # For figure regions, store relative path instead of base64
+                    figure_results.append(
+                        {
+                            "label": label,
+                            "text": f"![Figure](figures/{figure_filename})",
+                            "figure_path": f"figures/{figure_filename}",
+                            "bbox": [orig_x1, orig_y1, orig_x2, orig_y2],
+                            "reading_order": reading_order,
+                        }
+                    )
+                else:
+                    # Prepare element for parsing
+                    pil_crop = Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))
+                    element_info = {
+                        "crop": pil_crop,
+                        "label": label,
+                        "bbox": [orig_x1, orig_y1, orig_x2, orig_y2],
+                        "reading_order": reading_order,
+                    }
+                    
+                    # Group by type
+                    if label == "tab":
+                        table_elements.append(element_info)
+                    else:  # Text elements
+                        text_elements.append(element_info)
+
+            reading_order += 1
+
+        except Exception as e:
+            print(f"Error processing bbox with label {label}: {str(e)}")
+            continue
+
+    # Initialize results list
+    recognition_results = figure_results.copy()
+    
+    # Process text elements (in batches)
+    if text_elements:
+        text_results = process_element_batch(text_elements, model, "Read text in the image.", max_batch_size)
+        recognition_results.extend(text_results)
+    
+    # Process table elements (in batches)
+    if table_elements:
+        table_results = process_element_batch(table_elements, model, "Parse the table in the image.", max_batch_size)
+        recognition_results.extend(table_results)
+
+    # Sort elements by reading order
+    recognition_results.sort(key=lambda x: x.get("reading_order", 0))
+
+    return recognition_results
+
+
+def process_element_batch(elements, model, prompt, max_batch_size=None):
+    """Process elements of the same type in batches"""
+    results = []
+    
+    # Determine batch size
+    batch_size = len(elements)
+    if max_batch_size is not None and max_batch_size > 0:
+        batch_size = min(batch_size, max_batch_size)
+    
+    # Process in batches
+    for i in range(0, len(elements), batch_size):
+        batch_elements = elements[i:i+batch_size]
+        crops_list = [elem["crop"] for elem in batch_elements]
+        
+        # Use the same prompt for all elements in the batch
+        prompts_list = [prompt] * len(crops_list)
+        
+        # Batch inference
+        batch_results = model.chat(prompts_list, crops_list)
+        
+        # Add results
+        for j, result in enumerate(batch_results):
+            elem = batch_elements[j]
+            results.append({
+                "label": elem["label"],
+                "bbox": elem["bbox"],
+                "text": result.strip(),
+                "reading_order": elem["reading_order"],
+            })
+    
+    return results
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Document parsing based on DOLPHIN")
+    parser.add_argument("--model_path", default="./hf_model", help="Path to Hugging Face model")
+    parser.add_argument("--input_path", type=str, default="./demo", help="Path to input image/PDF or directory of files")
+    parser.add_argument(
+        "--save_dir",
+        type=str,
+        default=None,
+        help="Directory to save parsing results (default: same as input directory)",
+    )
+    parser.add_argument(
+        "--max_batch_size",
+        type=int,
+        default=16,
+        help="Maximum number of document elements to parse in a single batch (default: 16)",
+    )
+    args = parser.parse_args()
+
+    # Load Model
+    model = DOLPHIN(args.model_path)
+
+    # Collect Document Files (images and PDFs)
+    if os.path.isdir(args.input_path):
+        # Support both image and PDF files
+        file_extensions = [".jpg", ".jpeg", ".png", ".JPG", ".JPEG", ".PNG", ".pdf", ".PDF"]
+        
+        document_files = []
+        for ext in file_extensions:
+            document_files.extend(glob.glob(os.path.join(args.input_path, f"*{ext}")))
+        document_files = sorted(document_files)
+    else:
+        if not os.path.exists(args.input_path):
+            raise FileNotFoundError(f"Input path {args.input_path} does not exist")
+        
+        # Check if it's a supported file type
+        file_ext = os.path.splitext(args.input_path)[1].lower()
+        supported_exts = ['.jpg', '.jpeg', '.png', '.pdf']
+        
+        if file_ext not in supported_exts:
+            raise ValueError(f"Unsupported file type: {file_ext}. Supported types: {supported_exts}")
+        
+        document_files = [args.input_path]
+
+    save_dir = args.save_dir or (
+        args.input_path if os.path.isdir(args.input_path) else os.path.dirname(args.input_path)
+    )
+    setup_output_dirs(save_dir)
+
+    total_samples = len(document_files)
+    print(f"\nTotal files to process: {total_samples}")
+
+    # Process All Document Files
+    for file_path in document_files:
+        print(f"\nProcessing {file_path}")
+        try:
+            json_path, recognition_results = process_document(
+                document_path=file_path,
+                model=model,
+                save_dir=save_dir,
+                max_batch_size=args.max_batch_size,
+            )
+
+            print(f"Processing completed. Results saved to {save_dir}")
+
+        except Exception as e:
+            print(f"Error processing {file_path}: {str(e)}")
+            continue
+
+
+if __name__ == "__main__":
+    main()

deployment/ReadMe.md

@@ -0,0 +1,12 @@
+<h1 align="center">
+🚀 Dolphin Inference/Serving
+</h1>
+
+## vLLM
+> [Doc](./vllm/ReadMe.md)
+
+## TensorRT-LLM
+> [Doc](./tensorrt_llm/ReadMe.md)
+
+## Others
+

deployment/tensorrt_llm/ReadMe.md

@@ -0,0 +1,89 @@
+<h1 align="center">
+🚀 Dolphin TensorRT-LLM Demo
+</h1>
+
+## ✅ Introduction
+The Dolphin model employs a **Swin Encoder + MBart Decoder** architecture. In the HuggingFace Transformers [Config](https://huggingface.co/ByteDance/Dolphin/blob/main/config.json), 
+its architectures field is specified as "VisionEncoderDecoderModel". **Dolphin**, **[Nougat](https://huggingface.co/docs/transformers/model_doc/nougat)**, and **[Donut](https://huggingface.co/docs/transformers/model_doc/donut)** share the same model architecture. TensorRT-LLM has already supported the Nougat model. 
+Following Nougat's conversion script, we have successfully implemented Dolphin on TensorRT-LLM. 
+
+**Note:** [prompt_ids](./dolphin_runner.py#L120) MUST be of **int32** type, otherwise TensorRT-LLM will produce incorrect results.
+
+## 🛠️ Installation
+> We only test TensorRT-LLM 0.18.1 on Linux.
+
+https://nvidia.github.io/TensorRT-LLM/0.18.1/installation/linux.html
+
+
+## ⚡ Offline Inference
+```
+export MODEL_NAME="Dolphin"
+
+# predict elements reading order
+python run_dolphin.py \
+    --batch_size 1 \
+    --hf_model_dir tmp/hf_models/${MODEL_NAME} \
+    --visual_engine_dir tmp/trt_engines/${MODEL_NAME}/vision_encoder \
+    --llm_engine_dir tmp/trt_engines/${MODEL_NAME}/1-gpu/bfloat16 \
+    --max_new_tokens 4096 \
+    --repetition_penalty 1.0 \
+    --input_text "Parse the reading order of this document." \
+    --image_path "../../demo/page_imgs/page_1.jpeg"
+
+# recognize text/latex
+python run_dolphin.py \
+    --batch_size 1 \
+    --hf_model_dir tmp/hf_models/${MODEL_NAME} \
+    --visual_engine_dir tmp/trt_engines/${MODEL_NAME}/vision_encoder \
+    --llm_engine_dir tmp/trt_engines/${MODEL_NAME}/1-gpu/bfloat16 \
+    --max_new_tokens 4096 \
+    --repetition_penalty 1.0 \
+    --input_text "Read text in the image." \
+    --image_path "../../demo/element_imgs/block_formula.jpeg"
+
+
+python run_dolphin.py \
+    --batch_size 1 \
+    --hf_model_dir tmp/hf_models/${MODEL_NAME} \
+    --visual_engine_dir tmp/trt_engines/${MODEL_NAME}/vision_encoder \
+    --llm_engine_dir tmp/trt_engines/${MODEL_NAME}/1-gpu/bfloat16 \
+    --max_new_tokens 4096 \
+    --repetition_penalty 1.0 \
+    --input_text "Read text in the image." \
+    --image_path "../../demo/element_imgs/para_1.jpg"
+
+# recognize table
+python run_dolphin.py \
+    --batch_size 1 \
+    --hf_model_dir tmp/hf_models/${MODEL_NAME} \
+    --visual_engine_dir tmp/trt_engines/${MODEL_NAME}/vision_encoder \
+    --llm_engine_dir tmp/trt_engines/${MODEL_NAME}/1-gpu/bfloat16 \
+    --max_new_tokens 4096 \
+    --repetition_penalty 1.0 \
+    --input_text "Parse the table in the image." \
+    --image_path "../../demo/element_imgs/table_1.jpeg"
+```
+
+
+## ⚡ Online Inference
+```
+# 1. Start Api Server
+export MODEL_NAME="Dolphin"
+
+python api_server.py \
+    --hf_model_dir tmp/hf_models/${MODEL_NAME} \
+    --visual_engine_dir tmp/trt_engines/${MODEL_NAME}/vision_encoder \
+    --llm_engine_dir tmp/trt_engines/${MODEL_NAME}/1-gpu/bfloat16 \
+    --max_batch_size 16
+
+# 2. Predict
+# predict elements reading order
+python deployment/tensorrt_llm/api_client.py --image_path ./demo/page_imgs/page_1.jpeg --prompt "Parse the reading order of this document."
+
+# recognize text/latex
+python deployment/tensorrt_llm/api_client.py --image_path ./demo/element_imgs/block_formula.jpeg --prompt "Read text in the image."
+python deployment/tensorrt_llm/api_client.py --image_path ./demo/element_imgs/para_1.jpg --prompt "Read text in the image."
+
+# recognize table
+python deployment/tensorrt_llm/api_client.py --image_path ./demo/element_imgs/table_1.jpeg --prompt "Parse the table in the image."
+```

deployment/tensorrt_llm/api_client.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Example Python client for `vllm.entrypoints.api_server`
+Start the demo server:
+    python -m vllm.entrypoints.api_server --model <model_name>
+
+NOTE: The API server is used only for demonstration and simple performance
+benchmarks. It is not intended for production use.
+For production use, we recommend `vllm serve` and the OpenAI client API.
+"""
+
+import argparse
+import base64
+import json
+from argparse import Namespace
+from collections.abc import Iterable
+
+import requests
+
+
+def clear_line(n: int = 1) -> None:
+    LINE_UP = "\033[1A"
+    LINE_CLEAR = "\x1b[2K"
+    for _ in range(n):
+        print(LINE_UP, end=LINE_CLEAR, flush=True)
+
+
+def encode_image_base64(image_path: str) -> str:
+    """Encode local image to base64 format."""
+
+    with open(image_path, "rb") as f:
+        image_data = f.read()
+        result = base64.b64encode(image_data).decode("utf-8")
+
+    return result
+
+
+def post_http_request(
+        prompt: str, image_path: str, api_url: str, stream: bool = False
+) -> requests.Response:
+    headers = {"User-Agent": "Test Client"}
+    pload = {
+        "prompt": prompt,
+        "image_base64": encode_image_base64(image_path),
+    }
+    response = requests.post(api_url, headers=headers, json=pload, stream=stream)
+    return response
+
+
+def get_streaming_response(response: requests.Response) -> Iterable[list[str]]:
+    for chunk in response.iter_lines(
+            chunk_size=8192, decode_unicode=False, delimiter=b"\n"
+    ):
+        if chunk:
+            data = json.loads(chunk.decode("utf-8"))
+            output = data["text"]
+            yield output
+
+
+def get_response(response: requests.Response) -> list[str]:
+    data = json.loads(response.content)
+    output = data["text"]
+    return output
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument("--prompt", type=str, default="Parse the reading order of this document.")
+    parser.add_argument("--image_path", type=str, default="./demo/page_imgs/page_1.jpeg")
+    parser.add_argument("--stream", action="store_true")
+    return parser.parse_args()
+
+
+def main(args: Namespace):
+    prompt = args.prompt
+    image_path = args.image_path
+    api_url = f"http://{args.host}:{args.port}/generate"
+    stream = args.stream
+
+    print(f"Prompt: {prompt!r}\n", flush=True)
+    response = post_http_request(prompt, image_path, api_url, stream)
+
+    if stream:
+        num_printed_lines = 0
+        for h in get_streaming_response(response):
+            clear_line(num_printed_lines)
+            num_printed_lines = 0
+            for i, line in enumerate(h):
+                num_printed_lines += 1
+                print(f"Response {i}: {line!r}", flush=True)
+    else:
+        output = get_response(response)
+        print(f"Response: {output!r}", flush=True)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

deployment/tensorrt_llm/api_server.py

@@ -0,0 +1,112 @@
+# copied from: https://github.com/NVIDIA/TensorRT-LLM/blob/v0.18.1/examples/apps/fastapi_server.py
+
+#!/usr/bin/env python
+import asyncio
+import base64
+import io
+import logging
+import signal
+from http import HTTPStatus
+from PIL import Image
+from typing import Optional
+
+import click
+import uvicorn
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse, Response
+
+from tensorrt_llm.executor import CppExecutorError, RequestError
+from dolphin_runner import DolphinRunner, InferenceConfig
+
+TIMEOUT_KEEP_ALIVE = 5  # seconds.
+
+
+async def decode_image(image_base64: str) -> Image.Image:
+    image_data = base64.b64decode(image_base64)
+    image = Image.open(io.BytesIO(image_data))
+    return image
+
+
+class LlmServer:
+    def __init__(self, runner: DolphinRunner):
+        self.runner = runner
+        self.app = FastAPI()
+        self.register_routes()
+
+    def register_routes(self):
+        self.app.add_api_route("/health", self.health, methods=["GET"])
+        self.app.add_api_route("/generate", self.generate, methods=["POST"])
+
+    async def health(self) -> Response:
+        return Response(status_code=200)
+
+    async def generate(self, request: Request) -> Response:
+        """ Generate completion for the request.
+
+        The request should be a JSON object with the following fields:
+        - prompt: the prompt to use for the generation.
+        - image_base64: the image to use for the generation.
+        """
+        request_dict = await request.json()
+
+        prompt = request_dict.pop("prompt", "")
+        logging.info(f"request prompt: {prompt}")
+        image_base64 = request_dict.pop("image_base64", "")
+        image = await decode_image(image_base64)
+
+        try:
+            output_texts = self.runner.run([prompt], [image], 4024)
+            output_texts = [texts[0] for texts in output_texts]
+            return JSONResponse({"text": output_texts[0]})
+        except RequestError as e:
+            return JSONResponse(content=str(e),
+                                status_code=HTTPStatus.BAD_REQUEST)
+        except CppExecutorError:
+            # If internal executor error is raised, shutdown the server
+            signal.raise_signal(signal.SIGINT)
+
+    async def __call__(self, host, port):
+        config = uvicorn.Config(self.app,
+                                host=host,
+                                port=port,
+                                log_level="info",
+                                timeout_keep_alive=TIMEOUT_KEEP_ALIVE)
+        await uvicorn.Server(config).serve()
+
+
+@click.command()
+@click.option("--hf_model_dir", type=str, required=True)
+@click.option("--visual_engine_dir", type=str, required=True)
+@click.option("--llm_engine_dir", type=str, required=True)
+@click.option("--max_batch_size", type=int, default=16)
+@click.option("--max_new_tokens", type=int, default=4024)
+@click.option("--host", type=str, default=None)
+@click.option("--port", type=int, default=8000)
+def entrypoint(hf_model_dir: str,
+               visual_engine_dir: str,
+               llm_engine_dir: str,
+               max_batch_size: int,
+               max_new_tokens: int,
+               host: Optional[str] = None,
+               port: int = 8000):
+    host = host or "0.0.0.0"
+    port = port or 8000
+    logging.info(f"Starting server at {host}:{port}")
+
+    config = InferenceConfig(
+        max_new_tokens=max_new_tokens,
+        batch_size=max_batch_size,
+        log_level="info",
+        hf_model_dir=hf_model_dir,
+        visual_engine_dir=visual_engine_dir,
+        llm_engine_dir=llm_engine_dir,
+    )
+
+    dolphin_runner = DolphinRunner(config)
+    server = LlmServer(runner=dolphin_runner)
+
+    asyncio.run(server(host, port))
+
+
+if __name__ == "__main__":
+    entrypoint()