Full pytorch support

experiments/extended_benchmarks/run_timesfm.py

@@ -11,7 +11,6 @@
 # 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.
-
 """Evaluation script for timesfm."""
 
 import os
@@ -21,12 +20,10 @@ import time
 from absl import flags
 import numpy as np
 import pandas as pd
-from paxml import checkpoints
 import timesfm
 
 from .utils import ExperimentHandler
 
-
 dataset_names = [
     "m1_monthly",
     "m1_quarterly",
@@ -74,35 +71,27 @@ context_dict = {
     "m4_yearly": 64,
 }
 
-_MODEL_PATH = flags.DEFINE_string(
-    "model_path", "/home/timesfm_q10_20240501", "Path to model"
-)
+_MODEL_PATH = flags.DEFINE_string("model_path", "google/timesfm-1.0-200m",
+                                  "Path to model")
 _BATCH_SIZE = flags.DEFINE_integer("batch_size", 64, "Batch size")
 _HORIZON = flags.DEFINE_integer("horizon", 128, "Horizon")
 _BACKEND = flags.DEFINE_string("backend", "gpu", "Backend")
 _NUM_JOBS = flags.DEFINE_integer("num_jobs", 1, "Number of jobs")
 _SAVE_DIR = flags.DEFINE_string("save_dir", "./results", "Save directory")
 
-
 QUANTILES = list(np.arange(1, 10) / 10.0)
 
 
 def main():
   results_list = []
   tfm = timesfm.TimesFm(
-      context_len=512,
-      horizon_len=_HORIZON.value,
-      input_patch_len=32,
-      output_patch_len=128,
-      num_layers=20,
-      model_dims=1280,
-      backend=_BACKEND.value,
-      per_core_batch_size=_BATCH_SIZE.value,
-      quantiles=QUANTILES,
-  )
-  tfm.load_from_checkpoint(
-      _MODEL_PATH.value,
-      checkpoint_type=checkpoints.CheckpointType.FLAX,
+      hparams=timesfm.TimesFmHparams(
+          backend=_BACKEND.value,
+          per_core_batch_size=_BATCH_SIZE.value,
+          horizon_len=_HORIZON.value,
+      ),
+      checkpoint=timesfm.TimesFmCheckpoint(
+          huggingface_repo_id=_MODEL_PATH.value),
   )
   run_id = np.random.randint(100000)
   model_name = "timesfm"
@@ -127,9 +116,9 @@ def main():
     )
     total_time = time.time() - init_time
     time_df = pd.DataFrame({"time": [total_time], "model": model_name})
-    results = exp.evaluate_from_predictions(
-        models=[model_name], fcsts_df=fcsts_df, times_df=time_df
-    )
+    results = exp.evaluate_from_predictions(models=[model_name],
+                                            fcsts_df=fcsts_df,
+                                            times_df=time_df)
     print(results, flush=True)
     results_list.append(results)
     results_full = pd.concat(results_list)