Vision-LLMs for Spatiotemporal Traffic Forecasting

Environment Requirements

To use st-vision-llm, you will need the following packages:

• python>=3.10
• torch==2.4.1
• transformers==4.57.1
• accelerate==1.11.0
• ms-swift==3.10.1

Install Dependencies

pip install -r requirements.txt

Data Preparation

1. Split the Milan dataset and place the files in the structure below

   data/raw/
     Milan/
       train/*.txt
       val/*.txt
       test/*.txt
     Trentino/
       train/*.txt
       val/*.txt
       test/*.txt
   

2. Dataset preprocessing

   python dataset_prepare.py

   • Each input directory produces one .pt data file saved to data/processed/<City>_<split>.pt.
   • The file contains:
     • tensors: dict whose keys are channels (sms_in, internet, etc.) and values are [T, H, W] float32 tensors.
     • all_times: ISO timestamps in the Europe/Rome timezone.
     • Metadata such as cache_version, created_time, etc.

3. Visualize the results

   python dataset_prepare.py --check data/processed/Milan_train.pt

   This command randomly samples a channel and previews all frames as an animation.

NumToken Pipeline

1. Generate training corpora

   cd train_numtoken
   python num_encode_math.py

   This script automatically produces:

   • newtokens.txt: every <|FPm/e|> special token.
   • stage1_train.jsonl: numeric alignment samples.
   • stage2_train.jsonl / stage2_val.jsonl: numeric reasoning and basic arithmetic tasks.

2. Two-stage SFT

   bash train_numtoken/train_sft.sh

   • Stage 1 trains only embed_tokens and lm_head to expand the vocabulary, saving to train_numtoken/output_stage1.
   • Stage 2 automatically loads the latest Stage 1 checkpoint and continues full-parameter SFT on the stage2_* data, outputting to train_numtoken/output_stage2.
   • Change the MODEL variable at the top of train_sft.sh to swap the base model.

SFT Dataset and Training

1. Script entry point

   bash train_sft/train_sft.sh \
     <dataset_type> <ratio> <data_type> \
     <input_len> <output_len> \
     <x_start> <x_end> <y_start> <y_end> \
     <base_model>

   • dataset_type: Milan or Trentino; the script maps this to data/processed/<City>_{train,val}.pt.
   • ratio: sampling ratio for time slices in the training set.
   • data_type: channel name inside tensors.
   • input_len / output_len: number of historical frames vs. predicted frames.
   • (x_start, x_end, y_start, y_end): spatial coordinate range.
   • base_model: path or Hugging Face name of the NumToken-extended base model.

2. Example

   bash train_sft/train_sft.sh \
     Milan 0.02 internet \
     12 20 \
     45 55 45 55 \
     Saves/qwenvl-7B-numtoken-math

GRPO Dataset and Training

1. Usage

   bash train_grpo/train_grpo.sh \
     <dataset_type> <ratio> <data_type> \
     <input_len> <output_len> \
     <x_start> <x_end> <y_start> <y_end> \
     <base_model>

   Parameters are identical to the SFT script.

2. Example

   bash train_grpo/train_grpo.sh \
     Milan 0.02 internet \
     12 20 \
     45 55 45 55 \
     train_sft/sft/checkpoint-600