test_spec_multi.sh

#!/bin/bash
set -euo pipefail
export PYTHONNOUSERSITE=1
unset PYTHONPATH
# SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
# cd "$SCRIPT_DIR/../.."
# echo "CWD=$(pwd)"
export CUDA_VISIBLE_DEVICES=2
export CUBLAS_WORKSPACE_CONFIG=:4096:8
export NCCL_P2P_DISABLE=1
export NCCL_IB_DISABLE=1
export TORCH_DISTRIBUTED_DEBUG=DETAIL

# ===== zero-shot Chronos2 + RAG inference =====
model_name=chronos2
task_name=chronos2_rag_fut_spec

seq_len=288 # 历史长度可以调节
test_seq_len=$seq_len
test_pred_len=96

val_batch_size=16
num_workers=20

# 数据输入设置
data=TsFileDatasetMultiCV



data_name=xiexin_pvpf

pretrain_model_path=/data/qiuyunzhong/CKPT/Chronos/chronos2
checkpoints=./checkpoint
resume_dir=chronos-2

data_list_path=configs/PVPF/PVPF_tsfile.json

if [[ ! -d "$pretrain_model_path" ]]; then
    echo "[ERROR] pretrain_model_path not found: $pretrain_model_path" >&2
    echo "Please set pretrain_model_path to your local Chronos2 checkpoint directory." >&2
    exit 1
fi

if [[ ! -f "$data_list_path" ]]; then
    echo "[ERROR] data_list_path not found: $data_list_path" >&2
    exit 1
fi

# ===== RAG params =====
rag_k_mode=fixed
rag_max_k=30
rag_k=5
rag_strategy=coarse_wcos_fine_weuc
rag_max_kb=20000
rag_kb_stride=1

# rag_alignment: 1 to enable mean alignment between retrieved context & query
rag_alignment=0

# IMPORTANT: if seq_len > model context_length (often 2048), set these to keep some retrieved context.
rag_query_keep=$seq_len
des=spec_${data_name}_${model_name}_rag_fut_rmod${rag_k_mode}
run_ts=$(date +"%m%d%H%M%S")
filename=rmax${rag_max_k}_rk${rag_k}_sl${test_seq_len}_pl${test_pred_len}_${run_ts}
model_id=${filename}

mkdir -p ./logs/$des/${filename}
log_file=./logs/$des/${filename}/info.log
touch "$log_file"

{
    echo "============================================================"
    echo "run_timestamp: $run_ts"
    echo "data_list_path: $data_list_path"
    python -c 'import json,sys; data=json.load(open(sys.argv[1], "r", encoding="utf-8")); item=data[0] if isinstance(data, list) and data else {}; print("data_folder_path:", item.get("data_folder_path")); print("path:", item.get("path"))' "$data_list_path"
    echo "============================================================"
} | tee -a "$log_file"

alignment_flag=""
if [[ "$rag_alignment" -eq 1 ]]; then
    alignment_flag="--rag_alignment"
fi

python run.py \
    --task_name $task_name \
    --is_training 0 \
    --checkpoints $checkpoints \
    --resume_dir $resume_dir \
    --pretrain_model_path $pretrain_model_path \
    --model_id $model_id \
    --model $model_name \
    --seq_len $seq_len \
    --test_seq_len $test_seq_len \
    --test_pred_len $test_pred_len \
    --val_batch_size $val_batch_size \
    --gpu 0 \
    --num_workers $num_workers \
    --des $des \
    --data $data \
    --data_list_path $data_list_path \
    --rag_k $rag_k \
    --rag_strategy $rag_strategy \
    --rag_max_kb $rag_max_kb \
    --rag_kb_stride $rag_kb_stride \
    --rag_query_keep $rag_query_keep \
    --rag_k_mode $rag_k_mode \
    --rag_max_k $rag_max_k \
    --plt_rag \
    $alignment_flag \
    2>&1 | tee -a "$log_file"

echo "Finished zero-shot inference: logs saved in $log_file"