change conversion

app.py

@@ -19,25 +19,27 @@ from pytorch_forecasting import (
 from PIL import Image
 
 ## Functions 
-def raw_preds_to_df(raw,quantiles = None):
+def raw_preds_to_df(raw, idx_offset, quantiles = None):
     """
     raw is output of model.predict with return_index=True
     quantiles can be provided like [0.1,0.5,0.9] to get interpretable quantiles
     in the output, time_idx is the first prediction time index (one step after knowledge cutoff)
     pred_idx the index of the predicted date i.e. time_idx + h - 1
     """
-    index = raw[2]
-    preds = raw[0].prediction
-    dec_len = preds.shape[1]
+    index = raw.index
+    preds = raw.output.prediction
+    dec_len = raw.output.prediction.shape[1]
     n_quantiles = preds.shape[-1]
     preds_df = pd.DataFrame(index.values.repeat(dec_len * n_quantiles, axis=0),columns=index.columns)
-    preds_df = preds_df.assign(h=np.tile(np.repeat(np.arange(1,1+dec_len),n_quantiles),len(preds_df)//(dec_len*n_quantiles)))
-    preds_df = preds_df.assign(q=np.tile(np.arange(n_quantiles),len(preds_df)//n_quantiles))
-    preds_df = preds_df.assign(pred=preds.flatten().numpy())
+    preds_df = preds_df.assign(Horizon=np.tile(np.repeat(np.arange(1,1+dec_len),n_quantiles),len(preds_df)//(dec_len*n_quantiles)))
+    preds_df = preds_df.assign(Quantile=np.tile(np.arange(n_quantiles),len(preds_df)//n_quantiles))
+    preds_df = preds_df.assign(Prediction=preds.flatten().cpu().numpy())
     if quantiles is not None:
-        preds_df['q'] = preds_df['q'].map({i:q for i,q in enumerate(quantiles)})
+        preds_df['Quantile'] = preds_df['Quantile'].map({i:q for i,q in enumerate(quantiles)})
 
-    preds_df['pred_idx'] = preds_df['time_idx'] + preds_df['h'] - 1
+    preds_df['pred_idx'] = preds_df['time_idx'] + preds_df['Horizon'] - 1
+    preds_df['Date'] = pd.to_datetime(idx_offset)
+    preds_df['Date'] = preds_df['Date'] + preds_df['pred_idx'].apply(pd.DateOffset)
     return preds_df
 
 def prepare_dataset(_parameters, df, rain, temperature, datepicker, mapping):
@@ -57,7 +59,8 @@ def prepare_dataset(_parameters, df, rain, temperature, datepicker, mapping):
 
 def predict(_model, _dataloader, datepicker):
     out = _model.predict(_dataloader, mode="raw", return_x=True, return_index=True)
-    preds = raw_preds_to_df(out, quantiles = None)
+    first_date = min(df["Date"])
+    preds = raw_preds_to_df(out, first_date)
     return preds[["pred_idx", "Group", "pred"]]
 
 def adjust_data_for_plot(df, preds):
@@ -96,7 +99,7 @@ def generate_plot(df):
 
 @st.cache_data
 def load_data():
-    with open('data/parameters.pkl', 'rb') as f:
+    with open('data/parameters_q.pkl', 'rb') as f:
         parameters = pickle.load(f)
     df = pd.read_pickle('data/test_data.pkl')
     df = df.loc[(df["Branch"] == "15") & (df["Group"].isin(["6","7","4","1"]))]
@@ -104,7 +107,7 @@ def load_data():
 
 @st.cache_resource
 def init_model():
-    model = TemporalFusionTransformer.load_from_checkpoint('model/tft_check.ckpt', map_location=torch.device('cpu'))
+    model = TemporalFusionTransformer.load_from_checkpoint('model/tft_check_q.ckpt', map_location=torch.device('cpu'))
     return model
     
 def main():