update benchmark

Author: LongxingTan

.github/workflows/lint.yml

@@ -18,7 +18,7 @@ jobs:
           with:
             python-version: 3.8
         - name: Install Black
-          run: pip install black[jupyter]
+          run: pip install "black[jupyter]==24.8.0"
         - name: Run Black
           run: black --check .
 

.github/workflows/pypi_release.yml

@@ -23,11 +23,11 @@ jobs:
       - name: Install poetry
         shell: bash
         run: |
-          curl -sSL https://install.python-poetry.org | python3 -
+          curl -sSL https://install.python-poetry.org | python3 - --version 1.8.5
           python -m pip install poetry-dynamic-versioning[plugin]
 
       - name: Set poetry path variable
-        run: echo "/Users/runner/.local/bin" >> $GITHUB_PATH
+        run: echo "$HOME/.local/bin" >> "$GITHUB_PATH"
 
       - name: Build
         run: |

.github/workflows/test.yml

@@ -21,6 +21,8 @@ jobs:
         exclude:
           - python-version: 3.9
             tf-version: 2.13.1
+          - os: macOS-latest
+            tf-version: 2.13.1
 
     steps:
       - uses: actions/checkout@v4
@@ -37,9 +39,9 @@ jobs:
       - name: Install poetry
         shell: bash
         run: |
-          curl -sSL https://install.python-poetry.org | python3 -
+          curl -sSL https://install.python-poetry.org | python3 - --version 1.8.5
       - name: Set poetry path variable
-        run: echo "/Users/runner/.local/bin" >> $GITHUB_PATH
+        run: echo "$HOME/.local/bin" >> "$GITHUB_PATH"
 
       - name: Configure poetry
         shell: bash

.gitignore

@@ -96,3 +96,4 @@ coverage.xml
 !/weights/.gitkeep
 CLAUDE.md
 temp/
+*.keras

.pre-commit-config.yaml

@@ -19,7 +19,7 @@ repos:
     hooks:
     - id: isort
   - repo: https://github.com/psf/black
-    rev: 22.3.0
+    rev: 24.8.0
     hooks:
     - id: black
   - repo: https://github.com/nbQA-dev/nbQA

benchmark/README.md

@@ -0,0 +1,112 @@
+# TFTS Benchmark System
+
+A flexible benchmarking framework for evaluating TFTS models across multiple datasets with multiple metrics and multiple runs. Designed for reproducibility and paper-quality results.
+
+## Usage
+
+### Command-Line
+
+```bash
+python -m benchmark.cli \
+    --models rnn transformer dlinear \
+    --datasets sine air_passengers \
+    --metrics mae rmse mape \
+    --runs 3 \
+    --epochs 50 \
+    --output-dir results/
+```
+
+### Python API
+
+```python
+from benchmark import BenchmarkConfig, BenchmarkRunner
+
+config = BenchmarkConfig(
+    models=["rnn", "transformer", "dlinear"],
+    datasets=["sine", "air_passengers"],
+    metrics=["mae", "rmse", "mape"],
+    runs=3,
+    epochs=50,
+    output_dir="benchmark_results",
+)
+
+runner = BenchmarkRunner(config)
+results = runner.run()
+
+results.print_table()          # console
+results.to_csv("results.csv")   # CSV
+results.to_latex("results.tex") # LaTeX for papers
+```
+
+## Adding a New Dataset
+
+```python
+from benchmark import Dataset
+import pandas as pd
+
+class MyDataset(Dataset):
+    name = "my_dataset"
+    description = "Description of my dataset"
+    train_length = 24
+    predict_sequence_length = 8
+
+    def prepare_data(self, **kwargs):
+        # Load your data from any source (CSV, DB, API, etc.)
+        x, y = ...
+        return x, y
+
+    def get_train_valid_split(self, **kwargs):
+        x, y = self.prepare_data(**kwargs)
+        # split into train/valid
+        return (x_train, y_train), (x_valid, y_valid)
+```
+
+Then register it:
+
+```python
+from benchmark import BenchmarkRunner, DatasetRegistry
+
+registry = DatasetRegistry()
+registry.register("my_dataset", MyDataset)
+
+config = BenchmarkConfig(datasets=["my_dataset"], ...)
+runner = BenchmarkRunner(config, dataset_registry=registry)
+results = runner.run()
+```
+
+## Architecture
+
+- **BenchmarkRunner**: Orchestrates running models on datasets, collecting results.
+- **Dataset**: Abstract base; each dataset subclass implements `prepare_data()` and returns standardized format.
+- **DatasetRegistry**: Maintains a registry of all available datasets.
+- **ModelRegistry**: Wraps existing tfts model mapping.
+- **BenchmarkMetrics**: Computes standard time-series metrics (MAE, MSE, RMSE, MAPE, etc.)
+- **BenchmarkResults**: Formats and exports results (CSV, JSON, LaTeX, console table).
+
+## Metrics
+
+Available metrics:
+- `mae`: Mean Absolute Error
+- `mse`: Mean Squared Error
+- `rmse`: Root Mean Squared Error
+- `mape`: Mean Absolute Percentage Error
+- `smape`: Symmetric MAPE
+- `r2`: R-squared
+
+## Migrated Example Benchmarks
+
+The previous `examples/benchmarks` tasks are available as registered datasets:
+
+- `forecasting_sticker_sales`
+- `CMI_detect_sleep_states`
+
+Both support `data_path` overrides through `per_dataset_config`. If the source
+CSV is not available, they generate deterministic placeholder data so the
+benchmark runner and CLI remain usable without Kaggle downloads.
+
+## Output Files
+
+After running, the following files are generated in `output_dir`:
+- `results.json`: Raw results for each run.
+- `results.csv`: Averaged results (mean/std per model-dataset).
+- `results.tex`: LaTeX table for papers.

