File size: 11,223 Bytes
60b21d3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 | #!/usr/bin/env python3
# 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
"""
Test script for the PAMAP2 CoT loader.
"""
import unittest
from opentslm.logger import get_logger, set_global_verbose
from opentslm.time_series_datasets.pamap2.BalancedBatchSampler import BalancedBatchSampler
# Import and set up global logger with verbose mode
from opentslm.logger import get_logger, set_global_verbose
from opentslm.time_series_datasets.pamap2.BalancedBatchSampler import BalancedBatchSampler
class TestPAMAP2CoTLoader(unittest.TestCase):
"""
Unit tests for the PAMAP2 CoT loader functions.
"""
def setUp(self):
# Set up global logger with verbose mode for detailed output
set_global_verbose(True)
self.logger = get_logger()
from opentslm.time_series_datasets.pamap2.pamap2_cot_loader import load_pamap2_cot_splits
self.load_pamap2_cot_splits = load_pamap2_cot_splits
self.logger.loading("Loading PAMAP2 CoT dataset splits...")
self.train, self.val, self.test = self.load_pamap2_cot_splits()
self.logger.success(f"Dataset loaded successfully: Train={len(self.train)}, Val={len(self.val)}, Test={len(self.test)}")
def test_dataset_sizes(self):
"""Test that the datasets are non-empty and splits are correct."""
self.logger.info("Testing dataset sizes...")
self.assertGreater(len(self.train), 0)
self.assertGreater(len(self.val), 0)
self.assertGreater(len(self.test), 0)
self.logger.success("Dataset size tests passed")
def test_sample_keys(self):
"""Test that a sample contains all required keys."""
self.logger.info("Testing sample keys...")
sample = self.train[0]
required_keys = {"x_axis", "y_axis", "z_axis", "label"}
self.assertTrue(required_keys.issubset(sample.keys()))
self.logger.success("Sample keys test passed")
def test_axis_content(self):
"""Test that the axis data are lists and non-empty."""
self.logger.info("Testing axis content...")
sample = self.train[0]
for axis in ["x_axis", "y_axis", "z_axis"]:
self.assertIsInstance(sample[axis], list)
self.assertGreater(len(sample[axis]), 0)
self.logger.debug(f"{axis}: length={len(sample[axis])}")
self.logger.success("Axis content tests passed")
def test_label_is_string(self):
"""Test that the label is a string and non-empty."""
self.logger.info("Testing label format...")
sample = self.train[0]
self.assertIsInstance(sample["label"], str)
self.assertGreater(len(sample["label"]), 0)
self.logger.success(f"Label test passed: '{sample['label']}'")
def test_example_data(self):
"""Print example data to show what the dataset looks like."""
sample = self.train[0]
self.logger.info("="*80)
self.logger.info("EXAMPLE PAMAP2 COT DATASET SAMPLE")
self.logger.info("="*80)
self.logger.info(f"Label: '{sample['label']}'")
if 'rationale' in sample:
self.logger.info(f"Rationale: '{sample['rationale']}'")
for axis in ["x_axis", "y_axis", "z_axis"]:
self.logger.info(f"{axis}: length={len(sample[axis])}, first 5: {sample[axis][:5]}")
self.logger.info("="*80)
class TestPAMAP2CoTQADataset(unittest.TestCase):
"""
Unit tests for the PAMAP2CoTQADataset class.
"""
def setUp(self):
# Set up global logger with verbose mode for detailed output
set_global_verbose(True)
self.logger = get_logger()
from opentslm.time_series_datasets.pamap2.PAMAP2CoTQADataset import PAMAP2CoTQADataset
self.PAMAP2CoTQADataset = PAMAP2CoTQADataset
self.logger.loading("Initializing PAMAP2CoTQADataset...")
self.train_dataset = self.PAMAP2CoTQADataset(split="train", EOS_TOKEN="")
self.val_dataset = self.PAMAP2CoTQADataset(split="validation", EOS_TOKEN="")
self.test_dataset = self.PAMAP2CoTQADataset(split="test", EOS_TOKEN="")
self.logger.success(f"Datasets initialized: Train={len(self.train_dataset)}, Val={len(self.val_dataset)}, Test={len(self.test_dataset)}")
def test_dataset_sizes(self):
"""Test that the datasets are non-empty and splits are correct."""
self.logger.info("Testing QA dataset sizes...")
self.assertGreater(len(self.train_dataset), 0)
self.assertGreater(len(self.val_dataset), 0)
self.assertGreater(len(self.test_dataset), 0)
self.logger.success("QA dataset size tests passed")
def test_sample_keys(self):
"""Test that a sample contains all required keys."""
self.logger.info("Testing QA sample keys...")
sample = self.train_dataset[0]
required_keys = {"answer", "pre_prompt", "post_prompt", "time_series", "time_series_text"}
self.assertTrue(required_keys.issubset(sample.keys()))
self.logger.success("QA sample keys test passed")
def test_answer_is_rationale(self):
"""Test that the answer is a string (rationale)."""
self.logger.info("Testing answer format...")
sample = self.train_dataset[0]
self.assertIsInstance(sample["answer"], str)
self.assertGreater(len(sample["answer"]), 0)
self.logger.success(f"Answer test passed: length={len(sample['answer'])}")
def test_time_series_content(self):
"""Test that the time series and text are present and valid."""
self.logger.info("Testing time series content...")
sample = self.train_dataset[0]
self.assertIsInstance(sample["time_series"], list)
self.assertIsInstance(sample["time_series_text"], list)
self.assertGreater(len(sample["time_series"][0]), 0)
self.assertIsInstance(sample["time_series_text"][0], str)
self.logger.success(f"Time series test passed: {len(sample['time_series'])} series")
def test_time_series_text_includes_mean_std(self):
"""Test that each time_series_text includes 'mean' and 'std', and both are followed by a number."""
import re
self.logger.info("Testing time series text format...")
sample = self.train_dataset[0]
for i, text in enumerate(sample['time_series_text']):
self.logger.debug(f"Testing text {i}: {text[:100]}...")
self.assertIn('mean', text)
self.assertIn('std', text)
# Allow for any whitespace after 'mean' and 'std'
mean_match = re.search(r"mean\s+(-?\d+\.\d+)", text)
if not mean_match:
self.logger.error(f"DEBUG: {repr(text)}")
self.assertIsNotNone(mean_match, f"No mean value found in: {repr(text)}")
std_match = re.search(r"std\s+(-?\d+\.\d+)", text)
if not std_match:
self.logger.error(f"DEBUG: {repr(text)}")
self.assertIsNotNone(std_match, f"No std value found in: {repr(text)}")
self.logger.success("Time series text format tests passed")
def test_example_data_QA(self):
"""Print example data for PAMAP2CoTQADataset, showing all time series and text."""
sample = self.train_dataset[0]
self.logger.info("="*80)
self.logger.info("EXAMPLE PAMAP2CoTQADataset SAMPLE")
self.logger.info("="*80)
self.logger.info(f"Pre-prompt: '{sample['pre_prompt']}'")
self.logger.info(f"Post-prompt: '{sample['post_prompt']}'")
self.logger.info(f"Answer (rationale): '{sample['answer']}'")
self.logger.info(f"Number of time series: {len(sample['time_series'])}")
for i, (ts, ts_text) in enumerate(zip(sample['time_series'], sample['time_series_text'])):
self.logger.info(f"Time series {i} text: '{ts_text}'")
self.logger.info(f"Time series {i} length: {len(ts)}")
self.logger.info(f"First 10 values: {ts[:10]}")
self.logger.info(f"Last 10 values: {ts[-10:]}")
self.logger.info("="*80)
class TestBalancedBatchSampler(unittest.TestCase):
def test_balanced_batches(self):
"""
Test that BalancedBatchSampler produces batches where each class is equally represented.
This test creates an imbalanced label list (10 'a', 4 'b', 6 'c') and sets batch_size=6 (2 samples per class per batch).
For each batch yielded by the sampler, we check that every class appears exactly 2 times.
This ensures the sampler yields perfectly balanced mini-batches, even when the dataset is imbalanced.
"""
# Create a toy label list with imbalance
labels = ['a'] * 10 + ['b'] * 4 + ['c'] * 6
batch_size = 6 # 3 classes, so 2 samples per class per batch
sampler = BalancedBatchSampler(labels, batch_size)
for batch in sampler:
batch_labels = [labels[idx] for idx in batch]
counts = {l: batch_labels.count(l) for l in set(batch_labels)}
print(f"Batch labels: {batch_labels}")
print(f"Class counts in batch: {counts}")
# Each class should appear exactly 2 times per batch
for count in counts.values():
self.assertEqual(count, 2, f"Expected 2 samples per class per batch, got {count}")
def test_balanced_batches_pamap2cot(self):
"""
Test that BalancedBatchSampler produces balanced batches on the real PAMAP2CoTQADataset training split.
Prints batch labels and class counts for each batch.
"""
from opentslm.time_series_datasets.pamap2.PAMAP2CoTQADataset import PAMAP2CoTQADataset
from opentslm.time_series_datasets.pamap2.BalancedBatchSampler import BalancedBatchSampler
# Helper to extract label from answer string
def extract_label_from_answer(answer: str) -> str:
# Assumes answer ends with 'Answer: <label>' or 'Answer: <label>.'
if 'Answer:' in answer:
label = answer.split('Answer:')[-1].strip()
# Remove trailing period if present
if label.endswith('.'):
label = label[:-1]
return label.strip()
return ''
# Load the real dataset
dataset = PAMAP2CoTQADataset(split="train", EOS_TOKEN="")
labels = [extract_label_from_answer(row["answer"]) for row in dataset]
num_classes = len(set(labels))
batch_size = num_classes * 2 # 2 samples per class per batch
sampler = BalancedBatchSampler(labels, batch_size)
for i, batch in enumerate(sampler):
batch_labels = [labels[idx] for idx in batch]
counts = {l: batch_labels.count(l) for l in set(batch_labels)}
print(f"Batch {i} labels: {batch_labels}")
print(f"Batch {i} class counts: {counts}")
for l, count in counts.items():
self.assertEqual(count, 2, f"Expected 2 samples for class {l} per batch, got {count}")
if __name__ == "__main__":
unittest.main() |