Step 1. Build Dependencies. Our code is tested with Python 3.10.14. To run the codes, you should first install the following packages:

conda create -n geox python=3.10
conda activate geox
pip install --upgrade pip
pip install -r requirements.txt
pip install flash-attn==2.5.9.post1 --no-build-isolation

Step 1. Download and Prepare Data.

1. Follow the instructions here and download full dataset for GeoX.
2. To train the model, you are required to organize the files into the following folders:

./data/

  alignment/
    images/
    unified_formal_annotations.json

  geoqa/
    images/
    geoqa_train.json
    geoqa_test.json

  unigeo/
    images/
    unigeo_train.json
    unigeo_test.json

  geometry3k/
    images/
    geometry3k_train.json
    geometry3k_test.json

  pgps9k/
    images/
    pgps9k_train.json
    pgps9k_test.json

# Pretrain Geo-VIT
BASE_DIR="/path/to/your/base/directory"  # Modify this path as necessary
OUTPUT_DIR="/path/to/your/output/directory"  # Modify this path as necessary

python ${BASE_DIR}/pretrain/pretrain_encoder.py \
    --job_dir ${OUTPUT_DIR}/checkpoint/mae \
    --nodes 1 \
    --ngpus 8 \
    --accum_iter 16 \
    --batch_size 256 \
    --use_volta32 \
    --model mae_vit_base_patch16 \
    --mask_ratio 0.75 \
    --epochs 800 \
    --warmup_epochs 40 \
    --blr 1.5e-4 \
    --weight_decay 0.05 \
    --data_path ${BASE_DIR}/data  # Ensure the data path is correctly parameterized

# Pretrain Geo-LLM
DATA_FILE="/path/to/your/training/data"  # Modify this path as necessary
OUTPUT_DIR="/path/to/your/output/directory"  # Modify this path as necessary
MODEL_DIR="/path/to/LLEMMA/directory"  # Modify this path as necessary
LOG_FILE="${OUTPUT_DIR}/train.log"

if [ ! -d "${OUTPUT_DIR}" ]; then  
    mkdir -p "${OUTPUT_DIR}"
fi

GPU_DEVICES=""

deepspeed --include=localhost:${GPU_DEVICES} \
    main/train_llm.py \
    --config_name "${MODEL_DIR}/config.json" \
    --tokenizer_name "${MODEL_DIR}" \
    --model_name_or_path "${MODEL_DIR}" \
    --train_files "${DATA_FILE}" \
    --per_device_train_batch_size 64 \
    --per_device_eval_batch_size 32 \
    --do_train \
    --output_dir "${OUTPUT_DIR}" \
    --evaluation_strategy steps \
    --use_fast_tokenizer false \
    --max_eval_samples 0 \
    --learning_rate 1e-6 \
    --gradient_accumulation_steps 4 \
    --num_train_epochs 10 \
    --warmup_ratio 0.1 \
    --logging_dir "${OUTPUT_DIR}/logs" \
    --logging_strategy steps \
    --logging_steps 50 \
    --save_strategy steps \
    --preprocessing_num_workers 10 \
    --save_steps 20000000 \
    --eval_steps 500000000 \
    --save_total_limit 2000 \
    --seed 42 \
    --disable_tqdm false \
    --ddp_find_unused_parameters false \
    --block_size 1024 \
    --overwrite_output_dir \
    --report_to tensorboard \
    --run_name llm_pretrain \
    --bf16 \
    --bf16_full_eval \
    --gradient_checkpointing \
    --deepspeed configs/models/zero3.json \
    --ignore_data_skip true \
    --ddp_timeout 18000000 \
    | tee -a "${LOG_FILE}"

MODEL_DIR="/path/to/your/model/directory"  # Modify this path as necessary
Text_FILE="/path/to/your/training/data"  # Modify this path as necessary
IMAGE_FOLDER="/path/to/your/image/folder"  # Modify this path as necessary
OUTPUT_DIR="/path/to/your/output/directory"  # Modify this path as necessary
LOG_FILE="${OUTPUT_DIR}/train.log"
GPU_DEVICES=""

if [ ! -d "${OUTPUT_DIR}" ]; then  
    mkdir -p "${OUTPUT_DIR}"
fi
export MASTER_PORT=20728

deepspeed --include=localhost:${GPU_DEVICES} --master_port=$MASTER_PORT main/train_geox.py \
    --deepspeed ./configs/models/zero2.json \
    --model_name_or_path "${MODEL_DIR}" \
    --version geo_v1 \
    --data_path "${Text_FILE}" \
    --image_folder "${IMAGE_FOLDER}" \
    --vision_tower "${MODEL_DIR}/geo-vit.pth" \
    --gsformer_path "${MODEL_DIR}/gsformer.pth" \
    --mm_use_im_start_end False \
    --mm_use_im_patch_token False \
    --bf16 True \
    --output_dir "${OUTPUT_DIR}" \
    --num_train_epochs 100 \
    --per_device_train_batch_size 16 \
    --per_device_eval_batch_size 4 \
    --gradient_accumulation_steps 4 \
    --evaluation_strategy "no" \
    --save_strategy "steps" \
    --save_steps 1000 \
    --save_total_limit 100 \
    --learning_rate 3e-5 \
    --weight_decay 0. \
    --warmup_ratio 0.05 \
    --lr_scheduler_type "cosine" \
    --logging_steps 1 \
    --tf32 True \
    --model_max_length 2048 \
    --gradient_checkpointing True \
    --dataloader_num_workers 0 \
    --lazy_preprocess True \
    | tee -a "${LOG_FILE}"