Adding v2.0 support

notebooks/covariates.ipynb

@@ -9,7 +9,8 @@
     "This toturial notebook demonstrates how to utilize exogenous covariates with TimesFM when making forecasts. Before running this notebook, make sure:\n",
     "\n",
     "- You've read through the README of TimesFM.\n",
-    "- A local kernel with Python 3.10 is up and running."
+    "- A local kernel with Python 3.10 is up and running, for the jax version.\n",
+    "- Install the JAX version following the installation instructions."
    ]
   },
   {
@@ -19,36 +20,15 @@
     "## Setup the environment and install TimesFM."
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import os\n",
-    "os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'\n",
-    "os.environ['JAX_PMAP_USE_TENSORSTORE'] = 'false'"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "!pip install timesfm\n",
-    "import timesfm"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "## Load the checkpoint\n",
     "\n",
-    "**Notice:** Please set up the backend as per your machine (\"cpu\", \"gpu\" or \"tpu\"). This notebook will run by default on CPU.\n",
+    "**Notice:** Please set up the backend as per your machine (\"cpu\", \"gpu\" or \"tpu\"). This notebook will run by default on GPU.\n",
     "\n",
-    "We load the 1.0-200m model checkpoint from HuggingFace."
+    "We load the 2.0-500m model checkpoint from HuggingFace."
    ]
   },
   {
@@ -57,23 +37,21 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "timesfm_backend = \"cpu\"  # @param\n",
-    "\n",
-    "from jax._src import config\n",
-    "config.update(\n",
-    "    \"jax_platforms\", {\"cpu\": \"cpu\", \"gpu\": \"cuda\", \"tpu\": \"\"}[timesfm_backend]\n",
-    ")\n",
+    "import timesfm\n",
+    "timesfm_backend = \"gpu\"  # @param\n",
     "\n",
     "model = timesfm.TimesFm(\n",
-    "    context_len=512,\n",
-    "    horizon_len=128,\n",
-    "    input_patch_len=32,\n",
-    "    output_patch_len=128,\n",
-    "    num_layers=20,\n",
-    "    model_dims=1280,\n",
-    "    backend=timesfm_backend,\n",
-    ")\n",
-    "model.load_from_checkpoint(repo_id=\"google/timesfm-1.0-200m\")"
+    "      hparams=timesfm.TimesFmHparams(\n",
+    "          backend=timesfm_backend,\n",
+    "          per_core_batch_size=32,\n",
+    "          horizon_len=128,\n",
+    "          num_layers=50,\n",
+    "          use_positional_embedding=False,\n",
+    "          context_len=2048,\n",
+    "      ),\n",
+    "      checkpoint=timesfm.TimesFmCheckpoint(\n",
+    "          huggingface_repo_id=\"google/timesfm-2.0-500m-jax\"),\n",
+    "  )"
    ]
   },
   {
@@ -140,7 +118,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -176,7 +154,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -208,7 +186,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -239,13 +217,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import time\n",
+    "\n",
     "# Benchmark\n",
     "batch_size = 128\n",
     "context_len = 120\n",
     "horizon_len = 24\n",
     "input_data = get_batched_data_fn(batch_size = 128)\n",
     "metrics = defaultdict(list)\n",
-    "import time\n",
+    "\n",
     "\n",
     "for i, example in enumerate(input_data()):\n",
     "  raw_forecast, _ = model.forecast(\n",
@@ -307,15 +287,15 @@
    "source": [
     "You should see results close to \n",
     "```\n",
-    "eval_mae_timesfm: 6.762283045916956\n",
-    "eval_mae_xreg_timesfm: 5.39219617611074\n",
-    "eval_mae_xreg: 37.15275842572484\n",
-    "eval_mse_timesfm: 166.7771466306823\n",
-    "eval_mse_xreg_timesfm: 120.64757721021306\n",
-    "eval_mse_xreg: 1672.2116821201796\n",
+    "eval_mae_timesfm: 6.729583250571446\n",
+    "eval_mae_xreg_timesfm: 5.3375301110158\n",
+    "eval_mae_xreg: 37.152760709266\n",
+    "eval_mse_timesfm: 162.3132151851567\n",
+    "eval_mse_xreg_timesfm: 120.9900627409689\n",
+    "eval_mse_xreg: 1672.208769045399\n",
     "```\n",
     "\n",
-    "With the covariates, the TimesFM forecast Mean Absolute Error improves from 6.76 to 5.39, and Mean Squred Error from 166.78 to 120.65. The results of purely fitting the linear model are also provided for reference."
+    "With the covariates, the TimesFM forecast Mean Absolute Error improves from 6.73 to 5.34, and Mean Squred Error from 162.31 to 120.99. The results of purely fitting the linear model are also provided for reference."
    ]
   },
   {
@@ -381,9 +361,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "cuda-gpt",
+   "display_name": "chronos-v2",
    "language": "python",
-   "name": "cuda"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -395,7 +375,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.14"
+   "version": "3.10.15"
   }
  },
  "nbformat": 4,