Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,21 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+code/OpenTSLM/assets/cdhi_logo.png filter=lfs diff=lfs merge=lfs -text
+code/OpenTSLM/assets/ecg_rationale.png filter=lfs diff=lfs merge=lfs -text
+code/OpenTSLM/assets/eth_cdhi_logo.png filter=lfs diff=lfs merge=lfs -text
+code/OpenTSLM/assets/har_rationale.png filter=lfs diff=lfs merge=lfs -text
+code/OpenTSLM/assets/m4_caption.png filter=lfs diff=lfs merge=lfs -text
+code/OpenTSLM/assets/schematic_overview_2.png filter=lfs diff=lfs merge=lfs -text
+code/OpenTSLM/assets/schematic_overview_3.png filter=lfs diff=lfs merge=lfs -text
+code/OpenTSLM/assets/sleep_rationale.png filter=lfs diff=lfs merge=lfs -text
+code/OpenTSLM/src/opentslm/time_series_datasets/ecg_qa/dataset_distribution.png filter=lfs diff=lfs merge=lfs -text
+code/OpenTSLM/src/opentslm/time_series_datasets/ecg_qa/ecg_example_real.png filter=lfs diff=lfs merge=lfs -text
+data/ecg_qa_cot/ecg_qa_cot_test.csv filter=lfs diff=lfs merge=lfs -text
+data/ecg_qa_cot/ecg_qa_cot_train.csv filter=lfs diff=lfs merge=lfs -text
+data/ecg_qa_cot/ecg_qa_cot_val.csv filter=lfs diff=lfs merge=lfs -text
+results/Llama_3_2_1B/OpenTSLMSP/stage2_captioning/results/test_predictions.jsonl filter=lfs diff=lfs merge=lfs -text
+results/Llama_3_2_1B/OpenTSLMSP/stage2_captioning/results/test_predictions_rank_0.jsonl filter=lfs diff=lfs merge=lfs -text
+results/Llama_3_2_1B/OpenTSLMSP/stage3_cot/results/test_predictions.jsonl filter=lfs diff=lfs merge=lfs -text
+results/Llama_3_2_1B/OpenTSLMSP/stage3_cot/results/test_predictions_rank_0.jsonl filter=lfs diff=lfs merge=lfs -text
+results/Llama_3_2_1B/OpenTSLMSP/stage5_ecg_cot/results/test_predictions_rank_0.jsonl filter=lfs diff=lfs merge=lfs -text

RUN.md

@@ -0,0 +1,49 @@
+# TimeAgent — Llama-3.2-1B baseline: finish the ECG eval
+
+This bundle lets another server complete the ONLY remaining step for the
+Llama-3.2-1B OpenTSLM-SoftPrompt **baseline**: the stage-5 ECG-QA-CoT evaluation.
+Stages 1-4 (TSQA, M4, HAR, Sleep) are fully complete, and ECG **training** is
+done (early-stopped @ epoch 14). Only the ECG generation eval remains — it kept
+OOM-crashing on our shared/contended GPUs.
+
+## Contents
+- `results/Llama_3_2_1B/` — all stage checkpoints + metrics (stages 1-4 have
+  metrics.json so they're skipped; stage5_ecg_cot has the trained checkpoint,
+  no metrics.json → the eval will run).
+- `data/ecg_qa_cot/` — ECG-QA-CoT CSV splits (train/val/test).
+- `data/ptbxl.tar.gz` — PTB-XL records, tarred. On target: `cd data && tar xzf ptbxl.tar.gz` before running.
+- `code/OpenTSLM/` — the code (includes our fixes: greedy eval `do_sample=False`,
+  sdpa attention, left-padded batched generation, `EVAL_BATCH_SIZE` env).
+- `requirements-key.txt` — key package versions we ran with.
+
+## Setup on the target server
+```bash
+# 1) python env (Python 3.13; a venv or conda env)
+python3 -m venv venv && source venv/bin/activate
+pip install torch transformers peft wfdb datasets accelerate huggingface_hub numpy scipy pandas einops open_flamingo
+#   (match requirements-key.txt if you hit version issues)
+
+# 2) wire the checkpoints + data into the code tree (symlinks)
+cd code/OpenTSLM
+ln -sfn "$(realpath ../../results)"      results
+ln -sfn "$(realpath ../../data)"         src/data
+
+# 3) HF token (Llama-3.2-1B is gated; the base model auto-downloads)
+export HF_TOKEN=<your_hf_token_with_llama_access>
+export HF_HOME=/path/with/space/hf_cache
+```
+
+## Run the ECG eval
+```bash
+cd code/OpenTSLM
+# On a CLEAN 95GB GPU (no other tenants) you can batch the eval for speed:
+EVAL_BATCH_SIZE=4 python3 -u curriculum_learning.py --model OpenTSLMSP --llm_id meta-llama/Llama-3.2-1B
+#   If it OOMs, drop to EVAL_BATCH_SIZE=2, or omit it for batch=1 (slow but always fits).
+```
+It will: skip stages 1-4 (metrics exist) → load the stage5 ECG checkpoint →
+skip training (done) → run the ECG eval → write
+`results/Llama_3_2_1B/OpenTSLMSP/stage5_ecg_cot/results/metrics.json`.
+When that file appears, the 5-stage baseline is COMPLETE.
+
+Note: eval is greedy (deterministic). ECG sequences are ~7.7k tokens, so batched
+generation is memory-heavy — batch=4 needs a mostly-free 95GB GPU.

code/OpenTSLM/.github/workflows/publish.yml

@@ -0,0 +1,62 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+
+name: "Publish to PyPI"
+
+on:
+  workflow_dispatch:
+    inputs:
+      target:
+        description: 'Target'
+        required: true
+        default: 'PyPI'
+        type: choice
+        options:
+        - PyPI
+        - TestPyPI
+  # push:
+  #   tags:
+  #     # Publish on any tag starting with a `v`, e.g., v0.1.0
+  #     - v*
+
+env:
+  BRANCH_NAME: ${{ github.head_ref || github.ref_name }}
+
+run-name: Publish to ${{ inputs.target }}
+
+jobs:
+  run:
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+    permissions:
+      id-token: write
+      contents: read
+    steps:
+      - name: Verify branch
+        run: |
+          echo "Branch name is $BRANCH_NAME"
+          if [ "${{ inputs.target }}" == "PyPI" ] && [ "$BRANCH_NAME" != "main" ]; then
+            echo "::error title=Invalid branch::Only the main branch can be published to PyPI."
+            exit 1
+          fi
+      - name: Checkout
+        uses: actions/checkout@v5
+      - name: Install uv
+        uses: astral-sh/setup-uv@v7
+      - name: Build
+        run: uv build
+      # Check that basic features work and we didn't miss to include crucial files
+      - name: Smoke test (wheel)
+        run: uv run --isolated --no-project --with dist/*.whl test/smoke_test.py
+      - name: Smoke test (source distribution)
+        run: uv run --isolated --no-project --with dist/*.tar.gz test/smoke_test.py
+      - name: Publish
+        run: uv publish ${{ inputs.target == 'TestPyPI' && '--index testpypi' || '' }}
+      - name: Summary
+        run: |
+          echo "### Published OpenTSLM to ${{ inputs.target }} :rocket:" >> $GITHUB_STEP_SUMMARY
+          echo "Version: `$(uv version --short)`" >> $GITHUB_STEP_SUMMARY
+          echo "URL: https://${{ inputs.target == 'TestPyPI' && 'test.' || '' }}pypi.org/project/opentslm/$(uv version --short)/" >> $GITHUB_STEP_SUMMARY

code/OpenTSLM/.github/workflows/static-analysis.yml

@@ -0,0 +1,30 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+
+name: Static Analysis
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+  workflow_dispatch:
+  workflow_call:
+
+concurrency:
+  group: Static-Analysis-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  reuse_action:
+    name: REUSE Compliance Check
+    uses: StanfordBDHG/.github/.github/workflows/reuse.yml@v2
+    permissions:
+      contents: read
+  markdownlinkcheck:
+    name: Markdown Link Check
+    uses: StanfordBDHG/.github/.github/workflows/markdown-link-check.yml@v2
+    permissions:
+      contents: read

code/OpenTSLM/.gitignore

@@ -0,0 +1,22 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+
+.venv
+.vscode
+__pycache__
+.DS_STORE
+
+raw_data
+
+**/data/*
+!**/data/.gitkeep
+
+*.ts
+*.zip
+upload_to_huggingface.py
+
+dist/
+
+*.license

code/OpenTSLM/.linkspector.yml

@@ -0,0 +1,12 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+
+dirs:
+  - .
+useGitIgnore: true
+ignorePatterns:
+  - pattern: "doc:/"
+  - pattern: "http://localhost"
+  - pattern: "https://doi.org"

code/OpenTSLM/.python-version

@@ -0,0 +1,7 @@
+#
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+#
+# SPDX-License-Identifier: MIT
+#
+
+3.12

code/OpenTSLM/CITATION.cff

@@ -0,0 +1,69 @@
+#
+# This source file is part of the OpenTSLM open-source project
+#
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+#
+# SPDX-License-Identifier: MIT
+#
+
+message: "If you use this software or model, please cite it as below."
+title: "OpenTSLM: Time-Series Language Models for Reasoning over Multivariate Medical Text- and Time-Series Data"
+authors:
+  - family-names: "Langer"
+    given-names: "Patrick"
+    affiliation: "ETH Zurich"
+  - family-names: "Kaar"
+    given-names: "Thomas"
+    affiliation: "Stanford University"
+  - family-names: "Rosenblattl"
+    given-names: "Max"
+    affiliation: "Technical University of Munich"
+  - family-names: "Xu"
+    given-names: "Maxwell A."
+    affiliation: "University of Illinois Urbana-Champaign / Google Research"
+  - family-names: "Chow"
+    given-names: "Winnie"
+    affiliation: "Stanford University"
+  - family-names: "Maritsch"
+    given-names: "Martin"
+    affiliation: "Amazon"
+  - family-names: "Verma"
+    given-names: "Aradhana"
+    affiliation: "Stanford University"
+  - family-names: "Han"
+    given-names: "Brian"
+    affiliation: "Stanford University"
+  - family-names: "Kim"
+    given-names: "Daniel Seung"
+    affiliation: "University of Washington"
+  - family-names: "Chubb"
+    given-names: "Henry"
+    affiliation: "Stanford University"
+  - family-names: "Ceresnak"
+    given-names: "Scott"
+    affiliation: "Stanford University"
+  - family-names: "Zahedivash"
+    given-names: "Aydin"
+    affiliation: "Stanford University"
+  - family-names: "Sandhu"
+    given-names: "Alexander Tarlochan Singh"
+    affiliation: "Stanford University"
+  - family-names: "Rodriguez"
+    given-names: "Fatima"
+    affiliation: "Stanford University"
+  - family-names: "McDuff"
+    given-names: "Daniel"
+    affiliation: "University of Washington / Google Research"
+  - family-names: "Fleisch"
+    given-names: "Elgar"
+    affiliation: "ETH Zurich / University of St. Gallen"
+  - family-names: "Aalami"
+    given-names: "Oliver"
+    affiliation: "Stanford University"
+  - family-names: "Barata"
+    given-names: "Filipe"
+    affiliation: "ETH Zurich"
+  - family-names: "Schmiedmayer"
+    given-names: "Paul"
+    affiliation: "Stanford University"
+date-released: "2025-10-01"

code/OpenTSLM/CONTRIBUTORS.md

