GOT-OCR2.0 Documentation

This codebase supports post-training (stage-2/stage-3) using DeepSpeed.

Important Requirements:

• The <image> token in the conversations -> human -> value field is mandatory in your training samples.
• You must update data information in GOT/utils/constant.py and the data_name in GOT/data/conversation_dataset_qwen.py (line 37) to match your dataset.

Note: For stage-1 training, use the Vary-tiny-600k repository instead.

deepspeed   /GOT-OCR-2.0-master/GOT/train/train_GOT.py \
 --deepspeed /GOT-OCR-2.0-master/zero_config/zero2.json    --model_name_or_path /GOT_weights/ \
 --use_im_start_end True   \
 --bf16 True   \
 --gradient_accumulation_steps 2    \
 --evaluation_strategy "no"   \
 --save_strategy "steps"  \
 --save_steps 200   \
 --save_total_limit 1   \
 --weight_decay 0.    \
 --warmup_ratio 0.001     \
 --lr_scheduler_type "cosine"    \
 --logging_steps 1    \
 --tf32 True     \
 --model_max_length 8192    \
 --gradient_checkpointing True   \
 --dataloader_num_workers 8    \
 --report_to none  \
 --per_device_train_batch_size 2  \
 --num_train_epochs 1  \
 --learning_rate 2e-5   \
 --datasets pdf-ocr+scence \
 --output_dir /your/output/path

You can perform quick fine-tuning using the ms-swift library. By default, this performs SFT on the LLM and projector while freezing the vision encoder.

Custom Dataset Format: Your .jsonl files should follow one of these formats:

• Single image: {"query": "<image>55555", "response": "66666", "images": ["image_path"]}
• Multi-image: {"query": "<image><image>eeeee", "response": "fffff", "history": [], "images": ["image_path1", "image_path2"]}
• Text-only/History: {"query": "EEEEE", "response": "FFFFF", "history": [["query1", "response1"], ... ]}

Commands:

• Use --sft_type lora for LoRA fine-tuning.
• Use --deepspeed default-zero2 for distributed training with DeepSpeed ZeRO2.

# Setup ms-swift
git clone https://github.com/modelscope/ms-swift.git
cd ms-swift
pip install -e .[llm]

# Default LoRA fine-tune
CUDA_VISIBLE_DEVICES=0 swift sft\
--model_type got-ocr2 \
--model_id_or_path stepfun-ai/GOT-OCR2_0 \
--sft_type lora \
--dataset latex-ocr-print#5000

# LoRA fine-tune with DeepSpeed ZeRO2
NPROC_PER_NODE=4 \
CUDA_VISIBLE_DEVICES=0,1,2,3 swift sft \
--model_type got-ocr2 \
--model_id_or_path stepfun-ai/GOT-OCR2_0 \
--sft_type lora \
--dataset latex-ocr-print#5000 \
--deepspeed default-zero2

# Fine-tune with your own data
# --dataset train.jsonl --val_dataset val.jsonl