From 7b9574f666a92f5a4d99460b6e5c58fc64d14e9e Mon Sep 17 00:00:00 2001
From: Oleksandr Shchur <shchuro@amazon.com>
Date: Fri, 22 Mar 2024 17:24:44 +0000
Subject: [PATCH] Add item_id to ERCOT and ETT datasets

---
 src/gluonts/dataset/repository/_ercot.py     | 2 ++
 src/gluonts/dataset/repository/_ett_small.py | 2 ++
 2 files changed, 4 insertions(+)

diff --git a/src/gluonts/dataset/repository/_ercot.py b/src/gluonts/dataset/repository/_ercot.py
index 58f79910b4..d013a09c9f 100644
--- a/src/gluonts/dataset/repository/_ercot.py
+++ b/src/gluonts/dataset/repository/_ercot.py
@@ -34,6 +34,7 @@ def generate_ercot_dataset(dataset_path: Path, dataset_writer: DatasetWriter):
     test = [
         {
             "start": start,
+            "item_id": region,
             "target": df[region].to_numpy(dtype=np.float64),
         }
         for region in regions
@@ -42,6 +43,7 @@ def generate_ercot_dataset(dataset_path: Path, dataset_writer: DatasetWriter):
     train = [
         {
             "start": start,
+            "item_id": region,
             "target": df[region].to_numpy(dtype=np.float64)[
                 :-prediction_length
             ],
diff --git a/src/gluonts/dataset/repository/_ett_small.py b/src/gluonts/dataset/repository/_ett_small.py
index bf021a9659..757f595d96 100644
--- a/src/gluonts/dataset/repository/_ett_small.py
+++ b/src/gluonts/dataset/repository/_ett_small.py
@@ -47,6 +47,7 @@ def generate_ett_small_dataset(
             test.append(
                 {
                     "start": start,
+                    "item_id": col,
                     "target": df[col].values,
                 }
             )
@@ -60,6 +61,7 @@ def generate_ett_small_dataset(
             train.append(
                 {
                     "start": start,
+                    "item_id": col,
                     "target": df[col].values[:-prediction_length],
                 }
             )