Spaces:
Running
Running
| import os | |
| import gc | |
| import numpy as np | |
| import ray | |
| from ray.experimental import tqdm_ray | |
| import pandas as pd | |
| from pandas.tseries.frequencies import to_offset | |
| from pathlib import Path | |
| from tqdm import tqdm | |
| from collections import defaultdict | |
| from functools import cached_property | |
| from gluonts.time_feature import norm_freq_str | |
| from gift_eval.data import Dataset | |
| from .features import get_ts_features | |
| from .utils import persist_analysis | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| MAX_CONTEXT_LEN = 500 | |
| runtime_env = { | |
| 'env_vars': { | |
| "RAY_memory_usage_threshold": "0.85" | |
| } | |
| } | |
| if not os.getenv("NUM_CPUS"): | |
| print("NUM_CPUS environment variable not found. Setting to 1. Set NUM_CPUS to speed up processing.") | |
| NUM_CPUS = int(os.getenv("NUM_CPUS", "1")) | |
| def process_instance(self, test_input, test_label, dataset_freq): | |
| """ | |
| Process a single time series instance to compute features. | |
| Parameters: | |
| - self: Reference to the calling object. | |
| - test_input: Dictionary containing the input time series data. | |
| - test_label: Dictionary containing the label time series data. | |
| - dataset_freq: Frequency of the dataset. | |
| Returns: | |
| - DataFrame containing the computed features for the time series instance. | |
| """ | |
| np_inp = np.array(test_input["target"]) | |
| np_label = np.array(test_label["target"]) | |
| # Check if the input is 2D and trim to MAX_CONTEXT_LEN if necessary | |
| if len(np_inp.shape) == 2: | |
| if np_inp.shape[1] > MAX_CONTEXT_LEN: | |
| np_inp = np_inp[:, -MAX_CONTEXT_LEN:] | |
| np_instance = np.concatenate((np_inp, np_label), axis=1) | |
| else: | |
| if len(np_inp) > MAX_CONTEXT_LEN: | |
| np_inp = np_inp[-MAX_CONTEXT_LEN:] | |
| np_instance = np.concatenate((np_inp, np_label)) | |
| # Compute time series features | |
| window_features_df = get_ts_features( | |
| np_instance, norm_freq_str(to_offset(dataset_freq).name)) | |
| self.pbar.update.remote(1) | |
| return window_features_df | |
| def process_dataset(self, dataset, output_dir): | |
| """ | |
| Process an entire dataset to compute features for each time series instance. | |
| Parameters: | |
| - self: Reference to the calling object. | |
| - dataset: The dataset to be processed. | |
| - output_dir: Directory where the processed data will be saved. | |
| Returns: | |
| - None, but updates the progress bar and persists the analysis results. | |
| """ | |
| # Determine the directory for the dataset based on its term and name | |
| if str(dataset.term) == "Term.SHORT": | |
| dataset_dir = Path(os.path.join(os.path.dirname( | |
| output_dir), f"datasets/{dataset.name}")) | |
| else: | |
| if "/" in dataset.name: | |
| dataset_name, dataset_freq = dataset.name.split("/") | |
| dataset_name = f"{dataset_name}:{dataset.term}/{dataset_freq}" | |
| dataset_dir = Path(os.path.join(os.path.dirname( | |
| output_dir), f"datasets/{dataset_name}")) | |
| else: | |
| dataset_dir = Path(os.path.join(os.path.dirname( | |
| output_dir), f"datasets/{dataset.name}:{dataset.term}")) | |
| # Create the directory if it doesn't exist | |
| if not dataset_dir.exists(): | |
| dataset_dir.mkdir(parents=True, exist_ok=True) | |
| print("Directory created:", dataset_dir) | |
| else: | |
| print("Directory already exists:", dataset_dir) | |
| # Assume dataset has already been processed | |
| self.pbar.update.remote(dataset.hf_dataset.num_rows * dataset.windows) | |
| return None | |
| all_features_list = [] | |
| test_data = dataset.test_data | |
| # Process each instance in the dataset | |
| features = [process_instance.remote( | |
| self, test_input, test_label, dataset.freq) for test_input, test_label in test_data] | |
| for feature in features: | |
| try: | |
| # Retrieve the result with a timeout | |
| result = ray.get(feature, timeout=300) # 300 seconds timeout | |
| all_features_list.append(result) | |
| except ray.exceptions.GetTimeoutError: | |
| print("A task timed out and will be skipped.") | |
| continue # Skip this particular instance | |
| except Exception as e: | |
| print(f"An error occurred while processing: {e}") | |
| continue | |
| gc.collect() | |
| # Concatenate all features and persist the analysis | |
| all_features_df = pd.concat(all_features_list) | |
| persist_analysis(all_features_df, dataset_dir) | |
| class Analyzer(): | |
| """ | |
| Analyzer class to manage the analysis of multiple datasets, including feature computation and frequency distribution analysis. | |
| """ | |
| def __init__(self, datasets: list[Dataset]): | |
| """ | |
| Initialize the Analyzer with a list of datasets. | |
| Parameters: | |
| - datasets: List of Dataset objects to be analyzed. | |
| """ | |
| self.datasets = datasets | |
| ray.init(runtime_env=runtime_env) | |
| remote_tqdm = ray.remote(tqdm_ray.tqdm) | |
| self.pbar = remote_tqdm.remote(total=self._sum_windows_count) | |
| def print_datasets(self): | |
| """Print the names of all datasets.""" | |
| for dataset in self.datasets: | |
| print(dataset.name) | |
| def _sum_series_count(self) -> int: | |
| """Calculate the total number of series across all datasets.""" | |
| total_count = 0 | |
| for dataset in self.datasets: | |
| total_count += dataset.hf_dataset.num_rows | |
| return total_count | |
| def _sum_windows_count(self) -> int: | |
| """Calculate the total number of windows across all datasets.""" | |
| total_count = 0 | |
| for dataset in self.datasets: | |
| total_count += dataset.hf_dataset.num_rows * dataset.windows | |
| return total_count | |
| def freq_distribution_by_dataset(self): | |
| """Compute the frequency distribution by dataset.""" | |
| freqs = [norm_freq_str(to_offset(dataset.freq).name) | |
| for dataset in self.datasets] | |
| freq_counts = {freq: freqs.count(freq) for freq in set(freqs)} | |
| return freq_counts | |
| def freq_distribution_by_ts(self): | |
| """Compute the frequency distribution by time series count.""" | |
| freq_ts_counts = defaultdict(lambda: 0) | |
| for dataset in self.datasets: | |
| freq_ts_counts[norm_freq_str( | |
| to_offset(dataset.freq).name)] += dataset.hf_dataset.num_rows | |
| return freq_ts_counts | |
| def freq_distribution_by_ts_length(self): | |
| """Compute the frequency distribution by time series length.""" | |
| freq_dp_counts = defaultdict(lambda: 0) | |
| for dataset in self.datasets: | |
| freq_dp_counts[norm_freq_str( | |
| to_offset(dataset.freq).name)] += dataset.sum_series_length | |
| return freq_dp_counts | |
| def freq_distribution_by_window(self): | |
| """Compute the frequency distribution by window count.""" | |
| freq_window_counts = defaultdict(lambda: 0) | |
| for dataset in self.datasets: | |
| freq_window_counts[norm_freq_str( | |
| to_offset(dataset.freq).name)] += dataset.hf_dataset.num_rows * dataset.windows | |
| return freq_window_counts | |
| def features_by_window(self, output_dir): | |
| """ | |
| Create and persist features of each window for each dataset. | |
| Parameters: | |
| - output_dir: Directory where the features will be saved. | |
| """ | |
| ray.get([process_dataset.remote(self, dataset, output_dir) | |
| for dataset in self.datasets]) | |
| self.pbar.close.remote() | |
| all_datasets_df = [] | |
| # Aggregate the characteristics for each dataset | |
| with tqdm(total=len(self.datasets), desc="Computing ts features for whole benchmark") as pbar: | |
| for dataset in self.datasets: | |
| pbar.set_description(f"Loading ts features | {dataset.name}") | |
| if str(dataset.term) == "Term.SHORT": | |
| dataset_df_path = os.path.join(os.path.dirname( | |
| output_dir), f"datasets/{dataset.name}/features.csv") | |
| else: | |
| if "/" in dataset.name: | |
| dataset_name, dataset_freq = dataset.name.split("/") | |
| dataset_name = f"{dataset_name}:{dataset.term}/{dataset_freq}" | |
| dataset_df_path = Path(os.path.join(os.path.dirname( | |
| output_dir), f"datasets/{dataset_name}/features.csv")) | |
| else: | |
| dataset_df_path = Path(os.path.join(os.path.dirname( | |
| output_dir), f"datasets/{dataset.name}:{dataset.term}/features.csv")) | |
| df = pd.read_csv(dataset_df_path) | |
| all_datasets_df.append(df) | |
| pbar.update(1) | |
| # Concatenate all dataset features and persist the analysis | |
| all_features_df = pd.concat(all_datasets_df, ignore_index=True) | |
| persist_analysis(all_features_df, output_dir) | |
| return None | |