@@ -0,0 +1,17 @@
+<!--
+SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+
+SPDX-License-Identifier: MIT
+-->
+
+# OpenTSLM Contributors
+
+* [Patrick Langer](https://github.com/RealLast)
+* [Thomas Kaar](https://github.com/ThomasKaar)
+* [Max Rosenblattl](https://github.com/max-rosenblattl)
+* [Maxwell A. Xu](https://github.com/maxxu05)
+* [Winnie Chow](https://github.com/winnie0617)
+* [Martin Maritsch](https://github.com/masquare)
+* [Ning Wang](https://github.com/numericunderflow06)
+* [Paul Schmiedmayer](https://github.com/PSchmiedmayer)

code/OpenTSLM/LICENSE.md

@@ -0,0 +1,9 @@
+MIT License
+
+Copyright (c) 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

code/OpenTSLM/LICENSES/MIT.txt

@@ -0,0 +1,9 @@
+MIT License
+
+Copyright (c) 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

code/OpenTSLM/README.md

@@ -0,0 +1,385 @@
+<!--
+SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+
+SPDX-License-Identifier: MIT
+-->
+
+# OpenTSLM: Time-Series Language Models for Reasoning over Multivariate Medical Text- and Time-Series Data
+[![PyPI - Version](https://img.shields.io/pypi/v/opentslm)](https://pypi.org/project/opentslm)
+[![DOI](https://img.shields.io/badge/DOI-10.48550/arXiv.2510.02410-blue.svg)](https://doi.org/10.48550/arXiv.2510.02410)
+[![Static Analysis](https://github.com/StanfordBDHG/OpenTSLM/actions/workflows/static-analysis.yml/badge.svg)](https://github.com/StanfordBDHG/OpenTSLM/actions/workflows/static-analysis.yml)
+
+
+Large Language Models (LLMs) have emerged as powerful tools for interpreting multimodal data (e.g., images, audio, text), often surpassing specialized models. In medicine, they hold particular promise for synthesizing large volumes of clinical information into actionable insights and patient-facing digital health applications.  Yet, a major limitation remains their inability to handle time series data. To overcome this gap, we present OpenTSLM, a family of Time Series Language Models (TSLMs) created by integrating time series as a native modality to pretrained Large Language Models, enabling natural-language prompting and reasoning over multiple time series of any length [...] **[🔗 Read the full paper](https://doi.org/10.48550/arXiv.2510.02410)**  
+
+<p align="center">
+   <img src="https://raw.githubusercontent.com/StanfordBDHG/OpenTSLM/main/assets/schematic_overview_3.png" alt="Schematic Overview" width="100%">
+</p>
+
+
+## Examples
+
+OpenTSLM models can reason over multiple time series of any length at once, generating findings, captions, and rationales in natural language. We tested these models across a wide range of tasks spanning Human Activity Recognition (HAR) from 3-axis acceleration data, sleep staging from EEG readings, 12-lead ECG question answering, and time series captioning. Some examples are shown below, more are available in the paper.
+
+<p align="center">
+   <img src="https://raw.githubusercontent.com/StanfordBDHG/OpenTSLM/main/assets/ecg_rationale.png" alt="ECG Rationale" width="32%">
+   <img src="https://raw.githubusercontent.com/StanfordBDHG/OpenTSLM/main/assets/har_rationale.png" alt="HAR Rationale" width="32%">
+      <img src="https://raw.githubusercontent.com/StanfordBDHG/OpenTSLM/main/assets/m4_caption.png" alt="M4 Caption" width="34%">
+
+</p>
+
+## Installation
+
+```bash
+pip install opentslm
+```
+
+
+## LLM Setup
+
+OpenTSLM is designed to work with Llama and Gemma models, with Llama 3.2 1B as the default. These models are stored in Hugging Face repositories which may require access permissions. Follow these steps to gain access and download:
+
+1. **Request Access (for Llama models)**  
+    Visit the Llama model repository (e.g., https://huggingface.co/meta-llama/Llama-3.2-1B) or Gemma models repository (https://huggingface.co/google/gemma-3-270m) and request access from Meta.
+
+2. **Authenticate with Hugging Face**  
+    Log in to your Hugging Face account and configure the CLI:
+
+    ```bash
+    huggingface-cli login
+    ```
+
+3. **Create an API Token**
+    - Go to your Hugging Face settings: https://huggingface.co/settings/tokens
+    - Generate a new token with `read` scope.
+    - Copy the token for CLI login.
+
+### Supported Models
+
+OpenTSLM has been tested and works with the following models:
+
+**Llama Models:**
+
+- **meta-llama/Llama-3.2-1B** (default)
+- **meta-llama/Llama-3.2-3B**
+
+**Gemma Models:**
+
+- **google/gemma-3-270m**
+- **google/gemma-3-1b-pt**
+
+Other variants may work but have not been extensively tested.
+
+
+
+## 🤗 Quickstart with pretrained models on HuggingFace
+
+A factory class called `OpenTSLM` for easily loading pre-trained models from Hugging Face Hub. The `load_pretrained` method automatically detects the model type and returns the appropriate model instance.
+
+There are [demo scripts](demo/huggingface/) available which use the following minimal code. If you want to create your own applications, create a new file in **this repo folder** and use the following code as start:
+
+```python
+from opentslm import OpenTSLM
+from opentslm.time_series_datasets.TSQADataset import TSQADataset
+from opentslm.time_series_datasets.util import extend_time_series_to_match_patch_size_and_aggregate
+from torch.utils.data import DataLoader
+from opentslm.model_config import PATCH_SIZE
+
+REPO_ID = "OpenTSLM/llama-3.2-1b-tsqa-sp"
+
+# Use CPU or CUDA for inference. MPS does NOT work for pretrained HF checkpoints.
+model = OpenTSLM.load_pretrained(REPO_ID, device="cuda" if torch.cuda.is_available() else "cpu")
+test_dataset = TSQADataset("test", EOS_TOKEN=model.get_eos_token())
+
+test_loader = DataLoader(
+    test_dataset,
+    shuffle=False,
+    batch_size=1,
+    collate_fn=lambda batch: extend_time_series_to_match_patch_size_and_aggregate(
+         batch, patch_size=PATCH_SIZE
+    ),
+)
+
+for i, batch in enumerate(test_loader):
+    predictions = model.generate(batch, max_new_tokens=200)
+    for sample, pred in zip(batch, predictions):
+         print("Question:", sample.get("pre_prompt", "N/A"))
+         print("Answer:", sample.get("answer", "N/A"))
+         print("Output:", pred)
+    if i >= 4:
+         break
+```
+
+## Building and finetuning your own models
+
+To run the demos and use finetuning scripts **clone the repository** and set up all dependencies. We recommend using [uv](https://docs.astral.sh/uv/) to set up the environment, but you can also use pip:
+
+```bash
+git clone https://github.com/StanfordBDHG/OpenTSLM.git
+
+
+# uv environment management (recommended). Installs uv if it does not exist and creates the virtual environment
+command uv > /dev/null || curl -LsSf https://astral.sh/uv/install.sh | sh
+uv sync --all-groups
+source .venv/bin/activate
+
+
+# or alternatively install via pip:
+pip install -r requirements.txt
+
+```
+
+### HuggingFace Demo Scripts
+
+We provide ready-to-use demo scripts in the `demo/huggingface/` directory that demonstrate how to load pretrained models from HuggingFace Hub and run inference on the evaluation sets for each task:
+
+- **`01_test_hf_tsqa.py`** - Test TSQA (Time Series Question Answering) models
+- **`02_test_hf_m4.py`** - Test M4 (Time Series Captioning) models  
+- **`03_test_hf_har_cot.py`** - Test HAR CoT (Human Activity Recognition Chain-of-Thought) models
+- **`04_test_hf_sleep_cot.py`** - Test Sleep CoT (Sleep Stage Classification) models
+- **`05_test_hf_ecg_qa_cot.py`** - Test ECG QA CoT (ECG Question Answering) models
+
+Each script:
+1. Downloads the model checkpoint from HuggingFace Hub automatically (change repo id as neededs)
+2. Loads the corresponding test dataset
+3. Runs inference on the evaluation set
+4. Prints model outputs with sample information
+
+**Note:** The scripts above use the OpenTSLM-SP models except for ECG-QA, as they require less VRAM and should run on most hardware. Change the model checkpoints as needed in each file. 
+
+**Usage:**
+
+```bash
+# Run any of the demo scripts
+python demo/huggingface/01_test_hf_tsqa.py
+python demo/huggingface/02_test_hf_m4.py
+python demo/huggingface/03_test_hf_har_cot.py
+python demo/huggingface/04_test_hf_sleep_cot.py
+python demo/huggingface/05_test_hf_ecg_qa_cot.py
+```
+
+**Customizing the model:**
+
+Edit the `REPO_ID` variable at the top of each script to test different model variants. For example:
+
+```python
+# In 01_test_hf_tsqa.py
+REPO_ID = "OpenTSLM/llama-3.2-1b-tsqa-sp"  # Soft Prompt model
+# or
+REPO_ID = "OpenTSLM/llama-3.2-1b-tsqa-flamingo"  # Flamingo model
+```
+
+**Available models on HuggingFace:**
+
+All pretrained models are available under the `OpenTSLM` organization on HuggingFace Hub. Model names follow the pattern:
+- `OpenTSLM/{base_model}-{dataset}-{model_type}`
+   - `base_model`: `llama-3.2-1b`, `llama-3.2-3b`, `gemma-3-1b-pt`, `gemma-3-270m`
+   - `dataset`: `tsqa`, `m4`, `har`, `sleep`, `ecg`
+   - `model_type`: `sp` (Soft Prompt) or `flamingo` (Flamingo)
+
+Example: `OpenTSLM/llama-3.2-1b-ecg-flamingo`
+
+## Training: Multi-stage training (Curriculum)
+
+OpenTSLM uses curriculum learning with progressive training stages:
+
+### Training Stages
+
+1. **Stage 1 (MCQ)**: Multiple choice questions on time series data (TSQA dataset)
+2. **Stage 2 (Captioning)**: Generate detailed captions for time series (M4 dataset)
+3. **Stage 3 (CoT)**: Chain-of-thought reasoning on human activity recognition (HAR dataset)
+4. **Stage 4 (Sleep CoT)**: Chain-of-thought reasoning on sleep stage classification (SleepEDF dataset)
+5. **Stage 5 (ECG CoT)**: Chain-of-thought reasoning on ECG question answering (ECG QA dataset)
+
+> **⚠️ MPS/CUDA Compatibility Warning:**
+>
+> If you are using Apple's MPS (Metal Performance Shaders) backend (e.g., on Mac with Apple Silicon), you may encounter issues with training or inference. **Checkpoints trained with CUDA (NVIDIA GPUs) may not yield good results or may not be fully compatible when loaded and run on MPS.** For best results, use the same device type (CUDA or MPS) for both training and inference. CUDA is preferred in general.
+
+### Quick Start
+
+```bash
+# Run full curriculum with OpenTSLMFlamingo
+python curriculum_learning.py --model OpenTSLMSP
+
+# Run full curriculum with OpenTSLMSP
+python curriculum_learning.py --model OpenTSLMFlamingo
+
+# Run specific stages
+python curriculum_learning.py --model OpenTSLMFlamingo --stages stage1_mcq
+python curriculum_learning.py --model OpenTSLMFlamingo --stages stage2_captioning
+python curriculum_learning.py --model OpenTSLMFlamingo --stages stage3_cot
+python curriculum_learning.py --model OpenTSLMFlamingo --stages stage4_sleep_cot
+python curriculum_learning.py --model OpenTSLMFlamingo --stages stage5_ecg_cot
+
+# Run multiple stages
+python curriculum_learning.py --model OpenTSLMFlamingo --stages stage1_mcq stage2_captioning stage3_cot
+
+# Specify device
+python curriculum_learning.py --model OpenTSLMFlamingo --device cuda
+
+# Use different models
+python curriculum_learning.py --model OpenTSLMFlamingo --llm_id meta-llama/Llama-3.2-1B
+python curriculum_learning.py --model OpenTSLMFlamingo --llm_id google/gemma-3-270m
+
+# Run only evaluation
+python curriculum_learning.py --model OpenTSLMFlamingo --eval_only
+```
+
+### Command Line Arguments
+
+- `--model`: Model type (`OpenTSLMSP` or `OpenTSLMFlamingo`)
+- `--stages`: Stages to run (any combination of: `stage1_mcq`, `stage2_captioning`, `stage3_cot`, `stage4_sleep_cot`, `stage5_ecg_cot`)
+- `--device`: Device to use (`cuda`, `mps`, `cpu`)
+- `--eval_only`: Run evaluation only (requires an existing checkpoint for the stage)
+- `--llm_id`: Model ID (default: `meta-llama/Llama-3.2-1B`, supports Llama and Gemma models)
+- `--batch_size`: Batch size for training
+- `--gradient_checkpointing`: Enable gradient checkpointing for memory efficiency
+- `--verbose`: Enable verbose logging
+
+### Helper Scripts
+
+Helper scripts for analysis, testing, and batch processing are available in the `scripts/` directory:
+
+**Shell Scripts:**
+- **`run_all_memory.sh`** - Run comprehensive memory usage analysis across all stages
+- **`run_all_memory_missing.sh`** - Run memory analysis for missing stages only
+
+**Python Scripts:**
+- **`create_doctor_eval_dataset.py`** - Create evaluation dataset for doctor assessments
+- **`get_memory_use.py`** - Analyze and report memory usage across stages
+- **`plot_memory_usage.py`** - Visualize memory usage patterns
+- **`plot_memory_simulation.py`** - Simulate and plot memory requirements
+- **`plot_memory_simulation_per_length.py`** - Analyze memory usage by sequence length
+- **`hf_test.py`** - Test HuggingFace model loading and inference
+
+These scripts can be customized by editing the parameters directly or by passing command-line arguments.
+
+### Repository Naming Convention
+
+- Repository IDs ending with `-sp` will load and return `OpenTSLMSP` models
+- Repository IDs ending with `-flamingo` will load and return `OpenTSLMFlamingo` models
+
+
+## 📁 Results Structure
+
+During training, the script creates a structured results directory:
+
+```
+results/
+├── {llm_id}/
+│   ├── OpenTSLMSP/
+│   │   ├── stage1_mcq/
+│   │   │   ├── checkpoints/
+│   │   │   │   ├── best_model.pt
+│   │   │   │   └── loss_history.txt
+│   │   │   └── results/
+│   │   │       ├── test_predictions.jsonl
+│   │   │       └── metrics.json
+│   │   ├── stage2_captioning/
+│   │   │   ├── checkpoints/
+│   │   │   │   ├── best_model.pt
+│   │   │   │   └── loss_history.txt
+│   │   │   └── results/
+│   │   │       ├── test_predictions.jsonl
+│   │   │       └── metrics.json
+│   │   ├── stage3_cot/
+│   │   │   ├── checkpoints/
+│   │   │   │   ├── best_model.pt
+│   │   │   │   └── loss_history.txt
+│   │   │   └── results/
+│   │   │       ├── test_predictions.jsonl
+│   │   │       └── metrics.json
+│   │   ├── stage4_sleep_cot/
+│   │   │   ├── checkpoints/
+│   │   │   │   ├── best_model.pt
+│   │   │   │   └── loss_history.txt
+│   │   │   └── results/
+│   │   │       ├── test_predictions.jsonl
+│   │   │       └── metrics.json
+│   │   ├── stage5_ecg_cot/
+│   │   │   ├── checkpoints/
+│   │   │   │   ├── best_model.pt
+│   │   │   │   └── loss_history.txt
+│   │   │   └── results/
+│   │   │       ├── test_predictions.jsonl
+│   │   │       └── metrics.json
+│   │   └── curriculum_results.json
+│   └── OpenTSLMFlamingo/
+│       ├── stage1_mcq/
+│       ├── stage2_captioning/
+│       ├── stage3_cot/
+│       ├── stage4_sleep_cot/
+│       ├── stage5_ecg_cot/
+│       └── curriculum_results.json
+```
+
+Each stage automatically loads the best model from the previous stage, ensuring proper curriculum progression. Results are organized by model ID (sanitized), then by model type and stage. The `{llm_id}` directory name is derived from the `--llm_id` parameter (e.g., `meta-llama/Llama-3.2-1B` becomes `Llama3_2_1B`, `google/gemma-3-1b-pt` becomes `gemma_3_1b_pt`).
+
+
+## Authors
+
+This work was made possible through the collaborative efforts of an interdisciplinary team of researchers from computer science, medicine, and engineering. Thank you to all of the Co-authors of the TSLM publication:
+- Patrick Langer (Stanford University, ETH Zurich, ETH Agentic Systems Lab)
+- Thomas Kaar (Stanford University, TUM, ETH Agentic Systems Lab)
+- Max Rosenblattl (Stanford University, TUM, ETH Agentic Systems Lab)
+- Maxwell A. Xu (Google Research, University of Illinois Urbana-Champaign)
+- Winnie Chow (Stanford University)
+- Martin Maritsch (Amazon)
+- Robert Jakob (ETH Zurich, ETH Agentic Systems Lab)
+- Ning Wang (ETH Zurich, ETH Agentic Systems Lab)
+- Aradhana Verma (Stanford University)
+- Brian Han (Stanford University)
+- Daniel Seung Kim (University of Washington)
+- Henry Chubb (Stanford University)
+- Scott Ceresnak (Stanford University)
+- Aydin Zahedivash (Stanford University)
+- Alexander Tarlochan Singh Sandhu (Stanford University)
+- Fatima Rodriguez (Stanford University)
+- Daniel McDuff (Google Research, University of Washington)
+- Elgar Fleisch (ETH Zurich, University of St. Gallen)
+- Oliver Aalami (Stanford University)
+- Filipe Barata (ETH Zurich)
+- Paul Schmiedmayer (Stanford University)
+
+## Contributing
+
+Contributions to this project are welcome. Please make sure to read the [contribution guidelines](https://github.com/StanfordBDHG/.github/blob/main/CONTRIBUTING.md) and the [contributor covenant code of conduct](https://github.com/StanfordBDHG/.github/blob/main/CODE_OF_CONDUCT.md) first.
+
+You can find a list of all current contributors at [CONTRIBUTORS.md](CONTRIBUTORS.md).
+
+
+## Research Opportunities
+
+Are you a student interested in advancing the frontiers of time-series language models and digital health research? We welcome students to get involved in our research projects!
+
+Visit our [Student Research Opportunities page](http://bdh.stanford.edu/studentresearch) to learn more about current projects and how you can contribute to cutting-edge research at the intersection of AI and healthcare.
+
+For researchers and project partners interested in collaboration opportunities, please reach out to us at **digitalhealthresearch@stanford.edu**.
+
+
+## License
+
+This project is licensed under the MIT License.
+
+OpenTSLM uses [REUSE specification](https://reuse.software/spec/) to ensure consistent and machine-readable licensing across the repository.
+
+To add or update license headers, run:
+
+```bash
+reuse annotate --recursive \
+  --copyright "Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)" \
+  --copyright "This source file is part of the OpenTSLM open-source project." \
+  --license MIT \
+  --skip-unrecognized \
+  .
+```
+
+
+
+<div align="left">
+   <img src="https://raw.githubusercontent.com/StanfordBDHG/OpenTSLM/main/assets/stanford_biodesign_logo.png" alt="Stanford Biodesign" height="90">&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;
+   <img src="https://raw.githubusercontent.com/StanfordBDHG/OpenTSLM/main/assets/CDHI_white.svg" alt="ETH Centre for Digital Health Interventions" height="90">
+      <img src="https://raw.githubusercontent.com/StanfordBDHG/OpenTSLM/main/assets/ASLwhite.svg" alt="ETH Agentic Systems Lab" height="90">
+
+</div>

code/OpenTSLM/REUSE.toml

@@ -0,0 +1,6 @@
+version = 1
+
+[[annotations]]
+path = ["assets/**", "data/**", "**/*.png", "*.svg", "*.png", "**/*.pt", "**/*.jsonl", "**/*.json", ".gitignore", "**/uv.lock", "LICENSE.md", "**/requirements.txt"]
+SPDX-FileCopyrightText = "2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)"
+SPDX-License-Identifier = "MIT"

code/OpenTSLM/curriculum_learning.py

@@ -0,0 +1,1732 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+
+
+import os
+
+import json
+import os as _os
+import argparse
+from typing import List, Optional, Dict, Any, Callable
+from opentslm.time_series_datasets.TSQADataset import TSQADataset
+from opentslm.time_series_datasets.m4.M4QADataset import M4QADataset
+from opentslm.time_series_datasets.sleep.SleepEDFCoTQADataset import SleepEDFCoTQADataset
+from opentslm.time_series_datasets.har_cot.HARCoTQADataset import HARCoTQADataset
+from opentslm.time_series_datasets.ecg_qa.ECGQACoTQADataset import ECGQACoTQADataset
+from opentslm.time_series_datasets.util import (
+    extend_time_series_to_match_patch_size_and_aggregate,
+)
+import torch
+import torch.distributed as dist
+from torch.optim import AdamW
+from torch.nn.utils import clip_grad_norm_
+from torch.utils.data import ConcatDataset, DataLoader, Dataset
+from torch.utils.data.distributed import DistributedSampler
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.fsdp import (
+    CPUOffload,
+    MixedPrecision,
+    ShardingStrategy,
+    BackwardPrefetch,
+    FullStateDictConfig,
+    StateDictType,
+)
+from tqdm.auto import tqdm
+from transformers import get_linear_schedule_with_warmup
+
+from opentslm.model.encoder.TransformerCNNEncoder import TransformerCNNEncoder
+from opentslm.model.llm.OpenTSLMFlamingo import OpenTSLMFlamingo
+from opentslm.model.llm.OpenTSLMSP import OpenTSLMSP
+from opentslm.model.projector.MLPProjector import MLPProjector
+import datetime
+from opentslm.logger import get_logger, set_global_verbose
+
+from opentslm.model_config import (
+    BATCH_SIZE,
+    EARLY_STOP_PAT,
+    GRAD_CLIP_NORM,
+    LR_ENCODER,
+    LR_PROJECTOR,
+    NUM_EPOCHS,
+    PATCH_SIZE,
+    WARMUP_FRAC,
+    WEIGHT_DECAY,
+)
+
+
+# Global stage configuration - users can modify this to mix and match stages
+CURRICULUM_STAGES = [
+    "stage1_mcq",
+    "stage2_captioning",
+    "stage3_cot",
+    "stage4_sleep_cot",
+    "stage5_ecg_cot",
+]
+
+
+class CurriculumTrainer:
+    """
+    Curriculum learning trainer for OpenTSLM models.
+    Trains models stage by stage with shared training logic.
+    While this may look like a lot of code, it's actually quite modular.
+    We simply train either OpenTSLMSP or OpenTSLMFlamingo, both using the same training loop.
+    We train across different stages:
+    - stage1_mcq: Trains the model on a time-series MCQ dataset (TSQA)
+    - stage2_captioning: Trains the model on a time-series captioning dataset (M4 time series captioning)
+    - stage3_cot: Trains the model on a chain-of-thought reasoning dataset (HAR CoT)
+    - stage4_sleep_cot: Trains the model on sleep stage classification with chain-of-thought reasoning
+    - stage5_ecg_cot: Trains the model on ECG QA with chain-of-thought reasoning
+
+    Features:
+    - Automatic loss history tracking saved to loss_history.txt in each stage's checkpoints directory
+    - Loss history is appended to when resuming training, preserving all previous epochs
+    - Displays previous loss history when resuming training
+
+    If you run this script, you should be able to reproduce our results from the paper.
+    All datasets are automatically downloaded and processed.
+    """
+
+    def _sanitize_llm_id(self, llm_id: str) -> str:
+        """Sanitize llm_id for use in directory names (e.g., meta-llama/Llama-3.2-1B -> Llama3_2_1B)"""
+        if not llm_id:
+            return "unknown_llm"
+        # Take last part after /, replace . and - with _
+        name = llm_id.split("/")[-1]
+        name = name.replace(".", "_").replace("-", "_")
+        # Optionally, remove duplicate underscores
+        while "__" in name:
+            name = name.replace("__", "_")
+        return name
+
+    def __init__(
+        self,
+        model_type: str,
+        device: str = None,
+        gradient_checkpointing: bool = False,
+        dist_url: str = "env://",
+        dist_backend: str = "nccl",
+        local_rank: int = int(os.environ.get("LOCAL_RANK", 0)),
+        llm_id: str = None,
+    ):
+        """
+        Initialize the curriculum trainer.
+
+        Args:
+            model_type: Either 'OpenTSLMSP' or 'OpenTSLMFlamingo'
+            device: Device to use for training ('cuda', 'mps', or 'cpu')
+            gradient_checkpointing: Enable gradient checkpointing
+            dist_url: URL used to set up distributed training
+            dist_backend: Distributed backend
+            local_rank: Local GPU rank
+            llm_id: LLM model ID (e.g., 'google/medgemma-2b', 'meta-llama/Llama-3.2-1B')
+        """
+        self.model_type = model_type
+        self.device = device or self._get_device()
+        if self.device == "mps":
+            print(
+                "🚨 Warning: Using MPS, might not be fully compatible with the model. Use CUDA for best results."
+            )
+        self.llm_id = llm_id
+        self.llm_id_safe = self._sanitize_llm_id(llm_id)
+
+        # Distributed training parameters
+        self.gradient_checkpointing = gradient_checkpointing
+        self.dist_url = dist_url
+        self.dist_backend = dist_backend
+        self.local_rank = local_rank
+
+        # Initialize distributed training if needed
+        self.rank = 0
+        self.world_size = 1
+        if self._should_use_distributed():
+            self._init_distributed()
+
+        self.model = self._initialize_model()
+        self.results_dir = os.path.join("results", self.llm_id_safe, self.model_type)
+        self._create_results_dir()
+
+    def _get_device(self) -> str:
+        """Get the best available device."""
+        if torch.cuda.is_available():
+            return "cuda"
+        elif torch.backends.mps.is_available():
+            return "mps"
+        else:
+            return "cpu"
+
+    def _initialize_model(self):
+        """Initialize the specified model type."""
+        if self.model_type == "OpenTSLMSP":
+            model = OpenTSLMSP(llm_id=self.llm_id, device=self.device).to(self.device)
+
+        elif self.model_type == "OpenTSLMFlamingo":
+            model = OpenTSLMFlamingo(
+                cross_attn_every_n_layers=1,
+                gradient_checkpointing=self.gradient_checkpointing,
+                llm_id=self.llm_id,
+                device=self.device,
+            ).to(self.device)
+        else:
+            raise ValueError(f"Unknown model type: {self.model_type}")
+
+        # Use DDP for multi-GPU training (simpler and than FSDP)
+        if self.world_size > 1:
+            model = DDP(
+                model,
+                device_ids=[self.local_rank] if torch.cuda.is_available() else None,
+            )
+            if self.rank == 0:
+                print(f"Wrapped {self.model_type} with DDP for distributed training")
+
+        return model
+
+    def _get_cast_dtype(self, precision: str):
+        """Get cast dtype for mixed precision."""
+        if precision == "bf16":
+            return torch.bfloat16
+        elif precision == "fp16":
+            return torch.float16
+        else:
+            return None
+
+    def _create_results_dir(self):
+        """Create the results directory structure."""
+        os.makedirs(self.results_dir, exist_ok=True)
+        # model_dir now includes llm_id_safe
+        model_dir = self.results_dir
+        os.makedirs(model_dir, exist_ok=True)
+
+        # Create stage directories based on global configuration
+        for stage in CURRICULUM_STAGES:
+            stage_dir = os.path.join(model_dir, stage)
+            os.makedirs(stage_dir, exist_ok=True)
+            os.makedirs(os.path.join(stage_dir, "checkpoints"), exist_ok=True)
+            os.makedirs(os.path.join(stage_dir, "results"), exist_ok=True)
+
+    def _get_optimizer(
+        self,
+        batch_size: int = None,
+        lr_encoder: float = None,
+        lr_projector: float = None,
+        lr_base: float = None,
+    ):
+        """Get optimizer for the model with configurable learning rates."""
+        # Get the underlying model (handles DDP wrapping)
+        model = self._get_model()
+
+        if self.model_type == "OpenTSLMSP":
+            # Parameter groups with different learning rates for SP
+            enc_params = list(model.encoder.parameters())
+            proj_params = list(model.projector.projector.parameters())
+
+            # Use provided learning rates or defaults
+            encoder_lr = lr_encoder if lr_encoder is not None else LR_ENCODER
+            projector_lr = lr_projector if lr_projector is not None else LR_PROJECTOR
+
+            param_groups = [
+                {"params": enc_params, "lr": encoder_lr, "weight_decay": WEIGHT_DECAY},
+                {
+                    "params": proj_params,
+                    "lr": projector_lr,
+                    "weight_decay": WEIGHT_DECAY,
+                },
+            ]
+
+            # Add LoRA parameters if enabled
+            if hasattr(model, "lora_enabled") and model.lora_enabled:
+                lora_params = model.get_lora_parameters()
+                if lora_params:
+                    # Use projector LR for LoRA parameters (similar fine-tuning nature)
+                    param_groups.append(
+                        {
+                            "params": lora_params,
+                            "lr": projector_lr,
+                            "weight_decay": WEIGHT_DECAY,
+                        }
+                    )
+                    if self.rank == 0:
+                        print(f"📊 Learning rates for {self.model_type} (with LoRA):")
+                        print(f"   Encoder LR: {encoder_lr:.2e}")
+                        print(f"   Projector LR: {projector_lr:.2e}")
+                        print(
+                            f"   LoRA LR: {projector_lr:.2e} ({len(lora_params)} parameters)"
+                        )
+                else:
+                    raise RuntimeError(
+                        "LoRA is enabled but no trainable LoRA parameters found. This indicates a LoRA configuration issue."
+                    )
+            else:
+                if self.rank == 0:
+                    print(f"📊 Learning rates for {self.model_type}:")
+                    print(f"   Encoder LR: {encoder_lr:.2e}")
+                    print(f"   Projector LR: {projector_lr:.2e}")
+
+            return AdamW(param_groups)
+        else:
+            # For Flamingo, use grouped parameters
+            params_to_optimize = model.named_parameters()
+            params_to_optimize = list(
+                filter(
+                    lambda x: x[1].requires_grad
+                    and not getattr(x[1], "exclude_from_optimizer", False),
+                    params_to_optimize,
+                )
+            )
+
+            # Group parameters for weight decay
+            params_with_wd, params_without_wd = [], []
+            for n, p in params_to_optimize:
+                if "gated_cross_attn" in n:
+                    params_with_wd.append(p)
+                else:
+                    params_without_wd.append(p)
+
+            # Use provided base learning rate or default
+            base_lr = lr_base if lr_base is not None else 2e-4
+
+            if self.rank == 0:
+                print(f"📊 Learning rate for {self.model_type}:")
+                print(f"   Base LR: {base_lr:.2e}")
+
+            return torch.optim.AdamW(
+                [
+                    {"params": params_with_wd, "weight_decay": 0.1},
+                    {"params": params_without_wd, "weight_decay": 0.0},
+                ],
+                lr=base_lr,
+            )
+
+    def _merge_data_loaders(
+        self,
+        datasets: List[Dataset],
+        shuffle: bool,
+        batch_size: int,
+        patch_size: int,
+        distribute_data: bool = False,
+    ) -> DataLoader:
+        """Create a merged data loader from multiple datasets."""
+        merged_ds = ConcatDataset(datasets)
+
+        # Use distributed sampler if distributed training is enabled
+        if distribute_data and dist.is_initialized():
+            sampler = DistributedSampler(
+                merged_ds, num_replicas=self.world_size, rank=self.rank, shuffle=shuffle
+            )
+            return DataLoader(
+                merged_ds,
+                sampler=sampler,
+                batch_size=batch_size,
+                collate_fn=lambda batch: extend_time_series_to_match_patch_size_and_aggregate(
+                    batch, patch_size=patch_size
+                ),
+            )
+        else:
+            return DataLoader(
+                merged_ds,
+                shuffle=shuffle,
+                batch_size=batch_size,
+                collate_fn=lambda batch: extend_time_series_to_match_patch_size_and_aggregate(
+                    batch, patch_size=patch_size
+                ),
+            )
+
+    def _save_checkpoint(
+        self, stage: str, epoch: int, val_loss: float, optimizer, scheduler
+    ):
+        """Save model checkpoint for a specific stage."""
+        checkpoint_dir = os.path.join(self.results_dir, stage, "checkpoints")
+
+        # Only save on rank 0 for distributed training
+        if dist.is_initialized() and self.rank != 0:
+            return
+
+        # Get the underlying model (handles DDP wrapping)
+        model = self._get_model()
+
+        if self.model_type == "OpenTSLMSP":
+            checkpoint = {
+                "encoder_state": model.encoder.state_dict(),
+                "projector_state": model.projector.state_dict(),
+                "optimizer_state": optimizer.state_dict(),
+                "scheduler_state": scheduler.state_dict(),
+                "val_loss": val_loss,
+                "epoch": epoch,
+            }
+
+            # Add LoRA state to checkpoint
+            model.save_lora_state_to_checkpoint(checkpoint)
+        else:
+            # Handle DDP or single GPU case for OpenTSLMFlamingo
+            model_state = model.state_dict()
+            if hasattr(self.model, "module"):
+                # Remove 'module.' prefix for DDP
+                model_state = {
+                    k.replace("module.", ""): v for k, v in model_state.items()
+                }
+            checkpoint = {
+                "model_state": model_state,
+                "optimizer_state": optimizer.state_dict(),
+                "scheduler_state": scheduler.state_dict(),
+                "val_loss": val_loss,
+                "epoch": epoch,
+            }
+
+        checkpoint_path = os.path.join(checkpoint_dir, "best_model.pt")
+
+        # Check disk space before saving
+        if self.rank == 0:
+            import shutil
+
+            total, used, free = shutil.disk_usage(checkpoint_dir)
+            free_gb = free / (1024**3)
+            print(f"💾 Disk space: {free_gb:.2f} GB free in {checkpoint_dir}")
+
+            # Estimate checkpoint size (rough estimate)
+            estimated_size_gb = sum(
+                p.numel() * p.element_size() for p in self._get_model().parameters()
+            ) / (1024**3)
+            if (
+                free_gb < estimated_size_gb * 2
+            ):  # Need at least 2x the size for safe writing
+                print(
+                    f"⚠️  Warning: Low disk space. Need ~{estimated_size_gb:.2f} GB, have {free_gb:.2f} GB free"
+                )
+
+        # Try to save with error handling
+        try:
+            torch.save(checkpoint, checkpoint_path)
+        except Exception as e:
+            if self.rank == 0:
+                print(f"❌ Failed to save checkpoint: {e}")
+                print(f"   Checkpoint path: {checkpoint_path}")
+                print(
+                    f"   Checkpoint size: {sum(p.numel() * p.element_size() for p in self._get_model().parameters()) / 1024**3:.2f} GB"
+                )
+
+                raise RuntimeError(f"Failed to save checkpoint: {e}")
+
+    def _save_loss_history(
+        self, stage: str, epoch: int, train_loss: float, val_loss: float
+    ):
+        """Save loss history to a file for tracking training progress."""
+        if dist.is_initialized() and self.rank != 0:
+            return  # Only save on rank 0 for distributed training
+
+        checkpoint_dir = os.path.join(self.results_dir, stage, "checkpoints")
+        loss_history_file = os.path.join(checkpoint_dir, "loss_history.txt")
+
+        # Ensure the directory exists
+        os.makedirs(checkpoint_dir, exist_ok=True)
+
+        # Create the file with header if it doesn't exist
+        if not os.path.exists(loss_history_file):
+            with open(loss_history_file, "w") as f:
+                f.write("Epoch\tTrain_Loss\tVal_Loss\n")
+                f.write("-" * 30 + "\n")
+
+        # Append the current epoch's losses
+        with open(loss_history_file, "a") as f:
+            f.write(f"{epoch}\t{train_loss:.6f}\t{val_loss:.6f}\n")
+
+    def _display_loss_history(self, stage: str):
+        """Display the loss history for a stage if available."""
+        if dist.is_initialized() and self.rank != 0:
+            return  # Only display on rank 0 for distributed training
+
+        checkpoint_dir = os.path.join(self.results_dir, stage, "checkpoints")
+        loss_history_file = os.path.join(checkpoint_dir, "loss_history.txt")
+
+        if os.path.exists(loss_history_file):
+            try:
+                with open(loss_history_file, "r") as f:
+                    lines = f.readlines()
+
+                if len(lines) > 2:  # More than just header
+                    print(f"📊 Previous loss history for {stage}:")
+                    print("   Epoch\tTrain_Loss\tVal_Loss")
+                    print("   " + "-" * 30)
+
+                    # Show last 5 epochs (or all if less than 5)
+                    start_idx = max(2, len(lines) - 5)  # Skip header lines
+                    for line in lines[start_idx:]:
+                        if line.strip() and not line.startswith("-"):
+                            parts = line.strip().split("\t")
+                            if len(parts) == 3:
+                                epoch, train_loss, val_loss = parts
+                                print(f"   {epoch}\t{train_loss}\t{val_loss}")
+
+                    if len(lines) > 7:  # More than 5 epochs
+                        print(f"   ... and {len(lines) - 7} more epochs")
+                    print()
+            except Exception as e:
+                print(f"⚠️  Could not read loss history: {e}")
+
+    def _load_checkpoint(
+        self, stage: str, optimizer, scheduler, eval_only: bool = False
+    ):
+        """Load model checkpoint for a specific stage."""
+        checkpoint_path = os.path.join(
+            self.results_dir, stage, "checkpoints", "best_model.pt"
+        )
+
+        if os.path.exists(checkpoint_path):
+            # Always load checkpoint to CPU first to avoid GPU OOM spikes
+            checkpoint = torch.load(
+                checkpoint_path, map_location="cpu", weights_only=False
+            )
+
+            # Get the underlying model (handles DDP wrapping)
+            model = self._get_model()
+
+            if self.model_type == "OpenTSLMSP":
+                model.encoder.load_state_dict(checkpoint["encoder_state"])
+                model.projector.load_state_dict(checkpoint["projector_state"])
+
+                # Load LoRA state using the OpenTSLMSP method (allow missing for backward compatibility)
+                try:
+                    model.load_lora_state_from_checkpoint(
+                        checkpoint, allow_missing=True
+                    )
+                except RuntimeError as e:
+                    if self.rank == 0:
+                        print(f"❌ Failed to load LoRA state from checkpoint: {e}")
+                    raise
+
+                # Only load optimizer state when training
+                if (
+                    not eval_only
+                    and optimizer is not None
+                    and "optimizer_state" in checkpoint
+                ):
+                    optimizer.load_state_dict(checkpoint["optimizer_state"])
+            else:
+                # Handle DDP or single GPU case for OpenTSLMFlamingo
+                model_state = checkpoint["model_state"]
+                if hasattr(self.model, "module"):
+                    # Add 'module.' prefix for DDP
+                    model_state = {f"module.{k}": v for k, v in model_state.items()}
+
+                # Load state dict with strict=False to handle missing keys
+                try:
+                    missing_keys, unexpected_keys = self.model.load_state_dict(
+                        model_state, strict=False
+                    )
+                    if missing_keys and self.rank == 0:
+                        print(
+                            f"⚠️  Warning: Missing keys when loading checkpoint for {stage}:"
+                        )
+                        for key in missing_keys[:10]:  # Show first 10 missing keys
+                            print(f"   - {key}")
+                        if len(missing_keys) > 10:
+                            print(f"   ... and {len(missing_keys) - 10} more keys")
+                    if unexpected_keys and self.rank == 0:
+                        print(
+                            f"⚠️  Warning: Unexpected keys when loading checkpoint for {stage}:"
+                        )
+                        for key in unexpected_keys[
+                            :10
+                        ]:  # Show first 10 unexpected keys
+                            print(f"   - {key}")
+                        if len(unexpected_keys) > 10:
+                            print(f"   ... and {len(unexpected_keys) - 10} more keys")
+                except Exception as e:
+                    raise RuntimeError(
+                        f"Failed to load model state from checkpoint for {stage}: {e}"
+                    )
+
+                # Only load optimizer state when training
+                if (
+                    not eval_only
+                    and optimizer is not None
+                    and "optimizer_state" in checkpoint
+                ):
+                    optimizer.load_state_dict(checkpoint["optimizer_state"])
+
+            # Only load scheduler state when training
+            if (
+                not eval_only
+                and scheduler is not None
+                and "scheduler_state" in checkpoint
+            ):
+                scheduler.load_state_dict(checkpoint["scheduler_state"])
+
+            return checkpoint.get("epoch", "?"), checkpoint.get(
+                "val_loss", float("inf")
+            )
+        return None, float("inf")
+
+    def _load_previous_stage_model(
+        self, current_stage: str
+    ) -> Optional[Dict[str, Any]]:
+        """Load the best model from the previous stage and return its metrics."""
+        try:
+            current_idx = CURRICULUM_STAGES.index(current_stage)
+            if current_idx == 0:
+                # First stage, no previous model to load
+                return None
+            previous_stage = CURRICULUM_STAGES[current_idx - 1]
+            metrics_file = os.path.join(
+                self.results_dir, previous_stage, "results", "metrics.json"
+            )
+            if not os.path.exists(metrics_file):
+                # PATCH: If running stage2_captioning and previous stage metrics are missing, skip loading
+                if current_stage == "stage2_captioning":
+                    if self.rank == 0:
+                        print(
+                            f"⚠️  Skipping previous stage {previous_stage} because metrics file not found: {metrics_file}"
+                        )
+                    return None
+                raise RuntimeError(
+                    f"Previous stage {previous_stage} metrics file not found: {metrics_file}"
+                )
+            # Be robust to malformed JSON (e.g., concurrent writes or concatenated JSON)
+            try:
+                with open(metrics_file, "r") as f:
+                    metrics = json.load(f)
+            except Exception as e:
+                if self.rank == 0:
+                    print(
+                        f"⚠️  Warning: Could not parse metrics file for {previous_stage} ({metrics_file}): {e}"
+                    )
+                    print("   Proceeding without previous metrics.")
+                metrics = {}
+            # Load the model weights from previous stage
+            checkpoint_path = os.path.join(
+                self.results_dir, previous_stage, "checkpoints", "best_model.pt"
+            )
+            if not os.path.exists(checkpoint_path):
+                # PATCH: If running stage2_captioning and previous stage checkpoint is missing, skip loading
+                if current_stage == "stage2_captioning":
+                    if self.rank == 0:
+                        print(
+                            f"⚠️  Skipping previous stage {previous_stage} because checkpoint not found: {checkpoint_path}"
+                        )
+                    return None
+                raise RuntimeError(
+                    f"Previous stage {previous_stage} checkpoint not found: {checkpoint_path}"
+                )
+            print(
+                "Loading checkpoint from previous stage: ",
+                checkpoint_path,
+                "and model type: ",
+                self.model_type,
+                "and llm_id: ",
+                self.llm_id,
+            )
+            print("This might take a while")
+            checkpoint = torch.load(
+                checkpoint_path, map_location="cpu", weights_only=False
+            )
+            # Get the underlying model (handles DDP wrapping)
+            model = self._get_model()
+            if self.model_type == "OpenTSLMSP":
+                model.encoder.load_state_dict(checkpoint["encoder_state"])
+                model.projector.load_state_dict(checkpoint["projector_state"])
+
+                # Load LoRA state from previous stage (allow missing for stage transitions)
+                try:
+                    loaded_count = model.load_lora_state_from_checkpoint(
+                        checkpoint, allow_missing=True
+                    )
+                    if loaded_count > 0 and self.rank == 0:
+                        print(
+                            f"📥 Loaded LoRA adapters from previous stage: {loaded_count} parameters"
+                        )
+                except RuntimeError as e:
+                    if self.rank == 0:
+                        print(f"❌ Failed to load LoRA state from previous stage: {e}")
+                    # For previous stage loading, we can be more tolerant of LoRA mismatches
+                    # as stages might have different LoRA configurations
+            else:
+                # Handle OpenTSLMFlamingo with graceful loading
+                model_state = checkpoint["model_state"]
+                if hasattr(self.model, "module"):
+                    # Add 'module.' prefix for DDP
+                    model_state = {f"module.{k}": v for k, v in model_state.items()}
+                # Load state dict with strict=False to handle missing keys
+                try:
+                    missing_keys, unexpected_keys = self.model.load_state_dict(
+                        model_state, strict=False
+                    )
+                    if missing_keys and self.rank == 0:
+                        print(
+                            f"⚠️  Warning: Missing keys when loading previous stage {previous_stage}:"
+                        )
+                        for key in missing_keys[:5]:  # Show first 5 missing keys
+                            print(f"   - {key}")
+                        if len(missing_keys) > 5:
+                            print(f"   ... and {len(missing_keys) - 5} more keys")
+                        print(
+                            f"   This is normal when transitioning between stages with different model configurations."
+                        )
+                    if unexpected_keys and self.rank == 0:
+                        print(
+                            f"⚠️  Warning: Unexpected keys when loading previous stage {previous_stage}:"
+                        )
+                        for key in unexpected_keys[:5]:  # Show first 5 unexpected keys
+                            print(f"   - {key}")
+                        if len(unexpected_keys) > 5:
+                            print(f"   ... and {len(unexpected_keys) - 5} more keys")
+                except Exception as e:
+                    raise RuntimeError(
+                        f"Failed to load model state from previous stage {previous_stage}: {e}"
+                    )
+            return {
+                "stage": previous_stage,
+                "metrics": metrics,
+                "epoch": checkpoint.get("epoch", "?"),
+                "val_loss": checkpoint.get("val_loss", "?"),
+            }
+        except Exception as e:
+            raise RuntimeError(f"Failed to load previous stage model: {e}")
+
+    def _calculate_accuracy(
+        self, predictions: List[str], gold_answers: List[str]
+    ) -> float:
+        """Calculate accuracy for MCQ tasks."""
+        correct = 0
+        total = len(predictions)
+
+        for pred, gold in zip(predictions, gold_answers):
+            # Clean up predictions and gold answers
+            pred_clean = pred.strip()
+            gold_clean = gold.strip()
+
+            # Check if gold starts with the cleaned prediction (more robust matching)
+            if gold_clean.startswith(pred_clean) or pred_clean == gold_clean:
+                correct += 1
+
+        return correct / total if total > 0 else 0.0
+
+    def _evaluate_stage(
+        self,
+        stage: str,
+        test_loader: DataLoader,
+        stage_name: str,
+        metric_func: Callable = None,
+        epoch: int = None,
+    ) -> Dict[str, Any]:
+        """Evaluate model on test set for a specific stage."""
+        # Enable eval mode for all ranks
+        self.model.eval()
+        results = []
+        test_loss = 0.0
+
+        # Generation budget: MCQ answers are a few tokens and captions are short
+        # (M4 reference ~180 tokens, max ~220), so cap these stages at 256. This
+        # stops verbose backbones (e.g. Gemma) from generating to the full budget
+        # without truncating any real output. CoT stages keep 2000 for long rationales.
+        if "captioning" in stage_name or "mcq" in stage_name:
+            max_new_tokens = 256
+        else:
+            max_new_tokens = 2000
+
+        # Prepare per-rank streaming writer for test predictions
+        results_file_rank = os.path.join(
+            self.results_dir,
+            stage_name,
+            "results",
+            f"test_predictions_rank_{self.rank if dist.is_initialized() else 0}.jsonl",
+        )
+        final_results_file = os.path.join(
+            self.results_dir, stage_name, "results", "test_predictions.jsonl"
+        )
+        results_fp = None
+        # Ensure directory exists (defensive)
+        os.makedirs(os.path.dirname(results_file_rank), exist_ok=True)
+        if self.rank == 0:
+            print(f"[Eval] rank={self.rank}, world_size={self.world_size}")
+            print(f"Saving per-rank test predictions to: {results_file_rank}")
+            if dist.is_initialized():
+                print(
+                    f"Final merged predictions will be saved to: {final_results_file}"
+                )
+        # Open per-rank file in write mode to start fresh, then append per-sample
+        results_fp = open(results_file_rank, "w", encoding="utf-8")
+        if not results_fp:
+            raise RuntimeError(
+                f"Failed to open per-rank results file: {results_file_rank}"
+            )
+        try:
+            with torch.no_grad():
+                for batch in tqdm(
+                    test_loader, desc=f"Evaluating {stage_name}", disable=self.rank != 0
+                ):
+                    # Generate predictions with higher max_tokens (skip separate loss computation)
+                    predictions = self._get_model().generate(
+                        batch, max_new_tokens=max_new_tokens
+                    )
+
+                    # Collect results
+                    for sample, pred in zip(batch, predictions):
+                        result = {
+                            "pre_prompt": sample["pre_prompt"],
+                            "time_series_text": sample["time_series_text"],
+                            "post_prompt": sample["post_prompt"],
+                            "generated": pred,
+                            "gold": sample["answer"],
+                        }
+
+                        # Add time series ID for stage2 captioning
+                        if stage == "stage2_captioning" and "id" in sample:
+                            result["time_series_id"] = sample["id"]
+
+                        # Add template_id and ecg_id for stage5_ecg_cot
+                        if stage == "stage5_ecg_cot":
+                            if "template_id" in sample:
+                                result["template_id"] = sample["template_id"]
+                            if "ecg_id" in sample:
+                                result["ecg_id"] = sample["ecg_id"]
+                            if "correct_answer" in sample:
+                                result["correct_answer"] = sample["correct_answer"]
+                        results.append(result)
+                        # Stream write each result immediately to per-rank file
+                        results_fp.write(json.dumps(result, ensure_ascii=False) + "\n")
+                        results_fp.flush()
+                        try:
+                            os.fsync(results_fp.fileno())
+                        except Exception:
+                            pass
+        finally:
+            if results_fp is not None:
+                results_fp.close()
+
+        # Synchronize all ranks before merging
+        if dist.is_initialized():
+            dist.barrier()
+
+        # Rank 0 merges per-rank files into final results file
+        if (not dist.is_initialized()) or (self.rank == 0):
+            try:
+                # Overwrite final file each evaluation
+                with open(final_results_file, "w", encoding="utf-8") as merged_fp:
+                    if dist.is_initialized():
+                        num_ranks = self.world_size
+                    else:
+                        num_ranks = 1
+                    for r in range(num_ranks):
+                        part_file = os.path.join(
+                            self.results_dir,
+                            stage_name,
+                            "results",
+                            f"test_predictions_rank_{r}.jsonl",
+                        )
+                        if os.path.exists(part_file):
+                            with open(part_file, "r", encoding="utf-8") as pf:
+                                for line in pf:
+                                    merged_fp.write(line)
+                if self.rank == 0:
+                    print(f"Merged per-rank predictions into: {final_results_file}")
+            finally:
+                pass
+        avg_test_loss = float("nan")
+        # Calculate stage-specific metrics
+        metrics = {"test_loss": avg_test_loss}
+        if epoch is not None:
+            metrics["epoch"] = epoch
+        if metric_func:
+            # Compute metrics on rank 0 after merging, else minimal metrics
+            if (not dist.is_initialized()) or (self.rank == 0):
+                predictions = []
+                gold_answers = []
+                # Read from final merged file
+                merged_path = final_results_file
+                with open(merged_path, "r", encoding="utf-8") as f:
+                    for line in f:
+                        try:
+                            obj = json.loads(line)
+                            predictions.append(obj.get("generated", ""))
+                            gold_answers.append(obj.get("gold", ""))
+                        except Exception:
+                            continue
+                additional_metrics = metric_func(predictions, gold_answers)
+                metrics.update(additional_metrics)
+        # Save results only on rank 0 (or when not distributed)
+        if (not dist.is_initialized()) or (self.rank == 0):
+            # Save metrics
+            metrics_file = os.path.join(
+                self.results_dir, stage_name, "results", "metrics.json"
+            )
+            with open(metrics_file, "w") as f:
+                json.dump(metrics, f, indent=2)
+
+            print(f"✅ {stage_name} evaluation complete:")
+            print(f"   Test predictions saved to: {final_results_file}")
+            print(f"   Metrics saved to: {metrics_file}")
+            print(f"   Max tokens used for generation: {max_new_tokens}")
+            for metric, value in metrics.items():
+                if isinstance(value, (int, float)):
+                    print(f"   {metric}: {value:.4f}")
+                else:
+                    print(f"   {metric}: {value}")
+
+        # Signal other ranks that evaluation is complete
+        if dist.is_initialized():
+            dist.barrier()
+
+        return metrics
+
+    def _is_evaluation_completed(self, stage: str) -> bool:
+        """Check if evaluation was completed for a stage by looking for test predictions file."""
+        test_predictions_file = os.path.join(
+            self.results_dir, stage, "results", "test_predictions.jsonl"
+        )
+        metrics_file = os.path.join(self.results_dir, stage, "results", "metrics.json")
+
+        # Check if both files exist
+        if not os.path.exists(test_predictions_file) or not os.path.exists(
+            metrics_file
+        ):
+            return False
+
+        # Also check if metrics file has evaluation results
+        try:
+            with open(metrics_file, "r") as f:
+                metrics = json.load(f)
+            return "test_loss" in metrics
+        except:
+            return False
+
+    def _train_stage(
+        self,
+        stage_name: str,
+        dataset_class,
+        num_epochs: int,
+        lr_encoder: float,
+        lr_projector: float,
+        lr_base: float,
+        metric_func: Callable = None,
+        batch_size: int = None,
+        eval_only: bool = False,
+        sampler=None,
+    ) -> Dict[str, Any]:
+        """Generic training function for any stage."""
+        epoch = None
+        # Use provided batch_size or default to global BATCH_SIZE
+        if batch_size is None:
+            batch_size = BATCH_SIZE
+
+        if self.rank == 0:
+            print(f"\n🚀 Starting {stage_name} Training with {self.model_type}")
+            if eval_only:
+                print("🔍 EVAL-ONLY MODE: Skipping training, only running evaluation")
+            print("=" * 60)
+            print(f"📊 Stage Configuration:")
+            print(f"   Epochs: {num_epochs}")
+            if self.model_type == "OpenTSLMSP":
+                print(f"   Encoder LR: {lr_encoder:.2e}")
+                print(f"   Projector LR: {lr_projector:.2e}")
+            else:
+                print(f"   Base LR: {lr_base:.2e}")
+            print(f"   Batch size per GPU: {batch_size}")
+            if self.world_size > 1:
+                print(f"   Effective batch size: {batch_size * self.world_size}")
+            print()
+
+        # Check if checkpoint exists when in eval_only mode
+        if eval_only and not self._checkpoint_exists(stage_name):
+            raise RuntimeError(
+                f"Eval-only mode requires a checkpoint for {stage_name}, but none found at {os.path.join(self.results_dir, stage_name, 'checkpoints', 'best_model.pt')}"
+            )
+
+        # Load previous stage model and display metrics
+        try:
+            previous_stage_info = self._load_previous_stage_model(stage_name)
+            if previous_stage_info:
+                if self.rank == 0:
+                    print(f"📂 Loading best model from {previous_stage_info['stage']}:")
+                    print(f"   Achieved at epoch: {previous_stage_info['epoch']}")
+                    val_loss = previous_stage_info["val_loss"]
+                    if isinstance(val_loss, (int, float)):
+                        print(f"   Validation loss: {val_loss:.4f}")
+                    else:
+                        print(f"   Validation loss: {val_loss}")
+                    for metric, value in previous_stage_info["metrics"].items():
+                        if isinstance(value, (int, float)):
+                            print(f"   {metric}: {value:.4f}")
+                        else:
+                            print(f"   {metric}: {value}")
+                    print()
+            else:
+                # Only allow fresh model for first stage
+                if stage_name != CURRICULUM_STAGES[0]:
+                    raise RuntimeError(
+                        f"Cannot start {stage_name} with fresh model. Previous stage {CURRICULUM_STAGES[CURRICULUM_STAGES.index(stage_name) - 1]} must be completed first."
+                    )
+                if self.rank == 0:
+                    print("🆕 Starting with fresh model (first stage)")
+                    print()
+        except Exception as e:
+            if self.rank == 0:
+                print(f"❌ Error loading previous stage: {e}")
+            raise Exception(f"Error loading previous stage: {e}")
+
+        # Check if evaluation was already completed
+        evaluation_completed = self._is_evaluation_completed(stage_name)
+        if evaluation_completed and self.rank == 0:
+            print(
+                f"✅ Evaluation already completed for {stage_name}, skipping training and evaluation"
+            )
+            print(f"📂 Loading existing metrics...")
+
+            # Load and return existing metrics
+            metrics_file = os.path.join(
+                self.results_dir, stage_name, "results", "metrics.json"
+            )
+            with open(metrics_file, "r") as f:
+                metrics = json.load(f)
+
+            print(f"📊 Existing results for {stage_name}:")
+            for metric, value in metrics.items():
+                if isinstance(value, (int, float)):
+                    print(f"   {metric}: {value:.4f}")
+                else:
+                    print(f"   {metric}: {value}")
+
+            return metrics
+
+        # Enable LoRA if needed for this stage
+        self._enable_lora_if_needed(stage_name)
+
+        # Initialize optimizer and scheduler
+        optimizer = self._get_optimizer(batch_size, lr_encoder, lr_projector, lr_base)
+
+        # Create data loaders
+        if sampler is not None:
+            if self.world_size > 1:
+                get_logger().warning(
+                    "BalancedBatchSampler was provided, but distributed training (DDP) is enabled. BalancedBatchSampler will NOT be used. Data will be sharded using DistributedSampler instead. Typically for stage3_cot it is better to use BalancedBatchSampler, if dataset is imbalanced."
+                )
+                train_loader = self._merge_data_loaders(
+                    [
+                        dataset_class(
+                            "train", EOS_TOKEN=self._get_model().get_eos_token()
+                        )
+                    ],
+                    shuffle=True,
+                    batch_size=batch_size,
+                    patch_size=PATCH_SIZE,
+                    distribute_data=True,
+                )
+            else:
+                train_dataset = dataset_class(
+                    "train", EOS_TOKEN=self._get_model().get_eos_token()
+                )
+                train_loader = DataLoader(
+                    train_dataset,
+                    batch_sampler=sampler,
+                    collate_fn=lambda batch: extend_time_series_to_match_patch_size_and_aggregate(
+                        batch, patch_size=PATCH_SIZE
+                    ),
+                )
+        else:
+            train_loader = self._merge_data_loaders(
+                [dataset_class("train", EOS_TOKEN=self._get_model().get_eos_token())],
+                shuffle=True,
+                batch_size=batch_size,
+                patch_size=PATCH_SIZE,
+                distribute_data=self.world_size > 1,
+            )
+
+        val_loader = self._merge_data_loaders(
+            [dataset_class("validation", EOS_TOKEN=self._get_model().get_eos_token())],
+            shuffle=False,
+            batch_size=1,
+            patch_size=PATCH_SIZE,
+            distribute_data=False,  # Don't distribute validation
+        )
+
+        test_loader = self._merge_data_loaders(
+            [dataset_class("test", EOS_TOKEN=self._get_model().get_eos_token())],
+            shuffle=False,
+            # Test-set generation eval batch size. Default 1 (original behavior).
+            # Set EVAL_BATCH_SIZE>1 to batch generation (left-padded in OpenTSLMSP.generate)
+            # for a large speedup; validated to match batch=1 greedy output.
+            batch_size=int(os.environ.get("EVAL_BATCH_SIZE", "1")),
+            patch_size=PATCH_SIZE,
+            distribute_data=self.world_size > 1,
+        )
+
+        # Scheduler
+        total_steps = num_epochs * len(train_loader)
+        warmup_steps = int(WARMUP_FRAC * total_steps)
+        scheduler = get_linear_schedule_with_warmup(
+            optimizer,
+            num_warmup_steps=warmup_steps,
+            num_training_steps=total_steps,
+        )
+
+        if self.rank == 0:
+            print(f"📈 Total training steps: {total_steps}")
+            print(f"🔥 Warmup steps: {warmup_steps}")
+
+        # Load previous checkpoint if exists (for resuming current stage)
+        best_epoch, best_val_loss = self._load_checkpoint(
+            stage_name, optimizer, scheduler, eval_only=eval_only
+        )
+        if best_epoch is not None:
+            print(
+                f"📂 Resuming {stage_name} from epoch {best_epoch} (val_loss: {best_val_loss:.4f})"
+            )
+            # Display previous loss history if available
+            self._display_loss_history(stage_name)
+        else:
+            print(f"🆕 Starting fresh training for {stage_name}")
+            best_val_loss = float("inf")  # Ensure proper initialization
+
+        # Skip training loop if eval_only is True
+        if eval_only:
+            if self.rank == 0:
+                print(f"⏭️  Skipping training loop (eval_only mode)")
+                print(f"📂 Using existing checkpoint for evaluation")
+            epoch = best_epoch
+            epochs_no_improve = 0
+        else:
+            # Training loop
+            epochs_no_improve = 0
+            start_epoch = best_epoch + 1 if best_epoch is not None else 1
+            for epoch in range(start_epoch, num_epochs + 1):
+                # Set epoch for distributed sampler
+                if hasattr(train_loader.sampler, "set_epoch"):
+                    train_loader.sampler.set_epoch(epoch)
+
+                # Training
+                self.model.train()
+                running_loss = 0.0
+                prog = tqdm(
+                    train_loader,
+                    desc=f"Epoch {epoch}/{num_epochs}",
+                    disable=self.rank != 0,
+                )
+                for i, batch in enumerate(prog):
+                    # DEBUG PRINT: Only for the first batch of the first epoch
+                    if epoch == start_epoch and i == 0:
+                        print(f"[DEBUG] Batch {i} - batch size: {len(batch)}")
+                        if isinstance(batch, list) and isinstance(batch[0], dict):
+                            for k, v in batch[0].items():
+                                if hasattr(v, "shape"):
+                                    print(f"[DEBUG] Sample key '{k}' shape: {v.shape}")
+                                elif isinstance(v, list):
+                                    print(
+                                        f"[DEBUG] Sample key '{k}' list length: {len(v)}"
+                                    )
+                        import torch
+
+                        print(
+                            torch.cuda.memory_summary()
+                            if torch.cuda.is_available()
+                            else "No CUDA"
+                        )
+                    optimizer.zero_grad()
+                    loss = self._get_model().compute_loss(batch)
+                    loss.backward()
+
+                    # Handle gradient clipping for distributed training
+                    clip_grad_norm_(self._get_model().parameters(), GRAD_CLIP_NORM)
+
+                    optimizer.step()
+                    scheduler.step()
+
+                    running_loss += loss.item()
+                    if self.rank == 0:
+                        prog.set_postfix(
+                            loss=f"{loss.item():.4f}",
+                            lr=f"{scheduler.get_last_lr()[0]:.2e}",
+                        )
+
+                avg_train_loss = running_loss / len(train_loader)
+                if self.rank == 0:
+                    tqdm.write(f"Epoch {epoch} — train loss: {avg_train_loss:.4f}")
+
+                # Validation
+                val_loss = 0.0
+                self.model.eval()
+                with torch.no_grad():
+                    for batch in tqdm(
+                        val_loader,
+                        desc=f"Validating {stage_name}",
+                        disable=self.rank != 0,
+                    ):
+                        val_loss += self._get_model().compute_loss(batch).item()
+
+                avg_val_loss = val_loss / len(val_loader)
+
+                # Synchronize validation loss across all ranks
+                if dist.is_initialized():
+                    val_loss_tensor = torch.tensor(avg_val_loss, device=self.device)
+                    dist.all_reduce(val_loss_tensor, op=dist.ReduceOp.SUM)
+                    avg_val_loss = val_loss_tensor.item() / self.world_size
+
+                if self.rank == 0:
+                    tqdm.write(f"Epoch {epoch} — val   loss: {avg_val_loss:.4f}")
+                    tqdm.write(f"Epoch {epoch} — best  loss: {best_val_loss:.4f}")
+
+                # Save loss history for this epoch
+                self._save_loss_history(stage_name, epoch, avg_train_loss, avg_val_loss)
+
+                # Early stopping - all ranks need to make the same decision
+                should_save = avg_val_loss + 1e-4 < best_val_loss
+                if dist.is_initialized():
+                    save_tensor = torch.tensor(
+                        1 if should_save else 0, device=self.device
+                    )
+                    dist.all_reduce(save_tensor, op=dist.ReduceOp.SUM)
+                    should_save = (
+                        save_tensor.item() > 0
+                    )  # If any rank thinks we should save, we save
+
+                if should_save:
+                    best_val_loss = avg_val_loss
+                    epochs_no_improve = 0
+                    self._save_checkpoint(
+                        stage_name, epoch, avg_val_loss, optimizer, scheduler
+                    )
+                    if self.rank == 0:
+                        tqdm.write("✔️  New best model saved.\n")
+                else:
+                    epochs_no_improve += 1
+                    if self.rank == 0:
+                        tqdm.write(
+                            f"No improvement for {epochs_no_improve}/{EARLY_STOP_PAT} epochs.\n"
+                        )
+
+                    # Synchronize early stopping decision across all ranks
+                    if epochs_no_improve >= EARLY_STOP_PAT:
+                        if self.rank == 0:
+                            tqdm.write(
+                                f"\nEarly stopping triggered after {epoch} epochs."
+                            )
+                            tqdm.write(
+                                f"Final stats: best_val_loss={best_val_loss:.4f}, epochs_no_improve={epochs_no_improve}"
+                            )
+                        break
+
+                # Synchronize best_val_loss and epochs_no_improve across all ranks
+                if dist.is_initialized():
+                    best_loss_tensor = torch.tensor(best_val_loss, device=self.device)
+                    epochs_tensor = torch.tensor(epochs_no_improve, device=self.device)
+                    dist.broadcast(best_loss_tensor, src=0)
+                    dist.broadcast(epochs_tensor, src=0)
+                    best_val_loss = best_loss_tensor.item()
+                    epochs_no_improve = int(epochs_tensor.item())
+
+        # Load best model and evaluate
+        best_epoch, _ = self._load_checkpoint(stage_name, optimizer, scheduler)
+        if best_epoch is not None:
+            if self.rank == 0:
+                print(
+                    f"📂 Loaded best checkpoint from epoch {best_epoch} for evaluation."
+                )
+
+        if self.rank == 0:
+            if epoch is None:
+                epoch = best_epoch
+                print(f"🏁 Training completed for {stage_name}")
+                print(f"   Total epochs run: {epoch}")
+            else:
+                print(f"🏁 Training completed for {stage_name}")
+                print(f"   Total epochs run: {epoch}")
+                print(f"   Best validation loss: {best_val_loss:.4f}")
+                print(f"   Epochs without improvement: {epochs_no_improve}")
+
+        metrics = self._evaluate_stage(
+            stage_name, test_loader, stage_name, metric_func, best_epoch
+        )
+
+        return metrics
+
+    def stage1_mcq(
+        self, batch_size: int = None, eval_only: bool = False
+    ) -> Dict[str, Any]:
+        """Stage 1: Multiple Choice Question Answering (TSQA).
+
+        Configuration:
+        - Epochs: 20
+        - OpenTSLMSP: encoder_lr=2e-4, projector_lr=1e-4
+        - OpenTSLMFlamingo: base_lr=2e-4
+        - Metric: Accuracy
+        """
+        return self._train_stage(
+            stage_name="stage1_mcq",
+            dataset_class=TSQADataset,
+            num_epochs=30,
+            lr_encoder=2e-4,
+            lr_projector=1e-4,
+            lr_base=2e-4,
+            metric_func=lambda preds, golds: {
+                "accuracy": self._calculate_accuracy(preds, golds)
+            },
+            batch_size=batch_size,
+            eval_only=eval_only,
+        )
+
+    def stage2_captioning(
+        self, batch_size: int = None, eval_only: bool = False
+    ) -> Dict[str, Any]:
+        """Stage 2: Caption Generation (M4).
+
+        Configuration:
+        - Epochs: 15
+        - OpenTSLMSP: encoder_lr=1e-4, projector_lr=5e-5 (lower for fine-tuning)
+        - OpenTSLMFlamingo: base_lr=1e-4 (lower for fine-tuning)
+        - Metric: Test loss only
+        """
+        return self._train_stage(
+            stage_name="stage2_captioning",
+            dataset_class=M4QADataset,
+            num_epochs=20,
+            lr_encoder=2e-4,
+            lr_projector=1e-4,
+            lr_base=2e-4,
+            metric_func=None,  # Only test loss for captioning
+            batch_size=batch_size,
+            eval_only=eval_only,
+        )
+
+    def stage3_cot(
+        self, batch_size: int = None, eval_only: bool = False
+    ) -> Dict[str, Any]:
+        """Stage CoT: Chain-of-Thought Reasoning (HAR).
+
+        Configuration:
+        - Epochs: 100
+        - OpenTSLMSP: encoder_lr=2e-4, projector_lr=1e-4
+        - OpenTSLMFlamingo: base_lr=2e-4
+        - Metric: Test loss only (chain-of-thought reasoning)
+        """
+        sampler = None
+
+        return self._train_stage(
+            stage_name="stage3_cot",
+            dataset_class=HARCoTQADataset,
+            num_epochs=30,
+            lr_encoder=2e-4,
+            lr_projector=1e-4,
+            lr_base=2e-4,
+            metric_func=None,  # Only test loss for chain-of-thought reasoning
+            batch_size=batch_size,
+            eval_only=eval_only,
+            sampler=sampler,
+        )
+
+    def stage4_sleep_cot(
+        self, batch_size: int = None, eval_only: bool = False
+    ) -> Dict[str, Any]:
+        """Stage 4: Chain-of-Thought Reasoning (SleepEDF).
+
+        Configuration:
+        - Epochs: 60
+        - OpenTSLMSP: encoder_lr=2e-4, projector_lr=1e-4
+        - OpenTSLMFlamingo: base_lr=2e-4
+        - Metric: Test loss only (chain-of-thought reasoning)
+        """
+        sampler = None
+
+        return self._train_stage(
+            stage_name="stage4_sleep_cot",
+            dataset_class=SleepEDFCoTQADataset,
+            num_epochs=60,
+            lr_encoder=2e-4,
+            lr_projector=1e-4,
+            lr_base=2e-4,
+            metric_func=None,  # Only test loss for chain-of-thought reasoning
+            batch_size=batch_size,
+            eval_only=eval_only,
+            sampler=sampler,
+        )
+
+    def stage5_ecg_cot(
+        self, batch_size: int = None, eval_only: bool = False
+    ) -> Dict[str, Any]:
+        """Stage 5: Chain-of-Thought Reasoning (ECG QA CoT).
+
+        Configuration:
+        - Epochs: 60
+        - OpenTSLMSP: encoder_lr=2e-4, projector_lr=1e-4
+        - OpenTSLMFlamingo: base_lr=2e-4
+        - Metric: Test loss only (chain-of-thought reasoning)
+        """
+        sampler = None
+
+        return self._train_stage(
+            stage_name="stage5_ecg_cot",
+            dataset_class=ECGQACoTQADataset,
+            num_epochs=60,
+            lr_encoder=2e-4,
+            lr_projector=1e-4,
+            lr_base=2e-4,
+            metric_func=None,  # Only test loss for chain-of-thought reasoning
+            batch_size=batch_size,
+            eval_only=eval_only,
+            sampler=sampler,
+        )
+
+    def run_curriculum(
+        self, stages: List[str] = None, batch_size: int = None, eval_only: bool = False
+    ):
+        """Run the complete curriculum learning pipeline."""
+        if stages is None:
+            stages = CURRICULUM_STAGES
+
+        # Filter out completed stages
+        incomplete_stages = []
+        for stage in stages:
+            if self._is_stage_completed(stage):
+                if self.rank == 0:
+                    print(f"⏭️  Skipping completed stage: {stage}")
+            else:
+                incomplete_stages.append(stage)
+
+        if self.rank == 0:
+            print(f"🎓 Starting Curriculum Learning with {self.model_type}")
+            if eval_only:
+                print("🔍 EVAL-ONLY MODE: Will skip training and only run evaluation")
+            print(f"📊 All stages: {', '.join(stages)}")
+            print(f"🔄 Incomplete stages: {', '.join(incomplete_stages)}")
+            print(f"💻 Device: {self.device}")
+            if batch_size:
+                print(f"📦 Batch size: {batch_size}")
+            if self.world_size > 1:
+                print(f"🌐 Distributed training with {self.world_size} GPUs")
+            print("=" * 80)
+
+        results = {}
+
+        # Run only incomplete stages
+        for stage in incomplete_stages:
+            # Synchronize all ranks before starting each stage
+            if dist.is_initialized():
+                dist.barrier()
+
+            if stage == "stage1_mcq":
+                stage_results = self.stage1_mcq(
+                    batch_size=batch_size, eval_only=eval_only
+                )
+                results[stage] = stage_results
+                self._mark_stage_completed(stage, stage_results)
+            elif stage == "stage2_captioning":
+                stage_results = self.stage2_captioning(
+                    batch_size=batch_size, eval_only=eval_only
+                )
+                results[stage] = stage_results
+                self._mark_stage_completed(stage, stage_results)
+            elif stage == "stage3_cot":
+                stage_results = self.stage3_cot(
+                    batch_size=batch_size, eval_only=eval_only
+                )
+                results[stage] = stage_results
+                self._mark_stage_completed(stage, stage_results)
+            elif stage == "stage4_sleep_cot":
+                stage_results = self.stage4_sleep_cot(
+                    batch_size=batch_size, eval_only=eval_only
+                )
+                results[stage] = stage_results
+                self._mark_stage_completed(stage, stage_results)
+            elif stage == "stage5_ecg_cot":
+                stage_results = self.stage5_ecg_cot(
+                    batch_size=batch_size, eval_only=eval_only
+                )
+                results[stage] = stage_results
+                self._mark_stage_completed(stage, stage_results)
+            elif stage == "stage4_sleep_cot":
+                stage_results = self.stage4_sleep_cot(
+                    batch_size=batch_size, eval_only=eval_only
+                )
+                results[stage] = stage_results
+                self._mark_stage_completed(stage, stage_results)
+            elif stage == "stage5_ecg_cot":
+                stage_results = self.stage5_ecg_cot(
+                    batch_size=batch_size, eval_only=eval_only
+                )
+                results[stage] = stage_results
+                self._mark_stage_completed(stage, stage_results)
+            else:
+                if self.rank == 0:
+                    print(f"⚠️  Unknown stage: {stage}, skipping...")
+
+            # Synchronize all ranks after completing each stage
+            if dist.is_initialized():
+                dist.barrier()
+
+        # Save overall results only on rank 0
+        if self.rank == 0:
+            overall_results_file = os.path.join(
+                self.results_dir, "curriculum_results.json"
+            )
+            with open(overall_results_file, "w") as f:
+                json.dump(results, f, indent=2)
+
+            print(f"\n🎉 Curriculum Learning Complete!")
+            print(f"📁 All results saved to: {self.results_dir}/")
+            print(f"📊 Overall results: {overall_results_file}")
+
+        return results
+
+    def _should_use_distributed(self) -> bool:
+        """Check if distributed training should be used."""
+        return ("WORLD_SIZE" in os.environ and int(os.environ["WORLD_SIZE"]) > 1) or (
+            "LOCAL_RANK" in os.environ and int(os.environ["LOCAL_RANK"]) >= 0
+        )
+
+    def _init_distributed(self):
+        """Initialize distributed training."""
+        if "WORLD_SIZE" in os.environ:
+            self.world_size = int(os.environ["WORLD_SIZE"])
+        if "RANK" in os.environ:
+            self.rank = int(os.environ["RANK"])
+        elif "LOCAL_RANK" in os.environ:
+            self.rank = int(os.environ["LOCAL_RANK"])
+
+        # Initialize process group
+        dist.init_process_group(
+            backend=self.dist_backend,
+            init_method=self.dist_url,
+            world_size=self.world_size,
+            rank=self.rank,
+            timeout=datetime.timedelta(hours=999),
+        )
+
+        # Set device for this process
+        if torch.cuda.is_available():
+            torch.cuda.set_device(self.local_rank)
+            self.device = torch.device("cuda", self.local_rank)
+
+        if self.rank == 0:
+            print(f"Initialized distributed training with {self.world_size} GPUs")
+
+    def _is_stage_completed(self, stage: str) -> bool:
+        """Check if a stage is completed by verifying both training and evaluation were successful."""
+        metrics_file = os.path.join(self.results_dir, stage, "results", "metrics.json")
+
+        if not os.path.exists(metrics_file):
+            return False
+
+        try:
+            with open(metrics_file, "r") as f:
+                metrics = json.load(f)
+
+            # Check if the completion flag exists
+            if not metrics.get("completed", False):
+                return False
+
+            # Check if evaluation was actually completed by looking for test_loss
+            if "test_loss" not in metrics:
+                return False
+
+            # Check if test predictions file exists
+            test_predictions_file = os.path.join(
+                self.results_dir, stage, "results", "test_predictions.jsonl"
+            )
+            if not os.path.exists(test_predictions_file):
+                return False
+
+            return True
+
+        except:
+            return False
+
+    def _mark_stage_completed(self, stage: str, metrics: Dict[str, Any]):
+        """Mark a stage as completed by adding completion flag to metrics."""
+        metrics["completed"] = True
+        metrics["completion_epoch"] = metrics.get("epoch", "?")
+
+        metrics_file = os.path.join(self.results_dir, stage, "results", "metrics.json")
+        with open(metrics_file, "w") as f:
+            json.dump(metrics, f, indent=2)
+
+        if self.rank == 0:
+            print(f"✅ Stage {stage} marked as completed")
+
+    def _get_model(self):
+        """Get the underlying model (handles DDP wrapping)."""
+        if hasattr(self.model, "module"):
+            return self.model.module
+        return self.model
+
+    def _checkpoint_exists(self, stage: str) -> bool:
+        """Check if a checkpoint exists for a specific stage."""
+        checkpoint_path = os.path.join(
+            self.results_dir, stage, "checkpoints", "best_model.pt"
+        )
+        return os.path.exists(checkpoint_path)
+
+    def _enable_lora_if_needed(self, stage_name: str):
+        """Enable LoRA for OpenTSLMSP models in stages after stage2."""
+        if self.model_type != "OpenTSLMSP":
+            return  # LoRA only for OpenTSLMSP
+
+        # Get the underlying model (handles DDP wrapping)
+        model = self._get_model()
+
+        # Enable LoRA for stages after stage2_captioning
+        stages_with_lora = ["stage3_cot", "stage4_sleep_cot", "stage5_ecg_cot"]
+
+        if stage_name in stages_with_lora:
+            if not getattr(model, "lora_enabled", False):
+                if self.rank == 0:
+                    print(f"🔧 Enabling LoRA for {stage_name}")
+                try:
+                    model.enable_lora(lora_r=16, lora_alpha=32, lora_dropout=0.0)
+                    if self.rank == 0:
+                        print(f"✅ LoRA enabled for {stage_name}")
+                except Exception as e:
+                    if self.rank == 0:
+                        print(f"❌ Failed to enable LoRA for {stage_name}: {e}")
+                        print("   Continuing without LoRA...")
+            else:
+                if self.rank == 0:
+                    print(f"✅ LoRA already enabled for {stage_name}")
+        else:
+            if self.rank == 0:
+                if stage_name in ["stage1_mcq", "stage2_captioning"]:
+                    print(
+                        f"ℹ️  LoRA disabled for {stage_name} (only enabled for stages 3+)"
+                    )
+                else:
+                    print(f"ℹ️  LoRA not configured for {stage_name}")
+
+    def _enable_lora_if_needed(self, stage_name: str):
+        """Enable LoRA for OpenTSLMSP models in stages after stage2."""
+        if self.model_type != "OpenTSLMSP":
+            return  # LoRA only for OpenTSLMSP
+
+        # Get the underlying model (handles DDP wrapping)
+        model = self._get_model()
+
+        # Enable LoRA for stages after stage2_captioning
+        stages_with_lora = ["stage3_cot", "stage4_sleep_cot", "stage5_ecg_cot"]
+
+        if stage_name in stages_with_lora:
+            if not getattr(model, "lora_enabled", False):
+                if self.rank == 0:
+                    print(f"🔧 Enabling LoRA for {stage_name}")
+                try:
+                    model.enable_lora(lora_r=16, lora_alpha=32, lora_dropout=0.0)
+                    if self.rank == 0:
+                        print(f"✅ LoRA enabled for {stage_name}")
+                except Exception as e:
+                    if self.rank == 0:
+                        print(f"❌ Failed to enable LoRA for {stage_name}: {e}")
+                        print("   Continuing without LoRA...")
+            else:
+                if self.rank == 0:
+                    print(f"✅ LoRA already enabled for {stage_name}")
+        else:
+            if self.rank == 0:
+                if stage_name in ["stage1_mcq", "stage2_captioning"]:
+                    print(
+                        f"ℹ️  LoRA disabled for {stage_name} (only enabled for stages 3+)"
+                    )
+                else:
+                    print(f"ℹ️  LoRA not configured for {stage_name}")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Curriculum Learning for OpenTSLM Models"
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        choices=["OpenTSLMSP", "OpenTSLMFlamingo"],
+        required=True,
+        help="Model type to train",
+    )
+    parser.add_argument(
+        "--stages",
+        nargs="+",
+        choices=CURRICULUM_STAGES,
+        default=CURRICULUM_STAGES,
+        help="Stages to run (default: all stages)",
+    )
+    parser.add_argument(
+        "--device", type=str, default=None, help="Device to use (cuda, mps, cpu)"
+    )
+    parser.add_argument(
+        "--batch_size",
+        type=int,
+        default=None,
+        help="Batch size for training (default: use value from model_config.py)",
+    )
+
+    # Evaluation arguments
+    parser.add_argument(
+        "--eval_only",
+        default=False,
+        action="store_true",
+        help="Skip training and only run evaluation (requires existing checkpoint)",
+    )
+
+    # Model-specific arguments
+    parser.add_argument(
+        "--llm_id",
+        type=str,
+        default="meta-llama/Llama-3.2-1B",
+        help="LLM model ID for OpenTSLMFlamingo (e.g., 'google/medgemma-2b', 'meta-llama/Llama-3.2-1B')",
+    )
+
+    # Distributed training arguments
+    parser.add_argument(
+        "--gradient_checkpointing",
+        default=False,
+        action="store_true",
+        help="Enable gradient checkpointing",
+    )
+    parser.add_argument(
+        "--dist_url",
+        default="env://",
+        type=str,
+        help="URL used to set up distributed training",
+    )
+    parser.add_argument(
+        "--dist_backend", default="nccl", type=str, help="Distributed backend"
+    )
+    parser.add_argument(
+        "--local_rank",
+        type=int,
+        default=int(os.environ.get("LOCAL_RANK", 0)),
+        help="Local GPU rank",
+    )
+
+    # Logging arguments
+    parser.add_argument(
+        "--verbose", default=False, action="store_true", help="Enable verbose logging"
+    )
+
+    args = parser.parse_args()
+
+    # Set up global logging
+    set_global_verbose(args.verbose)
+    logger = get_logger(verbose=args.verbose)
+
+    # Initialize trainer
+    trainer = CurriculumTrainer(
+        args.model,
+        args.device,
+        gradient_checkpointing=args.gradient_checkpointing,
+        dist_url=args.dist_url,
+        dist_backend=args.dist_backend,
+        local_rank=args.local_rank,
+        llm_id=args.llm_id,
+    )
+
+    # Run curriculum
+    results = trainer.run_curriculum(args.stages, args.batch_size, args.eval_only)
+
+    # Print summary
+    logger.info("Final Results Summary:")
+    logger.info("=" * 40)
+    for stage, metrics in results.items():
+        logger.info(f"{stage.upper()}:")
+        for metric, value in metrics.items():
+            if isinstance(value, (int, float)):
+                logger.info(f"  {metric}: {value:.4f}")
+            else:
+                logger.info(f"  {metric}: {value}")
+
+
+if __name__ == "__main__":
+    main()

code/OpenTSLM/demo/huggingface/01_test_hf_tsqa.py

@@ -0,0 +1,72 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+"""
+Demo script for testing TSQA (Time Series Question Answering) model from HuggingFace.
+
+This script:
+1. Loads a pretrained model from HuggingFace Hub
+2. Loads the TSQA test dataset
+3. Generates predictions on the evaluation set
+4. Prints model outputs
+"""
+
+from opentslm.model.llm.OpenTSLM import OpenTSLM
+from opentslm.time_series_datasets.TSQADataset import TSQADataset
+from opentslm.time_series_datasets.util import extend_time_series_to_match_patch_size_and_aggregate
+from torch.utils.data import DataLoader
+from opentslm.model_config import PATCH_SIZE
+import torch
+# Model repository ID - change this to test different models
+REPO_ID = "OpenTSLM/llama-3.2-1b-tsqa-sp"
+
+def main():
+    print("=" * 60)
+    print("TSQA Model Demo")
+    print("=" * 60)
+    
+    # Load model from HuggingFace
+    print(f"\n📥 Loading model from {REPO_ID}...")
+    model = OpenTSLM.load_pretrained(REPO_ID, device="cuda" if torch.cuda.is_available() else "cpu")
+    
+    # Create dataset
+    print("\n📊 Loading TSQA test dataset...")
+    test_dataset = TSQADataset("test", EOS_TOKEN=model.get_eos_token())
+    
+    # Create data loader
+    test_loader = DataLoader(
+        test_dataset,
+        shuffle=False,
+        batch_size=1,
+        collate_fn=lambda batch: extend_time_series_to_match_patch_size_and_aggregate(
+            batch, patch_size=PATCH_SIZE
+        ),
+    )
+    
+    print(f"\n🔍 Running inference on {len(test_dataset)} test samples...")
+    print("=" * 60)
+    
+    # Iterate over evaluation set
+    for i, batch in enumerate(test_loader):
+        # Generate predictions
+        predictions = model.generate(batch, max_new_tokens=200)
+        
+        # Print results
+        for sample, pred in zip(batch, predictions):
+            print(f"\n📝 Sample {i + 1}:")
+            print(f"   Question: {sample.get('pre_prompt', 'N/A')}")
+            if 'time_series_text' in sample:
+                print(f"   Time Series Info: {sample['time_series_text'][:100]}...")
+            print(f"   Gold Answer: {sample.get('answer', 'N/A')}")
+            print(f"   Model Output: {pred}")
+            print("-" * 60)
+        
+        # Limit to first 5 samples for demo
+        if i >= 4:
+            print("\n✅ Demo complete! (Showing first 5 samples)")
+            break
+
+if __name__ == "__main__":
+    main()
+

code/OpenTSLM/demo/huggingface/02_test_hf_m4.py

@@ -0,0 +1,74 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+"""
+Demo script for testing M4 (Time Series Captioning) model from HuggingFace.
+
+This script:
+1. Loads a pretrained model from HuggingFace Hub
+2. Loads the M4 test dataset
+3. Generates predictions on the evaluation set
+4. Prints model outputs
+"""
+
+from opentslm.model.llm.OpenTSLM import OpenTSLM
+from opentslm.time_series_datasets.m4.M4QADataset import M4QADataset
+from opentslm.time_series_datasets.util import extend_time_series_to_match_patch_size_and_aggregate
+from torch.utils.data import DataLoader
+from opentslm.model_config import PATCH_SIZE
+import torch
+
+# Model repository ID - change this to test different models
+REPO_ID = "OpenTSLM/llama-3.2-1b-m4-sp"
+
+def main():
+    print("=" * 60)
+    print("M4 Captioning Model Demo")
+    print("=" * 60)
+    
+    # Load model from HuggingFace
+    print(f"\n📥 Loading model from {REPO_ID}...")
+    model = OpenTSLM.load_pretrained(REPO_ID, device="cuda" if torch.cuda.is_available() else "cpu")
+    
+    # Create dataset
+    print("\n📊 Loading M4 test dataset...")
+    test_dataset = M4QADataset("test", EOS_TOKEN=model.get_eos_token())
+    
+    # Create data loader
+    test_loader = DataLoader(
+        test_dataset,
+        shuffle=False,
+        batch_size=1,
+        collate_fn=lambda batch: extend_time_series_to_match_patch_size_and_aggregate(
+            batch, patch_size=PATCH_SIZE
+        ),
+    )
+    
+    print(f"\n🔍 Running inference on {len(test_dataset)} test samples...")
+    print("=" * 60)
+    
+    # Iterate over evaluation set
+    for i, batch in enumerate(test_loader):
+        # Generate predictions
+        predictions = model.generate(batch, max_new_tokens=200)
+        
+        # Print results
+        for sample, pred in zip(batch, predictions):
+            print(f"\n📝 Sample {i + 1}:")
+            if 'id' in sample:
+                print(f"   Time Series ID: {sample['id']}")
+            if 'pre_prompt' in sample:
+                print(f"   Prompt: {sample['pre_prompt']}")
+            print(f"   Gold Caption: {sample.get('answer', 'N/A')}")
+            print(f"   Model Output: {pred}")
+            print("-" * 60)
+        
+        # Limit to first 5 samples for demo
+        if i >= 9:
+            print("\n✅ Demo complete! (Showing first 10 samples)")
+            break
+
+if __name__ == "__main__":
+    main()
+

code/OpenTSLM/demo/huggingface/03_test_hf_har_cot.py

@@ -0,0 +1,77 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+"""
+Demo script for testing HAR CoT (Human Activity Recognition Chain-of-Thought) model from HuggingFace.
+
+This script:
+1. Loads a pretrained model from HuggingFace Hub
+2. Loads the HAR CoT test dataset
+3. Generates predictions on the evaluation set
+4. Prints model outputs
+"""
+
+from opentslm.model.llm.OpenTSLM import OpenTSLM
+from opentslm.time_series_datasets.har_cot.HARCoTQADataset import HARCoTQADataset
+from opentslm.time_series_datasets.util import extend_time_series_to_match_patch_size_and_aggregate
+from torch.utils.data import DataLoader
+from opentslm.model_config import PATCH_SIZE
+import torch
+
+# Model repository ID - change this to test different models
+REPO_ID = "OpenTSLM/llama-3.2-1b-har-sp"
+
+def main():
+    print("=" * 60)
+    print("HAR CoT Model Demo")
+    print("=" * 60)
+    
+    # Load model from HuggingFace
+    print(f"\n📥 Loading model from {REPO_ID}...")
+    enable_lora = False
+    if "-sp" in REPO_ID:
+        enable_lora = True
+    model = OpenTSLM.load_pretrained(REPO_ID, enable_lora=enable_lora, device="cuda" if torch.cuda.is_available() else "cpu")
+
+   
+    
+    # Create dataset
+    print("\n📊 Loading HAR CoT test dataset...")
+    test_dataset = HARCoTQADataset("test", EOS_TOKEN=model.get_eos_token())
+    
+    # Create data loader
+    test_loader = DataLoader(
+        test_dataset,
+        shuffle=False,
+        batch_size=1,
+        collate_fn=lambda batch: extend_time_series_to_match_patch_size_and_aggregate(
+            batch, patch_size=PATCH_SIZE
+        ),
+    )
+    
+    print(f"\n🔍 Running inference on {len(test_dataset)} test samples...")
+    print("=" * 60)
+    
+    # Iterate over evaluation set
+    for i, batch in enumerate(test_loader):
+        # Generate predictions
+        predictions = model.generate(batch, max_new_tokens=500)
+        
+        # Print results
+        for sample, pred in zip(batch, predictions):
+            print(f"\n📝 Sample {i + 1}:")
+            if 'pre_prompt' in sample:
+                print(f"   Question: {sample['pre_prompt']}")
+            print(f"   Gold Answer: {sample.get('answer', 'N/A')}")
+            print(f"   Model Output: {pred}")
+            print("-" * 60)
+        
+        # Limit to first 5 samples for demo
+        if i >= 4:
+            print("\n✅ Demo complete! (Showing first 5 samples)")
+            break
+
+if __name__ == "__main__":
+    main()
+

code/OpenTSLM/demo/huggingface/04_test_hf_sleep_cot.py

@@ -0,0 +1,75 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+"""
+Demo script for testing Sleep CoT (Sleep Stage Classification Chain-of-Thought) model from HuggingFace.
+
+This script:
+1. Loads a pretrained model from HuggingFace Hub
+2. Loads the Sleep CoT test dataset
+3. Generates predictions on the evaluation set
+4. Prints model outputs
+"""
+
+from opentslm.model.llm.OpenTSLM import OpenTSLM
+from opentslm.time_series_datasets.sleep.SleepEDFCoTQADataset import SleepEDFCoTQADataset
+from opentslm.time_series_datasets.util import extend_time_series_to_match_patch_size_and_aggregate
+from torch.utils.data import DataLoader
+from opentslm.model_config import PATCH_SIZE
+import torch
+
+# Model repository ID - change this to test different models
+REPO_ID = "OpenTSLM/llama-3.2-1b-sleep-sp"
+
+def main():
+    print("=" * 60)
+    print("Sleep CoT Model Demo")
+    print("=" * 60)
+    
+    # Load model from HuggingFace
+    print(f"\n📥 Loading model from {REPO_ID}...")
+    enable_lora = False
+    if "-sp" in REPO_ID:
+        enable_lora = True
+    model = OpenTSLM.load_pretrained(REPO_ID, enable_lora=enable_lora, device="cuda" if torch.cuda.is_available() else "cpu")
+    
+    # Create dataset
+    print("\n📊 Loading Sleep CoT test dataset...")
+    test_dataset = SleepEDFCoTQADataset("test", EOS_TOKEN=model.get_eos_token())
+    
+    # Create data loader
+    test_loader = DataLoader(
+        test_dataset,
+        shuffle=False,
+        batch_size=1,
+        collate_fn=lambda batch: extend_time_series_to_match_patch_size_and_aggregate(
+            batch, patch_size=PATCH_SIZE
+        ),
+    )
+    
+    print(f"\n🔍 Running inference on {len(test_dataset)} test samples...")
+    print("=" * 60)
+    
+    # Iterate over evaluation set
+    for i, batch in enumerate(test_loader):
+        # Generate predictions
+        predictions = model.generate(batch, max_new_tokens=500)
+        
+        # Print results
+        for sample, pred in zip(batch, predictions):
+            print(f"\n📝 Sample {i + 1}:")
+            if 'pre_prompt' in sample:
+                print(f"   Question: {sample['pre_prompt']}")
+            print(f"   Gold Answer: {sample.get('answer', 'N/A')}")
+            print(f"   Model Output: {pred}")
+            print("-" * 60)
+        
+        # Limit to first 5 samples for demo
+        if i >= 4:
+            print("\n✅ Demo complete! (Showing first 5 samples)")
+            break
+
+if __name__ == "__main__":
+    main()
+

code/OpenTSLM/demo/huggingface/05_test_hf_ecg_qa_cot.py

@@ -0,0 +1,79 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+"""
+Demo script for testing ECG QA CoT (ECG Question Answering Chain-of-Thought) model from HuggingFace.
+
+This script:
+1. Loads a pretrained model from HuggingFace Hub
+2. Loads the ECG QA CoT test dataset
+3. Generates predictions on the evaluation set
+4. Prints model outputs
+"""
+
+from opentslm.model.llm.OpenTSLM import OpenTSLM
+from opentslm.time_series_datasets.ecg_qa.ECGQACoTQADataset import ECGQACoTQADataset
+from opentslm.time_series_datasets.util import extend_time_series_to_match_patch_size_and_aggregate
+from torch.utils.data import DataLoader
+from opentslm.model_config import PATCH_SIZE
+import torch
+
+# Model repository ID - change this to test different models
+REPO_ID = "OpenTSLM/llama-3.2-1b-ecg-sp"
+
+def main():
+    print("=" * 60)
+    print("ECG QA CoT Model Demo")
+    print("=" * 60)
+    
+    # Load model from HuggingFace
+    print(f"\n📥 Loading model from {REPO_ID}...")
+    enable_lora = False
+    if "-sp" in REPO_ID:
+        enable_lora = True
+    model = OpenTSLM.load_pretrained(REPO_ID, enable_lora=enable_lora, device="cuda" if torch.cuda.is_available() else "cpu")
+    
+    # Create dataset
+    print("\n📊 Loading ECG QA CoT test dataset...")
+    test_dataset = ECGQACoTQADataset("test", EOS_TOKEN=model.get_eos_token())
+    
+    # Create data loader
+    test_loader = DataLoader(
+        test_dataset,
+        shuffle=False,
+        batch_size=1,
+        collate_fn=lambda batch: extend_time_series_to_match_patch_size_and_aggregate(
+            batch, patch_size=PATCH_SIZE
+        ),
+    )
+    
+    print(f"\n🔍 Running inference on {len(test_dataset)} test samples...")
+    print("=" * 60)
+    
+    # Iterate over evaluation set
+    for i, batch in enumerate(test_loader):
+        # Generate predictions
+        predictions = model.generate(batch, max_new_tokens=500)
+        
+        # Print results
+        for sample, pred in zip(batch, predictions):
+            print(f"\n📝 Sample {i + 1}:")
+            if 'pre_prompt' in sample:
+                print(f"   Question: {sample['pre_prompt']}")
+            if 'template_id' in sample:
+                print(f"   Template ID: {sample['template_id']}")
+            if 'ecg_id' in sample:
+                print(f"   ECG ID: {sample['ecg_id']}")
+            print(f"   Gold Answer: {sample.get('answer', 'N/A')}")
+            print(f"   Model Output: {pred}")
+            print("-" * 60)
+        
+        # Limit to first 5 samples for demo
+        if i >= 4:
+            print("\n✅ Demo complete! (Showing first 5 samples)")
+            break
+
+if __name__ == "__main__":
+    main()
+

code/OpenTSLM/evaluation/baseline/README.md

@@ -0,0 +1,213 @@
+<!--
+SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+
+SPDX-License-Identifier: MIT
+-->
+
+# Time Series LLM Evaluation Framework
+
+This directory contains a modular evaluation framework for testing LLMs on time series datasets. The system is designed to be easily extensible to new datasets and models.
+
+## Overview
+
+The evaluation framework consists of:
+
+1. **`common_evaluator.py`** - Core evaluation logic that can be reused across datasets
+2. **`evaluate_tsqa.py`** - TSQA-specific evaluation
+3. **`evaluate_pamap.py`** - PAMAP-specific evaluation  
+4. **`evaluate_all.py`** - Combined evaluation across all datasets
+5. **`test_baseline.py`** - Original baseline test (legacy)
+
+## Quick Start
+
+### Running Individual Dataset Evaluations
+
+```bash
+# Evaluate on TSQA dataset
+python evaluate_tsqa.py
+
+# Evaluate on PAMAP datasets
+python evaluate_pamap.py
+
+# Evaluate on all datasets
+python evaluate_all.py
+```
+
+### Adding New Models
+
+To add new models, simply update the `model_names` list in any evaluation script:
+
+```python
+model_names = [
+    "meta-llama/Llama-3.2-1B",
+    "google/gemma-3n-e2b",
+    "google/gemma-3n-e2b-it",
+    "microsoft/DialoGPT-medium",
+    "gpt2",
+]
+```
+
+### Adding New Datasets
+
+To add a new dataset:
+
+1. Create a new evaluation file (e.g., `evaluate_newdataset.py`)
+2. Define an evaluation function that takes `(ground_truth, prediction)` and returns metrics
+3. Import your dataset class and evaluation function in `evaluate_all.py`
+
+Example evaluation function:
+
+```python
+def evaluate_newdataset(ground_truth: str, prediction: str) -> Dict[str, Any]:
+    """Evaluate predictions for new dataset."""
+    gt_clean = ground_truth.lower().strip()
+    pred_clean = prediction.lower().strip()
+    
+    exact_match = gt_clean == pred_clean
+    similarity = SequenceMatcher(None, gt_clean, pred_clean).ratio()
+    
+    return {
+        "exact_match": int(exact_match),
+        "similarity": similarity,
+        "ground_truth": gt_clean,
+        "prediction": pred_clean,
+    }
+```
+
+## Output Files
+
+The evaluation system generates several output files:
+
+1. **Individual Results**: `evaluation_results_{model}_{dataset}.json` - Detailed results for each model-dataset combination
+2. **Summary CSV**: `evaluation_results_{timestamp}.csv` - Pandas DataFrame with all results
+3. **Console Output**: Real-time progress and summary statistics
+
+## Metrics
+
+The framework calculates several metrics for each prediction:
+
+- **exact_match**: Binary indicator for exact string match
+- **partial_match**: Binary indicator for partial string match
+- **contains_answer**: Binary indicator if prediction contains ground truth
+- **similarity**: Character-level similarity score (0-1)
+- **has_reasoning**: For CoT datasets, indicates if prediction contains reasoning
+
+## Configuration
+
+### Sample Limits
+
+For faster testing, you can limit the number of samples:
+
+```python
+max_samples=50  # Process only first 50 samples
+max_samples=None  # Process all samples
+```
+
+### Model Parameters
+
+You can customize model parameters:
+
+```python
+results_df = evaluator.evaluate_multiple_models(
+    model_names=model_names,
+    dataset_classes=dataset_classes,
+    evaluation_functions=evaluation_functions,
+    max_samples=50,
+    temperature=0.1,  # Model temperature
+    max_new_tokens=100,  # Maximum tokens to generate
+)
+```
+
+## Dataset-Specific Considerations
+
+### TSQADataset
+- Extracts answers after "Answer:" in predictions
+- Uses built-in time series formatting from dataset class
+- Calculates similarity metrics
+
+### PAMAP2AccQADataset
+- Activity classification from accelerometer data
+- Standard exact/partial match metrics
+
+### PAMAP2CoTQADataset
+- Chain-of-thought reasoning evaluation
+- Additional "has_reasoning" metric
+- More complex answer extraction
+
+## Extending the Framework
+
+### Adding Custom Metrics
+
+To add custom metrics, modify the evaluation function:
+
+```python
+def evaluate_custom(ground_truth: str, prediction: str) -> Dict[str, Any]:
+    # ... existing metrics ...
+    
+    # Add custom metric
+    custom_metric = calculate_custom_metric(ground_truth, prediction)
+    
+    return {
+        # ... existing metrics ...
+        "custom_metric": custom_metric,
+    }
+```
+
+### Adding New Dataset Classes
+
+1. Ensure your dataset class inherits from `QADataset`
+2. Implement the required abstract methods
+3. Create an evaluation function
+4. Add to the evaluation scripts
+
+### Batch Processing
+
+For large-scale evaluations, you can modify the framework to process models in batches or use distributed processing.
+
+## Troubleshooting
+
+### Common Issues
+
+1. **Model Loading Errors**: Check if the model name is correct and accessible
+2. **Memory Issues**: Reduce `max_samples` or use smaller models
+3. **Import Errors**: Ensure all dataset classes are properly imported
+
+### Debug Mode
+
+To enable detailed debugging, modify the evaluation scripts to print more information:
+
+```python
+# In common_evaluator.py, set debug=True
+if idx < 10:  # Print more samples for debugging
+    print(f"Sample {idx}: {input_text[:200]}...")
+```
+
+## Performance Tips
+
+1. **Use GPU**: The framework automatically detects and uses CUDA/MPS if available
+2. **Batch Processing**: For multiple models, consider running them in parallel
+3. **Sample Limiting**: Use `max_samples` for quick testing before full evaluation
+4. **Model Caching**: Models are loaded once per evaluation run 
+
+## Using OpenAI Models (ChatGPT, GPT-4, etc.)
+
+You can now evaluate OpenAI models (e.g., ChatGPT, GPT-4) using the same evaluation scripts. To do so:
+
+1. **Set your OpenAI API key** (required):
+   
+   ```bash
+   export OPENAI_API_KEY=sk-...
+   ```
+
+2. **Run the evaluation script with an OpenAI model name prefixed by `openai-`**:
+   
+   ```bash
+   python evaluate_tsqa.py openai-gpt-4
+   python evaluate_pamap.py openai-gpt-3.5-turbo
+   ```
+   This will use the OpenAI API instead of a local HuggingFace model.
+
+3. **Notes:**
+   - The model name after `openai-` should match the OpenAI API model name (e.g., `gpt-4`, `gpt-3.5-turbo`).
+   - You can adjust `max_new_tokens` and other parameters as usual. 

code/OpenTSLM/evaluation/baseline/common_evaluator.py

@@ -0,0 +1,722 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+
+import json
+import os
+import io
+import re
+import sys
+import base64
+from typing import Type, Callable, Dict, List, Any, Optional
+
+import numpy as np
+import pandas as pd
+import torch
+from torch.utils.data import Dataset
+from tqdm import tqdm
+from transformers.pipelines import pipeline
+import matplotlib.pyplot as plt
+from time import sleep
+
+from opentslm.logger import get_logger
+
+from .openai_pipeline import OpenAIPipeline
+
+
+class CommonEvaluator:
+    """
+    A common evaluation framework for testing LLMs on time series datasets.
+    """
+
+    def __init__(self, device: Optional[str] = None):
+        """
+        Initialize the evaluator.
+
+        Args:
+            device: Device to use for inference ('cuda', 'mps', 'cpu', or None for auto)
+        """
+        self.device = device or self._get_best_device()
+
+        if self.device == "mps":
+            print(
+                "⚠️ Warning!! MPS is available but not recommended for evaluation. Many LLMs do not produce reasonable output!"
+            )
+            print(
+                "⚠️ Warning!! MPS is available but not recommended for evaluation. Many LLMs do not produce reasonable output!"
+            )
+            print("⚠️ Better use CPU or CUDA for evaluation.")
+            sleep(10)
+
+    def _get_best_device(self) -> str:
+        """Get the best available device."""
+        if torch.cuda.is_available():
+            return "cuda"
+        elif torch.backends.mps.is_available():
+            return "mps"
+        else:
+            return "cpu"
+
+    def load_model(self, model_name: str, **pipeline_kwargs) -> pipeline:
+        """
+        Load a model using transformers pipeline or OpenAI API.
+        """
+        self.current_model_name = (
+            model_name  # Track the current model name for formatter selection
+        )
+        if model_name.startswith("openai-"):
+            # Use OpenAI API
+            openai_model = model_name.replace("openai-", "")
+            return OpenAIPipeline(model_name=openai_model, **pipeline_kwargs)
+        print(f"Loading model: {model_name}")
+        print(f"Using device: {self.device}")
+        # Default pipeline arguments
+        default_kwargs = {
+            "task": "text-generation",
+            "device": self.device,
+            "temperature": 0.1,
+        }
+        default_kwargs.update(pipeline_kwargs)
+        pipe = pipeline(model=model_name, **default_kwargs)
+        print(f"Model loaded successfully: {model_name}")
+        return pipe
+
+    def load_dataset(
+        self,
+        dataset_class: Type[Dataset],
+        split: str = "test",
+        format_sample_str: bool = True,
+        max_samples: Optional[int] = None,
+        **dataset_kwargs,
+    ) -> Dataset:
+        """
+        Load a dataset with proper formatting.
+        """
+        print(f"Loading dataset: {dataset_class.__name__}")
+
+        # Import the gruver formatters
+        from .gruver_llmtime_tokenizer import gpt_formatter, llama_formatter
+
+        # Choose formatter based on model type
+        model_name = getattr(self, "current_model_name", None)
+        if model_name is None and "model_name" in dataset_kwargs:
+            model_name = dataset_kwargs["model_name"]
+        if model_name is not None:
+            if model_name.startswith("openai-") or "gpt" in model_name.lower():
+                formatter = gpt_formatter
+                print(f"Using GPT formatter for model: {model_name}")
+            elif "llama" in model_name.lower():
+                formatter = llama_formatter
+                print(f"Using Llama formatter for model: {model_name}")
+            else:
+                print(f"Defaulting to Llama formatter for model: {model_name}")
+
+                formatter = llama_formatter
+        else:
+            formatter = llama_formatter
+
+        # Default dataset arguments
+        default_kwargs = {
+            "split": split,
+            "EOS_TOKEN": "",
+            "format_sample_str": format_sample_str,
+            "time_series_format_function": formatter,
+        }
+
+        # Add max_samples if provided
+        if max_samples is not None:
+            default_kwargs["max_samples"] = max_samples
+
+        # Update with provided kwargs
+        default_kwargs.update(dataset_kwargs)
+
+        dataset = dataset_class(**default_kwargs)
+        print(f"Loaded {len(dataset)} {split} samples")
+
+        return dataset
+
+    def evaluate_model_on_dataset(
+        self,
+        model_name: str,
+        dataset_class: Type[Dataset],
+        evaluation_function: Callable[[str, str], Dict[str, Any]],
+        max_samples: Optional[int] = None,
+        use_plot: bool = False,
+        **pipeline_kwargs,
+    ) -> Dict[str, Any]:
+        """
+        Evaluate a model on a dataset using a custom evaluation function.
+
+        Args:
+            model_name: Name of the model to evaluate
+            dataset_class: Dataset class to use
+            evaluation_function: Function that takes (ground_truth, prediction) and returns metrics
+            max_samples: Maximum number of samples to evaluate (None for all)
+            **pipeline_kwargs: Additional arguments for model pipeline
+
+        Returns:
+            Dictionary containing evaluation results
+        """
+        print(
+            f"Starting evaluation with model {model_name} on dataset {dataset_class.__name__}"
+        )
+        print("=" * 60)
+
+        # Load model
+        pipe = self.load_model(model_name, **pipeline_kwargs)
+
+        # Load dataset
+        dataset = self.load_dataset(dataset_class, max_samples=max_samples)
+
+        # Check for existing results to resume from
+        existing_count = self._get_existing_results_count(
+            model_name, dataset_class.__name__
+        )
+        start_idx = existing_count
+
+        # Limit samples if specified
+        if max_samples is not None:
+            dataset_size = min(len(dataset), max_samples)
+            print(f"Processing samples {start_idx} to {dataset_size}...")
+        else:
+            dataset_size = len(dataset)
+            print(f"Processing samples {start_idx} to {dataset_size}...")
+
+        if start_idx >= dataset_size:
+            print(f"✅ All {dataset_size} samples already processed!")
+            return self._consolidate_jsonl_results(model_name, dataset_class.__name__)
+
+        # Initialize tracking
+        total_samples = dataset_size
+        successful_inferences = existing_count  # Start with existing count
+        all_metrics = []
+        results = []
+        first_error_printed = False  # Track if we've printed the first error
+
+        print("\nRunning inference...")
+        print("=" * 80)
+
+        # Get max_new_tokens for generation (default 1000)
+        max_new_tokens = pipeline_kwargs.pop("max_new_tokens", 1000)
+
+        # Load existing metrics if resuming
+        if start_idx > 0:
+            print(f"📂 Loading existing results from {start_idx} completed samples...")
+            jsonl_file = self._get_jsonl_file_path(model_name, dataset_class.__name__)
+            if os.path.exists(jsonl_file):
+                with open(jsonl_file, "r") as f:
+                    for line in f:
+                        if line.strip():
+                            result = json.loads(line.strip())
+                            all_metrics.append(result["metrics"])
+                            results.append(result)
+
+        # Process each sample starting from where we left off
+        for idx in tqdm(range(start_idx, dataset_size), desc="Processing samples"):
+            try:
+                sample = dataset[idx]
+                plot_data = None
+
+                # Clean up prompt for TSQADataset (if needed)
+                if use_plot and hasattr(sample, "get") and sample.get("prompt"):
+                    plot_data = self.get_plot_from_prompt(sample["prompt"])
+                    pattern = r"The following is the accelerometer data on the [xyz]-axis\n([\-0-9, ]+)"
+                    sample["prompt"] = re.sub(pattern, "", sample["prompt"])
+
+                # Clean up prompt for TSQADataset (if needed)
+                if hasattr(sample, "get") and sample.get("prompt"):
+                    pattern = r"This is the time series, it has mean (-?\d+\.\d{4}) and std (-?\d+\.\d{4})\."
+                    replacement = "This is the time series:"
+                    sample["prompt"] = re.sub(pattern, replacement, sample["prompt"])
+
+                # Create input text
+                input_text = sample["prompt"]
+                target_answer = sample["answer"]
+
+                # Generate prediction
+                outputs = pipe(
+                    input_text,
+                    max_new_tokens=max_new_tokens,
+                    return_full_text=False,
+                    plot_data=plot_data,
+                )
+
+                # Extract generated text
+                if outputs and len(outputs) > 0:
+                    generated_text = outputs[0]["generated_text"].strip()
+                    successful_inferences += 1
+
+                    # Evaluate using custom function (optionally with sample)
+                    try:
+                        import inspect
+
+                        sig = inspect.signature(evaluation_function)
+                        if len(sig.parameters) >= 3:
+                            metrics = evaluation_function(
+                                target_answer, generated_text, sample
+                            )
+                        else:
+                            metrics = evaluation_function(target_answer, generated_text)
+                    except Exception:
+                        # Fallback to 2-arg call
+                        metrics = evaluation_function(target_answer, generated_text)
+                    all_metrics.append(metrics)
+
+                    # Store detailed results
+                    result = {
+                        "sample_idx": idx,
+                        "input_text": input_text,
+                        "target_answer": target_answer,
+                        "generated_answer": generated_text,
+                        "metrics": metrics,
+                    }
+                    # Include template_id if present in sample for downstream analysis
+                    if isinstance(sample, dict) and "template_id" in sample:
+                        result["template_id"] = sample["template_id"]
+                    results.append(result)
+
+                    # Save individual result immediately to prevent data loss
+                    self._save_individual_result(
+                        result, model_name, dataset_class.__name__
+                    )
+
+                    # Print progress for first few samples
+                    if idx < 10:
+                        print(f"\nSAMPLE {idx + 1}:")
+                        print(f"PROMPT: {input_text}...")
+                        print(f"TARGET: {target_answer}")
+                        print(f"PREDICTION: {generated_text}")
+                        print(f"METRICS: {metrics}")
+                        print("=" * 80)
+
+                    # Print first error for debugging
+                    if not first_error_printed and metrics.get("accuracy", 1) == 0:
+                        print(f"\n❌ FIRST ERROR (Sample {idx + 1}):")
+                        print(f"TARGET: {target_answer}")
+                        print(f"PREDICTION: {generated_text}")
+                        print("=" * 80)
+                        first_error_printed = True
+            except Exception as e:
+                print(f"Error processing sample {idx}: {e}")
+                continue
+
+        # Calculate aggregate metrics
+        if successful_inferences > 0:
+            # Aggregate metrics across all samples
+            aggregate_metrics = self._aggregate_metrics(all_metrics)
+
+            # Calculate success rate
+            success_rate = successful_inferences / total_samples
+
+            # Prepare final results
+            final_results = {
+                "model_name": model_name,
+                "dataset_name": dataset_class.__name__,
+                "total_samples": total_samples,
+                "successful_inferences": successful_inferences,
+                "success_rate": success_rate,
+                "metrics": aggregate_metrics,
+                "detailed_results": results,
+            }
+
+            # Print summary
+            self._print_summary(final_results)
+
+            # Consolidate JSONL results into final JSON file
+            consolidated_file = self._consolidate_jsonl_results(
+                model_name, dataset_class.__name__
+            )
+            if consolidated_file:
+                # Update the consolidated file with correct total_samples
+                with open(consolidated_file, "r") as f:
+                    consolidated_data = json.load(f)
+                consolidated_data["total_samples"] = total_samples
+                consolidated_data["success_rate"] = success_rate
+                with open(consolidated_file, "w") as f:
+                    json.dump(consolidated_data, f, indent=2)
+
+            return final_results
+        else:
+            print("❌ No successful inferences completed!")
+            return {
+                "model_name": model_name,
+                "dataset_name": dataset_class.__name__,
+                "total_samples": total_samples,
+                "successful_inferences": 0,
+                "success_rate": 0.0,
+                "metrics": {},
+                "detailed_results": [],
+            }
+
+    def _aggregate_metrics(self, metrics_list: List[Dict[str, Any]]) -> Dict[str, Any]:
+        """
+        Aggregate metrics across all samples.
+
+        Args:
+            metrics_list: List of metric dictionaries
+
+        Returns:
+            Aggregated metrics
+        """
+        if not metrics_list:
+            return {}
+
+        # Get all unique metric keys
+        all_keys = set()
+        for metrics in metrics_list:
+            all_keys.update(metrics.keys())
+
+        aggregated = {}
+        for key in all_keys:
+            values = [metrics.get(key, 0) for metrics in metrics_list]
+            if all(isinstance(v, (int, float)) for v in values):
+                # Calculate overall accuracy/percentage
+                accuracy = np.mean(values) * 100
+                aggregated[key] = accuracy
+            else:
+                # For non-numeric metrics, just count occurrences
+                aggregated[key] = {
+                    "values": values,
+                    "count": len(values),
+                }
+
+        return aggregated
+
+    def _print_summary(self, results: Dict[str, Any]):
+        """Print evaluation summary."""
+        print("\n" + "=" * 80)
+        print("EVALUATION RESULTS")
+        print("=" * 80)
+        print(f"Model: {results['model_name']}")
+        print(f"Dataset: {results['dataset_name']}")
+        print(f"Total samples processed: {results['total_samples']}")
+        print(f"Successful inferences: {results['successful_inferences']}")
+        print(f"Success rate: {results['success_rate']:.2%}")
+
+        if results["metrics"]:
+            print("\nAggregated Metrics:")
+            for metric_name, metric_values in results["metrics"].items():
+                if isinstance(metric_values, (int, float)):
+                    print(f"  {metric_name}: {metric_values:.1f}%")
+                else:
+                    print(f"  {metric_name}: {metric_values}")
+
+    def _save_results(self, results: Dict[str, Any]):
+        """Save detailed results to file."""
+        import os
+
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        detailed_dir = os.path.join(
+            current_dir, "..", "results", "baseline", "detailed"
+        )
+        os.makedirs(detailed_dir, exist_ok=True)
+        normalized_model_id = re.sub(r"[^a-z0-9]", "-", results["model_name"].lower())
+        normalized_dataset_name = re.sub(
+            r"[^a-z0-9]", "-", results["dataset_name"].lower()
+        )
+        results_file = os.path.join(
+            detailed_dir,
+            f"evaluation_results_{normalized_model_id}_{normalized_dataset_name}.json",
+        )
+
+        with open(results_file, "w") as f:
+            json.dump(results, f, indent=2)
+
+        print(f"\nDetailed results saved to: {results_file}")
+
+    def _save_individual_result(
+        self, result: Dict[str, Any], model_name: str, dataset_name: str
+    ):
+        """Save individual result incrementally to prevent data loss."""
+        import os
+
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        detailed_dir = os.path.join(
+            current_dir, "..", "results", "baseline", "detailed"
+        )
+        os.makedirs(detailed_dir, exist_ok=True)
+        normalized_model_id = re.sub(r"[^a-z0-9]", "-", model_name.lower())
+        normalized_dataset_name = re.sub(r"[^a-z0-9]", "-", dataset_name.lower())
+        results_file = os.path.join(
+            detailed_dir,
+            f"evaluation_results_{normalized_model_id}_{normalized_dataset_name}.jsonl",
+        )
+
+        # Append individual result as JSONL
+        with open(results_file, "a") as f:
+            json.dump(result, f)
+            f.write("\n")
+
+    def _consolidate_jsonl_results(self, model_name: str, dataset_name: str) -> str:
+        """Consolidate JSONL results into final JSON file."""
+        import os
+
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        detailed_dir = os.path.join(
+            current_dir, "..", "results", "baseline", "detailed"
+        )
+        normalized_model_id = re.sub(r"[^a-z0-9]", "-", model_name.lower())
+        normalized_dataset_name = re.sub(r"[^a-z0-9]", "-", dataset_name.lower())
+
+        jsonl_file = os.path.join(
+            detailed_dir,
+            f"evaluation_results_{normalized_model_id}_{normalized_dataset_name}.jsonl",
+        )
+        json_file = os.path.join(
+            detailed_dir,
+            f"evaluation_results_{normalized_model_id}_{normalized_dataset_name}.json",
+        )
+
+        # Read all JSONL results
+        individual_results = []
+        if os.path.exists(jsonl_file):
+            with open(jsonl_file, "r") as f:
+                for line in f:
+                    if line.strip():
+                        individual_results.append(json.loads(line.strip()))
+
+        # Create consolidated results structure
+        if individual_results:
+            # Calculate aggregate metrics
+            all_metrics = [result["metrics"] for result in individual_results]
+            aggregate_metrics = self._aggregate_metrics(all_metrics)
+
+            # Calculate success rate
+            successful_inferences = len(individual_results)
+            total_samples = len(individual_results)  # This will be updated by caller
+
+            consolidated_results = {
+                "model_name": model_name,
+                "dataset_name": dataset_name,
+                "total_samples": total_samples,
+                "successful_inferences": successful_inferences,
+                "success_rate": (
+                    successful_inferences / total_samples if total_samples > 0 else 0.0
+                ),
+                "metrics": aggregate_metrics,
+                "detailed_results": individual_results,
+            }
+
+            # Save consolidated results
+            with open(json_file, "w") as f:
+                json.dump(consolidated_results, f, indent=2)
+
+            print(f"\nConsolidated results saved to: {json_file}")
+            return json_file
+
+        return None
+
+    def _get_existing_results_count(self, model_name: str, dataset_name: str) -> int:
+        """Get count of existing results from JSONL file for resuming interrupted evaluations."""
+        import os
+
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        detailed_dir = os.path.join(
+            current_dir, "..", "results", "baseline", "detailed"
+        )
+        normalized_model_id = re.sub(r"[^a-z0-9]", "-", model_name.lower())
+        normalized_dataset_name = re.sub(r"[^a-z0-9]", "-", dataset_name.lower())
+        jsonl_file = os.path.join(
+            detailed_dir,
+            f"evaluation_results_{normalized_model_id}_{normalized_dataset_name}.jsonl",
+        )
+
+        if os.path.exists(jsonl_file):
+                    if line.strip():
+                        count += 1
+            return count
+        return 0
+
+    def _get_jsonl_file_path(self, model_name: str, dataset_name: str) -> str:
+        """Get the JSONL file path for a model-dataset combination."""
+        import os
+
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        detailed_dir = os.path.join(
+            current_dir, "..", "results", "baseline", "detailed"
+        )
+        normalized_model_id = re.sub(r"[^a-z0-9]", "-", model_name.lower())
+        normalized_dataset_name = re.sub(r"[^a-z0-9]", "-", dataset_name.lower())
+        return os.path.join(
+            detailed_dir,
+            f"evaluation_results_{normalized_model_id}_{normalized_dataset_name}.jsonl",
+        )
+
+    def evaluate_multiple_models(
+        self,
+        model_names: List[str],
+        dataset_classes: List[Type[Dataset]],
+        evaluation_functions: Dict[str, Callable[[str, str], Dict[str, Any]]],
+        max_samples: Optional[int] = None,
+        **pipeline_kwargs,
+    ) -> pd.DataFrame:
+        """
+        Evaluate multiple models on multiple datasets.
+
+        Args:
+            model_names: List of model names to evaluate
+            dataset_classes: List of dataset classes to evaluate on
+            evaluation_functions: Dictionary mapping dataset class names to evaluation functions
+            max_samples: Maximum number of samples per evaluation
+            **pipeline_kwargs: Additional arguments for model pipeline
+
+        Returns:
+            DataFrame with results for all model-dataset combinations
+        """
+        all_results = []
+
+        # Generate filename once at the beginning
+        import os
+
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        results_dir = os.path.join(current_dir, "..", "results", "baseline")
+        os.makedirs(results_dir, exist_ok=True)
+        df_filename = os.path.join(results_dir, "evaluation_results.csv")
+        print(f"Results will be saved to: {df_filename}")
+
+        # Load existing results if file exists
+        existing_df = None
+        if os.path.exists(df_filename):
+            try:
+                existing_df = pd.read_csv(df_filename)
+                print(f"Found existing results file with {len(existing_df)} entries")
+            except Exception as e:
+                print(f"Warning: Could not read existing results file: {e}")
+
+        for model_name in model_names:
+            for dataset_class in dataset_classes:
+                dataset_name = dataset_class.__name__
+
+                if dataset_name not in evaluation_functions:
+                    print(f"Warning: No evaluation function found for {dataset_name}")
+                    continue
+
+                # Check if this model-dataset combination already exists in results
+                if existing_df is not None:
+                    existing_result = existing_df[
+                        (existing_df["model"] == model_name)
+                        & (existing_df["dataset"] == dataset_name)
+                    ]
+                    if not existing_result.empty:
+                        print(
+                            f"⏭️  Skipping {model_name} on {dataset_name} (already evaluated)"
+                        )
+                        continue
+
+                evaluation_function = evaluation_functions[dataset_name]
+
+                print(f"\n{'=' * 80}")
+                print(f"Evaluating {model_name} on {dataset_name}")
+                print(f"{'=' * 80}")
+
+                try:
+                    results = self.evaluate_model_on_dataset(
+                        model_name=model_name,
+                        dataset_class=dataset_class,
+                        evaluation_function=evaluation_function,
+                        max_samples=max_samples,
+                        use_plot=False,
+                        **pipeline_kwargs,
+                    )
+
+                    # Extract key metrics for DataFrame
+                    row = {
+                        "model": model_name,
+                        "dataset": dataset_name,
+                        "total_samples": results["total_samples"],
+                        "successful_inferences": results["successful_inferences"],
+                        "success_rate": results["success_rate"],
+                    }
+
+                    # Add specific metrics
+                    if results["metrics"]:
+                        for metric_name, metric_values in results["metrics"].items():
+                            if isinstance(metric_values, (int, float)):
+                                row[metric_name] = metric_values
+                            else:
+                                row[metric_name] = str(metric_values)
+
+                    all_results.append(row)
+
+                    # Combine with existing results and save
+                    current_df = pd.DataFrame(all_results)
+                    if existing_df is not None:
+                        # Append new results
+                        final_df = pd.concat(
+                            [existing_df, current_df], ignore_index=True
+                        )
+                    else:
+                        final_df = current_df
+
+                    final_df.to_csv(df_filename, index=False)
+                    print(f"✅ Results updated: {df_filename}")
+
+                except Exception as e:
+                    print(f"Error evaluating {model_name} on {dataset_name}: {e}")
+                    all_results.append(
+                        {
+                            "model": model_name,
+                            "dataset": dataset_name,
+                            "status": "Failed",
+                        }
+                    )
+
+                    # Save DataFrame even after errors
+                    current_df = pd.DataFrame(all_results)
+                    if existing_df is not None:
+                        final_df = pd.concat(
+                            [existing_df, current_df], ignore_index=True
+                        )
+                    else:
+                        final_df = current_df
+                    final_df.to_csv(df_filename, index=False)
+                    print(f"⚠️  Results updated (with error): {df_filename}")
+
+        print(f"\nFinal results saved to: {df_filename}")
+        return final_df
+
+    def get_plot_from_prompt(self, prompt: str):
+        """
+        Parse time series data from the prompt and return a base64 image.
+        """
+        # Parse the time series data from the prompt
+        time_series_data = []
+
+        # Extract data for each axis using regex
+        axes = ["x-axis", "y-axis", "z-axis"]
+        for axis in axes:
+            pattern = f"accelerometer data on the {axis}\\n([\\-0-9, ]+)"
+            match = re.search(pattern, prompt.lower())
+            if match:
+                # Extract the data and convert to a list of integers
+                data_str = match.group(1).strip()
+                data_str = data_str.replace(" ", "")
+                data = [int(val.strip()) for val in data_str.split(",") if val.strip()]
+                time_series_data.append(data)
+
+        # Create the plot
+        num_series = len(time_series_data)
+        fig, axes = plt.subplots(
+            num_series, 1, figsize=(10, 4 * num_series), sharex=True
+        )
+        # If there's only one series, axes won't be an array
+        if num_series == 1:
+            axes = [axes]
+
+        # Plot each time series in its own subplot
+        axis_names = {0: "X-axis", 1: "Y-axis", 2: "Z-axis"}
+        for i, series in enumerate(time_series_data):
+            axes[i].plot(series, marker="o", linestyle="-", markersize=0)
+            axes[i].grid(True, alpha=0.3)
+            axes[i].set_title(f"Accelerometer - {axis_names.get(i)}")
+
+        plt.tight_layout()
+
+        # Convert plot to base64 image
+        img_buffer = io.BytesIO()
+        plt.savefig(img_buffer, format="png", bbox_inches="tight", dpi=100)
+        plt.close()
+        img_buffer.seek(0)
+        image_data = base64.b64encode(img_buffer.getvalue()).decode("utf-8")
+        return image_data

code/OpenTSLM/evaluation/baseline/common_evaluator_plot.py

@@ -0,0 +1,459 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+
+import io
+import base64
+from typing import Type, Callable, Dict, List, Any, Optional
+
+import torch
+from torch.utils.data import Dataset
+from tqdm import tqdm
+from transformers.pipelines import pipeline
+import matplotlib.pyplot as plt
+from time import sleep
+from PIL import Image
+import pandas as pd
+
+
+from openai_pipeline import OpenAIPipeline
+from common_evaluator import CommonEvaluator
+
+
+class CommonEvaluatorPlot(CommonEvaluator):
+    """
+    A common evaluation framework for testing LLMs on time series datasets with plot generation.
+    """
+
+    def load_model(self, model_name: str, **pipeline_kwargs) -> pipeline:
+        """
+        Load a model using transformers pipeline or OpenAI API.
+        """
+        self.current_model_name = (
+            model_name  # Track the current model name for formatter selection
+        )
+        if model_name.startswith("openai-"):
+            # Use OpenAI API
+            openai_model = model_name.replace("openai-", "")
+            return OpenAIPipeline(model_name=openai_model, **pipeline_kwargs)
+        print(f"Loading model: {model_name}")
+        print(f"Using device: {self.device}")
+        # Default pipeline arguments
+        default_kwargs = {
+            "task": "image-text-to-text",
+            "device": self.device,
+            "temperature": 0.1,
+        }
+        default_kwargs.update(pipeline_kwargs)
+        pipe = pipeline(model=model_name, **default_kwargs)
+        print(f"Model loaded successfully: {model_name}")
+        return pipe
+
+    def load_dataset(
+        self,
+        dataset_class: Type[Dataset],
+        split: str = "test",
+        format_sample_str: bool = True,
+        max_samples: Optional[int] = None,
+        **dataset_kwargs,
+    ) -> Dataset:
+        """
+        Load a dataset with proper formatting.
+        """
+        print(f"Loading dataset: {dataset_class.__name__}")
+        formatter = None
+
+        # Default dataset arguments
+        default_kwargs = {
+            "split": split,
+            "EOS_TOKEN": "",
+            "format_sample_str": format_sample_str,
+            "time_series_format_function": formatter,
+        }
+        # Add max_samples if provided
+        if max_samples is not None:
+            default_kwargs["max_samples"] = max_samples
+
+        # Update with provided kwargs
+        default_kwargs.update(dataset_kwargs)
+
+        dataset = dataset_class(**default_kwargs)
+        print(f"Loaded {len(dataset)} {split} samples")
+
+        return dataset
+
+    def evaluate_multiple_models(
+        self,
+        model_names: List[str],
+        dataset_classes: List[Type[Dataset]],
+        evaluation_functions: Dict[str, Callable[[str, str], Dict[str, Any]]],
+        plot_functions: Optional[Dict[str, Callable[[Any], str]]] = None,
+        max_samples: Optional[int] = None,
+        **pipeline_kwargs,
+    ) -> pd.DataFrame:
+        """
+        Evaluate multiple models on multiple datasets with plotting support.
+
+        Args:
+            model_names: List of model names to evaluate
+            dataset_classes: List of dataset classes to evaluate on
+            evaluation_functions: Dictionary mapping dataset class names to evaluation functions
+            plot_functions: Optional dict mapping dataset class names to plot functions
+            max_samples: Maximum number of samples per evaluation
+            **pipeline_kwargs: Additional arguments for model pipeline
+        Returns:
+            DataFrame with results for all model-dataset combinations
+        """
+        all_results = []
+
+        # Generate filename once at the beginning
+        import os
+
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        results_dir = os.path.join(current_dir, "..", "results", "baseline")
+        os.makedirs(results_dir, exist_ok=True)
+        df_filename = os.path.join(results_dir, "evaluation_results.csv")
+        print(f"Results will be saved to: {df_filename}")
+        # Load existing results if file exists
+        existing_df = None
+        if os.path.exists(df_filename):
+            try:
+                existing_df = pd.read_csv(df_filename)
+                print(f"Found existing results file with {len(existing_df)} entries")
+            except Exception as e:
+                print(f"Warning: Could not read existing results file: {e}")
+
+        for model_name in model_names:
+            for dataset_class in dataset_classes:
+                dataset_name = dataset_class.__name__
+
+                if dataset_name not in evaluation_functions:
+                    print(f"Warning: No evaluation function found for {dataset_name}")
+                    continue
+
+                # Check if this model-dataset combination already exists in results
+                if existing_df is not None:
+                    existing_result = existing_df[
+                        (existing_df["model"] == model_name)
+                        & (existing_df["dataset"] == dataset_name)
+                    ]
+                    if not existing_result.empty:
+                        print(
+                            f"⏭️  Skipping {model_name} on {dataset_name} (already evaluated)"
+                        )
+                        continue
+
+                evaluation_function = evaluation_functions[dataset_name]
+                plot_fn = None
+                if plot_functions is not None and dataset_name in plot_functions:
+                    plot_fn = plot_functions[dataset_name]
+                else:
+                    print(f"Warning: No plot function found for {dataset_name}")
+
+                print(f"\n{'=' * 80}")
+                print(f"Evaluating {model_name} on {dataset_name}")
+                print(f"{'=' * 80}")
+
+                try:
+                    results = self.evaluate_model_on_dataset(
+                        model_name=model_name,
+                        dataset_class=dataset_class,
+                        evaluation_function=evaluation_function,
+                        plot_function=plot_fn,
+                        max_samples=max_samples,
+                        **pipeline_kwargs,
+                    )
+
+                    # Extract key metrics for DataFrame
+                    row = {
+                        "model": model_name,
+                        "dataset": dataset_name,
+                        "total_samples": results["total_samples"],
+                        "successful_inferences": results["successful_inferences"],
+                        "success_rate": results["success_rate"],
+                    }
+                    # Add specific metrics
+                    if results["metrics"]:
+                        for metric_name, metric_values in results["metrics"].items():
+                            if isinstance(metric_values, (int, float)):
+                                row[metric_name] = metric_values
+                            else:
+                                row[metric_name] = str(metric_values)
+
+                    all_results.append(row)
+
+                    # Combine with existing results and save
+                    current_df = pd.DataFrame(all_results)
+                    if existing_df is not None:
+                        # Append new results
+                        final_df = pd.concat(
+                            [existing_df, current_df], ignore_index=True
+                        )
+                    else:
+                        final_df = current_df
+
+                    final_df.to_csv(df_filename, index=False)
+                    print(f"✅ Results updated: {df_filename}")
+
+                except Exception as e:
+                    print(f"Error evaluating {model_name} on {dataset_name}: {e}")
+                    all_results.append(
+                        {
+                            "model": model_name,
+                            "dataset": dataset_name,
+                            "status": "Failed",
+                        }
+                    )
+
+                    # Save DataFrame even after errors
+                    current_df = pd.DataFrame(all_results)
+                    if existing_df is not None:
+                        final_df = pd.concat(
+                            [existing_df, current_df], ignore_index=True
+                        )
+                    else:
+                        final_df = current_df
+                    final_df.to_csv(df_filename, index=False)
+                    print(f"⚠️  Results updated (with error): {df_filename}")
+        print(f"\nFinal results saved to: {df_filename}")
+        return final_df
+
+    def evaluate_model_on_dataset(
+        self,
+        model_name: str,
+        dataset_class: Type[Dataset],
+        evaluation_function: Callable[[str, str], Dict[str, Any]],
+        plot_function: Optional[Callable[[Any], str]] = None,
+        max_samples: Optional[int] = None,
+        **pipeline_kwargs,
+    ) -> Dict[str, Any]:
+        """
+        Evaluate a model on a dataset using a custom evaluation function.
+
+        Args:
+            model_name: Name of the model to evaluate
+            dataset_class: Dataset class to use
+            evaluation_function: Function that takes (ground_truth, prediction) and returns metrics
+            max_samples: Maximum number of samples to evaluate (None for all)
+            **pipeline_kwargs: Additional arguments for model pipeline
+
+        Returns:
+            Dictionary containing evaluation results
+        """
+        print(
+            f"Starting evaluation with model {model_name} on dataset {dataset_class.__name__}"
+        )
+        print("=" * 60)
+
+        # Load model
+        pipe = self.load_model(model_name, **pipeline_kwargs)
+
+        # Load dataset
+        dataset = self.load_dataset(
+            dataset_class, format_sample_str=False, max_samples=max_samples
+        )
+
+        # Limit samples if specified
+        if max_samples is not None:
+            dataset_size = min(len(dataset), max_samples)
+            print(f"Processing first {dataset_size} samples...")
+        else:
+            dataset_size = len(dataset)
+            print(f"Processing all {dataset_size} samples...")
+        # Initialize tracking
+        total_samples = 0
+        successful_inferences = 0
+        all_metrics = []
+        results = []
+        first_error_printed = False  # Track if we've printed the first error
+
+        print("\nRunning inference...")
+        print("=" * 80)
+
+        # Get max_new_tokens for generation (default 1000)
+        max_new_tokens = pipeline_kwargs.pop("max_new_tokens", 1000)
+
+        for idx in tqdm(range(dataset_size), desc="Processing samples"):
+            try:
+                sample = dataset[idx]
+                plot_data = None
+
+                if isinstance(sample, dict) and "time_series" in sample:
+                    plot_data = plot_function(sample["time_series"])
+                else:
+                    raise ValueError(
+                        f"Sample {sample} does not contain 'time_series' key"
+                    )
+
+                target_answer = sample["answer"]
+                input_text = sample["pre_prompt"] + sample["post_prompt"]
+
+                # Generate prediction
+                if isinstance(pipe, OpenAIPipeline):
+                    outputs = pipe(
+                        input_text,
+                        max_new_tokens=max_new_tokens,
+                        return_full_text=False,
+                        plot_data=plot_data,
+                    )
+                else:  # For Hugging Face pipelines, convert plot_data (base64) to PIL and pass via images
+                    try:
+                        img_bytes = base64.b64decode(plot_data)
+                        img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+
+                        # Check if this is a pretrained model (pt) vs instruction-tuned (it)
+                        if "pt" in model_name.lower() and "gemma" in model_name.lower():
+                            # For pretrained Gemma models, use model and processor directly
+                            # to avoid pipeline issues with chat templates
+
+                            try:
+                                outputs = pipe(
+                                    text=f"{sample['pre_prompt']} <start_of_image> {sample['post_prompt']}",
+                                    images=img,
+                                    max_new_tokens=max_new_tokens,
+                                    return_full_text=False,
+                                )
+                            except Exception as e:
+                                raise RuntimeError(
+                                    f"Failed to call pretrained pipeline: {e}"
+                                )
+                        else:
+                            # For instruction-tuned models, use chat template format
+                            messages = [
+                                {
+                                    "role": "user",
+                                    "content": [
+                                        {"type": "image", "image": img},
+                                        {"type": "text", "text": input_text},
+                                    ],
+                                }
+                            ]
+                            outputs = pipe(
+                                text=messages,
+                                max_new_tokens=max_new_tokens,
+                                return_full_text=False,
+                            )
+                    except Exception as e:
+                        raise RuntimeError(f"Failed to decode plot image: {e}")
+
+                # Extract generated text
+                if outputs and len(outputs) > 0:
+                    generated_text = outputs[0]["generated_text"].strip()
+                    successful_inferences += 1
+
+                    # Evaluate using custom function (optionally with sample)
+                    try:
+                        import inspect
+
+                        sig = inspect.signature(evaluation_function)
+                        if len(sig.parameters) >= 3:
+                            metrics = evaluation_function(
+                                target_answer, generated_text, sample
+                            )
+                        else:
+                            metrics = evaluation_function(target_answer, generated_text)
+                    except Exception:
+                        # Fallback to 2-arg call
+                        metrics = evaluation_function(target_answer, generated_text)
+                    all_metrics.append(metrics)
+                    # Store detailed results
+                    result = {
+                        "sample_idx": idx,
+                        "input_text": input_text,
+                        "target_answer": target_answer,
+                        "generated_answer": generated_text,
+                        "metrics": metrics,
+                    }
+                    results.append(result)
+                    # Print progress for first few samples
+                    if idx < 10:
+                        print(f"\nSAMPLE {idx + 1}:")
+                        print(f"PROMPT: {input_text}...")
+                        print(f"TARGET: {target_answer}")
+                        print(f"PREDICTION: {generated_text}")
+                        print(f"METRICS: {metrics}")
+                        print("=" * 80)
+
+                    # Print first error for debugging
+                    if not first_error_printed and metrics.get("accuracy", 1) == 0:
+                        print(f"\n❌ FIRST ERROR (Sample {idx + 1}):")
+                        print(f"TARGET: {target_answer}")
+                        print(f"PREDICTION: {generated_text}")
+                        print("=" * 80)
+                        first_error_printed = True
+                else:
+                    raise ValueError(f"Unexpectedly found empty outputs")
+
+                total_samples += 1
+
+                # Save results every 100 samples
+                if total_samples % 50 == 0:
+                    # Calculate current aggregate metrics
+                    current_aggregate_metrics = (
+                        self._aggregate_metrics(all_metrics) if all_metrics else {}
+                    )
+                    current_success_rate = (
+                        successful_inferences / total_samples
+                        if total_samples > 0
+                        else 0.0
+                    )
+
+                    # Prepare intermediate results
+                    intermediate_results = {
+                        "model_name": model_name,
+                        "dataset_name": dataset_class.__name__,
+                        "total_samples": total_samples,
+                        "successful_inferences": successful_inferences,
+                        "success_rate": current_success_rate,
+                        "metrics": current_aggregate_metrics,
+                        "detailed_results": results,
+                    }
+
+                    # Save intermediate results
+                    self._save_results(intermediate_results)
+                    print(
+                        f"💾 Intermediate results saved after {total_samples} samples"
+                    )
+
+            except Exception as e:
+                print(f"Error processing sample {idx}: {e}")
+                total_samples += 1
+                continue
+        # Calculate aggregate metrics
+        if successful_inferences > 0:
+            # Aggregate metrics across all samples
+            aggregate_metrics = self._aggregate_metrics(all_metrics)
+
+            # Calculate success rate
+            success_rate = successful_inferences / total_samples
+
+            # Prepare final results
+            final_results = {
+                "model_name": model_name,
+                "dataset_name": dataset_class.__name__,
+                "total_samples": total_samples,
+                "successful_inferences": successful_inferences,
+                "success_rate": success_rate,
+                "metrics": aggregate_metrics,
+                "detailed_results": results,
+            }
+
+            # Print summary
+            self._print_summary(final_results)
+
+            # Save results
+            self._save_results(final_results)
+
+            return final_results
+        else:
+            print("❌ No successful inferences completed!")
+            return {
+                "model_name": model_name,
+                "dataset_name": dataset_class.__name__,
+                "total_samples": total_samples,
+                "successful_inferences": 0,
+                "success_rate": 0.0,
+                "metrics": {},
+                "detailed_results": [],
+            }

code/OpenTSLM/evaluation/baseline/eval.sh

@@ -0,0 +1,38 @@
+#!/usr/bin/env bash
+
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+
+set -euo pipefail
+
+# List of models to evaluate
+models=(
+    "google/gemma-3-270m"
+    "google/gemma-3-1b-pt"
+    "meta-llama/Llama-3.2-1B"
+    "meta-llama/Llama-3.2-3B"
+    
+)
+
+#echo "Starting evaluation of ${#models[@]} models on TSQA dataset..."
+#for model in "${models[@]}"; do
+#    echo "Evaluating $model on TSQA..."
+#    python evaluate_tsqa.py "$model"
+#done
+
+
+echo "Starting evaluation of ${#models[@]} models on HAR dataset..."
+for model in "${models[@]}"; do
+    echo "Evaluating $model on PAMAP..."
+    python evaluate_har.py "$model"
+done
+
+#echo "Starting evaluation of ${#models[@]} models on SleepEDF CoT dataset..."
+#for model in "${models[@]}"; do
+#    echo "Evaluating $model on SleepEDF CoT..."
+#    python evaluate_sleep_cot.py "$model"
+#done
+
+echo "All evaluations completed!"

code/OpenTSLM/evaluation/baseline/eval_remaining.sh

@@ -0,0 +1,18 @@
+#!/usr/bin/env bash
+
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+
+python evaluate_tsqa.py "google/gemma-3-270m"
+python evaluate_sleep_cot.py "google/gemma-3-270m"
+
+python evaluate_tsqa.py "google/gemma-3-270m"
+python evaluate_sleep_cot.py "google/gemma-3-270m"
+python evaluate_har.py "google/gemma-3-1b-pt"
+
+python evaluate_ecg_qa.py "google/gemma-3-270m"
+python evaluate_ecg_qa.py "google/gemma-3-1b-pt"
+python evaluate_ecg_qa.py "meta-llama/Llama-3.2-1B"
+python evaluate_ecg_qa.py "meta-llama/Llama-3.2-3B"

code/OpenTSLM/evaluation/baseline/evaluate_all.py

@@ -0,0 +1,92 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+
+from typing import Dict, Any, Callable
+from common_evaluator import CommonEvaluator
+from opentslm.time_series_datasets.TSQADataset import TSQADataset
+from opentslm.time_series_datasets.pamap2.PAMAP2AccQADataset import PAMAP2AccQADataset
+from opentslm.time_series_datasets.pamap2.PAMAP2CoTQADataset import PAMAP2CoTQADataset
+
+# Import evaluation functions
+from evaluate_tsqa import evaluate_tsqa
+from evaluate_pamap import evaluate_pamap_acc, evaluate_pamap_cot
+
+
+def main():
+    """Main function to run comprehensive evaluation across all datasets."""
+    
+    # Define models to evaluate
+    model_names = [
+        "meta-llama/Llama-3.2-1B",
+        # Add more models as needed
+        # "google/gemma-3n-e2b",
+        # "google/gemma-3n-e2b-it",
+        # "microsoft/DialoGPT-medium",
+        # "gpt2",
+    ]
+    
+    # Define datasets to evaluate on
+    dataset_classes = [
+        TSQADataset,
+        PAMAP2AccQADataset,
+        PAMAP2CoTQADataset,
+    ]
+    
+    # Define evaluation functions
+    evaluation_functions = {
+        "TSQADataset": evaluate_tsqa,
+        "PAMAP2AccQADataset": evaluate_pamap_acc,
+        "PAMAP2CoTQADataset": evaluate_pamap_cot,
+    }
+    
+    # Initialize evaluator
+    evaluator = CommonEvaluator()
+    
+    # Run comprehensive evaluation
+    results_df = evaluator.evaluate_multiple_models(
+        model_names=model_names,
+        dataset_classes=dataset_classes,
+        evaluation_functions=evaluation_functions,
+        max_samples=50,  # Limit for faster testing, set to None for full evaluation
+    )
+    
+    print("\n" + "="*80)
+    print("COMPREHENSIVE EVALUATION RESULTS")
+    print("="*80)
+    print(results_df.to_string(index=False))
+    
+    # Print summary statistics
+    print("\n" + "="*80)
+    print("SUMMARY STATISTICS")
+    print("="*80)
+    
+    # Group by dataset and show average metrics
+    for dataset in results_df['dataset'].unique():
+        dataset_results = results_df[results_df['dataset'] == dataset]
+        print(f"\n{dataset}:")
+        print(f"  Number of models evaluated: {len(dataset_results)}")
+        
+        # Show average accuracy
+        if 'accuracy' in dataset_results.columns:
+            avg_accuracy = dataset_results['accuracy'].mean()
+            print(f"  Average accuracy: {avg_accuracy:.1f}%")
+    
+    # Group by model and show average metrics
+    print(f"\nBy Model:")
+    for model in results_df['model'].unique():
+        model_results = results_df[results_df['model'] == model]
+        print(f"\n{model}:")
+        print(f"  Number of datasets evaluated: {len(model_results)}")
+        
+        # Show average accuracy
+        if 'accuracy' in model_results.columns:
+            avg_accuracy = model_results['accuracy'].mean()
+            print(f"  Average accuracy: {avg_accuracy:.1f}%")
+    
+    return results_df
+
+
+if __name__ == "__main__":
+    main() 

code/OpenTSLM/evaluation/baseline/evaluate_ecg_qa.py

@@ -0,0 +1,293 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+
+import re
+import sys
+from typing import Dict, Any, List, Tuple
+
+
+from common_evaluator import CommonEvaluator
+from opentslm.time_series_datasets.ecg_qa.ECGQACoTQADataset import ECGQACoTQADataset
+
+
+def extract_answer(text: str) -> str:
+    """
+    Extract the final answer from model text, following the parser rules:
+    - If "Answer: " present, take substring after the last occurrence
+    - Strip any special end tokens like <|...|> or <eos>
+    - Trim trailing periods and whitespace
+    """
+    if text is None:
+        return ""
+    if "Answer: " not in text:
+        return text.strip()
+    answer = text.split("Answer: ")[-1].strip()
+    answer = re.sub(r"<\|.*?\|>|<eos>$", "", answer).strip()
+    answer = re.sub(r"\.$", "", answer).strip()
+    return answer
+
+
+def normalize_label(label: str) -> str:
+    """Lowercase, strip, and remove trailing punctuation to match parser behavior."""
+    if label is None:
+        return ""
+    return label.lower().strip().rstrip(".,!?;:")
+
+
+def evaluate_ecg_metrics(
+    ground_truth: str, prediction: str, sample: Dict[str, Any] | None = None
+) -> Dict[str, Any]:
+    """
+    Evaluate ECG-QA CoT predictions using per-template answers from CSV.
+    Normalization matches the parser in evaluation/opentslm/ecg_qa_cot/parse_ecg_qa_cot_data.py.
+    """
+    # Extract answers
+    pred_raw = extract_answer(prediction)
+    gt_raw = extract_answer(ground_truth)
+
+    # Normalize
+    pred_norm = normalize_label(pred_raw)
+    gt_norm = normalize_label(gt_raw)
+
+    # Per-template supported answers (strict)
+    if not isinstance(sample, dict):
+        print(f"DEBUG: Sample type: {type(sample)}")
+        print(f"DEBUG: Sample content: {sample}")
+        raise ValueError(
+            "Sample must be a dict containing 'template_id' for ECG-QA evaluation"
+        )
+
+    template_id = sample.get("template_id") or sample.get("cot_template_id")
+    if template_id is None:
+        print(f"DEBUG: Sample keys: {sample.keys()}")
+        print(f"DEBUG: Sample content: {sample}")
+        raise ValueError("Missing 'template_id' in sample for ECG-QA evaluation")
+
+    possible_answers = ECGQACoTQADataset.get_possible_answers_for_template(
+        int(template_id)
+    )
+    if not possible_answers:
+        raise ValueError(f"No possible answers found for template_id={template_id}")
+
+    possible_answers_lower = [a.lower().strip() for a in possible_answers]
+
+    # Supported flags: restrict to template answers strictly
+    pred_supported = pred_norm in possible_answers_lower
+    gt_supported = gt_norm in possible_answers_lower
+
+    # Exact match
+    is_correct = int(pred_norm == gt_norm)
+
+    # For single-label exact-match, precision=recall=F1=accuracy per-sample
+    f1 = float(is_correct)
+
+    return {
+        "accuracy": is_correct,
+        "f1_score": f1,
+        "precision": f1,
+        "recall": f1,
+        "prediction_normalized": pred_norm,
+        "ground_truth_normalized": gt_norm,
+        "prediction_supported": pred_supported,
+        "ground_truth_supported": gt_supported,
+        "template_id": template_id,
+        "possible_answers": possible_answers,
+    }
+
+
+# --- Parser-matching aggregation helpers ---
+
+
+def _calculate_template_f1_stats(data_points: List[Dict[str, Any]]) -> Dict[str, Any]:
+    if not data_points:
+        return {}
+
+    from collections import defaultdict
+
+    template_groups: Dict[int, List[Dict[str, Any]]] = defaultdict(list)
+    for point in data_points:
+        template_id = point.get("template_id")
+        if template_id is None:
+            raise ValueError(f"Missing template_id in data point: {point}")
+        template_groups[int(template_id)].append(point)
+
+    template_stats: Dict[int, Dict[str, Any]] = {}
+    total_samples = 0
+    total_correct = 0
+    total_f1_sum = 0.0
+
+    for template_id, points in template_groups.items():
+        if not points:
+            continue
+
+        possible_answers = points[0].get("possible_answers", [])
+        if not possible_answers:
+            raise ValueError(f"No possible answers found for template {template_id}")
+
+        # Initialize per-class counts
+        class_predictions: Dict[str, Dict[str, int]] = {}
+        for answer in possible_answers:
+            class_predictions[answer.lower()] = {"tp": 0, "fp": 0, "fn": 0}
+
+        # Count TP/FP/FN
+        for p in points:
+            gt_class = p.get("ground_truth_normalized", "")
+            pred_class = p.get("prediction_normalized", "")
+            pred_supported = p.get("prediction_supported", False)
+
+            if gt_class in class_predictions:
+                if pred_class == gt_class:
+                    class_predictions[gt_class]["tp"] += 1
+                else:
+                    class_predictions[gt_class]["fn"] += 1
+                    if pred_supported and pred_class in class_predictions:
+                        class_predictions[pred_class]["fp"] += 1
+
+        # Per-class and macro-F1
+        class_f1_scores: Dict[str, Dict[str, float]] = {}
+        template_f1_sum = 0.0
+        valid_classes = 0
+
+        for class_name, counts in class_predictions.items():
+            tp = counts["tp"]
+            fp = counts["fp"]
+            fn = counts["fn"]
+            precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
+            recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0
+            f1 = (
+                2 * (precision * recall) / (precision + recall)
+                if (precision + recall) > 0
+                else 0.0
+            )
+
+            class_f1_scores[class_name] = {
+                "f1": f1,
+                "precision": precision,
+                "recall": recall,
+                "tp": tp,
+                "fp": fp,
+                "fn": fn,
+            }
+            template_f1_sum += f1
+            valid_classes += 1
+
+        macro_f1 = template_f1_sum / valid_classes if valid_classes > 0 else 0.0
+
+        template_correct = sum(1 for p in points if p.get("accuracy", False))
+        template_accuracy = template_correct / len(points) if points else 0.0
+        template_avg_f1 = (
+            sum(p.get("f1_score", 0.0) for p in points) / len(points) if points else 0.0
+        )
+
+        template_stats[template_id] = {
+            "num_samples": len(points),
+            "accuracy": template_accuracy,
+            "average_f1": template_avg_f1,
+            "macro_f1": macro_f1,
+            "class_f1_scores": class_f1_scores,
+            "num_classes": valid_classes,
+            "correct_predictions": template_correct,
+        }
+
+        total_samples += len(points)
+        total_correct += template_correct
+        total_f1_sum += template_avg_f1 * len(points)
+
+    template_macro_f1s = [stats["macro_f1"] for stats in template_stats.values()]
+    overall_macro_f1 = (
+        sum(template_macro_f1s) / len(template_macro_f1s) if template_macro_f1s else 0.0
+    )
+    overall_accuracy = total_correct / total_samples if total_samples > 0 else 0.0
+    overall_avg_f1 = total_f1_sum / total_samples if total_samples > 0 else 0.0
+
+    return {
+        "overall": {
+            "total_samples": total_samples,
+            "total_templates": len(template_stats),
+            "accuracy": overall_accuracy,
+            "average_f1": overall_avg_f1,
+            "macro_f1": overall_macro_f1,
+        },
+        "per_template": template_stats,
+    }
+
+
+def _build_data_points_from_results(
+    detailed_results: List[Dict[str, Any]],
+) -> List[Dict[str, Any]]:
+    data_points: List[Dict[str, Any]] = []
+    for r in detailed_results:
+        m = r.get("metrics", {})
+        template_id = m.get("template_id", r.get("template_id"))
+        if template_id is None:
+            raise ValueError("Missing template_id in detailed result")
+        dp = {
+            "template_id": int(template_id),
+            "accuracy": m.get("accuracy", 0),
+            "f1_score": m.get("f1_score", 0.0),
+            "precision": m.get("precision", 0.0),
+            "recall": m.get("recall", 0.0),
+            "prediction_normalized": m.get("prediction_normalized", ""),
+            "ground_truth_normalized": m.get("ground_truth_normalized", ""),
+            "prediction_supported": m.get("prediction_supported", False),
+            "ground_truth_supported": m.get("ground_truth_supported", False),
+            "possible_answers": m.get("possible_answers", []),
+        }
+        if not dp["possible_answers"]:
+            raise ValueError(
+                f"No possible answers in metrics for template {template_id}"
+            )
+        data_points.append(dp)
+    return data_points
+
+
+def main():
+    """Main function to run ECG-QA CoT evaluation with parser-matching F1 aggregation."""
+    if len(sys.argv) != 2:
+        print("Usage: python evaluate_ecg_qa.py <model_name>")
+        print("Example: python evaluate_ecg_qa.py meta-llama/Llama-3.2-1B")
+        sys.exit(1)
+
+    model_name = sys.argv[1]
+
+    evaluator = CommonEvaluator()
+    # Run single evaluation to keep detailed results for F1 aggregation
+    results = evaluator.evaluate_model_on_dataset(
+        model_name=model_name,
+        dataset_class=ECGQACoTQADataset,
+        evaluation_function=evaluate_ecg_metrics,
+        max_samples=490,
+        max_new_tokens=400,
+    )
+
+    # Build data points and compute parser-matching F1 stats
+    detailed_results = results.get("detailed_results", [])
+    data_points = _build_data_points_from_results(detailed_results)
+    f1_stats = _calculate_template_f1_stats(data_points)
+
+    # Print parser-like summary
+    overall = f1_stats.get("overall", {})
+    print("\n" + "=" * 80)
+    print("FINAL RESULTS SUMMARY (Parser-matching)")
+    print("=" * 80)
+    print(f"Total templates: {overall.get('total_templates', 0)}")
+    print(f"Average F1 Score: {overall.get('average_f1', 0):.4f}")
+    print(f"Macro-F1 Score: {overall.get('macro_f1', 0):.4f}")
+
+    per_template = f1_stats.get("per_template", {})
+    if per_template:
+        print(f"\nPer-Template Statistics:")
+        for template_id, stats in sorted(per_template.items()):
+            print(f"  Template {template_id}:")
+            print(f"    Samples: {stats['num_samples']}")
+            print(f"    Accuracy: {stats['accuracy']:.4f}")
+            print(f"    Average F1: {stats['average_f1']:.4f}")
+            print(f"    Macro-F1: {stats['macro_f1']:.4f}")
+
+    return f1_stats
+
+
+if __name__ == "__main__":
+    main()

code/OpenTSLM/evaluation/baseline/evaluate_ecg_qa_plot.py

@@ -0,0 +1,355 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+
+import re
+import sys
+import io
+import base64
+import os
+from datetime import datetime
+from typing import Dict, Any, List, Optional
+
+import matplotlib.pyplot as plt
+
+from common_evaluator_plot import CommonEvaluatorPlot
+from opentslm.time_series_datasets.ecg_qa.ECGQACoTQADataset import ECGQACoTQADataset
+
+
+def extract_answer(text: str) -> str:
+    """
+    Extract the final answer from model text, following the parser rules:
+    - If "Answer: " present, take substring after the last occurrence
+    - Strip any special end tokens like <|...|> or <eos>
+    - Trim trailing periods and whitespace
+    """
+    if text is None:
+        return ""
+    if "Answer: " not in text:
+        return text.strip()
+    answer = text.split("Answer: ")[-1].strip()
+    answer = re.sub(r"<\|.*?\|>|<eos>$", "", answer).strip()
+    answer = re.sub(r"\.$", "", answer).strip()
+    return answer
+
+
+def normalize_label(label: str) -> str:
+    """Lowercase, strip, and remove trailing punctuation to match parser behavior."""
+    if label is None:
+        return ""
+    return label.lower().strip().rstrip(".,!?;:")
+
+
+def evaluate_ecg_metrics(
+    ground_truth: str, prediction: str, sample: Optional[Dict[str, Any]] = None
+) -> Dict[str, Any]:
+    """
+    Evaluate ECG-QA CoT predictions using per-template answers from CSV.
+    Normalization matches the parser used in evaluate_ecg_qa.py.
+    """
+    # Extract answers
+    pred_raw = extract_answer(prediction)
+    gt_raw = extract_answer(ground_truth)
+
+    # Normalize
+    pred_norm = normalize_label(pred_raw)
+    gt_norm = normalize_label(gt_raw)
+
+    # Per-template supported answers (strict)
+    if not isinstance(sample, dict):
+        raise ValueError(
+            "Sample must be a dict containing 'template_id' for ECG-QA evaluation"
+        )
+
+    template_id = sample.get("template_id") or sample.get("cot_template_id")
+    if template_id is None:
+        raise ValueError("Missing 'template_id' in sample for ECG-QA evaluation")
+
+    possible_answers = ECGQACoTQADataset.get_possible_answers_for_template(
+        int(template_id)
+    )
+    if not possible_answers:
+        raise ValueError(f"No possible answers found for template_id={template_id}")
+
+    possible_answers_lower = [a.lower().strip() for a in possible_answers]
+
+    # Supported flags: restrict to template answers strictly
+    pred_supported = pred_norm in possible_answers_lower
+    gt_supported = gt_norm in possible_answers_lower
+
+    # Exact match
+    is_correct = int(pred_norm == gt_norm)
+
+    # For single-label exact-match, precision=recall=F1=accuracy per-sample
+    f1 = float(is_correct)
+
+    return {
+        "accuracy": is_correct,
+        "f1_score": f1,
+        "precision": f1,
+        "recall": f1,
+        "prediction_normalized": pred_norm,
+        "ground_truth_normalized": gt_norm,
+        "prediction_supported": pred_supported,
+        "ground_truth_supported": gt_supported,
+        "template_id": template_id,
+        "possible_answers": possible_answers,
+    }
+
+
+def generate_ecg_plot(time_series: List[List[float]]) -> str:
+    """
+    Create a base64 PNG plot for multi-lead ECG time series.
+
+    - Accepts a list of 1D lists/arrays, one per ECG lead.
+    - Renders each lead as a separate subplot with grid and title.
+    """
+    if time_series is None:
+        return None
+
+    ts_list = list(time_series)
+    if not ts_list:
+        return None
+
+    num_series = len(ts_list)
+    lead_names = [
+        "I",
+        "II",
+        "III",
+        "aVR",
+        "aVL",
+        "aVF",
+        "V1",
+        "V2",
+        "V3",
+        "V4",
+        "V5",
+        "V6",
+    ]
+
+    # Limit to a reasonable number of subplots; if more, still plot all with generic names
+    fig_height = max(3, min(2 + 0.9 * num_series, 20))
+    fig, axes = plt.subplots(num_series, 1, figsize=(12, fig_height), sharex=True)
+    if num_series == 1:
+        axes = [axes]
+
+    for i, series in enumerate(ts_list):
+        axes[i].plot(series, linewidth=1.0)
+        axes[i].grid(True, alpha=0.3)
+        name = lead_names[i] if i < len(lead_names) else f"Lead {i + 1}"
+        axes[i].set_title(f"ECG Lead {name}")
+        axes[i].set_ylabel("mV")
+    axes[-1].set_xlabel("Time (samples)")
+
+    plt.tight_layout()
+    # Save plot to disk instead of showing it
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    results_dir = os.path.join(current_dir, "..", "results", "baseline", "plots")
+    os.makedirs(results_dir, exist_ok=True)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
+    filename = f"ecg_plot_{timestamp}.png"
+    file_path = os.path.join(results_dir, filename)
+    plt.savefig(file_path, format="png", bbox_inches="tight", dpi=110)
+
+    # Also save to buffer for returning base64 to caller
+    img_buffer = io.BytesIO()
+    plt.savefig(img_buffer, format="png", bbox_inches="tight", dpi=110)
+    plt.close()
+    img_buffer.seek(0)
+    image_data = base64.b64encode(img_buffer.getvalue()).decode("utf-8")
+    return image_data
+
+
+def _calculate_template_f1_stats(data_points: List[Dict[str, Any]]) -> Dict[str, Any]:
+    if not data_points:
+        return {}
+
+    from collections import defaultdict
+
+    template_groups: Dict[int, List[Dict[str, Any]]] = defaultdict(list)
+    for point in data_points:
+        template_id = point.get("template_id")
+        if template_id is None:
+            raise ValueError(f"Missing template_id in data point: {point}")
+        template_groups[int(template_id)].append(point)
+
+    template_stats: Dict[int, Dict[str, Any]] = {}
+    total_samples = 0
+    total_correct = 0
+    total_f1_sum = 0.0
+
+    for template_id, points in template_groups.items():
+        if not points:
+            continue
+
+        possible_answers = points[0].get("possible_answers", [])
+        if not possible_answers:
+            raise ValueError(f"No possible answers found for template {template_id}")
+
+        # Initialize per-class counts
+        class_predictions: Dict[str, Dict[str, int]] = {}
+        for answer in possible_answers:
+            class_predictions[answer.lower()] = {"tp": 0, "fp": 0, "fn": 0}
+
+        # Count TP/FP/FN
+        for p in points:
+            gt_class = p.get("ground_truth_normalized", "")
+            pred_class = p.get("prediction_normalized", "")
+            pred_supported = p.get("prediction_supported", False)
+
+            if gt_class in class_predictions:
+                if pred_class == gt_class:
+                    class_predictions[gt_class]["tp"] += 1
+                else:
+                    class_predictions[gt_class]["fn"] += 1
+                    if pred_supported and pred_class in class_predictions:
+                        class_predictions[pred_class]["fp"] += 1
+
+        # Per-class and macro-F1
+        class_f1_scores: Dict[str, Dict[str, float]] = {}
+        template_f1_sum = 0.0
+        valid_classes = 0
+
+        for class_name, counts in class_predictions.items():
+            tp = counts["tp"]
+            fp = counts["fp"]
+            fn = counts["fn"]
+            precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
+            recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0
+            f1 = (
+                2 * (precision * recall) / (precision + recall)
+                if (precision + recall) > 0
+                else 0.0
+            )
+
+            class_f1_scores[class_name] = {
+                "f1": f1,
+                "precision": precision,
+                "recall": recall,
+                "tp": tp,
+                "fp": fp,
+                "fn": fn,
+            }
+            template_f1_sum += f1
+            valid_classes += 1
+
+        macro_f1 = template_f1_sum / valid_classes if valid_classes > 0 else 0.0
+
+        template_correct = sum(1 for p in points if p.get("accuracy", False))
+        template_accuracy = template_correct / len(points) if points else 0.0
+        template_avg_f1 = (
+            sum(p.get("f1_score", 0.0) for p in points) / len(points) if points else 0.0
+        )
+
+        template_stats[template_id] = {
+            "num_samples": len(points),
+            "accuracy": template_accuracy,
+            "average_f1": template_avg_f1,
+            "macro_f1": macro_f1,
+            "class_f1_scores": class_f1_scores,
+            "num_classes": valid_classes,
+            "correct_predictions": template_correct,
+        }
+
+        total_samples += len(points)
+        total_correct += template_correct
+        total_f1_sum += template_avg_f1 * len(points)
+
+    template_macro_f1s = [stats["macro_f1"] for stats in template_stats.values()]
+    overall_macro_f1 = (
+        sum(template_macro_f1s) / len(template_macro_f1s) if template_macro_f1s else 0.0
+    )
+    overall_accuracy = total_correct / total_samples if total_samples > 0 else 0.0
+    overall_avg_f1 = total_f1_sum / total_samples if total_samples > 0 else 0.0
+
+    return {
+        "overall": {
+            "total_samples": total_samples,
+            "total_templates": len(template_stats),
+            "accuracy": overall_accuracy,
+            "average_f1": overall_avg_f1,
+            "macro_f1": overall_macro_f1,
+        },
+        "per_template": template_stats,
+    }
+
+
+def _build_data_points_from_results(
+    detailed_results: List[Dict[str, Any]],
+) -> List[Dict[str, Any]]:
+    data_points: List[Dict[str, Any]] = []
+    for r in detailed_results:
+        m = r.get("metrics", {})
+        template_id = m.get("template_id", r.get("template_id"))
+        if template_id is None:
+            raise ValueError("Missing template_id in detailed result")
+        dp = {
+            "template_id": int(template_id),
+            "accuracy": m.get("accuracy", 0),
+            "f1_score": m.get("f1_score", 0.0),
+            "precision": m.get("precision", 0.0),
+            "recall": m.get("recall", 0.0),
+            "prediction_normalized": m.get("prediction_normalized", ""),
+            "ground_truth_normalized": m.get("ground_truth_normalized", ""),
+            "prediction_supported": m.get("prediction_supported", False),
+            "ground_truth_supported": m.get("ground_truth_supported", False),
+            "possible_answers": m.get("possible_answers", []),
+        }
+        if not dp["possible_answers"]:
+            raise ValueError(
+                f"No possible answers in metrics for template {template_id}"
+            )
+        data_points.append(dp)
+    return data_points
+
+
+def main():
+    """Main function to run ECG-QA CoT evaluation with plotting."""
+    if len(sys.argv) != 2:
+        print("Usage: python evaluate_ecqqa_plot.py <model_name>")
+        print("Example: python evaluate_ecqqa_plot.py meta-llama/Llama-3.2-1B")
+        sys.exit(1)
+
+    model_name = sys.argv[1]
+
+    evaluator = CommonEvaluatorPlot()
+    # Run single evaluation to keep detailed results for F1 aggregation
+    results = evaluator.evaluate_model_on_dataset(
+        model_name=model_name,
+        dataset_class=ECGQACoTQADataset,
+        evaluation_function=evaluate_ecg_metrics,
+        plot_function=generate_ecg_plot,
+        max_samples=10000,
+        max_new_tokens=400,
+    )
+
+    # Build data points and compute parser-matching F1 stats
+    detailed_results = results.get("detailed_results", [])
+    data_points = _build_data_points_from_results(detailed_results)
+    f1_stats = _calculate_template_f1_stats(data_points)
+
+    # Print parser-like summary
+    overall = f1_stats.get("overall", {})
+    print("\n" + "=" * 80)
+    print("FINAL RESULTS SUMMARY (Parser-matching)")
+    print("=" * 80)
+    print(f"Total templates: {overall.get('total_templates', 0)}")
+    print(f"Average F1 Score: {overall.get('average_f1', 0):.4f}")
+    print(f"Macro-F1 Score: {overall.get('macro_f1', 0):.4f}")
+
+    per_template = f1_stats.get("per_template", {})
+    if per_template:
+        print(f"\nPer-Template Statistics:")
+        for template_id, stats in sorted(per_template.items()):
+            print(f"  Template {template_id}:")
+            print(f"    Samples: {stats['num_samples']}")
+            print(f"    Accuracy: {stats['accuracy']:.4f}")
+            print(f"    Average F1: {stats['average_f1']:.4f}")
+            print(f"    Macro-F1: {stats['macro_f1']:.4f}")
+
+    return f1_stats
+
+
+if __name__ == "__main__":
+    main()

code/OpenTSLM/evaluation/baseline/evaluate_har.py

@@ -0,0 +1,83 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+
+import re
+import sys
+from typing import Dict, Any
+
+from opentslm.time_series_datasets.har_cot.HARCoTQADataset import HARCoTQADataset
+
+
+def extract_label_from_prediction(prediction: str) -> str:
+    """
+    Extract the label from the model's prediction.
+    - If 'Answer:' is present, take everything after the last 'Answer:'
+    - Otherwise, take the last word
+    - Strips whitespace and punctuation
+    """
+    pred = prediction.strip()
+    # Find the last occurrence of 'Answer:' (case-insensitive)
+    match = list(re.finditer(r"answer:\s*", pred, re.IGNORECASE))
+    if match:
+        # Take everything after the last 'Answer:'
+        start = match[-1].end()
+        label = pred[start:].strip()
+    else:
+        # Take the last word
+        label = pred.split()[-1] if pred.split() else ""
+    # Remove trailing punctuation (e.g., period, comma)
+    label = re.sub(r"[\.,;:!?]+$", "", label)
+    return label.lower()
+
+
+def evaluate_har_acc(ground_truth: str, prediction: str) -> Dict[str, Any]:
+    """
+    <<<<<<< HEAD
+    <<<<<<< HEAD
+        Evaluate HARAccQADataset predictions against ground truth.
+    =======
+        Evaluate HARCoTQADataset predictions against ground truth.
+    >>>>>>> RealLast/ECG-QA-Integration
+    =======
+        Evaluate HARCoTQADataset predictions against ground truth.
+    >>>>>>> main
+        Extracts the label from the end of the model's output and compares to ground truth.
+    """
+    gt_clean = ground_truth.lower().strip()
+    pred_label = extract_label_from_prediction(prediction)
+    accuracy = int(gt_clean == pred_label)
+    return {"accuracy": accuracy}
+
+
+def main():
+    """Main function to run HAR evaluation."""
+    if len(sys.argv) != 2:
+        print("Usage: python evaluate_har.py <model_name>")
+        print("Example: python evaluate_har.py meta-llama/Llama-3.2-1B")
+        sys.exit(1)
+
+    model_name = sys.argv[1]
+
+    dataset_classes = [HARCoTQADataset]
+    evaluation_functions = {
+        "HARCoTQADataset": evaluate_har_acc,
+    }
+    evaluator = CommonEvaluator()
+    results_df = evaluator.evaluate_multiple_models(
+        model_names=[model_name],
+        dataset_classes=dataset_classes,
+        evaluation_functions=evaluation_functions,
+        max_samples=None,  # Limit for faster testing, set to None for full evaluation,
+        max_new_tokens=400,
+    )
+    print("\n" + "=" * 80)
+    print("FINAL RESULTS SUMMARY")
+    print("=" * 80)
+    print(results_df.to_string(index=False))
+    return results_df
+
+
+if __name__ == "__main__":
+    main()

code/OpenTSLM/evaluation/baseline/evaluate_har_plot.py

@@ -0,0 +1,113 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+
+import re
+import sys
+import argparse
+import io
+import base64
+from typing import Dict, Any
+
+import matplotlib.pyplot as plt
+
+from common_evaluator_plot import CommonEvaluatorPlot
+from opentslm.time_series_datasets.pamap2.PAMAP2AccQADataset import PAMAP2AccQADataset
+from opentslm.time_series_datasets.har_cot.HARAccQADataset import HARAccQADataset
+
+def extract_label_from_prediction(prediction: str) -> str:
+    """
+    Extract the label from the model's prediction.
+    - If 'Answer:' is present, take everything after the last 'Answer:'
+    - Otherwise, take the last word
+    - Strips whitespace and punctuation
+    """
+    pred = prediction.strip()
+    # Find the last occurrence of 'Answer:' (case-insensitive)
+    match = list(re.finditer(r'answer:\s*', pred, re.IGNORECASE))
+    if match:
+        # Take everything after the last 'Answer:'
+        start = match[-1].end()
+        label = pred[start:].strip()
+    else:
+        # Take the last word
+        label = pred.split()[-1] if pred.split() else ''
+    # Remove trailing punctuation (e.g., period, comma)
+    label = re.sub(r'[\.,;:!?]+$', '', label)
+    return label.lower()
+
+
+def evaluate_har_acc(ground_truth: str, prediction: str) -> Dict[str, Any]:
+    """
+    Evaluate HARAccQADataset predictions against ground truth.
+    Extracts the label from the end of the model's output and compares to ground truth.
+    """
+    gt_clean = ground_truth.lower().strip()
+    pred_label = extract_label_from_prediction(prediction)
+    accuracy = int(gt_clean == pred_label)
+    return {"accuracy": accuracy}
+
+
+def generate_time_series_plot(time_series) -> str:
+    """
+    Create a base64 PNG plot from a list/tuple of 1D numpy arrays (e.g., [x, y, z]).
+    """
+    if time_series is None:
+        return None
+    ts_list = list(time_series)
+
+    num_series = len(ts_list)
+    fig, axes = plt.subplots(num_series, 1, figsize=(10, 4 * num_series), sharex=True)
+    if num_series == 1:
+        axes = [axes]
+
+    axis_names = {0: 'X-axis', 1: 'Y-axis', 2: 'Z-axis'}
+    for i, series in enumerate(ts_list):
+        axes[i].plot(series, marker='o', linestyle='-', markersize=0)
+        axes[i].grid(True, alpha=0.3)
+        axes[i].set_title(f"Accelerometer - {axis_names.get(i, f'Axis {i+1}')}" )
+
+    plt.tight_layout()
+
+    img_buffer = io.BytesIO()
+    plt.savefig(img_buffer, format='png', bbox_inches='tight', dpi=100)
+    plt.close()
+    img_buffer.seek(0)
+    image_data = base64.b64encode(img_buffer.getvalue()).decode('utf-8')
+    return image_data
+
+
+def main():
+    """Main function to run HAR evaluation."""
+    if len(sys.argv) != 2:
+        print("Usage: python evaluate_har_plot.py <model_name>")
+        print("Example: python evaluate_har_plot.py meta-llama/Llama-3.2-1B")
+        sys.exit(1)
+    
+    model_name = sys.argv[1]
+
+    dataset_classes = [HARAccQADataset]
+    evaluation_functions = {
+        "HARAccQADataset": evaluate_har_acc,
+    }
+    evaluator = CommonEvaluatorPlot()
+    plot_functions = {
+        "HARAccQADataset": generate_time_series_plot,
+    }
+    results_df = evaluator.evaluate_multiple_models(
+        model_names=[model_name],
+        dataset_classes=dataset_classes,
+        evaluation_functions=evaluation_functions,
+        plot_functions=plot_functions,
+        max_samples=None,  # Limit for faster testing, set to None for full evaluation,
+        max_new_tokens=400,
+    )
+    print("\n" + "="*80)
+    print("FINAL RESULTS SUMMARY")
+    print("="*80)
+    print(results_df.to_string(index=False))
+    return results_df
+
+if __name__ == "__main__":
+    main()

code/OpenTSLM/evaluation/baseline/evaluate_pamap.py

@@ -0,0 +1,75 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+
+import re
+import sys
+from typing import Dict, Any
+
+from common_evaluator import CommonEvaluator
+from opentslm.time_series_datasets.pamap2.PAMAP2AccQADataset import PAMAP2AccQADataset
+
+
+def extract_label_from_prediction(prediction: str) -> str:
+    """
+    Extract the label from the model's prediction.
+    - If 'Answer:' is present, take everything after the last 'Answer:'
+    - Otherwise, take the last word
+    - Strips whitespace and punctuation
+    """
+    pred = prediction.strip()
+    # Find the last occurrence of 'Answer:' (case-insensitive)
+    match = list(re.finditer(r'answer:\s*', pred, re.IGNORECASE))
+    if match:
+        # Take everything after the last 'Answer:'
+        start = match[-1].end()
+        label = pred[start:].strip()
+    else:
+        # Take the last word
+        label = pred.split()[-1] if pred.split() else ''
+    # Remove trailing punctuation (e.g., period, comma)
+    label = re.sub(r'[\.,;:!?]+$', '', label)
+    return label.lower()
+
+
+def evaluate_pamap_acc(ground_truth: str, prediction: str) -> Dict[str, Any]:
+    """
+    Evaluate PAMAP2AccQADataset predictions against ground truth.
+    Extracts the label from the end of the model's output and compares to ground truth.
+    """
+    gt_clean = ground_truth.lower().strip()
+    pred_label = extract_label_from_prediction(prediction)
+    accuracy = int(gt_clean == pred_label)
+    return {"accuracy": accuracy}
+
+
+def main():
+    """Main function to run PAMAP evaluation."""
+    if len(sys.argv) != 2:
+        print("Usage: python evaluate_pamap.py <model_name>")
+        print("Example: python evaluate_pamap.py meta-llama/Llama-3.2-1B")
+        sys.exit(1)
+    
+    model_name = sys.argv[1]
+    
+    dataset_classes = [PAMAP2AccQADataset]
+    evaluation_functions = {
+        "PAMAP2AccQADataset": evaluate_pamap_acc,
+    }
+    evaluator = CommonEvaluator()
+    results_df = evaluator.evaluate_multiple_models(
+        model_names=[model_name],
+        dataset_classes=dataset_classes,
+        evaluation_functions=evaluation_functions,
+        max_samples=None,  # Limit for faster testing, set to None for full evaluation,
+        max_new_tokens=400,
+    )
+    print("\n" + "="*80)
+    print("FINAL RESULTS SUMMARY")
+    print("="*80)
+    print(results_df.to_string(index=False))
+    return results_df
+
+if __name__ == "__main__":
+    main() 

code/OpenTSLM/evaluation/baseline/evaluate_sleep_cot.py

@@ -0,0 +1,93 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+
+import re
+import sys
+from typing import Dict, Any
+
+from common_evaluator import CommonEvaluator
+from opentslm.time_series_datasets.sleep.SleepEDFCoTQADataset import SleepEDFCoTQADataset
+
+
+def extract_sleep_stage_from_prediction(prediction: str) -> str:
+    """
+    Extract the sleep stage label from the model's prediction.
+    - If 'Answer:' is present, take everything after the last 'Answer:'
+    - Otherwise, take the last word/phrase
+    - Strips whitespace and punctuation
+    """
+    pred = prediction.strip()
+    # Find the last occurrence of 'Answer:' (case-insensitive)
+    match = list(re.finditer(r'answer:\s*', pred, re.IGNORECASE))
+    if match:
+        # Take everything after the last 'Answer:'
+        start = match[-1].end()
+        label = pred[start:].strip()
+    else:
+        # Take the last word/phrase (sleep stages can be multi-word)
+        words = pred.split()
+        if len(words) >= 2 and words[-2].lower() in ['non-rem', 'rem']:
+            # Handle multi-word sleep stages like "Non-REM stage 1", "REM sleep"
+            label = ' '.join(words[-2:])
+        else:
+            # Single word stage
+            label = words[-1] if words else ''
+    # Look for "Answer: [answer]" pattern with more precise matching
+    answer_match = re.search(r'Answer:\s*(.+?)(?:\.|$)', pred, re.IGNORECASE)
+    if answer_match:
+        label = answer_match.group(1).strip()
+        # Remove trailing period if present
+        if label.endswith('.'):
+            label = label[:-1]
+        return label.strip().lower()
+    # Remove trailing punctuation (e.g., period, comma)
+    label = re.sub(r'[\.,;:!?]+$', '', label)
+    return label.strip().lower()
+
+
+def evaluate_sleep_cot(ground_truth: str, prediction: str) -> Dict[str, Any]:
+    """
+    Evaluate SleepEDFCoTQADataset predictions against ground truth.
+    Extracts the sleep stage label from the end of the model's output and compares to ground truth.
+    """
+    gt_clean = ground_truth.lower().strip()
+    pred_label = extract_sleep_stage_from_prediction(prediction)
+    accuracy = int(gt_clean == pred_label)
+    return {"accuracy": accuracy}
+
+
+def main():
+    """Main function to run SleepEDF CoT evaluation."""
+    if len(sys.argv) != 2:
+        print("Usage: python evaluate_sleep_cot.py <model_name>")
+        print("Example: python evaluate_sleep_cot.py meta-llama/Llama-3.2-1B")
+        sys.exit(1)
+    
+    model_name = sys.argv[1]
+    
+    dataset_classes = [SleepEDFCoTQADataset]
+    evaluation_functions = {
+        "SleepEDFCoTQADataset": evaluate_sleep_cot,
+    }
+    
+    evaluator = CommonEvaluator()
+    results_df = evaluator.evaluate_multiple_models(
+        model_names=[model_name],
+        dataset_classes=dataset_classes,
+        evaluation_functions=evaluation_functions,
+        max_samples=None,  # Set to None for full evaluation
+        max_new_tokens=1000,  # Sleep stage classification might need more tokens for reasoning
+    )
+    
+    print("\n" + "="*80)
+    print("FINAL RESULTS SUMMARY")
+    print("="*80)
+    print(results_df.to_string(index=False))
+    
+    return results_df
+
+
+if __name__ == "__main__":
+    main() 

code/OpenTSLM/evaluation/baseline/evaluate_sleep_plot.py

@@ -0,0 +1,120 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+
+import re
+import sys
+import io
+import base64
+from typing import Dict, Any
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from common_evaluator_plot import CommonEvaluatorPlot
+from opentslm.time_series_datasets.sleep.SleepEDFCoTQADataset import SleepEDFCoTQADataset
+
+
+def extract_label_from_text(text: str) -> str:
+    """
+    Extract the label from a free-form rationale or prediction text.
+    - If 'Answer:' is present (case-insensitive), take everything after the last 'Answer:'
+    - Otherwise, take the last word
+    - Strip whitespace and trailing punctuation
+    - Lowercase for comparison
+    """
+    if text is None:
+        return ""
+    pred = text.strip()
+    matches = list(re.finditer(r"answer:\s*", pred, re.IGNORECASE))
+    if matches:
+        start = matches[-1].end()
+        label = pred[start:].strip()
+    else:
+        label = pred.split()[-1] if pred.split() else ""
+    label = re.sub(r"[\.,;:!?]+$", "", label)
+    return label.lower()
+
+
+def evaluate_sleep_stage(
+    ground_truth_text: str, prediction_text: str
+) -> Dict[str, Any]:
+    """
+    Evaluate SleepEDFCoTQADataset predictions against ground truth.
+    For SleepEDF, the dataset's "answer" is a rationale ending with 'Answer: <label>'.
+    We therefore extract the label from BOTH ground truth and prediction and compare.
+    """
+    gt_label = extract_label_from_text(ground_truth_text)
+    pred_label = extract_label_from_text(prediction_text)
+    accuracy = int(gt_label == pred_label)
+    return {"accuracy": accuracy, "gt_label": gt_label, "pred_label": pred_label}
+
+
+def generate_time_series_plot(time_series) -> str:
+    """
+    Create a base64 PNG plot from one or more time series.
+    - Accepts a single 1D array/list or a collection of 1D arrays/lists.
+    - If a 2D numpy array is provided, each row is treated as a separate series.
+    """
+    if time_series is None:
+        return None
+    ts_list = list(time_series)
+
+    num_series = len(ts_list)
+    fig, axes = plt.subplots(num_series, 1, figsize=(10, 4 * num_series), sharex=True)
+    if num_series == 1:
+        axes = [axes]
+
+    axis_names = {0: "EEG", 1: "EOG", 2: "EMG"}
+    for i, series in enumerate(ts_list):
+        axes[i].plot(series, marker="o", linestyle="-", markersize=0)
+        axes[i].grid(True, alpha=0.3)
+        axes[i].set_title(f"{axis_names.get(i, f'Axis {i + 1}')}")
+
+    plt.tight_layout()
+
+    img_buffer = io.BytesIO()
+    plt.savefig(img_buffer, format="png", bbox_inches="tight", dpi=100)
+    plt.close()
+    img_buffer.seek(0)
+    image_data = base64.b64encode(img_buffer.getvalue()).decode("utf-8")
+    return image_data
+
+
+def main():
+    """Main function to run SleepEDF evaluation with plotting."""
+    if len(sys.argv) != 2:
+        print("Usage: python evaluate_sleep_plot.py <model_name>")
+        print("Example: python evaluate_sleep_plot.py openai-gpt-4o")
+        sys.exit(1)
+
+    model_name = sys.argv[1]
+
+    dataset_classes = [SleepEDFCoTQADataset]
+    evaluation_functions = {
+        "SleepEDFCoTQADataset": evaluate_sleep_stage,
+    }
+    evaluator = CommonEvaluatorPlot()
+    plot_functions = {
+        "SleepEDFCoTQADataset": generate_time_series_plot,
+    }
+
+    results_df = evaluator.evaluate_multiple_models(
+        model_names=[model_name],
+        dataset_classes=dataset_classes,
+        evaluation_functions=evaluation_functions,
+        plot_functions=plot_functions,
+        max_samples=None,  # Set to None for full evaluation
+        max_new_tokens=400,
+    )
+
+    print("\n" + "=" * 80)
+    print("FINAL RESULTS SUMMARY")
+    print("=" * 80)
+    print(results_df.to_string(index=False))
+    return results_df
+
+
+if __name__ == "__main__":
+    main()

code/OpenTSLM/evaluation/baseline/evaluate_tsqa.py

@@ -0,0 +1,89 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+
+import re
+import sys
+from typing import Dict, Any
+
+from common_evaluator import CommonEvaluator
+from opentslm.time_series_datasets.TSQADataset import TSQADataset
+
+
+def evaluate_tsqa(ground_truth: str, prediction: str) -> Dict[str, Any]:
+    """
+    Evaluate TSQA predictions against ground truth.
+    
+    Args:
+        ground_truth: The correct answer
+        prediction: The model's prediction
+        
+    Returns:
+        Dictionary containing evaluation metrics
+    """
+    # Clean up strings for comparison
+    gt_clean = ground_truth.lower().strip()
+    pred_clean = prediction.lower().strip()
+    
+    # Only compare the first 3 characters of the ground truth and prediction,
+    # because the ground truth answer is always in the format "(a)", "(b)", or "(c)",
+    # but the model might answer with the full answer, e.g. "(a) This time series has a constant volatility."
+    gt_clean, pred_clean = gt_clean[:3], pred_clean[:3]
+
+    # Extract the actual answer from the prediction (everything after "Answer:")
+    answer_match = re.search(r'answer:\s*(.+)', pred_clean, re.IGNORECASE)
+    if answer_match:
+        pred_answer = answer_match.group(1).strip()
+    else:
+        # If no "Answer:" found, use the entire prediction
+        pred_answer = pred_clean
+    
+    # Calculate accuracy (exact match)
+    accuracy = int(gt_clean == pred_answer)
+    
+    return {
+        "accuracy": accuracy,
+    }
+
+
+def main():
+    """Main function to run TSQA evaluation."""
+    
+    if len(sys.argv) != 2:
+        print("Usage: python evaluate_tsqa.py <model_name>")
+        print("Example: python evaluate_tsqa.py meta-llama/Llama-3.2-1B")
+        sys.exit(1)
+    
+    model_name = sys.argv[1]
+    
+    # Define datasets to evaluate on
+    dataset_classes = [TSQADataset]
+    
+    # Define evaluation functions
+    evaluation_functions = {
+        "TSQADataset": evaluate_tsqa,
+    }
+    
+    # Initialize evaluator
+    evaluator = CommonEvaluator()
+    
+    # Run evaluation
+    results_df = evaluator.evaluate_multiple_models(
+        model_names=[model_name],
+        dataset_classes=dataset_classes,
+        evaluation_functions=evaluation_functions,
+        max_samples=None,  # Limit for faster testing, set to None for full evaluation
+        max_new_tokens=40,
+    )
+    
+    print("\n" + "="*80)
+    print("FINAL RESULTS SUMMARY")
+    print("="*80)
+    print(results_df.to_string(index=False))
+    
+    return results_df
+
+
+if __name__ == "__main__":
+    main() 

code/OpenTSLM/evaluation/baseline/evaluate_tsqa_plot.py

@@ -0,0 +1,124 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+
+import re
+import sys
+import io
+import base64
+from typing import Dict, Any
+
+import matplotlib.pyplot as plt
+
+from common_evaluator_plot import CommonEvaluatorPlot
+from opentslm.time_series_datasets.TSQADataset import TSQADataset
+
+
+def evaluate_tsqa(ground_truth: str, prediction: str) -> Dict[str, Any]:
+    """
+    Evaluate TSQA predictions against ground truth.
+
+    - TSQA answers are in the format "(a)", "(b)", or "(c)".
+    - The model might answer with a longer explanation like
+      "(a) This time series has a constant volatility.", so we compare only
+      the first 3 characters of the answer (e.g., "(a)").
+    - If the prediction contains an "Answer:" tag, take everything after it.
+    """
+    # Clean up strings for comparison
+    gt_clean = ground_truth.lower().strip()
+    pred_clean = prediction.lower().strip()
+
+    # Only compare the first 3 characters of the ground truth and prediction
+    gt_clean, pred_clean = gt_clean[:3], pred_clean[:3]
+
+    # Extract the actual answer from the prediction (everything after "Answer:")
+    answer_match = re.search(r"answer:\s*(.+)", pred_clean, re.IGNORECASE)
+    if answer_match:
+        pred_answer = answer_match.group(1).strip()
+    else:
+        # If no "Answer:" found, use the entire prediction
+        pred_answer = pred_clean
+
+    # Calculate accuracy (exact match)
+    accuracy = int(gt_clean == pred_answer)
+
+    return {
+        "accuracy": accuracy,
+    }
+
+
+def generate_time_series_plot(time_series) -> str:
+    """
+    Create a base64 PNG plot from one or more time series.
+
+    - TSQA typically has a single 1D time series, but we handle multiple for robustness.
+    - Accepts either a single 1D list/array or a list of 1D lists/arrays.
+    """
+    if time_series is None:
+        return None
+
+    # Ensure we have a list of series
+    ts_list = list(time_series)
+    # If it's a flat list of numbers, wrap it as a single series
+    if ts_list and all(isinstance(x, (int, float)) for x in ts_list):
+        ts_list = [ts_list]
+
+    num_series = len(ts_list)
+    fig, axes = plt.subplots(num_series, 1, figsize=(10, 4 * num_series), sharex=True)
+    if num_series == 1:
+        axes = [axes]
+
+    for i, series in enumerate(ts_list):
+        axes[i].plot(series, marker="o", linestyle="-", markersize=0)
+        axes[i].grid(True, alpha=0.3)
+        axes[i].set_title(f"Time Series {i + 1}")
+        axes[i].set_ylabel("Value")
+    axes[-1].set_xlabel("Time")
+
+    plt.tight_layout()
+
+    img_buffer = io.BytesIO()
+    plt.savefig(img_buffer, format="png", bbox_inches="tight", dpi=100)
+    plt.close()
+    img_buffer.seek(0)
+    image_data = base64.b64encode(img_buffer.getvalue()).decode("utf-8")
+    return image_data
+
+
+def main():
+    """Main function to run TSQA evaluation with plotting."""
+    if len(sys.argv) != 2:
+        print("Usage: python evaluate_tsqa_plot.py <model_name>")
+        print("Example: python evaluate_tsqa_plot.py openai-gpt-4o")
+        sys.exit(1)
+
+    model_name = sys.argv[1]
+
+    dataset_classes = [TSQADataset]
+    evaluation_functions = {
+        "TSQADataset": evaluate_tsqa,
+    }
+    evaluator = CommonEvaluatorPlot()
+    plot_functions = {
+        "TSQADataset": generate_time_series_plot,
+    }
+
+    results_df = evaluator.evaluate_multiple_models(
+        model_names=[model_name],
+        dataset_classes=dataset_classes,
+        evaluation_functions=evaluation_functions,
+        plot_functions=plot_functions,
+        max_samples=None,  # Set to None for full evaluation
+        max_new_tokens=80,
+    )
+
+    print("\n" + "=" * 80)
+    print("FINAL RESULTS SUMMARY")
+    print("=" * 80)
+    print(results_df.to_string(index=False))
+    return results_df
+
+
+if __name__ == "__main__":
+    main()

code/OpenTSLM/evaluation/baseline/gruver_llmtime_tokenizer.py

@@ -0,0 +1,175 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+
+import numpy as np
+from functools import partial
+from dataclasses import dataclass
+
+# Please check the original code at https://github.com/ngruver/llmtime/blob/main/data/serialize.py
+@dataclass
+class SerializerSettings:
+    """
+    Settings for time series serialization and deserialization.
+    """
+    base: int = 10                  # Numeric base for representation
+    prec: int = 3                  # Number of digits after the 'decimal' point
+    signed: bool = True            # Whether to include a sign for positive values
+    fixed_length: bool = False     # If True, pad to a fixed number of digits
+    max_val: float = 1e7           # Maximum absolute value allowed
+    time_sep: str = ' ,'           # Separator between time steps
+    bit_sep: str = ' '             # Separator between individual digits
+    plus_sign: str = ''            # String to prepend for positive values
+    minus_sign: str = ' -'         # String to prepend for negative values
+    half_bin_correction: bool = True  # Apply half-bin correction on deserialization
+    decimal_point: str = ''        # Literal to mark decimal point in serialized string
+    missing_str: str = ' Nan'      # Representation for missing (NaN) values
+
+
+def vec_num2repr(val: np.ndarray, base: int, prec: int, max_val: float):
+    """
+    Convert an array of floats into sign and digit representation in the given base.
+    """
+    base = float(base)
+    sign = np.where(val >= 0, 1, -1)
+    mag = np.abs(val)
+    max_bit_pos = int(np.ceil(np.log(max_val) / np.log(base)).item())
+
+    # Compute digits before the "decimal"
+    before = []
+    rem = mag.copy()
+    for i in range(max_bit_pos):
+        power = base ** (max_bit_pos - i - 1)
+        digit = (rem / power).astype(int)
+        before.append(digit)
+        rem = rem - digit * power
+    before = np.stack(before, axis=-1)
+
+    # Compute digits after the "decimal"
+    if prec > 0:
+        after = []
+        for i in range(prec):
+            power = base ** (-(i + 1))
+            digit = (rem / power).astype(int)
+            after.append(digit)
+            rem = rem - digit * power
+        after = np.stack(after, axis=-1)
+        digits = np.concatenate([before, after], axis=-1)
+    else:
+        digits = before
+
+    return sign, digits
+
+
+def vec_repr2num(sign: np.ndarray, digits: np.ndarray, base: int, prec: int, half_bin_correction: bool = True) -> np.ndarray:
+    """
+    Convert sign and digit arrays back into floats.
+    """
+    base = float(base)
+    # Reverse the digit order for positional weights
+    flipped = np.flip(digits, axis=-1)
+    total_digits = digits.shape[1]
+    # Compute positional powers
+    powers = -np.arange(-prec, total_digits-prec)
+    mags = np.sum(flipped / (base ** powers), axis=-1)
+    if half_bin_correction:
+        mags = mags + (0.5 / (base ** prec))
+    return sign * mags
+
+
+def serialize_arr(arr: np.ndarray, settings: SerializerSettings) -> str:
+    """
+    Serialize a 1D numpy array into a digit-level string using the given settings.
+    """
+    # Validate range
+    clean = np.where(np.isnan(arr), 0.0, arr)
+    assert np.all(np.abs(clean) <= settings.max_val), \
+        f"Values must be within ±{settings.max_val}"
+
+    # Convert numbers to sign and digits
+    to_repr = partial(vec_num2repr, base=settings.base, prec=settings.prec, max_val=settings.max_val)
+    sign_arr, digits_arr = to_repr(clean)
+    is_nan = np.isnan(arr)
+
+    def format_digits(digits):
+        # Optionally strip leading zeros
+        if not settings.fixed_length:
+            nz = np.where(digits != 0)[0]
+            if nz.size > 0:
+                digits = digits[nz[0]:]
+            else:
+                digits = np.array([0], dtype=int)
+        # Insert decimal point if specified
+        if settings.decimal_point and settings.prec > 0:
+            point_idx = len(digits) - settings.prec
+            digits = np.concatenate([digits[:point_idx], [-1], digits[point_idx:]])
+        # Join with bit separator
+        return settings.bit_sep.join(str(d) for d in digits if d != -1)
+
+    parts = []
+    for s, digs, missing in zip(sign_arr, digits_arr, is_nan):
+        if missing:
+            parts.append(settings.missing_str)
+        else:
+            sign_str = settings.plus_sign if s > 0 else settings.minus_sign
+            digit_str = format_digits(digs)
+            parts.append(sign_str + digit_str)
+
+    result = settings.time_sep.join(parts) + settings.time_sep
+    return result
+
+
+def deserialize_str(bit_str: str, settings: SerializerSettings, ignore_last: bool = False, steps: int = None) -> np.ndarray:
+    """
+    Deserialize a serialized string back into a numpy array of floats.
+    """
+    tokens = [t for t in bit_str.split(settings.time_sep) if t]
+    if ignore_last:
+        tokens = tokens[:-1]
+    if steps is not None:
+        tokens = tokens[:steps]
+
+    sign_list = []
+    digit_list = []
+    for tok in tokens:
+        if tok == settings.missing_str.strip():
+            sign_list.append(1)
+            digit_list.append([0])
+            continue
+        # Determine sign
+        if settings.signed and tok.startswith(settings.minus_sign):
+            sign_list.append(-1)
+            tok = tok[len(settings.minus_sign):]
+        else:
+            sign_list.append(1)
+            if settings.signed and tok.startswith(settings.plus_sign):
+                tok = tok[len(settings.plus_sign):]
+        # Split digits
+        digs = [int(ch) for ch in tok.split(settings.bit_sep) if ch.isdigit()]
+        digit_list.append(digs)
+
+    # Pad to equal length
+    maxlen = max(len(d) for d in digit_list)
+    padded = np.array([([0]*(maxlen - len(d)) + d) for d in digit_list])
+
+    nums = vec_repr2num(
+        np.array(sign_list), padded,
+        base=settings.base, prec=settings.prec,
+        half_bin_correction=settings.half_bin_correction
+    )
+    return nums
+
+gpt_settings = SerializerSettings(base=10, prec=3, signed=True, time_sep=', ', bit_sep=' ', minus_sign='-')
+llama = SerializerSettings(base=10, prec=3, signed=True, half_bin_correction=True)
+
+# Then wrap serialize_arr so it only takes the array
+def gpt_formatter(arr: np.ndarray) -> str:
+    return serialize_arr(arr, gpt_settings)
+
+def llama_formatter(arr: np.ndarray) -> str:
+    return serialize_arr(arr, llama)
+
+# Backward compatibility: default to llama_formatter
+# (or you can set to gpt_formatter if you prefer)
+gruver_et_al_formatter = llama_formatter

code/OpenTSLM/evaluation/baseline/openai_pipeline.py

@@ -0,0 +1,46 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+
+import os
+from openai import OpenAI
+from typing import Optional
+
+class OpenAIPipeline:
+    def __init__(self, model_name: str, api_key: str = None, temperature: float = 0.1, max_new_tokens: int = 1000):
+        """
+        A small wrapper around the new OpenAI v1 client for chat completions.
+        """
+        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
+        if not self.api_key:
+            raise ValueError("API key must be provided via argument or OPENAI_API_KEY env var")
+        # instantiate the v1 client
+        self.client = OpenAI(api_key=self.api_key)
+        self.model_name = model_name
+        self.temperature = temperature
+
+    def __call__(self, prompt: str, max_new_tokens: int = 1000, return_full_text: bool = False, plot_data: Optional[str] = None):
+        """
+        Send a single-user-message chat completion request and return the generated text in HuggingFace pipeline format.
+        """
+
+        if plot_data:
+            content = [
+                {"type": "text", "text": prompt},
+                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{plot_data}", "detail": "high"}}
+            ]
+        else:
+            content = [{"type": "text", "text": prompt}]
+
+        resp = self.client.chat.completions.create(
+            model=self.model_name,
+            messages=[
+                {"role": "user", "content": content}
+            ],
+            temperature=self.temperature,
+            max_tokens=max_new_tokens,
+        )
+        # Return in HuggingFace pipeline format
+        return [{"generated_text": resp.choices[0].message.content.strip()}]
+