Utility plot fixes

[ajati]

This PR fixes a bug in the plotting utility function plot_preds in the file tsfm_public/models/tinytimemixer/utils/ttm_utils.py related to truncation of context_length for better visualization. It also makes the truncation optional (default: True). The PR also re-runs all the notebooks which use the plot_preds function, and deletes some unused imports.

[vijaye12]

Looks good @ajati Lets wait for @wgifford to review and merge it.

[wgifford]

Approving, but we need to relocate to make these more general purpose (eventually).

[wgifford]

Eventually these utilities should be moved out of model-specific folders.

[ajati]

Agree.

[SatishGune]

Have these utilities been moved out to some other location? The below code in ttm_tutorial.py is failing

from tsfm_public.models.tinytimemixer.utils import (
count_parameters,
plot_preds,
)

Here are the errors.
ImportError: cannot import name 'count_parameters' from 'tsfm_public.models.tinytimemixer.utils' (/content/tsfm/tsfm_public/models/tinytimemixer/utils/init.py)

ImportError: cannot import name 'plot_preds' from 'tsfm_public.models.tinytimemixer.utils' (/content/tsfm/tsfm_public/models/tinytimemixer/utils/init.py)

[wgifford]

@SatishGune They were moved -- let me double check.

[wgifford]

plot_preds is now plot_predictions. It can be imported with from tsfm_public.toolkit.visualization import plot_predictions. The signature of plot_predictions is now:

def plot_predictions(
    model: torch.nn.Module,
    dset: torch.utils.data.Dataset,
    plot_dir: str = None,
    num_plots: int = 10,
    plot_prefix: str = "valid",
    channel: int = -1,
    truncate_history: bool = True,
)

It takes a model as its first argument instead of a trainer. The ttm_tutorial.ipynb has been updated to reflect these changes.

[SatishGune]

Thanks @wgifford . Much appreciated.