#!/bin/bash

set -x

source env.sh

echo "args: $@"

# set the dataset dir via `DATADIR_QuarkGluon`
DATADIR=${DATADIR_QuarkGluon}
[[ -z $DATADIR ]] && DATADIR='./datasets/QuarkGluon'

# set a comment via `COMMENT`
suffix=${COMMENT}

model=$1
extraopts=""
modelopts="networks/example_ParticleTransformer.py"  # ParT

# /public/home/mashun/gaonengsuo/particle_transformer/training/QuarkGluon/ParT/20241121-093112_example_ParticleTransformer_ranger_lr0.001_batch512/net_best_epoch_state.pt

weaver --predict \
    --data-test "${DATADIR}/test_file_*.parquet" \
    --data-config data/QuarkGluon/qg_${FEATURE_TYPE}.yaml \
    --network-config $modelopts \
    --model-prefix  training/QuarkGluon/ParT/20241121-093112_example_ParticleTransformer_ranger_lr0.001_batch512/net_best_epoch_state.pt\
    --batch-size 512 \
    --gpus 0 \
    --predict-output pred.root \