[📂 Project Pages] [📖 Paper] [🤗 Weights] [🤗 Demo] [🚀 Quick Start]
- Introduction
- Installation
- Quick Start
- Streamlit Demo
- BPE Token Visualization
- Token Family
- Release Plans
We are excited to announce the release of TokenOCR, 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 TokenOCR,
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 TokenOCR to construct a document-level MLLM, TokenVL, for VQA-based document understanding tasks.
In summary:
(1) The first token-level image text dataset (TokenIT) is proposed;
(2) The first token-level text image foundation model, TokenOCR, is proposed to support downstream tasks.
(3) The image-as-text semantic capability inspires us to develop TokenVL, a VQA-based MLLM tailored for document perception, understanding, and reasoning.
conda create -n tokenocr python=3.9
conda activate tokenocr
pip install -r requirements.txt
Install flash-attn==2.3.6 (optional):
pip install flash-attn==2.3.6 --no-build-isolation
Alternatively you can compile from source:
git clone https://github.com/Dao-AILab/flash-attention.git
cd flash-attention
git checkout v2.3.6
python setup.py install
If you don't use flash-attn, please modify the configs of weights, referring to this
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/TokenOCR_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.pyWe are excited to present an interactive demo of our project using Streamlit. This demo allows users to explore the capabilities of our model——TokenOCR.
To run the Streamlit demo, you need to wrap the dependencies and then run:pip install requirement_app.txt
streamlit run app.py --server.port 8400
- Interactive Interface: Easily upload the image, enter the bpe you want to query, and click the RUN button to view the results of TokenOCR's processing.
- Real-time Results: Both models, based on internvl and resnet50, give users instant feedback in bpe.
- User-Friendly: Designed to be intuitive, even for users without a technical background.
- Access the Demo: [Link to your Streamlit demo]
- Upload a Document or Image: Use the interface to upload your files.
- Text input: Input your text related to the content of the images.
- View Results: See How models generate bpe visualizations in real time.
Then a simple Web-UI to interactive:
We welcome any feedback or suggestions to improve the demo. Please feel free to reach out via [contact information or GitHub issues].
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 |
| TokenOCR | token-level | TokenIT | 20M | 1.8B |
TokenOCR
An overview of the proposed TokenOCR, 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.
In the following table, we provide all models 🤗 link of the TokenOCR series. You can use prompt ' ' to get a highlight background.
| Model Name | Description |
|---|---|
| TokenOCR_2048_Bilingual_seg | Backbone is ViT；feature dimension is 2048; support interactive with English and Chinese texts. |
| TokenOCR_4096_English_seg | (We recommend đź‘Ť) Backbone is ViT; feature dimension is 4096; only supports interactive with English texts. |
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 TokenOCR. Please refer to our technical report for more details.
TokenVL
we employ the TokenOCR 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.
âś… Inference code and weights for TokenOCR
- Release Character-level Text Image Foundation Model (CharOCR)
- Code & model checkpoint for TokenVL
- Data for the Pre-training and Fine-tuning of TokenVL
- TokenIT data and script
This project is released under the MIT License.
If you find this project useful in your research, please consider citing:
@inproceedings{guan2025TokenOCR,
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},
journal={arXiv preprint arXiv:2503.02304},
year={2025}
}