#!/bin/bash
# Time-R1 baseline on ActivityForensics (single-span split-per-segment).
#
# Prereqs (run once from baselines/time_r1):
#   1. python data_forensics/build_forensics_json.py \
#          --annot_dir /ces/zt/activityforensics/annot \
#          --video_root /ces/zt \
#          --output_dir ./dataset/forensics/annotation
#   2. python data_forensics/link_cache.py \
#          --forensics_cache /path/to/forensics_grpo/preprocessed \
#          --annotation_json_dir ./dataset/forensics/annotation \
#          --output_dir ./dataset/forensics/preprocessed
#
# Time-R1's reward (single-span IoU + optional format) is kept UNCHANGED.
# Only data processing is adapted.

set -e

export WANDB_PROJECT=Forensics-Baselines
export WANDB_NAME=TimeR1_forensics_7B
export PYTHONPATH=".:$PYTHONPATH"

OUTDIR=./logs/finetune/${WANDB_NAME}
mkdir -p "${OUTDIR}"
export DEBUG_MODE="true"
export LOG_PATH="${OUTDIR}/${WANDB_NAME}.txt"

# Edit these paths for your machine.
MODEL_PATH=${MODEL_PATH:-./ckpts/Qwen2.5-VL-7B-Instruct}
ANNOT_DIR=${ANNOT_DIR:-./dataset/forensics/annotation}
PREPROC_DIR=${PREPROC_DIR:-./dataset/forensics/preprocessed}

torchrun --nproc_per_node="${NPROC:-4}" \
    --nnodes="${NNODES:-1}" \
    --node_rank="${NODE_RANK:-0}" \
    --master_addr="${MASTER_ADDR:-127.0.0.1}" \
    --master_port="${MASTER_PORT:-12370}" \
    finetune.py \
    --deepspeed scripts/zero3_offload.json \
    --output_dir "${OUTDIR}" \
    --model_name_or_path "${MODEL_PATH}" \
    --train_data_path "${ANNOT_DIR}/train.json" \
    --eval_data_path  "${ANNOT_DIR}/val.json" \
    --video_folder /unused \
    --preprocessed_data_path "${PREPROC_DIR}" \
    --dataset_name forensics \
    --max_prompt_length 8192 \
    --max_completion_length 1024 \
    --num_generations 8 \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 2 \
    --logging_steps 1 \
    --bf16 \
    --torch_dtype bfloat16 \
    --data_seed 42 \
    --gradient_checkpointing true \
    --attn_implementation flash_attention_2 \
    --fix_vit true \
    --slide_window false \
    --num_train_epochs 2 \
    --run_name "${WANDB_NAME}" \
    --report_to tensorboard \
    --reward_funcs iou format \
    --temperature 1.0 \
    --beta 0.0 \
    --prompt_type v1 \
    --is_curriculum_learning false \
    --logging_dir "${OUTDIR}" \
    --save_steps 100 \
    --use_grpo true \
    --save_only_model true