Delete starcaster_eval_pipeline.py

starcaster_eval_pipeline.py

@@ -1,223 +0,0 @@
-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]}")