benchmark/__init__.py

@@ -0,0 +1,38 @@
+"""TFTS Benchmark System.
+
+A flexible benchmarking framework for evaluating TFTS modelsacross multiple
+datasets with multiple metrics and multiple runs. Designed for reproducibility
+and paper-quality results.
+
+Example::
+
+    from benchmark import BenchmarkRunner, BenchmarkConfig
+    from benchmark.datasets import GrocerysalesDataset, RecruitRestaurantDataset
+
+    config = BenchmarkConfig(
+        models=["rnn", "transformer", "dlinear"],
+        datasets=["grocery_sales", "recruit_restaurant"],
+        metrics=["mae", "rmse", "mape"],
+        runs=5,
+        output_dir="benchmark_results",
+    )
+
+    runner = BenchmarkRunner(config)
+    results = runner.run()
+    results.to_latex("benchmark_results.tex")
+    results.to_csv("benchmark_results.csv")
+"""
+
+from benchmark.base import BenchmarkConfig, Dataset
+from benchmark.formatter import BenchmarkResults
+from benchmark.registry import DatasetRegistry, ModelRegistry
+from benchmark.runner import BenchmarkRunner
+
+__all__ = [
+    "BenchmarkRunner",
+    "BenchmarkConfig",
+    "BenchmarkResults",
+    "Dataset",
+    "DatasetRegistry",
+    "ModelRegistry",
+]

benchmark/base.py

