From 11d3cc66d520ddfe20cdd4420ecbe8025e0153a8 Mon Sep 17 00:00:00 2001 From: "@gabriel.pereira" Date: Thu, 26 Mar 2026 16:21:57 -0300 Subject: [PATCH] adding timesfm project and exceptions on gitignore --- .gitignore | 32 +++ timesfm-forecast/.gitignore | 23 +++ timesfm-forecast/LICENSE | 22 +++ timesfm-forecast/README.md | 130 ++++++++++++ timesfm-forecast/examples/sample.csv | 33 ++++ timesfm-forecast/pyproject.toml | 40 ++++ timesfm-forecast/src/timesfm_app/__init__.py | 16 ++ timesfm-forecast/src/timesfm_app/cli.py | 13 ++ timesfm-forecast/src/timesfm_app/core.py | 198 +++++++++++++++++++ timesfm-forecast/src/timesfm_app/ui/app.py | 178 +++++++++++++++++ timesfm-forecast/tests/test_pipeline.py | 49 +++++ 11 files changed, 734 insertions(+) create mode 100644 .gitignore create mode 100755 timesfm-forecast/.gitignore create mode 100755 timesfm-forecast/LICENSE create mode 100755 timesfm-forecast/README.md create mode 100755 timesfm-forecast/examples/sample.csv create mode 100755 timesfm-forecast/pyproject.toml create mode 100755 timesfm-forecast/src/timesfm_app/__init__.py create mode 100755 timesfm-forecast/src/timesfm_app/cli.py create mode 100755 timesfm-forecast/src/timesfm_app/core.py create mode 100755 timesfm-forecast/src/timesfm_app/ui/app.py create mode 100755 timesfm-forecast/tests/test_pipeline.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ce31563 --- /dev/null +++ b/.gitignore @@ -0,0 +1,32 @@ +# Byte-compiled / cache files +__pycache__/ +*.py[cod] +*.pyc +*pycache +*.so +*.pyo +*.pyd + +# Virtual environments +.env/ +.venv/ +env/ +venv/ + +# Distribution / packaging +*.egg-info/ +dist/ +build/ + +# OS files +.DS_Store + +# Streamlit +.streamlit/ +**/.streamlit/ + +# Lock files (optional, if not needed) +uv.lock + +# Data files +*.csv diff --git a/timesfm-forecast/.gitignore b/timesfm-forecast/.gitignore new file mode 100755 index 0000000..1733b80 --- /dev/null +++ b/timesfm-forecast/.gitignore @@ -0,0 +1,23 @@ + +# Python +__pycache__/ +*.py[cod] +*.egg-info/ +*.egg +*.pyo +*.pyd + +# Environments +.venv/ +.uv/ + +# Streamlit +.streamlit/ +**/.streamlit/ + +# OS +.DS_Store + +# Editors +.vscode/ +.idea/ diff --git a/timesfm-forecast/LICENSE b/timesfm-forecast/LICENSE new file mode 100755 index 0000000..2ea3942 --- /dev/null +++ b/timesfm-forecast/LICENSE @@ -0,0 +1,22 @@ + +MIT License + +Copyright (c) 2026 Gabriel Faria Pereira + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/timesfm-forecast/README.md b/timesfm-forecast/README.md new file mode 100755 index 0000000..ec542ff --- /dev/null +++ b/timesfm-forecast/README.md @@ -0,0 +1,130 @@ + +# TimesFM Forecast (Streamlit, OOP) + +A maintainable, OOP-based Streamlit app to forecast multiple time series (by `key`) using **TimesFM**. + +- Upload CSV with columns: `date`, `value`, `key` +- Select a **horizon range** (e.g., 1–12) +- App uses **all data as training** (no holdout) +- **Progress bar** while forecasting +- **Download** CSV with forecasts (`key, date, step, forecast`) +- Clean separation between UI and core logic + +--- + +## 🚀 Quickstart (using `uv`) + +> Requires Python **3.10+**. +> `uv` docs: https://docs.astral.sh/uv/ + +```bash +# 1) Clone +# (If you already downloaded this folder locally, cd into it and skip clone.) +# git clone https://github.com//timesfm-forecast.git +cd timesfm-forecast + +# 2) Create a virtualenv (managed by uv) +uv venv +source .venv/bin/activate # Windows: .venv\Scriptsctivate + +# 3) Install dependencies (editable mode) +uv pip install -e . + +# 4) Run the app +uv run timesfm-app +# or directly: +# uv run streamlit run src/timesfm_app/ui/app.py +``` + +Open the URL shown in your terminal (typically http://localhost:8501). + +--- + +## 📦 CSV Format + +Upload a CSV with **columns**: + +- `date` – a date per observation (the app prefers `dd/mm/yyyy`, but will try general parsing) +- `value` – numeric target +- `key` – series identifier (one forecast per `key`) + +Example: [examples/sample.csv](examples/sample.csv) + +--- + +## ⚙️ Configuration (in the UI) + +- **TimesFM model**: defaults to `google/timesfm-2.0-500m-pytorch` +- **Device**: `cuda` if available, otherwise `cpu` +- **Frequency code (TimesFM)**: defaults to **2 = monthly** (matches your original draft) +- **Fallback pandas frequency**: used to create future dates if per-key inference fails (`D,W,M,Q,Y`) +- **Batch size**: controls throughput vs memory +- **Horizon range**: inclusive range (e.g., 1..12). Output includes each step with its aligned date. + +> The app uses `pandas.infer_freq` to detect per-key frequency; if inference fails, it falls back to your selection. + +--- + +## 🧠 Design / OOP + +- `TimesFMService` – encapsulates model loading + batch inference. +- `ForecastPipeline` – orchestrates validation, batching, inference, and future index construction. Returns a **long** DataFrame: + - `key, date, step, forecast` +- `DataValidator` / `FrequencyHelper` – stateless utility classes. +- `ui/app.py` – Streamlit-only, thin UI. + +This structure makes it straightforward to add more backends (e.g., Prophet, Chronos) by introducing a new service class. + +--- + +## 🧪 Tests + +Install dev extras and run: + +```bash +uv pip install -e ".[dev]" +uv run pytest +``` + +We provide a minimal test (`tests/test_pipeline.py`) that injects a **fake service** to validate pipeline behavior without loading a real model. + +--- + +## 🖥️ GPU vs CPU (PyTorch) + +By default, this project depends on `torch` without a pinned wheel. If you need a **CUDA** build, install the wheel for your CUDA version. Examples: + +```bash +# CUDA 12.1 (example) +uv pip install torch --index-url https://download.pytorch.org/whl/cu121 + +# CPU-only (explicit) +uv pip install torch --index-url https://download.pytorch.org/whl/cpu +``` + +Then run the app and select **Device = cuda** in the sidebar. Make sure your NVIDIA drivers & CUDA runtime match the wheel. + +--- + +## 📝 Output + +You can download a CSV with columns: + +- `key` – series id +- `date` – predicted timestamp (aligned to step) +- `step` – horizon (1..H) +- `forecast` – mean prediction + +--- + +## 🧯 Troubleshooting + +- **Model download slow / blocked**: the first run downloads the model weights from Hugging Face. Ensure internet connectivity and retry. You can also pre-download the model to your HF cache. +- **Out-of-memory on GPU**: reduce `Batch size`, or switch **Device** to `cpu`. +- **Dates misaligned**: pick the correct fallback pandas frequency (e.g., `M` for monthly) if your data has irregular gaps that prevent inference. + +--- + +## 📄 License + +MIT diff --git a/timesfm-forecast/examples/sample.csv b/timesfm-forecast/examples/sample.csv new file mode 100755 index 0000000..ab89130 --- /dev/null +++ b/timesfm-forecast/examples/sample.csv @@ -0,0 +1,33 @@ + +date,value,key +2022-01-01,120,A +2022-02-01,118,A +2022-03-01,121,A +2022-04-01,123,A +2022-05-01,125,A +2022-06-01,127,A +2022-07-01,130,A +2022-08-01,129,A +2022-09-01,131,A +2022-10-01,133,A +2022-11-01,135,A +2022-12-01,138,A +2023-01-01,140,A +2023-02-01,142,A +2023-03-01,145,A + +2022-01-01,80,B +2022-02-01,82,B +2022-03-01,81,B +2022-04-01,83,B +2022-05-01,84,B +2022-06-01,85,B +2022-07-01,87,B +2022-08-01,88,B +2022-09-01,90,B +2022-10-01,89,B +2022-11-01,91,B +2022-12-01,92,B +2023-01-01,93,B +2023-02-01,95,B +2023-03-01,96,B diff --git a/timesfm-forecast/pyproject.toml b/timesfm-forecast/pyproject.toml new file mode 100755 index 0000000..116b0f5 --- /dev/null +++ b/timesfm-forecast/pyproject.toml @@ -0,0 +1,40 @@ + +[project] +name = "timesfm-forecast" +version = "0.1.0" +description = "Streamlit TimesFM forecast app (OOP, multi-key, horizon range, progress, CSV download)." +readme = "README.md" +requires-python = ">=3.10" +authors = [ + { name = "Gabriel Faria Pereira" } +] +license = { text = "MIT" } + +# Keep torch unpinned here (see README for CUDA/CPU wheels) +dependencies = [ + "streamlit>=1.31", + "pandas>=2.0", + "numpy>=1.24", + "transformers>=4.45.0", + "torch", + "accelerate>=0.33", +] + +[project.optional-dependencies] +dev = [ + "pytest>=7.0", + "ruff>=0.4", +] + +[project.scripts] +timesfm-app = "timesfm_app.cli:main" + +[build-system] +requires = ["setuptools>=68", "wheel"] +build-backend = "setuptools.build_meta" + +[tool.setuptools] +package-dir = {"" = "src"} + +[tool.setuptools.packages.find] +where = ["src"] diff --git a/timesfm-forecast/src/timesfm_app/__init__.py b/timesfm-forecast/src/timesfm_app/__init__.py new file mode 100755 index 0000000..b46d757 --- /dev/null +++ b/timesfm-forecast/src/timesfm_app/__init__.py @@ -0,0 +1,16 @@ + +__all__ = [ + "ForecastConfig", + "TimesFMService", + "ForecastPipeline", + "DataValidator", + "FrequencyHelper", +] + +from .core import ( + ForecastConfig, + TimesFMService, + ForecastPipeline, + DataValidator, + FrequencyHelper, +) diff --git a/timesfm-forecast/src/timesfm_app/cli.py b/timesfm-forecast/src/timesfm_app/cli.py new file mode 100755 index 0000000..13e420f --- /dev/null +++ b/timesfm-forecast/src/timesfm_app/cli.py @@ -0,0 +1,13 @@ + +import pathlib +import subprocess +import sys + +def main() -> int: + """ + Launch the Streamlit UI via CLI: + $ timesfm-app + """ + app_path = pathlib.Path(__file__).with_name("ui").joinpath("app.py") + cmd = [sys.executable, "-m", "streamlit", "run", str(app_path)] + return subprocess.call(cmd) diff --git a/timesfm-forecast/src/timesfm_app/core.py b/timesfm-forecast/src/timesfm_app/core.py new file mode 100755 index 0000000..858a72c --- /dev/null +++ b/timesfm-forecast/src/timesfm_app/core.py @@ -0,0 +1,198 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Callable, Iterable, List, Optional, Tuple +import math + +import numpy as np +import pandas as pd +import torch +from transformers import TimesFmModelForPrediction + + +@dataclass(frozen=True) +class ForecastConfig: + model_name: str = "google/timesfm-2.0-500m-pytorch" + device: str = "cpu" # "cpu" or "cuda" + batch_size: int = 16 + timesfm_freq_code: int = 2 # 2 = Monthly (as per your draft) + pandas_freq_fallback: str = "M" # for building future index + horizon_min: int = 1 + horizon_max: int = 12 # inclusive + + def validate(self) -> None: + if self.horizon_min < 1: + raise ValueError("horizon_min must be >= 1") + if self.horizon_max < self.horizon_min: + raise ValueError("horizon_max must be >= horizon_min") + if self.batch_size < 1: + raise ValueError("batch_size must be >= 1") + + +class DataValidator: + REQUIRED = {"date", "value", "key"} + + @staticmethod + def standardize_columns(df: pd.DataFrame) -> pd.DataFrame: + out = df.copy() + out.columns = out.columns.str.lower() + return out + + @staticmethod + def validate_schema(df: pd.DataFrame) -> None: + cols = set(df.columns) + if not DataValidator.REQUIRED.issubset(cols): + raise ValueError( + f"CSV must contain columns {sorted(DataValidator.REQUIRED)}, " + f"found {sorted(cols)}" + ) + + @staticmethod + def parse_dates(df: pd.DataFrame) -> pd.DataFrame: + out = df.copy() + sample = str(out["date"].iloc[0]) + if "/" in sample and sample.count("/") == 2: + out["date"] = pd.to_datetime(out["date"], dayfirst=True, errors="coerce") + else: + out["date"] = pd.to_datetime(out["date"], dayfirst=False, errors="coerce") + return out + + @staticmethod + def sanitize(df: pd.DataFrame) -> pd.DataFrame: + out = df.copy() + out = out.dropna(subset=["date", "value", "key"]) + out["value"] = pd.to_numeric(out["value"], errors="coerce") + out = out.dropna(subset=["value"]) + out = out.sort_values(["key", "date"]).reset_index(drop=True) + return out + + +class FrequencyHelper: + @staticmethod + def infer_pandas_freq(dates: pd.Series) -> str: + if dates.is_monotonic_increasing and dates.notna().all(): + try: + inferred = pd.infer_freq(pd.DatetimeIndex(dates)) + return inferred or "" + except Exception: + return "" + return "" + + @staticmethod + def future_dates( + last_date: pd.Timestamp, periods: int, pandas_freq: str + ) -> pd.DatetimeIndex: + if periods <= 0: + return pd.DatetimeIndex([]) + try: + offset = pd.tseries.frequencies.to_offset(pandas_freq) + except Exception: + offset = pd.tseries.frequencies.to_offset("M") # safe fallback + start = last_date + offset + return pd.date_range(start=start, periods=periods, freq=offset) + + +class TimesFMService: + def __init__(self, model_name: str, device: str = "cpu"): + self.model_name = model_name + self.device = device + self._model: Optional[TimesFmModelForPrediction] = None + + @property + def model(self) -> TimesFmModelForPrediction: + if self._model is None: + self._model = TimesFmModelForPrediction.from_pretrained( + self.model_name, + dtype=torch.float32, + device_map=self.device, # "cpu" or "cuda" + ) + return self._model + + @torch.no_grad() + def forecast_batch( + self, past_values: List[torch.Tensor], freq_code: int, max_h: int + ) -> np.ndarray: + if max_h < 1: + raise ValueError("max_h must be >= 1") + freq_tensor = torch.tensor([freq_code] * len(past_values), dtype=torch.long) + out = self.model( + past_values=past_values, + freq=freq_tensor, + return_dict=True, + ) + preds = out.mean_predictions[:, :max_h].detach().cpu().numpy() + return preds + + +class ForecastPipeline: + def __init__(self, svc: TimesFMService, config: ForecastConfig): + config.validate() + self.svc = svc + self.cfg = config + + def _chunked(self, seq: List, n: int) -> Iterable[List]: + for i in range(0, len(seq), n): + yield seq[i : i + n] + + def run( + self, raw_df: pd.DataFrame, progress_cb: Optional[callable] = None + ) -> pd.DataFrame: + df = DataValidator.standardize_columns(raw_df) + DataValidator.validate_schema(df) + df = DataValidator.parse_dates(df) + df = DataValidator.sanitize(df) + + keys = sorted(df["key"].unique().tolist()) + n_keys = len(keys) + if n_keys == 0: + return pd.DataFrame(columns=["key", "date", "step", "forecast"]) + + series_payload: List[torch.Tensor] = [] + meta: List[Tuple[str, pd.Timestamp, str]] = [] + + for k in keys: + sub = df[df["key"] == k] + values = sub["value"].astype(float).values + series_payload.append(torch.tensor(values, dtype=torch.float32)) + last_dt = sub["date"].iloc[-1] + inferred = FrequencyHelper.infer_pandas_freq(sub["date"]) + meta.append((k, last_dt, inferred)) + + bsz = self.cfg.batch_size + total_batches = math.ceil(n_keys / bsz) + results: List[pd.DataFrame] = [] + + horizon_min, horizon_max = self.cfg.horizon_min, self.cfg.horizon_max + + for batch_idx, (payload_chunk, meta_chunk) in enumerate( + zip(self._chunked(series_payload, bsz), self._chunked(meta, bsz)), start=1 + ): + if progress_cb: + progress_cb(batch_idx, total_batches, "Forecasting…") + + preds = self.svc.forecast_batch( + past_values=payload_chunk, + freq_code=self.cfg.timesfm_freq_code, + max_h=horizon_max, + ) + + for i, (k, last_dt, inferred) in enumerate(meta_chunk): + pfreq = inferred or self.cfg.pandas_freq_fallback + future_idx = FrequencyHelper.future_dates(last_dt, horizon_max, pfreq) + + steps = np.arange(horizon_min, horizon_max + 1, dtype=int) + dates = future_idx[steps - 1] + fc_vals = preds[i, steps - 1] + + res = pd.DataFrame( + { + "key": k, + "date": dates, + "step": steps, + "forecast": fc_vals, + } + ) + results.append(res) + + out_df = pd.concat(results, ignore_index=True) + return out_df diff --git a/timesfm-forecast/src/timesfm_app/ui/app.py b/timesfm-forecast/src/timesfm_app/ui/app.py new file mode 100755 index 0000000..94f2538 --- /dev/null +++ b/timesfm-forecast/src/timesfm_app/ui/app.py @@ -0,0 +1,178 @@ +import io + +import pandas as pd +import streamlit as st +import torch + +from timesfm_app.core import ( + ForecastConfig, + TimesFMService, + ForecastPipeline, + DataValidator, +) + +st.set_page_config(page_title="TimesFM Forecast", layout="wide") + + +def horizon_range_slider(label: str, default_min: int = 1, default_max: int = 12): + return st.slider( + label, + min_value=1, + max_value=120, + value=(default_min, default_max), + step=1, + help="Select an inclusive range of horizons. The app outputs each step separately.", + ) + + +@st.cache_resource(show_spinner="Preparing TimesFM service…") +def get_service(model_name: str, device: str) -> TimesFMService: + return TimesFMService(model_name=model_name, device=device) + + +def main(): + st.title("📈 TimesFM Forecast") + + st.markdown( + """ +Upload a CSV with columns **`date, value, key`**. + +**Features** +- Uses **all rows** as training (no holdout) +- Choose a **range of horizons** (e.g., 1–12) +- Multi-series (**by key**) batching +- **Progress bar** during inference +- **Download** forecasts as CSV + """ + ) + + with st.sidebar: + st.header("Settings") + + model_name = st.text_input( + "TimesFM model", + value="google/timesfm-2.0-500m-pytorch", + ) + + cuda_available = torch.cuda.is_available() + device = st.selectbox( + "Device", + options=(["cuda", "cpu"] if cuda_available else ["cpu"]), + index=(0 if cuda_available else 0), + ) + + timesfm_freq_code = st.number_input( + "TimesFM frequency code", + min_value=0, + max_value=99, + value=2, # monthly by default + step=1, + help="Default 2 = monthly (matches your draft).", + ) + + pandas_freq_fallback = st.selectbox( + "Fallback pandas frequency for future dates", + options=["D", "W", "M", "Q", "Y"], + index=2, # "M" + help="Used when per-key frequency inference fails.", + ) + + batch_size = st.slider("Batch size", 1, 64, 16, 1) + h_min, h_max = horizon_range_slider("Horizon range", 1, 12) + + uploaded = st.file_uploader("Upload CSV", type=["csv"]) + + if "forecast" not in st.session_state: + st.session_state["forecast"] = None + if "history" not in st.session_state: + st.session_state["history"] = None + + if uploaded is None: + st.info("👆 Upload a CSV to begin.") + return + + raw_df = pd.read_csv(uploaded) + st.dataframe(raw_df.head(10), use_container_width=True) + + run = st.button("▶️ Run forecast", type="primary") + + if run: + cfg = ForecastConfig( + model_name=model_name, + device=device, + batch_size=batch_size, + timesfm_freq_code=timesfm_freq_code, + pandas_freq_fallback=pandas_freq_fallback, + horizon_min=h_min, + horizon_max=h_max, + ) + + svc = get_service(model_name=cfg.model_name, device=cfg.device) + pipe = ForecastPipeline(svc=svc, config=cfg) + + progress = st.progress(0.0) + info = st.empty() + + def progress_cb(batch_idx: int, total_batches: int, message: str): + pct = min(1.0, batch_idx / max(1, total_batches)) + progress.progress(pct) + info.write(f"{message} ({batch_idx}/{total_batches})") + + st.info("Running TimesFM… This may take a while for large files.") + + try: + forecast_df = pipe.run(raw_df, progress_cb=progress_cb) + finally: + info.empty() + progress.empty() + + hist = DataValidator.standardize_columns(raw_df) + hist = DataValidator.parse_dates(hist) + hist = DataValidator.sanitize(hist) + + st.session_state["forecast"] = forecast_df + st.session_state["history"] = hist + + st.success("✅ Forecast complete!") + + forecast_df = st.session_state["forecast"] + hist = st.session_state["history"] + + if forecast_df is None or hist is None: + return + + st.subheader("Forecast preview (long format)") + st.dataframe(forecast_df.head(20), use_container_width=True) + + st.subheader("Visualize one key") + keys_sorted = sorted(hist["key"].unique().tolist()) + sel_key = st.selectbox("Select key", options=keys_sorted) + + hist_k = hist[hist["key"] == sel_key][["date", "value"]].rename( + columns={"value": "y"} + ) + hist_k = hist_k.assign(series="history") + + fcast_k = forecast_df[forecast_df["key"] == sel_key][["date", "forecast"]] + fcast_k = fcast_k.rename(columns={"forecast": "y"}).assign(series="forecast") + + display_df = pd.concat([hist_k, fcast_k], ignore_index=True).sort_values("date") + st.line_chart(data=display_df, x="date", y="y", color="series", height=340) + + st.subheader("Download") + buf = io.StringIO() + forecast_df.to_csv(buf, index=False) + st.download_button( + label="💾 Download forecasts (CSV)", + data=buf.getvalue().encode("utf-8"), + file_name="timesfm_forecasts_long.csv", + mime="text/csv", + ) + + st.caption( + "Output columns: key, date, step, forecast. Each 'step' is the horizon (1..H)." + ) + + +if __name__ == "__main__": + main() diff --git a/timesfm-forecast/tests/test_pipeline.py b/timesfm-forecast/tests/test_pipeline.py new file mode 100755 index 0000000..e146711 --- /dev/null +++ b/timesfm-forecast/tests/test_pipeline.py @@ -0,0 +1,49 @@ +import numpy as np +import pandas as pd + +from timesfm_app.core import ForecastConfig, ForecastPipeline, TimesFMService + + +class FakeTimesFMService(TimesFMService): + """ + Replace the model call with deterministic values for tests. + Returns predictions: [1, 2, ..., max_h] for each series in the batch. + """ + + def __init__(self): + pass + + def forecast_batch(self, past_values, freq_code: int, max_h: int) -> np.ndarray: + base = np.arange(1, max_h + 1, dtype=float) + return np.tile(base, (len(past_values), 1)) + + +def test_pipeline_long_output_shape(): + df = pd.DataFrame( + { + "date": pd.date_range("2023-01-01", periods=5, freq="ME").strftime( + "%Y-%m-%d" + ), + "value": [1, 2, 3, 4, 5], + "key": ["A"] * 5, + } + ) + + cfg = ForecastConfig( + model_name="dummy", + device="cpu", + batch_size=2, + timesfm_freq_code=2, + pandas_freq_fallback="M", + horizon_min=1, + horizon_max=3, + ) + + svc = FakeTimesFMService() + pipe = ForecastPipeline(svc=svc, config=cfg) + out = pipe.run(df) + + assert len(out) == 3 + assert out["key"].unique().tolist() == ["A"] + assert out["step"].tolist() == [1, 2, 3] + assert np.allclose(out["forecast"].values, [1.0, 2.0, 3.0])