OpenTSLab/SciTS

1import json2import random3from pathlib import Path4from typing import Sequence, Callable5 6import fire7import h5py8import numpy as np9from tqdm import tqdm10 11from infer_eval_utils import (12    read_time_series_data,13    GENERATION_TASK_IDS, 14    CLASSIFICATION_TASK_IDS, 15    EVENT_DETECTION_TASK_IDS, 16    ANOMALY_DETECTION_TASK_IDS,17    MCQ_TASK_IDS, 18    IMPUTATION_TASK_IDS19)20 21 22def read_raw_data(path: str | None) -> Sequence:23    if path is None:24        return []25 26    return read_time_series_data(path)27 28 29def transform_raw_data_single_channel(raw_data: list | np.ndarray,30                                      dataset_id: str) -> str:31 32    if dataset_id in GENERATION_TASK_IDS:33        data_str_list = []34        for x in raw_data:35            if x != "X":36                data_str_list.append(f"{float(x):.3f}")37            else:38                data_str_list.append(x)39        data_text = " ".join(data_str_list)40    else:41        data_text = " ".join([f"{float(x):.3f}" for x in raw_data])42 43    return data_text44 45 46def transform_raw_data_to_text(raw_data: list | np.ndarray, dataset_id: str,47                               channel_detail: list[str]) -> str:48 49    if isinstance(raw_data, np.ndarray):50        if raw_data.ndim > 1 and raw_data.shape[1] > 1 and len(51                channel_detail) == 0:52            channel_detail = [f"channel {i}" for i in range(raw_data.shape[1])]53 54        if len(channel_detail) <= 1:55            data_text = transform_raw_data_single_channel(raw_data, dataset_id)56        else:57            data_text = ""58            for channel_idx, channel_name in enumerate(channel_detail):59                channel_data = raw_data[:, channel_idx]60                channel_data_text = transform_raw_data_single_channel(61                    channel_data, dataset_id)62                data_text += f"{channel_name}: {channel_data_text} "63    else:64        data_text = transform_raw_data_single_channel(raw_data, dataset_id)65    return data_text66 67 68def transform_gt_data_to_text(gt_data: list | np.ndarray,69                              dataset_id: str) -> str:70    gt_data = np.array(gt_data)71    if gt_data.ndim == 1:72        data_text = transform_raw_data_single_channel(gt_data, dataset_id)73    else:74        data_text = ""75        for channel_idx in range(gt_data.shape[1]):76            channel_data = gt_data[:, channel_idx]77            channel_data_text = transform_raw_data_single_channel(78                channel_data, dataset_id)79            data_text += f"{channel_data_text}\n"80 81    return data_text82 83 84def get_extra_instruction(dataset_id: str, ) -> str:85    extra_instruction = ""86    if dataset_id == "ASU01_ASG02":87        extra_instruction = "Answer yes or no in the first line. If the Gravitational Wave is detected, answer the index of the starting time point in the second line."88    elif dataset_id == "EAU01_EAG02":89        extra_instruction = "Answer yes or no in the first line. If an Earthquake event is detected, answer the starting time point index of the P-wave in the second line, " \90                            "answer the starting time point index of the S-wave in the third line."91    elif dataset_id == "MFU01_MFU02":92        extra_instruction = "Output the diameter in the first line, and the position in the second line."93    elif dataset_id == "PHU01":94        extra_instruction = "Give each answer in a line. For example, if the answer is ['NORM', 'MI'], you should output: NORM\nMI."95    elif dataset_id == "MAG01":96        extra_instruction = "Give answer of each channel in a line so the number of predicted time points in each line should match the given one. For example, if " \97            "it is required to predict the next 5 time points, and the predicted x0, x1, x2 are [[0.1, 0.2, 0.3, 0.4, 0.5], [0.4, 0.5, 0.6, 0.7, 0.8], [0.7, 0.8, 0.9, 0.1, 0.2]], " \98            "you should output: 0.1 0.2 0.3 0.4 0.5\n0.4 0.5 0.6 0.7 0.8\n0.7 0.8 0.9 0.1 0.2."99    elif dataset_id in ANOMALY_DETECTION_TASK_IDS:100        extra_instruction = "Answer yes if anomaly points are detected, and no if there are only normal points."101    elif dataset_id in GENERATION_TASK_IDS:102        extra_instruction = "Output the values separated by spaces."103    return extra_instruction104 105 106def extract_gt(data: dict, dataset_id: str) -> str | dict | Path:107    if dataset_id in CLASSIFICATION_TASK_IDS:108        gt = data["gt_result"]["gt_class"]109        if isinstance(gt, dict) and len(gt) == 1:110            gt = gt["default"]111        if isinstance(gt, list) and len(gt) == 1:112            gt = gt[0]113    elif dataset_id in GENERATION_TASK_IDS:114        gt = data["gt_ts"]["path"].strip("/")115    elif dataset_id in EVENT_DETECTION_TASK_IDS:116        gt = data["gt_result"]117    elif dataset_id in ANOMALY_DETECTION_TASK_IDS:118        gt = data["gt_result"]["contain"]119    elif dataset_id in MCQ_TASK_IDS:120        gt = data["gt_result"]["answer"]121    else:122        raise ValueError(f"Unsupported dataset id: {dataset_id}")123    return gt124 125 126def initialize_model() -> Callable:127    """128    Initialize the model here. The model can be called by:129    130    ```python131    response = model(prompt) 132    # or133    response = model(prompt, max_tokens=max_tokens)  # to limit the response length134    ```135    """136    pass137 138 139def infer_dataset(model: Callable, dataset_data: list, scits_dir: Path,140                  dataset_id: str, output_path: Path):141    print(f"Inferring {dataset_id}")142 143    if dataset_id in GENERATION_TASK_IDS:144        ext = "h5"145    else:146        ext = "jsonl"147 148    output_path = Path(output_path) / f"{dataset_id}.{ext}"149    output_path.parent.mkdir(parents=True, exist_ok=True)150    completed_ids = []151 152    if str(output_path).endswith(".jsonl"):153        has_metadata = False154 155    if output_path.exists():156        if str(output_path).endswith(".jsonl"):157            with open(output_path, 'r') as f:158                for line in f.readlines():159                    data = json.loads(line)160                    if "id" in data:161                        completed_ids.append(data["id"])162                    else:163                        has_metadata = True164        elif str(output_path).endswith(".h5"):165            with h5py.File(output_path, 'r') as f:166                completed_ids = list(f.keys())167 168    completed_ids = set(completed_ids)169    random.shuffle(dataset_data)170    dataset_data = dataset_data[:10]171 172    try:173        seq_length = dataset_data[0]["input_ts"]["length"]174    except:175        seq_length = None176 177    for sample in tqdm(dataset_data):178        id = sample["id"].replace(179            "/", "%2F")  # to avoid errors related to "/" in hdf5180 181        if id in completed_ids:182            continue183 184        # Load raw data185        if sample["input_ts"] is None:186            raw_data_path = None187            channel_detail = None188        else:189            raw_data_path = scits_dir / sample["input_ts"]["path"].strip("/")190            channel_detail = sample["input_ts"]["channel_detail"]191 192        raw_data = read_raw_data(raw_data_path)193        raw_data_text = transform_raw_data_to_text(raw_data, dataset_id,194                                                   channel_detail)195 196        gt = extract_gt(sample, dataset_id)197        extra_instruction = get_extra_instruction(dataset_id)198 199        if dataset_id in GENERATION_TASK_IDS:200            # give max_tokens to save cost for generation tasks201            gt_data = read_time_series_data(scits_dir / gt)202            gt_data_text = transform_gt_data_to_text(203                gt_data, dataset_id)204            max_tokens = len(gt_data_text)205        else:206            max_tokens = None207 208        prompt_text = f'{sample["input_text"]} {extra_instruction} Give me the answer directly, ' \209                    f'without any other extra content (including punctuation). ' \210                    f'{raw_data_text}'211        output_text = model(text=prompt_text, max_tokens=max_tokens)212        # print(f"output_text: {output_text}")213 214        if dataset_id not in GENERATION_TASK_IDS:215            with open(output_path, 'a') as writer:216                if not has_metadata:217                    metadata = {}218                    if dataset_id in EVENT_DETECTION_TASK_IDS:219                        metadata["seq_length"] = seq_length220                    writer.write(json.dumps(metadata) + "\n")221                    has_metadata = True222 223                pred_results = output_text224                if "class_list" in sample["gt_result"] and isinstance(225                        sample["gt_result"]["class_list"], dict) and len(226                            sample["gt_result"]["class_list"]) > 1:227                    if len(output_text.split("\n")) != len(228                            sample["gt_result"]["class_list"]):229                        pred_results = "NA"230                    else:231                        pred_results = {}232                        for class_name, pred_result in zip(233                                sample["gt_result"]["class_list"],234                                output_text.split("\n")):235                            pred_results[class_name] = pred_result236                writer.write(237                    json.dumps({238                        "id": id,239                        "output": pred_results,240                        "ground_truth": gt241                    }) + "\n")242        else:243            if "\n" not in output_text:244                pred_result = np.fromstring(output_text.strip(),245                                            dtype=np.float32,246                                            sep=' ')247            else:248                try:249                    pred_result = np.vstack([250                        np.fromstring(x.strip(), dtype=np.float32, sep=' ')251                        for x in output_text.split("\n")252                    ]).transpose()253                except ValueError:254                    pred_result = np.array([])255 256            with h5py.File(output_path, 'a') as writer:257                writer[f"{id}/pred_result"] = pred_result258                writer[f"{id}/gt_path"] = gt.__str__().encode("utf-8")259                if dataset_id in IMPUTATION_TASK_IDS:260                    writer[f"{id}/input_ts_path"] = sample["input_ts"][261                        "path"].strip("/").encode("utf-8")262                if "base_path" not in writer:263                    writer["base_path"] = scits_dir.__str__().encode("utf-8")264 265 266 267def infer(268    scits_dir: str,269    output_dir: str,270):271 272    # Initialize caller273    model: Callable = initialize_model()274 275    scits_dir = Path(scits_dir)276    output_dir = Path(output_dir)277    dataset_data = []278    prev_dataset_id = None279    with open(scits_dir / "meta_data.jsonl", 'r') as f:280        for line in f.readlines():281            sample = json.loads(line)282            dataset_id = "_".join(sample["task_id"])283 284            if dataset_id != prev_dataset_id:285                if prev_dataset_id is not None:286                    infer_dataset(model, dataset_data, scits_dir, prev_dataset_id,287                                  output_dir)288                dataset_data = []289                prev_dataset_id = dataset_id290            dataset_data.append(sample)291 292    infer_dataset(model, dataset_data, scits_dir, prev_dataset_id, output_dir)293 294 295if __name__ == '__main__':296    fire.Fire(infer)297