AWildHippo
commited on
Commit
•
f9e4730
1
Parent(s):
d11b9f4
Upload starcaster_eval_pipeline.py with huggingface_hub
Browse files- starcaster_eval_pipeline.py +223 -0
starcaster_eval_pipeline.py
ADDED
@@ -0,0 +1,223 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import warnings
|
2 |
+
import logging
|
3 |
+
import json
|
4 |
+
import sys
|
5 |
+
|
6 |
+
sys.path.append("/lustre/orion/csc605/scratch/rolandriachi/starcaster/Time-LLM/")
|
7 |
+
sys.path.append("/lustre/orion/csc605/scratch/rolandriachi/starcaster/UniTime/")
|
8 |
+
sys.path.append("/lustre/orion/csc605/scratch/rolandriachi/starcaster/Time-LLM/models")
|
9 |
+
sys.path.append("/lustre/orion/csc605/scratch/rolandriachi/starcaster/UniTime/models")
|
10 |
+
|
11 |
+
import torch
|
12 |
+
import torch.nn as nn
|
13 |
+
import torch.nn.functional as F
|
14 |
+
import pandas as pd
|
15 |
+
|
16 |
+
from TimeLLM import Model as TimeLLMModel
|
17 |
+
from unitime import UniTime as UniTimeModel
|
18 |
+
|
19 |
+
IMPLEMENTED_BASELINES = [TimeLLMModel, UniTimeModel]
|
20 |
+
|
21 |
+
from typing import Optional, Union, Dict, Callable, Iterable
|
22 |
+
|
23 |
+
def truncate_mse_loss(future_time, future_pred):
|
24 |
+
# Assumes future_time.shape == (B, T1) and future_pred.shape == (B, T2)
|
25 |
+
min_length = min(future_time.shape[-1], future_pred.shape[-1])
|
26 |
+
return F.mse_loss(future_time[...,:min_length], future_pred[...,:min_length])
|
27 |
+
|
28 |
+
def truncate_mae_loss(future_time, future_pred):
|
29 |
+
# Assumes future_time.shape == (B, T1) and future_pred.shape == (B, T2)
|
30 |
+
min_length = min(future_time.shape[-1], future_pred.shape[-1])
|
31 |
+
return F.l1_loss(future_time[...,:min_length], future_pred[...,:min_length])
|
32 |
+
|
33 |
+
class DotDict(dict):
|
34 |
+
"""dot.notation access to dictionary attributes"""
|
35 |
+
__getattr__ = dict.get
|
36 |
+
__setattr__ = dict.__setitem__
|
37 |
+
__delattr__ = dict.__delitem__
|
38 |
+
|
39 |
+
def find_pred_len_from_path(path: str) -> int:
|
40 |
+
if "pl_96" or "pl96" in path: pred_len = 96
|
41 |
+
elif "pl_192" or "pl192" in path: pred_len = 192
|
42 |
+
elif "pl_336" or "pl336" in path: pred_len = 336
|
43 |
+
elif "pl720" or "pl720" in path: pred_lent = 720
|
44 |
+
else:
|
45 |
+
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}}'.")
|
46 |
+
|
47 |
+
return pred_len
|
48 |
+
|
49 |
+
def find_model_name_from_path(path: str) -> str:
|
50 |
+
path = path.lower()
|
51 |
+
if "time-llm" in path or "timellm" in path: model_name = "time-llm"
|
52 |
+
elif "unitime" in path: model_name = "unitime"
|
53 |
+
else:
|
54 |
+
raise ValueError(f"Could not determine model name from path {path}. Expected path to contain either 'time-llm', 'timellm', or 'unitime'.")
|
55 |
+
|
56 |
+
return model_name
|
57 |
+
|
58 |
+
TIME_LLM_CONFIGS = DotDict({
|
59 |
+
"task_name" : "long_term_forecast", "seq_len" : 512, "enc_in" : 7, "d_model" : 32, "d_ff" : 128, "llm_layers" : 32, "llm_dim" : 4096,
|
60 |
+
"patch_len" : 16, "stride" : 8, "llm_model" : "LLAMA", "llm_layers" : 32, "prompt_domain" : 1, "content" : None, "dropout" : 0.1,
|
61 |
+
"d_model" : 32, "n_heads" : 8, "enc_in" : 7
|
62 |
+
})
|
63 |
+
|
64 |
+
logger = logging.getLogger(__name__)
|
65 |
+
logger.setLevel(logging.INFO)
|
66 |
+
UNITIME_CONFIGS = DotDict({
|
67 |
+
"max_token_num" : 17, "mask_rate" : 0.5, "patch_len" : 16, "max_backcast_len" : 96, "max_forecast_len" : 720, "logger" : logger,
|
68 |
+
"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,
|
69 |
+
})
|
70 |
+
|
71 |
+
class TimeLLMStarCasterWrapper(nn.Module):
|
72 |
+
|
73 |
+
def __init__(self, time_llm_model):
|
74 |
+
super().__init__()
|
75 |
+
|
76 |
+
assert isinstance(time_llm_model, TimeLLMModel), f"TimeLLMStarCasterWrapper can only wrap a model of class TimeLLM.Model but got {type(time_llm_model)}"
|
77 |
+
self.base_model = time_llm_model
|
78 |
+
|
79 |
+
def forward(self, past_time, context):
|
80 |
+
self.base_model.description = context
|
81 |
+
return self.base_model(x_enc=past_time.unsqueeze(-1), x_mark_enc=None, x_dec=None, x_mark_dec=None).squeeze(-1)
|
82 |
+
|
83 |
+
class UniTimeStarCasterWrapper(nn.Module):
|
84 |
+
|
85 |
+
def __init__(self, unitime_model):
|
86 |
+
super().__init__()
|
87 |
+
|
88 |
+
assert isinstance(unitime_model, UniTimeModel), f"UniTimeStarCasterWrapper can only wrap a model of class TimeLLM.Model but got {type(unitime_model)}"
|
89 |
+
self.base_model = unitime_model
|
90 |
+
|
91 |
+
def forward(self, past_time, context):
|
92 |
+
past_time = past_time.unsqueeze(-1)
|
93 |
+
mask = torch.ones_like(past_time)
|
94 |
+
data_id = -1
|
95 |
+
seq_len = 96
|
96 |
+
stride = 16
|
97 |
+
|
98 |
+
info = (data_id, seq_len, stride, context[:17])
|
99 |
+
return self.base_model(info=info, x_inp=past_time, mask=mask).squeeze(-1)
|
100 |
+
|
101 |
+
class StarCasterBaseline(nn.Module):
|
102 |
+
|
103 |
+
def __init__(self, model):
|
104 |
+
super().__init__()
|
105 |
+
|
106 |
+
# TODO: Make this more extendable
|
107 |
+
if type(model) not in IMPLEMENTED_BASELINES:
|
108 |
+
raise NotImplementedError(f"StarCasterBaseline currently only handles models of type {IMPLEMENTED_BASELINES}.")
|
109 |
+
|
110 |
+
self.base_model = model
|
111 |
+
if isinstance(self.base_model, TimeLLMModel):
|
112 |
+
self.wrapped_model = TimeLLMStarCasterWrapper(self.base_model)
|
113 |
+
if isinstance(self.base_model, UniTimeModel):
|
114 |
+
self.wrapped_model = UniTimeStarCasterWrapper(self.base_model)
|
115 |
+
|
116 |
+
def forward(self, past_time, context):
|
117 |
+
return self.wrapped_model(past_time, context)
|
118 |
+
|
119 |
+
def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False):
|
120 |
+
return self.base_model.load_state_dict(state_dict, strict, assign)
|
121 |
+
|
122 |
+
|
123 |
+
class EvaluationPipeline:
|
124 |
+
|
125 |
+
def __init__(
|
126 |
+
self,
|
127 |
+
dataset: Iterable,
|
128 |
+
model: TimeLLMModel,
|
129 |
+
metrics: Optional[Union[Callable, Dict[str, Callable]]] = None
|
130 |
+
):
|
131 |
+
self.dataset = dataset
|
132 |
+
self.metrics = metrics if metrics is not None else {"mse_loss" : truncate_mse_loss}
|
133 |
+
|
134 |
+
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
135 |
+
if self.device == "cpu":
|
136 |
+
warnings.warn("Warning: No CUDA device detected, proceeding with EvaluationPipeline on CPU .....")
|
137 |
+
|
138 |
+
self.model = StarCasterBaseline(model).to(self.device)
|
139 |
+
|
140 |
+
# TODO: This method needs to be replaced to handle actual StarCaster benchmark
|
141 |
+
def get_evaluation_loader(self) -> Iterable:
|
142 |
+
samples = []
|
143 |
+
for sample in self.dataset.values():
|
144 |
+
past_time = torch.from_numpy(sample["past_time"].to_numpy().T).float().to(self.device)
|
145 |
+
future_time = torch.from_numpy(sample["future_time"].to_numpy().T).float().to(self.device)
|
146 |
+
context = sample["context"]
|
147 |
+
|
148 |
+
samples.append([past_time, future_time, context])
|
149 |
+
|
150 |
+
return samples
|
151 |
+
|
152 |
+
def compute_loss(self, future_time, future_pred):
|
153 |
+
return {m_name : m(future_time, future_pred) for m_name, m in self.metrics.items()}
|
154 |
+
|
155 |
+
def evaluation_step(self, past_time, future_time, context):
|
156 |
+
with torch.no_grad():
|
157 |
+
future_pred = self.model(past_time, context)
|
158 |
+
loss = self.compute_loss(future_time, future_pred)
|
159 |
+
return loss, future_pred
|
160 |
+
|
161 |
+
@torch.no_grad()
|
162 |
+
def eval(self):
|
163 |
+
model.eval()
|
164 |
+
infer_dataloader = self.get_evaluation_loader()
|
165 |
+
losses, predictions = {m_name : [] for m_name in self.metrics.keys()}, []
|
166 |
+
for past_time, future_time, context in infer_dataloader:
|
167 |
+
loss_dict, preds = self.evaluation_step(past_time, future_time, context)
|
168 |
+
|
169 |
+
for m_name, loss in loss_dict.items(): losses[m_name].append(loss)
|
170 |
+
predictions.append(preds)
|
171 |
+
|
172 |
+
model.train()
|
173 |
+
return losses, predictions
|
174 |
+
|
175 |
+
if __name__ == "__main__":
|
176 |
+
# from argparse import ArgumentParser
|
177 |
+
|
178 |
+
# parser = ArgumentParser()
|
179 |
+
|
180 |
+
# parser.add_argument("--data_path", type=str, required=True)
|
181 |
+
# parser.add_argument("--ckpt_path", type=str, default=None)
|
182 |
+
|
183 |
+
# args = parser.parse_args()
|
184 |
+
|
185 |
+
# args = TIME_LLM_CONFIGS
|
186 |
+
args = DotDict(dict())
|
187 |
+
|
188 |
+
# 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"
|
189 |
+
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"
|
190 |
+
args.data_path = "./example_data_dict_simple_dtypes.pkl"
|
191 |
+
|
192 |
+
dataset = pd.read_pickle(args.data_path)
|
193 |
+
# args.pred_len = find_pred_len_from_path(args.ckpt_path)
|
194 |
+
# args.model_name = find_model_name_from_path(args.ckpt_path)
|
195 |
+
args.pred_len = 96
|
196 |
+
args.model_name = "unitime" # "time-llm"
|
197 |
+
|
198 |
+
if args.model_name == "time-llm":
|
199 |
+
args.update(TIME_LLM_CONFIGS)
|
200 |
+
elif args.model_name == "unitime":
|
201 |
+
args.update(UNITIME_CONFIGS)
|
202 |
+
|
203 |
+
print(f"Initializing model from config:\n{args} .....")
|
204 |
+
|
205 |
+
if args.model_name == "time-llm":
|
206 |
+
model = TimeLLMModel(args)
|
207 |
+
elif args.model_name == "unitime":
|
208 |
+
model = UniTimeModel(args)
|
209 |
+
|
210 |
+
if args.ckpt_path is not None:
|
211 |
+
print(f"Loading model checkpoint from path {args.ckpt_path} .....")
|
212 |
+
ckpt = torch.load(args.ckpt_path)
|
213 |
+
if args.model_name == "time-llm":
|
214 |
+
model.load_state_dict(ckpt["module"]) # TODO: Change this to not be specific to the Time-LLM checkpoint
|
215 |
+
elif args.model_name == "unitime":
|
216 |
+
model.load_state_dict(ckpt)
|
217 |
+
|
218 |
+
pipeline = EvaluationPipeline(dataset, model, metrics={"mse_loss" : truncate_mse_loss, "mae_loss" : truncate_mae_loss})
|
219 |
+
|
220 |
+
print(f"Evaluating .....")
|
221 |
+
losses, predictions = pipeline.eval()
|
222 |
+
print(f"Got losses: {losses}")
|
223 |
+
print(f"Predictions has shape: {[pred.shape for pred in predictions]}")
|