File size: 2,899 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 | #!/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 M4QADataset with caption generation.
"""
import unittest
from opentslm.time_series_datasets.m4.M4QADataset import M4QADataset as _M4QADataset
class TestM4QADataset(unittest.TestCase):
"""
Unit tests for the M4QADataset class and loader.
"""
def setUp(self):
self.M4QADataset = _M4QADataset
self.train_dataset = self.M4QADataset("train", "")
self.val_dataset = self.M4QADataset("validation", "")
self.test_dataset = self.M4QADataset("test", "")
def test_dataset_sizes(self):
"""Test that the datasets are non-empty and splits are correct."""
self.assertGreater(len(self.train_dataset), 0)
self.assertGreater(len(self.val_dataset), 0)
self.assertGreater(len(self.test_dataset), 0)
def test_sample_keys(self):
"""Test that a sample contains all required keys."""
sample = self.train_dataset[0]
required_keys = {"answer", "post_prompt", "pre_prompt", "time_series", "time_series_text"}
self.assertTrue(required_keys.issubset(sample.keys()))
def test_time_series_content(self):
"""Test that the time series and text are present and valid."""
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)
def test_caption_is_answer(self):
"""Test that the answer is a string (caption)."""
sample = self.train_dataset[0]
self.assertIsInstance(sample["answer"], str)
self.assertGreater(len(sample["answer"]), 0)
def test_example_data(self):
"""Print example data to show what the dataset looks like."""
sample = self.train_dataset[0]
print("\n" + "="*80)
print("EXAMPLE M4 DATASET SAMPLE")
print("="*80)
print(f"Pre-prompt: '{sample['pre_prompt']}'")
print(f"Post-prompt: '{sample['post_prompt']}'")
print(f"Answer (caption): '{sample['answer'][:200]}...'")
print(f"Number of time series: {len(sample['time_series'])}")
if sample['time_series']:
ts = sample['time_series'][0]
ts_text = sample['time_series_text'][0]
print(f"Time series text: '{ts_text}'")
print(f"Time series length: {len(ts)}")
print(f"First 10 time series values: {ts[:10]}")
print(f"Last 10 time series values: {ts[-10:]}")
print("="*80)
if __name__ == "__main__":
unittest.main()
|