Adding v2.0 support

notebooks/finetuning.ipynb

@@ -9,7 +9,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -34,7 +34,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -61,14 +61,22 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "timesfm_backend = \"gpu\"  # @param\n",
+    "\n",
     "tfm = timesfm.TimesFm(\n",
     "      hparams=timesfm.TimesFmHparams(\n",
-    "          backend=\"gpu\",\n",
+    "          backend=timesfm_backend,\n",
     "          per_core_batch_size=32,\n",
     "          horizon_len=128,\n",
+    "          num_layers=50,\n",
+    "          # Se this to True for v1.0 checkpoints\n",
+    "          use_positional_embedding=False,\n",
+    "          # Note that we could set this to as high as 2048 but keeping it 512 here so that\n",
+    "          # both v1.0 and 2.0 checkpoints work\n",
+    "          context_len=512,\n",
     "      ),\n",
     "      checkpoint=timesfm.TimesFmCheckpoint(\n",
-    "          huggingface_repo_id=\"google/timesfm-1.0-200m\"),\n",
+    "          huggingface_repo_id=\"google/timesfm-2.0-500m-jax\"),\n",
     "  )"
    ]
   },
@@ -81,7 +89,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -126,7 +134,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -147,7 +155,7 @@
     "pred_len = 96\n",
     "\n",
     "num_ts = len(ts_cols)\n",
-    "batch_size = 16\n",
+    "batch_size = 8\n",
     "\n",
     "dtl = data_loader.TimeSeriesdata(\n",
     "      data_path=data_path,\n",
@@ -187,7 +195,7 @@
    "outputs": [],
    "source": [
     "for tbatch in tqdm(train_batches.as_numpy_iterator()):\n",
-    "    pass\n",
+    "    break\n",
     "print(tbatch[0].shape)"
    ]
   },
@@ -224,7 +232,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -248,7 +256,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 10,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -276,7 +284,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -296,7 +304,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 12,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -325,7 +333,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 13,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -416,7 +424,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 19,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -445,7 +453,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 20,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -455,7 +463,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 21,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -470,7 +478,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 22,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -583,13 +591,13 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## There is around a __9%__ reduction in MAE from finetuning."
+    "## There is around a __7%__ reduction in MAE from finetuning."
    ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "tfm_env_v3",
+   "display_name": "chronos-v2",
    "language": "python",
    "name": "python3"
   },
@@ -603,7 +611,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.14"
+   "version": "3.10.15"
   }
  },
  "nbformat": 4,