@@ -0,0 +1,136 @@
+"""Base classes for the TFTS benchmark system."""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+import logging
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import pandas as pd
+import tensorflow as tf
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class BenchmarkConfig:
+    """Configuration for a benchmark run.
+
+    Attributes:
+        models: List of model names to evaluate (e.g., ["rnn", "transformer"]).
+            Use ``["all"]`` to run all registered models.
+        datasets: List of dataset names to evaluate (e.g., ["sine", "grocery_sales"]).
+            Use ``["all"]`` to run all registered datasets.
+        metrics: List of metric names to compute.
+            Available: ``"mae"``, ``"mse"``, ``"rmse"``, ``"mape"``, ``"smape"``, ``"r2"``.
+        runs: Number of runs per model-dataset pair (for statistical significance).
+        epochs: Number of training epochs per run.
+        batch_size: Batch size for training.
+        learning_rate: Learning rate for the optimizer.
+        train_length: Lookback window length. If None, each dataset provides its own.
+        predict_sequence_length: Forecast horizon. If None, each dataset provides its own.
+        seed: Base seed. Each run uses ``seed + run_idx`` for reproducibility.
+        output_dir: Directory to save results.
+        save_models: Whether to save trained model weights.
+        verbose: Verbosity level (0=silent, 1=progress, 2=detailed).
+        device: Device to run on (e.g., ``"/gpu:0"`` or ``"/cpu:0"``).
+        per_dataset_config: Optional per-dataset configuration overrides.
+            Keys are dataset names, values are dicts with keys like ``train_length``,
+            ``predict_sequence_length``, ``epochs``, etc.
+    """
+
+    models: List[str] = field(default_factory=lambda: ["all"])
+    datasets: List[str] = field(default_factory=lambda: ["all"])
+    metrics: List[str] = field(default_factory=lambda: ["mae", "rmse", "mape"])
+    runs: int = 1
+    epochs: int = 50
+    batch_size: int = 32
+    learning_rate: float = 1e-3
+    train_length: Optional[int] = None
+    predict_sequence_length: Optional[int] = None
+    seed: int = 42
+    output_dir: str = "benchmark_results"
+    save_models: bool = False
+    verbose: int = 1
+    device: str = ""
+    per_dataset_config: Dict[str, Dict[str, Any]] = field(default_factory=dict)
+
+    def __post_init__(self):
+        if self.runs < 1:
+            raise ValueError("runs must be >= 1")
+        if self.epochs < 1:
+            raise ValueError("epochs must be >= 1")
+
+    def get_dataset_config(self, dataset_name: str) -> Dict[str, Any]:
+        """Get configuration for a specific dataset, merging with defaults."""
+        config = {
+            "epochs": self.epochs,
+            "batch_size": self.batch_size,
+            "learning_rate": self.learning_rate,
+            "train_length": self.train_length,
+            "predict_sequence_length": self.predict_sequence_length,
+        }
+        if dataset_name in self.per_dataset_config:
+            config.update(self.per_dataset_config[dataset_name])
+        return config
+
+
+class Dataset(ABC):
+    """Abstract base class for benchmark datasets.
+
+        Each dataset subclass implements ``prepare_data()`` to load and format
+    the data. The class also provides metadata about the dataset.
+
+        Attributes:
+            name: Unique identifier for the dataset.
+            description: Human-readable description.
+            train_length: Default lookback window length (can be overridden by config).
+            predict_sequence_length: Default forecast horizon (can be overridden by config).
+    """
+
+    name: str = ""
+    description: str = ""
+    train_length: int = 24
+    predict_sequence_length: int = 8
+    num_features: int = 1
+    is_multivariate: bool = False
+    is_grouped: bool = False
+    target_column: str = "target"
+    time_column: str = "time"
+    group_column: Optional[str] = None
+
+    @abstractmethod
+    def prepare_data(self, **kwargs) -> Union[
+        Tuple[np.ndarray, np.ndarray],
+        Tuple[Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray]],
+        tf.data.Dataset,
+    ]:
+        """Prepare and return the dataset.
+
+        Returns:
+            Either (x_train, y_train), ((x_train, y_train), (x_valid, y_valid)),
+            or a tf.data.Dataset.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_train_valid_split(self, **kwargs) -> Tuple[Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray]]:
+        """Return train and validation splits.
+
+        Returns:
+            Tuple of (x_train, y_train), (x_valid, y_valid).
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def list_params(cls) -> Dict[str, Any]:
+        """Return dataset parameters for display."""
+        return {
+            "name": cls.name,
+            "description": cls.description,
+            "train_length": cls.train_length,
+            "predict_sequence_length": cls.predict_sequence_length,
+            "num_features": cls.num_features,
+            "is_multivariate": cls.is_multivariate,
+            "is_grouped": cls.is_grouped,
+        }