Update README.md

README.md

@@ -52,23 +52,28 @@ TTMs that can cater to many common forecasting settings in practice.
 Each pre-trained model will be released in a different branch name in this model card. Kindly access the required model using our 
 getting started [notebook](https://github.com/IBM/tsfm/blob/main/notebooks/hfdemo/ttm_getting_started.ipynb) mentioning the branch name.
 
-
+ 
 ## Model Releases (along with the branch name where the models are stored):
 
-- **512-96-r2**: Given the last 512 time-points (i.e. context length), this model can forecast up to next 96 time-points (i.e. forecast length)
+
+
+- **512-96-r2**: Given the last 512 time-points (i.e. context length), this model can forecast up to the next 96 time-points (i.e. forecast length)
     in future. (branch name: main)
   
-- **1024-96-r2**: Given the last 1024 time-points (i.e. context length), this model can forecast up to next 96 time-points (i.e. forecast length)
+- **1024-96-r2**: Given the last 1024 time-points (i.e. context length), this model can forecast up to the next 96 time-points (i.e. forecast length)
     in future.  (branch name: 1024-96-r2) [[Benchmarks]]
 
-- **1536-96-r2**: Given the last 1536 time-points (i.e. context length), this model can forecast up to next 96 time-points (i.e. forecast length)
+- **1536-96-r2**: Given the last 1536 time-points (i.e. context length), this model can forecast up to the next 96 time-points (i.e. forecast length)
     in future. (branch name: 1536-96-r2)
 
-- Likewise, we have models released for forecast lengths up to 720 timepoints. The branch names for these are as follows: 512-192-r2, 1024-192-r2, 1536-192-r2, 512-336-r2,
-  512-336-r2, 1024-336-r2, 1536-336-r2, 512-720-r2, 1024-720-r2, 1536-720-r2
+- Likewise, we have models released for forecast lengths up to 720 timepoints. The branch names for these are as follows: `512-192-r2`, `1024-192-r2`, `1536-192-r2`, `512-336-r2`,
+  `512-336-r2`, `1024-336-r2`, `1536-336-r2`, `512-720-r2`, `1024-720-r2`, `1536-720-r2`
 
 - Please use the [[get_model]](https://github.com/ibm-granite/granite-tsfm/blob/main/tsfm_public/toolkit/get_model.py) utility to automatically select the required model based on your input context length and forecast length requirement.
 
+- We currently allow 3 context lengths (512, 1024 and 1536) and 4 forecast lengths (96, 192, 336, 720). Users need to provide the exact context length as input.
+but can provide any forecast lengths up to 720 in get_model().
+
   
 
     
@@ -143,12 +148,63 @@ In addition, TTM also supports exogenous infusion and categorical data infusion.
   
 ## Uses
 
+
+Automatic Model selection
+```
+def get_model(
+    model_path,
+    model_name: str = "ttm",
+    context_length: int = None,
+    prediction_length: int = None,
+    freq_prefix_tuning: bool = None,
+    **kwargs,
+):
+    
+    TTM Model card offers a suite of models with varying context_length and forecast_length combinations.
+    This wrapper automatically selects the right model based on the given input context_length and prediction_length abstracting away the internal
+    complexity.
+
+    Args:
+        model_path (str):
+            HF model card path or local model path (Ex. ibm-granite/granite-timeseries-ttm-r1)
+        model_name (*optional*, str)
+            model name to use. Allowed values: ttm
+        context_length (int):
+            Input Context length. For ibm-granite/granite-timeseries-ttm-r1, we allow 512 and 1024.
+            For ibm-granite/granite-timeseries-ttm-r2 and  ibm/ttm-research-r2, we allow 512, 1024 and 1536
+        prediction_length (int):
+            Forecast length to predict. For ibm-granite/granite-timeseries-ttm-r1, we can forecast upto 96.
+            For ibm-granite/granite-timeseries-ttm-r2 and  ibm/ttm-research-r2, we can forecast upto 720.
+            Model is trained for fixed forecast lengths (96,192,336,720) and this model add required `prediction_filter_length` to the model instance for required pruning.
+            For Ex. if we need to forecast 150 timepoints given last 512 timepoints using model_path = ibm-granite/granite-timeseries-ttm-r2, then get_model will select the
+            model from 512_192_r2 branch and applies prediction_filter_length = 150 to prune the forecasts from 192 to 150. prediction_filter_length also applies loss
+            only to the pruned forecasts during finetuning.
+        freq_prefix_tuning (*optional*, bool):
+            Future use. Currently do no use this parameter.
+        kwargs:
+            Pass all the extra fine-tuning model parameters intended to be passed in the from_pretrained call to update model configuration.
+
+
+```
+
 ```
 # Load Model from HF Model Hub mentioning the branch name in revision field
 
+
 model = TinyTimeMixerForPrediction.from_pretrained(
-                "https://huggingface.co/ibm/TTM", revision="main"
-            ) 
+                "https://huggingface.co/ibm-granite/granite-timeseries-ttm-r2", revision="main"
+            )
+
+or
+
+from tsfm_public.toolkit.get_model import get_model
+model = get_model(
+            model_path="https://huggingface.co/ibm-granite/granite-timeseries-ttm-r2",
+            context_length=512,
+            prediction_length=96
+        )
+
+
 
 # Do zeroshot
 zeroshot_trainer = Trainer(
@@ -166,6 +222,14 @@ zeroshot_output = zeroshot_trainer.evaluate(dset_test)
 for param in model.backbone.parameters():
   param.requires_grad = False
 
+finetune_model = get_model(
+            model_path="https://huggingface.co/ibm-granite/granite-timeseries-ttm-r2",
+            context_length=512,
+            prediction_length=96,
+            # pass other finetune params of decoder or head
+            head_dropout = 0.2
+        )
+
 finetune_forecast_trainer = Trainer(
         model=model,
         args=finetune_forecast_args,