Upload starcaster_eval_pipeline.py with huggingface_hub

starcaster_eval_pipeline.py

@@ -0,0 +1,223 @@
+import warnings
+import logging
+import json
+import sys
+
+sys.path.append("/lustre/orion/csc605/scratch/rolandriachi/starcaster/Time-LLM/")
+sys.path.append("/lustre/orion/csc605/scratch/rolandriachi/starcaster/UniTime/")
+sys.path.append("/lustre/orion/csc605/scratch/rolandriachi/starcaster/Time-LLM/models")
+sys.path.append("/lustre/orion/csc605/scratch/rolandriachi/starcaster/UniTime/models")
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import pandas as pd
+
+from TimeLLM import Model as TimeLLMModel
+from unitime import UniTime as UniTimeModel
+
+IMPLEMENTED_BASELINES = [TimeLLMModel, UniTimeModel]
+
+from typing import Optional, Union, Dict, Callable, Iterable
+
+def truncate_mse_loss(future_time, future_pred):
+    # Assumes future_time.shape == (B, T1) and future_pred.shape == (B, T2)
+    min_length = min(future_time.shape[-1], future_pred.shape[-1])
+    return F.mse_loss(future_time[...,:min_length], future_pred[...,:min_length])
+
+def truncate_mae_loss(future_time, future_pred):
+    # Assumes future_time.shape == (B, T1) and future_pred.shape == (B, T2)
+    min_length = min(future_time.shape[-1], future_pred.shape[-1])
+    return F.l1_loss(future_time[...,:min_length], future_pred[...,:min_length])
+
+class DotDict(dict):
+    """dot.notation access to dictionary attributes"""
+    __getattr__ = dict.get
+    __setattr__ = dict.__setitem__
+    __delattr__ = dict.__delitem__
+
+def find_pred_len_from_path(path: str) -> int:
+    if "pl_96" or "pl96" in path: pred_len = 96
+    elif "pl_192" or "pl192" in path: pred_len = 192
+    elif "pl_336" or "pl336" in path: pred_len = 336
+    elif "pl720" or "pl720" in path: pred_lent = 720
+    else:
+        raise ValueError(f"Could not determine prediction length of model from path {path}. Expected path to contain a substring of the form 'pl_{{pred_len}}' or 'pl{{pred_len}}'.")
+
+    return pred_len
+
+def find_model_name_from_path(path: str) -> str:
+    path = path.lower()
+    if "time-llm" in path or "timellm" in path: model_name = "time-llm"
+    elif "unitime" in path: model_name = "unitime"
+    else:
+        raise ValueError(f"Could not determine model name from path {path}. Expected path to contain either 'time-llm', 'timellm', or 'unitime'.")
+
+    return model_name
+
+TIME_LLM_CONFIGS = DotDict({
+    "task_name" : "long_term_forecast", "seq_len" : 512, "enc_in" : 7, "d_model" : 32, "d_ff" : 128, "llm_layers" : 32, "llm_dim" : 4096, 
+    "patch_len" : 16, "stride" : 8, "llm_model" : "LLAMA", "llm_layers" : 32, "prompt_domain" : 1, "content" : None, "dropout" : 0.1, 
+    "d_model" : 32, "n_heads" : 8, "enc_in" : 7
+    })
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+UNITIME_CONFIGS = DotDict({
+    "max_token_num" : 17, "mask_rate" : 0.5, "patch_len" : 16, "max_backcast_len" : 96, "max_forecast_len" : 720, "logger" : logger,
+    "model_path" : "gpt2", "lm_layer_num" : 6, "lm_ft_type" : "freeze", "ts_embed_dropout" : 0.3, "dec_trans_layer_num" : 2, "dec_head_dropout" : 0.1,
+})
+
+class TimeLLMStarCasterWrapper(nn.Module):
+
+    def __init__(self, time_llm_model):
+        super().__init__()
+
+        assert isinstance(time_llm_model, TimeLLMModel), f"TimeLLMStarCasterWrapper can only wrap a model of class TimeLLM.Model but got {type(time_llm_model)}"
+        self.base_model = time_llm_model
+    
+    def forward(self, past_time, context):
+        self.base_model.description = context
+        return self.base_model(x_enc=past_time.unsqueeze(-1), x_mark_enc=None, x_dec=None, x_mark_dec=None).squeeze(-1)
+    
+class UniTimeStarCasterWrapper(nn.Module):
+
+    def __init__(self, unitime_model):
+        super().__init__()
+
+        assert isinstance(unitime_model, UniTimeModel), f"UniTimeStarCasterWrapper can only wrap a model of class TimeLLM.Model but got {type(unitime_model)}"
+        self.base_model = unitime_model
+
+    def forward(self, past_time, context):
+        past_time = past_time.unsqueeze(-1)
+        mask = torch.ones_like(past_time)
+        data_id = -1
+        seq_len = 96
+        stride = 16
+
+        info = (data_id, seq_len, stride, context[:17])
+        return self.base_model(info=info, x_inp=past_time, mask=mask).squeeze(-1)
+
+class StarCasterBaseline(nn.Module):
+
+    def __init__(self, model):
+        super().__init__()
+
+        # TODO: Make this more extendable
+        if type(model) not in IMPLEMENTED_BASELINES:
+            raise NotImplementedError(f"StarCasterBaseline currently only handles models of type {IMPLEMENTED_BASELINES}.")
+        
+        self.base_model = model
+        if isinstance(self.base_model, TimeLLMModel):
+            self.wrapped_model = TimeLLMStarCasterWrapper(self.base_model)
+        if isinstance(self.base_model, UniTimeModel):
+            self.wrapped_model = UniTimeStarCasterWrapper(self.base_model)
+    
+    def forward(self, past_time, context):
+        return self.wrapped_model(past_time, context)
+    
+    def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False):
+        return self.base_model.load_state_dict(state_dict, strict, assign)
+
+
+class EvaluationPipeline:
+
+    def __init__(
+        self, 
+        dataset: Iterable,
+        model: TimeLLMModel, 
+        metrics: Optional[Union[Callable, Dict[str, Callable]]] = None
+    ):
+        self.dataset = dataset
+        self.metrics = metrics if metrics is not None else {"mse_loss" : truncate_mse_loss}
+
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        if self.device == "cpu":
+            warnings.warn("Warning: No CUDA device detected, proceeding with EvaluationPipeline on CPU .....")
+
+        self.model = StarCasterBaseline(model).to(self.device)
+    
+    # TODO: This method needs to be replaced to handle actual StarCaster benchmark
+    def get_evaluation_loader(self) -> Iterable:
+        samples = []
+        for sample in self.dataset.values():
+            past_time = torch.from_numpy(sample["past_time"].to_numpy().T).float().to(self.device)
+            future_time = torch.from_numpy(sample["future_time"].to_numpy().T).float().to(self.device)
+            context = sample["context"]
+
+            samples.append([past_time, future_time, context])
+
+        return samples
+
+    def compute_loss(self, future_time, future_pred):
+        return {m_name : m(future_time, future_pred) for m_name, m in self.metrics.items()}
+    
+    def evaluation_step(self, past_time, future_time, context):
+        with torch.no_grad():
+            future_pred = self.model(past_time, context)
+            loss = self.compute_loss(future_time, future_pred)
+        return loss, future_pred
+            
+    @torch.no_grad()
+    def eval(self):
+        model.eval()
+        infer_dataloader = self.get_evaluation_loader()
+        losses, predictions = {m_name : [] for m_name in self.metrics.keys()}, []
+        for past_time, future_time, context in infer_dataloader:
+            loss_dict, preds = self.evaluation_step(past_time, future_time, context)
+
+            for m_name, loss in loss_dict.items(): losses[m_name].append(loss)
+            predictions.append(preds)
+
+        model.train()
+        return losses, predictions
+
+if __name__ == "__main__":
+    # from argparse import ArgumentParser
+
+    # parser = ArgumentParser()
+
+    # parser.add_argument("--data_path", type=str, required=True)
+    # parser.add_argument("--ckpt_path", type=str, default=None)
+    
+    # args = parser.parse_args()
+
+    # args = TIME_LLM_CONFIGS
+    args = DotDict(dict())
+
+    # args.ckpt_path = "./Time-LLM/checkpoints/long_term_forecast_ETTh1_512_96_TimeLLM_ETTh1_ftM_sl512_ll48_pl96_dm32_nh8_el2_dl1_df128_fc3_ebtimeF_Exp_0-TimeLLM-ETTh1/best_checkpoint/pytorch_model/mp_rank_00_model_states.pt"
+    args.ckpt_path = "/lustre/orion/csc605/scratch/rolandriachi/starcaster/UniTime/outputs/checkpoint_gpt2-small_full_etth1-96_instruct_6_2_0.5_96/model_s2036.pth"
+    args.data_path = "./example_data_dict_simple_dtypes.pkl"
+
+    dataset = pd.read_pickle(args.data_path)
+    # args.pred_len = find_pred_len_from_path(args.ckpt_path)
+    # args.model_name = find_model_name_from_path(args.ckpt_path)
+    args.pred_len = 96
+    args.model_name = "unitime" # "time-llm"
+    
+    if args.model_name == "time-llm":
+        args.update(TIME_LLM_CONFIGS)
+    elif args.model_name == "unitime":
+        args.update(UNITIME_CONFIGS)
+
+    print(f"Initializing model from config:\n{args} .....")
+
+    if args.model_name == "time-llm":
+        model = TimeLLMModel(args)
+    elif args.model_name == "unitime":
+        model = UniTimeModel(args)
+
+    if args.ckpt_path is not None:
+        print(f"Loading model checkpoint from path {args.ckpt_path} .....")
+        ckpt = torch.load(args.ckpt_path)
+        if args.model_name == "time-llm":
+            model.load_state_dict(ckpt["module"]) # TODO: Change this to not be specific to the Time-LLM checkpoint
+        elif args.model_name == "unitime":
+            model.load_state_dict(ckpt)
+
+    pipeline = EvaluationPipeline(dataset, model, metrics={"mse_loss" : truncate_mse_loss, "mae_loss" : truncate_mae_loss})
+
+    print(f"Evaluating .....")
+    losses, predictions = pipeline.eval()
+    print(f"Got losses: {losses}")
+    print(f"Predictions has shape: {[pred.shape for pred in predictions]}")