#!/bin/bash
# ============================================================================
# OpenTSLM A0 BASELINE — full curriculum on THEIR original datasets
# (stage1 TSQA -> stage2 M4 -> stage3 HAR -> stage4 Sleep -> stage5 ECG)
# Usage:  bash run_baseline.sh <llm_id> <gpu>
#   e.g.  bash run_baseline.sh google/gemma-3-270m 0
#         bash run_baseline.sh meta-llama/Llama-3.2-1B 4
# ============================================================================

# ---- HF TOKEN (needs access to gated Llama/Gemma models) -------------------
export HF_TOKEN=<HF_TOKEN_REMOVED_SET_YOUR_OWN>
# --------------------------------------------------------------------------

export HF_HOME=/mnt/nvme2/adinath/timeagent/hf_cache   # model + TSQA cache (nvme2, fresh disk)
export TMPDIR=/mnt/nvme2/adinath/timeagent/tmp         # keep temp off the full root/nvme0
export TOKENIZERS_PARALLELISM=false
# Reclaim reserved-but-unallocated memory / reduce fragmentation on the long
# 12-lead ECG soft-prompt sequences (stage5) that otherwise OOM under GPU contention.
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

LLM_ID="${1:-google/gemma-3-270m}"     # backbone (arg 1, default gemma-3-270m)
GPU="${2:-0}"                           # GPU (arg 2)
export CUDA_VISIBLE_DEVICES=$GPU

SAFE=$(echo "$LLM_ID" | sed 's#.*/##; s/[.-]/_/g')
LOG=/mnt/nvme2/adinath/timeagent/logs/opentslm_baseline_${SAFE}.log

# Call the venv python by absolute path: `source activate` lands in conda's
# base env here (conda auto-activation hijacks PATH). venv is on nvme2 now.
VENV_PY=/mnt/nvme2/timeagent/venv/bin/python3
export PYTHONPATH=/home/mbz-imran/Adinath/TimeAgent/OpenTSLM/src
cd /home/mbz-imran/Adinath/TimeAgent/OpenTSLM            # results/ -> nvme symlink

echo "=== OpenTSLM baseline | llm=$LLM_ID | GPU=$GPU | log=$LOG ==="
"$VENV_PY" -u curriculum_learning.py \
    --model OpenTSLMSP \
    --llm_id "$LLM_ID" \
    2>&1 | tee "$LOG"