File size: 2,695 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 | # 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 List, Literal, Optional, Tuple
from datasets import Dataset
from opentslm.prompt.text_time_series_prompt import TextTimeSeriesPrompt
from opentslm.time_series_datasets.monash.MonashDataset import MonashDataset
from opentslm.time_series_datasets.QADataset import QADataset
from opentslm.time_series_datasets.util import (
extend_time_series_to_match_patch_size_and_aggregate,
)
import torch
from torch.utils.data import DataLoader, random_split
from tqdm.auto import tqdm
TIME_SERIS_LABELS = ["The following is PPG data", "The following is ECG data"]
class MonashSPO2QADataset(QADataset):
def _load_splits(self) -> Tuple[Dataset, Dataset, Dataset]:
train = MonashDataset(
_data_dir="monash_datasets",
data_name="BIDMC32SpO2/BIDMC32SpO2_TRAIN",
)
test = MonashDataset(
_data_dir="monash_datasets",
data_name="BIDMC32SpO2/BIDMC32SpO2_TEST",
)
train_size = int(len(train) * 0.9)
val_size = len(train) - train_size
train, val = random_split(
train, [train_size, val_size], generator=torch.Generator().manual_seed(42)
)
return train, val, test
def _get_answer(self, row) -> str:
return str(round(row["answer"], 2))
def _get_pre_prompt(self, _row) -> str:
return "You are given PPG and ECG data. Your task is to predict the average blood oxygen saturation over the given the 32 second period."
def _get_post_prompt(self, _row) -> str:
return "Answer:"
def _get_text_time_series_prompt_list(self, row) -> List[TextTimeSeriesPrompt]:
if len(row["time_series"][0]) != len(TIME_SERIS_LABELS):
raise RuntimeError(
"question labels and time series from the data must be of the same length"
)
# TODO normalize
return [
TextTimeSeriesPrompt(time_series_label, time_series)
for time_series_label, time_series in zip(
TIME_SERIS_LABELS, row["time_series"][0]
)
]
if __name__ == "__main__":
dataset = MonashSPO2QADataset(split="train", EOS_TOKEN="")
dataloader = DataLoader(
dataset,
batch_size=4,
shuffle=True,
collate_fn=lambda batch: extend_time_series_to_match_patch_size_and_aggregate(
batch, patch_size=4
),
)
for batch in tqdm(dataloader):
print(batch)
|