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

Introduction

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.

Token Family

TokenIT

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.

The comparisons with other visual foundation models:

VFM	Granularity	Dataset	#Image	#Pairs
CLIP	image-level	WIT400M	400M	0.4B
DINO	image-level	ImageNet	14M	-
SAM	pixel-level	SA1B	11M	1.1B
TokenFD	token-level	TokenIT	20M	1.8B

TokenFD

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.

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

image segmentation

visual question answering

TokenVL

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.

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

Document Understanding Results

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}
}

TongkunGuan/TokenFD

A Token-level Text Image Foundation Model for Document Understanding

Model Cards

Quick Start