File size: 6,665 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 | # 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
import torch
from typing import List, Tuple
from datasets import Dataset
from opentslm.prompt.text_time_series_prompt import TextTimeSeriesPrompt
from opentslm.time_series_datasets.QADataset import QADataset
DATASET_SIZE = 200
class SimulationQADataset(QADataset):
def __init__(
self,
split,
EOS_TOKEN,
length: int = 100,
num_series: int = 1,
format_sample_str: bool = False,
time_series_format_function=None,
):
"""
Initialize SimulationQADataset with one or more time series of variable length.
Args:
split: Dataset split (train/test/validation) - all return the same single item
EOS_TOKEN: End-of-sequence token
length: Length of the generated time series (default: 100)
num_series: Number of time series to generate (default: 1)
format_sample_str: Whether to format as string
time_series_format_function: Optional time series formatting function
"""
self.length = length
self.num_series = num_series
super().__init__(
split, EOS_TOKEN, format_sample_str, time_series_format_function
)
def _load_splits(self) -> Tuple[Dataset, Dataset, Dataset]:
"""
Creates a dataset with 200 items, each with random time series data.
Each item will have num_series time series of length elements.
"""
all_items = []
for _ in range(DATASET_SIZE):
# Generate random time series for this item
time_series_data = {}
for i in range(self.num_series):
series = torch.tensor(np.random.randn(self.length), dtype=torch.float32)
# Normalize the series
mean_val = series.mean().item()
std_val = max(series.std().item(), 1e-6)
normalized_series = (series - mean_val) / std_val
time_series_data[f"series_{i}"] = normalized_series.tolist()
time_series_data[f"series_text_{i}"] = (
f"This is a time series with mean {mean_val:.4f} and std {std_val:.4f}."
)
item_data = {
**time_series_data,
"Question": f"You are given different time series. All have the same length of {self.length} data points. What is the pattern of the time series?",
"Answer": "This is a random pattern.",
}
all_items.append(item_data)
# Convert to HuggingFace Dataset format
dataset_dict = {}
for key in all_items[0].keys():
dataset_dict[key] = [item[key] for item in all_items]
dataset = Dataset.from_dict(dataset_dict)
# Return the same dataset for all splits
return dataset, dataset, dataset
def _get_answer(self, row) -> str:
"""Get the answer from the data row."""
return row["Answer"]
def _get_pre_prompt(self, row) -> str:
"""Get the question/pre-prompt from the data row."""
return row["Question"]
def _get_post_prompt(self, row) -> str:
"""Get the post-prompt from the data row."""
return "Predict the pattern of the time series. Answer:"
def _get_text_time_series_prompt_list(self, row) -> List[TextTimeSeriesPrompt]:
"""
Convert the time series data from the current row to TextTimeSeriesPrompt format.
Each row now contains its own random time series data.
"""
prompts = []
for i in range(self.num_series):
series_key = f"series_{i}"
series_text_key = f"series_text_{i}"
if series_key in row and series_text_key in row:
series_data = row[series_key]
text_description = row[series_text_key]
prompts.append(TextTimeSeriesPrompt(text_description, series_data))
return prompts
if __name__ == "__main__":
# Example usage - Single time series
print("=== Single Time Series Dataset (10,000 items) ===")
dataset_single = SimulationQADataset("train", "", length=50, num_series=1)
print(f"Dataset length: {len(dataset_single)}")
sample_single = dataset_single[0]
print(f"Sample keys: {sample_single.keys()}")
print(f"Question: {sample_single['pre_prompt'][:100]}...")
print(f"Answer: {sample_single['answer']}")
print(f"Number of time series prompts: {len(sample_single['time_series_prompts'])}")
print("\n=== Multiple Time Series Dataset (10,000 items) ===")
# Example usage - Multiple time series
dataset_multi = SimulationQADataset("train", "", length=50, num_series=3)
print(f"Dataset length: {len(dataset_multi)}")
sample_multi = dataset_multi[0]
print(f"Sample keys: {sample_multi.keys()}")
print(f"Question: {sample_multi['pre_prompt'][:100]}...")
print(f"Answer: {sample_multi['answer']}")
print(f"Number of time series prompts: {len(sample_multi['time_series_prompts'])}")
# Show time series prompt details
for i, ts_prompt in enumerate(sample_multi["time_series_prompts"]):
print(f"Time series {i}: {ts_prompt.text[:50]}...")
# Test different splits (should all be the same)
print("\n=== Testing Different Splits ===")
train_dataset = SimulationQADataset("train", "", length=50, num_series=2)
val_dataset = SimulationQADataset("validation", "", length=50, num_series=2)
test_dataset = SimulationQADataset("test", "", length=50, num_series=2)
print(f"Train length: {len(train_dataset)}")
print(f"Val length: {len(val_dataset)}")
print(f"Test length: {len(test_dataset)}")
# Test that different items have different data
print("\n=== Testing Randomness ===")
item_0 = dataset_single[0]
item_100 = dataset_single[100]
print(f"Item 0 Series length: {len(item_0['time_series_prompts'][0].time_series)}")
print(
f"Item 100 Series length: {len(item_100['time_series_prompts'][0].time_series)}"
)
print(
f"First 5 values of item 0: {item_0['time_series_prompts'][0].time_series[:5]}"
)
print(
f"First 5 values of item 100: {item_100['time_series_prompts'][0].time_series[:5]}"
)
print(
"Values are different:",
item_0["time_series_prompts"][0].time_series[:5]
!= item_100["time_series_prompts"][0].time_series[:5],
) |