| --- |
| license: mit |
| pipeline_tag: image-to-text |
| base_model: TokenFD |
| base_model_relation: finetune |
| --- |
| |
| <center> |
|
|
| <h1 style="color: black;">A Token-level Text Image Foundation Model for Document Understanding</h1> |
|
|
|
|
| [\[📂 GitHub\]](https://github.com/Token-family/TokenFD) [\[📖 Paper\]](https://arxiv.org/pdf/2503.02304) [\[🆕 Project Pages\]](https://token-family.github.io/project_page/) [\[🤗 HF Demo\]](https://huggingface.co/spaces/TongkunGuan/Token-level_Text_Image_Foundation_Model) |
| |
| </center> |
| |
| <!-- <div align="center"> |
| <img width="500" alt="image" src="https://cdn-uploads.huggingface.co/production/uploads/650d4a36cbd0c7d550d3b41b/dQ_JfK_I91WXzIq52D015.png"> |
| </div> --> |
| |
| <center> |
| |
| |
| <!-- ### Model Cards --> |
| <h2 style="color: #4CAF50;">Model Cards</h2> |
| |
| <!-- In the following table, we provide all models [🤗 link] of the series. --> |
| |
| <!-- | Model Name | Description | |
| | :-----------------------: | :-------------------------------------------------------------------: | |
| | [R50](https://huggingface.co/TongkunGuan/R50) |Backbone is ResNet-50；feature dimension is 2048; support interactive with English and Chinese texts. | |
| | [TokenFD-2048-Bilingual-seg](https://huggingface.co/TongkunGuan/TokenFD_2048_Bilingual_seg) |Backbone is ViT；feature dimension is 2048; support interactive with English and Chinese texts. | |
| | [TokenFD-4096-English-seg](https://huggingface.co/TongkunGuan/TokenFD_4096_English_seg) |(We recommend 👍) Backbone is ViT; feature dimension is 4096; only supports interactive with English texts. | |
| --> |
| 我们在当前开源的基座模型上进行了定制适配，模型地址： |
| |
| [TokenFD-ResNet50-bilingual](https://huggingface.co/TongkunGuan/R50) |
| |
| [TokenFD-InternViT2.5-bilingual](https://huggingface.co/TongkunGuan/TokenFD_2048_Bilingual_seg) |
| |
| [TokenFD-InternViT2.5-english](https://huggingface.co/TongkunGuan/TokenFD_4096_English_seg) |
| |
| [TokenFD-QwenViT2.5-bilingual](https://huggingface.co/TongkunGuan/TokenFD_no_dis) |
| |
| </center> |
| |
| <div align="center"> |
| <img width="800" alt="image" src="https://cdn-uploads.huggingface.co/production/uploads/650d4a36cbd0c7d550d3b41b/MHfn1TgL1DkiEhNmD8rnK.gif"> |
| </div> |
| |
| ### Quick Start |
| |
| > \[!Warning\] |
| > 🚨 Note: Since there are fewer Chinese images in public data than English, we recommend you use the **`TokenFD-4096-English-seg`** version. |
| |
| ```python |
| import os |
| import torch |
| from transformers import AutoTokenizer |
| from internvl.model.internvl_chat import InternVLChatModel |
| from utils import post_process, generate_similiarity_map, load_image |
|
|
| checkpoint = '/mnt/dolphinfs/hdd_pool/docker/user/hadoop-mt-ocr/guantongkun/VFM_try/processed_models/TokenFD_4096_English_seg' |
| image_path = './demo_images/0000000.png' |
| input_query = '11/12/2020' |
| out_dir = 'results' |
|
|
| if not os.path.exists(out_dir): |
| os.makedirs(out_dir, exist_ok=True) |
| |
| """loading model, tokenizer, tok_embeddings """ |
| tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True, use_fast=False) |
| model = InternVLChatModel.from_pretrained(checkpoint, low_cpu_mem_usage=True, torch_dtype=torch.bfloat16).eval() |
| model = model.cuda() |
| |
| """loading image """ |
| pixel_values, images, target_aspect_ratio = load_image(image_path) |
| |
|
|
| """loading query texts """ |
| if input_query[0] in '!"#$%&\'()*+,-./0123456789:;<=>?@^_{|}~0123456789': |
| input_ids = tokenizer(input_query)['input_ids'][1:] |
| else: |
| input_ids = tokenizer(' '+input_query)['input_ids'][1:] |
| input_ids = torch.Tensor(input_ids).long().to(model.device) |
| input_embeds = model.tok_embeddings(input_ids).clone() |
| all_bpe_strings = [tokenizer.decode(input_id) for input_id in input_ids] |
| |
|
|
| """Obtaining similarity """ |
| with torch.no_grad(): |
| vit_embeds, _ = model.forward_tokenocr(pixel_values.to(model.device)) #(vit_batch_size, 16*16, 2048) |
| vit_embeds_local, resized_size = post_process(vit_embeds, target_aspect_ratio) |
| token_features = vit_embeds_local / vit_embeds_local.norm(dim=-1, keepdim=True) |
| input_embedings = input_embeds / input_embeds.norm(dim=-1, keepdim=True) |
| similarity = input_embedings @ token_features.t() |
| attn_map = similarity.reshape(len(input_embedings), resized_size[0], resized_size[1]) |
| |
| """generate map locally """ |
| generate_similiarity_map(images, attn_map, all_bpe_strings, out_dir, target_aspect_ratio) |
| |
| |
| """user command """ |
| # python quick_start.py |
| ``` |
| <center> |
| |
| <!-- # Introduction --> |
| <h2 style="color: #4CAF50;">Introduction</h2> |
| |
| </center> |
| |
| We are excited to announce the release of **`TokenFD`**, the first token-level visual foundation model specifically tailored for text-image-related tasks, |
| designed to support a variety of traditional downstream applications. To facilitate the pretraining of TokenFD, |
| we also devise a high-quality data production pipeline that constructs the first token-level image text dataset, |
| **`TokenIT`**, comprising 20 million images and 1.8 billion token-mask pairs. |
| Furthermore, leveraging this foundation with exceptional image-as-text capability, |
| we seamlessly replace previous VFMs with TokenFD to construct a document-level MLLM, **`TokenVL`**, for VQA-based document understanding tasks. |
| |
| <center> |
| |
| <!-- # Token Family --> |
| <h2 style="color: #4CAF50;">Token Family</h2> |
| |
| </center> |
| |
| <!-- ## TokenIT --> |
| <h2 style="color: #4CAF50;">TokenIT</h2> |
| |
| In the following picture, we provide an overview of the self-constructed token-level **TokenIT** dataset, comprising 20 million images and 1.8 billion |
| text-mask pairs. |
| |
| As depicted in Figure 2 (a), each sample in this dataset includes a raw image, a mask image, and a JSON file. |
| The JSON file provides the question-answer pairs and several BPE tokens randomly selected from the answer, along with |
| the ordinal number of each BPE token in the answer and its corresponding pixel value on the mask image. Consequently, |
| **each BPE token corresponds one-to-one with a pixel-level mask**. |
| The data ratios are summarized in Figure 2 (b). Figure 2 (c) and (d) further provide the number distribution |
| of tokens per image type and a word cloud of the top 100 tokens, respectively. |
| |
| <div align="center"> |
| <img width="1000" alt="image" src="https://cdn-uploads.huggingface.co/production/uploads/650d4a36cbd0c7d550d3b41b/WcQwU3-xjyT5Vm-pZhACo.png"> |
| </div> |
| |
| <!--  --> |
| |
| The comparisons with other visual foundation models: |
| |
| | VFM | Granularity | Dataset | #Image | #Pairs | |
| |:-------------------|:------------|:---------|:------:|:------:| |
| | [CLIP](https://github.com/openai/CLIP) | image-level | WIT400M | 400M | 0.4B | |
| | [DINO](https://github.com/facebookresearch/dino) | image-level | ImageNet | 14M | - | |
| | [SAM](https://github.com/facebookresearch/SAM) | pixel-level | SA1B | 11M | 1.1B | |
| | **TokenFD** | **token-level** | **TokenIT** | **20M** | **1.8B** | |
| |
| |
| <!-- ## TokenFD |
| --> |
| <h2 style="color: #4CAF50;">TokenFD</h2> |
| |
| ### Model Architecture |
| |
| An overview of the proposed TokenFD, where the token-level image features and token-level language |
| features are aligned within the same semantic space. This “image-as-text” alignment seamlessly facilitates user-interactive |
| applications, including text segmentation, retrieval, and visual question answering. |
| |
| <div align="center"> |
| <img width="1000" alt="image" src="https://cdn-uploads.huggingface.co/production/uploads/650d4a36cbd0c7d550d3b41b/6vNkEPzolBWVM--beoxLI.png"> |
| </div> |
| |
| |
| |
| ### Evaluation on Vision Capability |
| |
| We present a comprehensive evaluation of the vision encoder’s performance across various domains and tasks. |
| The evaluation is divided into two key categories: |
| |
| (1) text retrial; |
| (2) image segmentation; |
| (3) visual question answering; |
| |
| This approach allows us to assess the representation quality of TokenFD. |
| Please refer to our technical report for more details. |
| |
| #### text retrial |
| |
| <div align="left"> |
| <img width="500" alt="image" src="https://cdn-uploads.huggingface.co/production/uploads/650d4a36cbd0c7d550d3b41b/wlLcdB0hpC666PrEQSDaM.png"> |
| </div> |
| |
| <!--  --> |
| |
| #### image segmentation |
| |
| <div align="left"> |
| <img width="500" alt="image" src="https://cdn-uploads.huggingface.co/production/uploads/650d4a36cbd0c7d550d3b41b/0HqFXP8OC2tLH4d7scdMt.png"> |
| </div> |
| |
| <!--  --> |
| |
| #### visual question answering |
| |
| <div align="left"> |
| <img width="500" alt="image" src="https://cdn-uploads.huggingface.co/production/uploads/650d4a36cbd0c7d550d3b41b/3PBP0akDiMbupu_Gr7lzP.png"> |
| </div> |
| |
| <!--  |
| --> |
| |
| <!-- ## TokenVL --> |
| <h2 style="color: #4CAF50;">TokenVL</h2> |
| |
| we employ the TokenFD as the visual foundation model and further develop an MLLM, named TokenVL, tailored for document understanding. |
| Following the previous training paradigm, TokenVL also includes two stages: |
| |
| **Stage 1: LLM-guided Token Alignment Training for text parsing tasks.** |
| |
| <div align="center"> |
| <img width="500" alt="image" src="https://cdn-uploads.huggingface.co/production/uploads/650d4a36cbd0c7d550d3b41b/5ZCzz1tYy0bnIIZFxgTPN.png"> |
| </div> |
| |
| |
| The framework of LLM-guided Token Alignment Training. Existing MLLMs primarily enhance spatial-wise text perception capabilities by integrating localization prompts to predict coordinates. However, this implicit |
| method makes it difficult for these models to have a precise understanding. |
| In contrast, the proposed token alignment uses BPE token masks to directly and explicitly align text with corresponding pixels in the input image, enhancing the MLLM’s localization awareness. |
| |
| **Stage 2: Supervised Instruction Tuning for VQA tasks.** |
| |
| During the Supervised Instruction Tuning stage, we cancel the token alignment branch as answers may not appear in the image for some reasoning tasks |
| (e.g., How much taller is the red bar compared to the green bar?). This also ensures no computational overhead during inference to improve the document understanding capability. Finally, we inherit the |
| remaining weights from the LLM-guided Token Alignment and unfreeze all parameters to facilitate comprehensive parameter updates. |
| |
| ### OCRBench Results |
| |
| <div align="center"> |
| <img width="1300" alt="image" src="https://cdn-uploads.huggingface.co/production/uploads/650d4a36cbd0c7d550d3b41b/DZej5Ogpho3wpZC4KVAMO.png"> |
| </div> |
| |
| ### Document Understanding Results |
| |
| <div align="center"> |
| <img width="1300" alt="image" src="https://cdn-uploads.huggingface.co/production/uploads/650d4a36cbd0c7d550d3b41b/Msfs1YkDQHq2-djhm6QqD.png"> |
| </div> |
| |
| |
| |
| ## License |
| |
| This project is released under the MIT License. |
| |
| ## Citation |
| |
| If you find this project useful in your research, please consider citing: |
| |
| ```BibTeX |
| @inproceedings{guan2025TokenFD, |
| title={A Token-level Text Image Foundation Model for Document Understanding}, |
| author={Tongkun Guan, Zining Wang, Pei Fu, Zhentao Guo, Wei Shen, Kai zhou, Tiezhu Yue, Chen Duan, Hao Sun, Qianyi Jiang, Junfeng Luo, Xiaokang Yang}, |
| booktitle={Arxiv}, |
| year={2025} |
| } |
| ``` |
| |
