Merge branch 'master' into feature/lora

src/timesfm/data_loader.py

@@ -11,13 +11,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 """TF dataloaders for general timeseries datasets.
 
 The expected input format is csv file with a datetime index.
 """
 
-
 from absl import logging
 import numpy as np
 import pandas as pd
@@ -79,9 +77,8 @@ class TimeSeriesdata(object):
       self.data_df['ccol'] = np.zeros(self.data_df.shape[0])
       cat_cov_cols = ['ccol']
     self.data_df.fillna(0, inplace=True)
-    self.data_df.set_index(
-        pd.DatetimeIndex(self.data_df[datetime_col]), inplace=True
-    )
+    self.data_df.set_index(pd.DatetimeIndex(self.data_df[datetime_col]),
+                           inplace=True)
     self.num_cov_cols = num_cov_cols
     self.cat_cov_cols = cat_cov_cols
     self.ts_cols = ts_cols
@@ -94,18 +91,16 @@ class TimeSeriesdata(object):
             data_df_idx[-1] + pd.Timedelta(1, freq=freq),
             periods=pred_len + 1,
             freq=freq,
-        )
-    )
+        ))
     self.time_df = time_features.TimeCovariates(
-        date_index, holiday=holiday
-    ).get_covariates()
+        date_index, holiday=holiday).get_covariates()
     self.hist_len = hist_len
     self.pred_len = pred_len
     self.batch_size = batch_size
     self.freq = freq
     self.normalize = normalize
     self.data_mat = self.data_df[self.ts_cols].to_numpy().transpose()
-    self.data_mat = self.data_mat[:, 0 : self.test_range[1]]
+    self.data_mat = self.data_mat[:, 0:self.test_range[1]]
     self.time_mat = self.time_df.to_numpy().transpose()
     self.num_feat_mat = self.data_df[num_cov_cols].to_numpy().transpose()
     self.cat_feat_mat, self.cat_sizes = self._get_cat_cols(cat_cov_cols)
@@ -135,7 +130,7 @@ class TimeSeriesdata(object):
 
   def _normalize_data(self):
     self.scaler = StandardScaler()
-    train_mat = self.data_mat[:, self.train_range[0] : self.train_range[1]]
+    train_mat = self.data_mat[:, 0:self.train_range[1]]
     self.scaler = self.scaler.fit(train_mat.transpose())
     self.data_mat = self.scaler.transform(self.data_mat.transpose()).transpose()
 
@@ -253,9 +248,8 @@ class TimeSeriesdata(object):
       gen_fn = self.train_gen
     else:
       gen_fn = lambda: self.test_val_gen(mode, shift)
-    output_types = tuple(
-        [tf.float32] * 2 + [tf.int32] + [tf.float32] * 2 + [tf.int32] * 2
-    )
+    output_types = tuple([tf.float32] * 2 + [tf.int32] + [tf.float32] * 2 +
+                         [tf.int32] * 2)
     dataset = tf.data.Dataset.from_generator(gen_fn, output_types)
     dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
     return dataset

src/timesfm/timesfm.py

@@ -783,6 +783,7 @@ class TimesFm:
       model_name: str = "timesfm",
       window_size: int | None = None,
       num_jobs: int = 1,
+      verbose: bool = True,
   ) -> pd.DataFrame:
     """Forecasts on a list of time series.
 
@@ -800,6 +801,7 @@ class TimesFm:
       window_size: window size of trend + residual decomposition. If None then
         we do not do decomposition.
       num_jobs: number of parallel processes to use for dataframe processing.
+      verbose: output model states in terminal.
 
     Returns:
       Future forecasts dataframe.
@@ -819,7 +821,8 @@ class TimesFm:
     new_inputs = []
     uids = []
     if num_jobs == 1:
-      print("Processing dataframe with single process.")
+      if verbose:
+        print("Processing dataframe with single process.")
       for key, group in df_sorted.groupby("unique_id"):
         inp, uid = process_group(
             key,
@@ -832,7 +835,8 @@ class TimesFm:
     else:
       if num_jobs == -1:
         num_jobs = multiprocessing.cpu_count()
-      print("Processing dataframe with multiple processes.")
+      if verbose:
+        print("Processing dataframe with multiple processes.")
       with multiprocessing.Pool(processes=num_jobs) as pool:
         results = pool.starmap(
             process_group,
@@ -842,12 +846,14 @@ class TimesFm:
             ],
         )
       new_inputs, uids = zip(*results)
-    print("Finished preprocessing dataframe.")
+    if verbose:
+        print("Finished preprocessing dataframe.")
     freq_inps = [freq_map(freq)] * len(new_inputs)
     _, full_forecast = self.forecast(
         new_inputs, freq=freq_inps, window_size=window_size
     )
-    print("Finished forecasting.")
+    if verbose:
+        print("Finished forecasting.")
     fcst_df = make_future_dataframe(
         uids=uids,
         last_times=df_sorted.groupby("unique_id")["ds"].tail(1),