added a verbose parameter to TimesFM.forecast_on_df(), in order to make it possible to desactivate print outputs.

This commit is contained in:
periodLeo
2024-07-09 20:16:43 +09:00
parent 1facd4eaab
commit c0b93d4b6a
+10 -4
View File
@@ -518,6 +518,7 @@ class TimesFm:
model_name: str = "timesfm", model_name: str = "timesfm",
window_size: int | None = None, window_size: int | None = None,
num_jobs: int = 1, num_jobs: int = 1,
verbose: bool = True,
) -> pd.DataFrame: ) -> pd.DataFrame:
"""Forecasts on a list of time series. """Forecasts on a list of time series.
@@ -535,6 +536,7 @@ class TimesFm:
window_size: window size of trend + residual decomposition. If None then window_size: window size of trend + residual decomposition. If None then
we do not do decomposition. we do not do decomposition.
num_jobs: number of parallel processes to use for dataframe processing. num_jobs: number of parallel processes to use for dataframe processing.
verbose: output model states in terminal.
Returns: Returns:
Future forecasts dataframe. Future forecasts dataframe.
@@ -554,7 +556,8 @@ class TimesFm:
new_inputs = [] new_inputs = []
uids = [] uids = []
if num_jobs == 1: 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"): for key, group in df_sorted.groupby("unique_id"):
inp, uid = process_group( inp, uid = process_group(
key, key,
@@ -567,7 +570,8 @@ class TimesFm:
else: else:
if num_jobs == -1: if num_jobs == -1:
num_jobs = multiprocessing.cpu_count() 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: with multiprocessing.Pool(processes=num_jobs) as pool:
results = pool.starmap( results = pool.starmap(
process_group, process_group,
@@ -577,12 +581,14 @@ class TimesFm:
], ],
) )
new_inputs, uids = zip(*results) new_inputs, uids = zip(*results)
print("Finished preprocessing dataframe.") if verbose:
print("Finished preprocessing dataframe.")
freq_inps = [freq_map(freq)] * len(new_inputs) freq_inps = [freq_map(freq)] * len(new_inputs)
_, full_forecast = self.forecast( _, full_forecast = self.forecast(
new_inputs, freq=freq_inps, window_size=window_size new_inputs, freq=freq_inps, window_size=window_size
) )
print("Finished forecasting.") if verbose:
print("Finished forecasting.")
fcst_df = make_future_dataframe( fcst_df = make_future_dataframe(
uids=uids, uids=uids,
last_times=df_sorted.groupby("unique_id")["ds"].tail(1), last_times=df_sorted.groupby("unique_id")["ds"].tail(1),