Added notebook to sample from predictions and load model

src/data/dataset.py

@@ -1,3 +1,4 @@
+import datetime
 import torch
 from torch.utils.data import Dataset, DataLoader
 import pandas as pd
@@ -29,7 +30,6 @@ class NrvDataset(Dataset):
             range(len(dataframe) - self.sequence_length - self.predict_sequence_length)
         )
         self.valid_indices = sorted(list(total_indices - set(self.samples_to_skip)))
-        print(len(self.valid_indices))
 
         self.history_features = []
         if self.data_config.LOAD_HISTORY:
@@ -69,7 +69,9 @@ class NrvDataset(Dataset):
             self.time_feature = None
     
         self.nrv = torch.tensor(dataframe["nrv"].values).float().reshape(-1)
+        self.datetime = dataframe["datetime"]
 
+        print(dataframe.columns)
         self.history_features, self.forecast_features = self.preprocess_data(dataframe)
 
     def skip_samples(self, dataframe):
@@ -203,3 +205,6 @@ class NrvDataset(Dataset):
             return None, torch.stack(targets)
 
         return torch.stack(features), torch.stack(targets)
+
+    def get_idx_for_date(self, date: datetime.date):
+        return self.datetime[self.datetime.dt.date == date].index[0]