Kaggle competition: Jane Street Market Prediction
<SOS> metric: JaneStreetPnl Kaggle data source: jane-street-market-prediction<define_variables>
!mkdir cache
DATA_PATH = '.. /input/prostate-cancer-grade-assessment'
MODELS_PATH = '.'
IMG_SIZE = 128
SEQ_LEN = 36
TARGET_SIZE = [528, 528]
BATCH_SIZE = 16
MDL_VERSION = 'v0'
SEED = 80
LEVEL = 1<train_model>
class PurgedGroupTimeSeriesSplit(_BaseKFold):

 @_deprecate_positional_args
def __init__(self,
n_splits=5,
*,
max_train_group_size=np.inf,
max_test_group_size=np.inf,
group_gap=None,
verbose=False 
):
super().__init__(n_splits, shuffle=False, random_state=None)
 self.max_train_group_size = max_train_group_size
self.group_gap = group_gap
self.max_test_group_size = max_test_group_size
self.verbose = verbose
def split(self, X, y=None, groups=None):

 if groups is None:
raise ValueError(
 "The 'groups' parameter should not be None")
 X, y, groups = indexable(X, y, groups)
 n_samples = _num_samples(X)
 n_splits = self.n_splits
group_gap = self.group_gap
max_test_group_size = self.max_test_group_size
max_train_group_size = self.max_train_group_size
n_folds = n_splits + 1
group_dict = {}
u, ind = np.unique(groups, return_index=True)
 unique_groups = u[np.argsort(ind)]
n_samples = _num_samples(X)
 n_groups = _num_samples(unique_groups)
 for idx in np.arange(n_samples):
if(groups[idx] in group_dict):
group_dict[groups[idx]].append(idx)
 else:
group_dict[groups[idx]] = [idx]
if n_folds > n_groups:
raise ValueError(
("Cannot have number of folds={0} greater than"
" the number of groups={1}" ).format(n_folds,
n_groups))
group_test_size = min(n_groups // n_folds, max_test_group_size)
 group_test_starts = range(n_groups - n_splits * group_test_size,
n_groups, group_test_size)
 for group_test_start in group_test_starts:
train_array = []
test_array = []
group_st = max(0, group_test_start - group_gap - max_train_group_size)
 for train_group_idx in unique_groups[group_st:(group_test_start - group_gap)]:
train_array_tmp = group_dict[train_group_idx]
train_array = np.sort(np.unique(
 np.concatenate(( train_array,
train_array_tmp)) ,
axis=None), axis=None)
 train_end = train_array.size
for test_group_idx in unique_groups[group_test_start:
group_test_start +
group_test_size]:
test_array_tmp = group_dict[test_group_idx]
test_array = np.sort(np.unique(
 np.concatenate(( test_array,
test_array_tmp)) ,
axis=None), axis=None)
 test_array = test_array[group_gap:]
if self.verbose > 0:
pass
yield [int(i)for i in train_array], [int(i)for i in test_array]
def get_axis_max_min(array, axis=0):
one_axis = list(( array != 255 ).sum(axis=tuple([x for x in(0, 1, 2)if x != axis])))
 axis_min = next(( i for i, x in enumerate(one_axis)if x), 0)
 axis_max = len(one_axis)- next(( i for i, x in enumerate(one_axis[::-1])if x), 0)
 return axis_min, axis_max
class DataGenPanda(Sequence):
def __init__(self, imgs_path, df, batch_size=32,
mode='fit', shuffle=False, aug=None,
seq_len=SEQ_LEN, img_size=IMG_SIZE, n_classes=6):
self.imgs_path = imgs_path
self.df = df
self.shuffle = shuffle
self.mode = mode
self.aug = aug
self.batch_size = batch_size
self.img_size = img_size
self.seq_len = seq_len
self.n_classes = n_classes
self.side = int(seq_len **.5)
 self.on_epoch_end()
def __len__(self):
return int(np.floor(len(self.df)/ self.batch_size))
def on_epoch_end(self):
self.indexes = np.arange(len(self.df))
if self.shuffle:
np.random.shuffle(self.indexes)
 def __getitem__(self, index):
X = np.zeros(( self.batch_size, self.side * self.img_size, self.side * self.img_size, 3), dtype=np.float32)
 imgs_batch = self.df[index * self.batch_size :(index + 1)* self.batch_size]['image_id'].values
for i, img_name in enumerate(imgs_batch):
img_path = '{}/{}.tiff'.format(self.imgs_path, img_name)
 img_patches = self.get_patches(img_path)
 X[i, ] = self.glue_to_one(img_patches)
 X = tf.image.convert_image_dtype(X, tf.float32)
 X = tf.image.resize(X, size = TARGET_SIZE)
 if self.mode == 'fit':
y = np.zeros(( self.batch_size, self.n_classes), dtype=np.float32)
 lbls_batch = self.df[index * self.batch_size :(index + 1)* self.batch_size]['isup_grade'].values
for i in range(self.batch_size):
y[i, lbls_batch[i]] = 1
return X, y
elif self.mode == 'predict':
return X
else:
raise AttributeError('mode parameter error')
 def get_patches(self, img_path):
num_patches = self.seq_len
p_size = self.img_size
img = skimage.io.MultiImage(img_path)[LEVEL]
a0min, a0max = get_axis_max_min(img, axis=0)
 a1min, a1max = get_axis_max_min(img, axis=1)
 img = img[a0min:a0max, a1min:a1max, :].astype(np.float32)/ 255
if self.aug:
img = self.aug(image=img)['image']
pad0, pad1 =(p_size - img.shape[0] % p_size)% p_size,(p_size - img.shape[1] % p_size)% p_size
img = np.pad(
 img,
[
[pad0 // 2, pad0 - pad0 // 2],
[pad1 // 2, pad1 - pad1 // 2],
[0, 0]
],
constant_values=1 
)
 img = img.reshape(img.shape[0] // p_size, p_size, img.shape[1] // p_size, p_size, 3)
 img = img.transpose(0, 2, 1, 3, 4 ).reshape(-1, p_size, p_size, 3)
 if len(img)< num_patches:
img = np.pad(
 img,
[
[0, num_patches - len(img)],
[0, 0],
[0, 0],
[0, 0]
],
constant_values=1 
)
 idxs = np.argsort(img.reshape(img.shape[0], -1 ).sum(-1)) [:num_patches]
return np.array(img[idxs])
 def glue_to_one(self, imgs_seq):
img_glue = np.zeros(( self.img_size * self.side, self.img_size * self.side, 3), dtype=np.float32)
 for i, ptch in enumerate(imgs_seq):
x = i // self.side
y = i % self.side
img_glue[x * self.img_size :(x + 1)* self.img_size,
y * self.img_size :(y + 1)* self.img_size, :] = ptch
return img_glue<split>
class FinData(Dataset):
def __init__(self, data, target, date, mode='train', transform=None, cache_dir=None, multi=False):
self.data = data
self.target = target
self.mode = mode
self.transform = transform
self.cache_dir = cache_dir
self.date = date
self.multi = multi
def __getitem__(self, index):
if torch.is_tensor(index):
index.to_list()
if self.transform:
return self.transform(self.data.iloc[index].values)
 else:
if type(index)is list:
if self.multi == False:
sample = {
'target': torch.Tensor(self.target.iloc[index].values),
'data': torch.FloatTensor(self.data[index]),
'date': torch.Tensor(self.date.iloc[index].values)
 }
elif self.multi == True:
sample = {
'target': torch.Tensor(self.target[index]),
'data': torch.FloatTensor(self.data[index]),
'date': torch.Tensor(self.date.iloc[index].values)
 }
else:
if self.multi == False:
sample = {
'target': torch.Tensor(self.target.iloc[index]),
'data': torch.FloatTensor(self.data[index]),
'date': torch.Tensor(self.date.iloc[index])
 }
elif self.multi == True:
sample = {
'target': torch.Tensor(self.target[index]),
'data': torch.FloatTensor(self.data[index]),
'date': torch.Tensor(self.date.iloc[index])
 }
return sample
def __len__(self):
return len(self.data)
 def load_data(root_dir, mode, overide=None):
if overide:
data = dt.fread(overide ).to_pandas()
elif mode == 'train':
data = dt.fread(root_dir + 'train.csv' ).to_pandas()
elif mode == 'test':
data = dt.fread(root_dir + 'example_test.csv' ).to_pandas()
elif mode == 'sub':
data = dt.fread(root_dir + 'example_sample_submission.csv' ).to_pandas()
return data
def preprocess_data(data: pd.DataFrame, scale: bool = False, nn: bool = False,
action: str = 'weight'):

 data = data.query('weight > 0' ).reset_index(drop=True)
 data = data.query('date > 85' ).reset_index(drop=True)
 if action == 'weight':
data['action'] =(
(data['weight'].values * data['resp'].values)> 0 ).astype('float32')
 if action == 'combined':
data['action'] =(
(data['resp'].values > 0)and(data['resp_1'].values > 0)and(data['resp_2'].values > 0)and(
 data['resp_3'].values > 0)and(data['resp_4'].values > 0)).astype('float32')
 if action == 'multi':
resp_cols = ['resp', 'resp_1', 'resp_2', 'resp_3', 'resp_4']
for i in range(len(resp_cols)) :
data['action_' + str(i)] =(data['weight'] * data[resp_cols[i]] > 0 ).astype('int')

features = [col for col in data.columns if 'feature' in col] + ['weight']
date = data['date']
if action == 'multi':
target = np.array([data['action_' + str(i)]
for i in range(len(resp_cols)) ] ).T
else:
target = data['action']
data = data[features]
if scale:
scaler = StandardScaler()
data = scaler.fit_transform(data)
 if not scale and nn:
data = data.values
return data, target, features, date
def calc_data_mean(array, cache_dir=None, fold=None, train=True, mode='mean'):
if train:
if mode == 'mean':
f_mean = np.nanmean(array, axis=0)
 if cache_dir and fold:
np.save(f'{cache_dir}/f_{fold}_mean.npy', f_mean)
 elif cache_dir:
np.save(f'{cache_dir}/f_mean.npy', f_mean)
 array = np.nan_to_num(array)+ np.isnan(array)* f_mean
if mode == 'median':
f_med = np.nanmedian(array, axis=0)
 if cache_dir and fold:
np.save(f'{cache_dir}/f_{fold}_median.npy', f_med)
 elif cache_dir:
np.save(f'{cache_dir}/f_median.npy', f_med)
 array = np.nan_to_num(array)+ np.isnan(array)* f_med
if mode == 'zero':
array = np.nan_to_num(array)+ np.isnan(array)* 0
if not train:
if mode == 'mean':
f_mean = np.load(f'{cache_dir}/f_mean.npy')
 array = np.nan_to_num(array)+ np.isnan(array)* f_mean
if mode == 'median':
f_med = np.load(f'{cache_dir}/f_med.npy')
 array = np.nan_to_num(array)+ np.isnan(array)* f_med
if mode == 'zero':
array = np.nan_to_num(array)+ np.isnan(array)* 0
return array
def weighted_mean(scores, sizes):
largest = np.max(sizes)
 weights = [size / largest for size in sizes]
return np.average(scores, weights=weights)
 def create_dataloaders(dataset: Dataset, indexes: dict, batch_size):
train_idx = indexes.get('train', None)
 val_idx = indexes.get('val', None)
 test_idx = indexes.get('test', None)
 dataloaders = {}
if train_idx:
train_set = Subset(
 dataset, train_idx)
 train_sampler = BatchSampler(SequentialSampler(
 train_set), batch_size=batch_size, drop_last=False)
 dataloaders['train'] = DataLoader(
 dataset, sampler=train_sampler, num_workers=10, pin_memory=True)
 if val_idx:
val_set = Subset(dataset, val_idx)
 val_sampler = BatchSampler(SequentialSampler(
 val_set), batch_size=batch_size, drop_last=False)
 dataloaders['val'] = DataLoader(
 dataset, sampler=val_sampler, num_workers=10, pin_memory=True)
 if test_idx:
test_set = Subset(dataset, test_idx)
 test_sampler = BatchSampler(SequentialSampler(
 test_set), batch_size=batch_size, drop_last=False)
 dataloaders['test'] = DataLoader(
 dataset, sampler=test_sampler, num_workers=10, pin_memory=True)
 return dataloaders
def seed_everything(seed):
random.seed(seed)
 torch.manual_seed(seed)
 np.random.seed(seed)
 os.environ['PYTHONHASHSEED'] = str(seed)
 torch.cuda.manual_seed(seed)
 torch.backends.cudnn.deterministic = True
def load_model(path, input_size, output_size, p, pl_lightning):
if os.path.isdir(path):
models = []
for file in os.listdir(path):
if pl_lightning:
model = Regressor.load_from_checkpoint(checkpoint_path=file, input_size=input_size,
output_size=output_size, params=p)
 else:
model = Regressor(input_size, output_size, params=p)
 model.load_state_dict(torch.load(f'{path}/{file}'))
models.append(model)
 return models
elif os.path.isfile(path):
if pl_lightning:
return Regressor.load_from_checkpoint(checkpoint_path=path, input_size=input_size,
output_size=output_size, params=p)
 else:
model = Regressor(input_size, output_size, params=p)
 model.load_state_dict(torch.load(path))
return model
train = pd.read_csv('{}/train.csv'.format(DATA_PATH))
print('train: ', train.shape, '| unique ids:', sum(train['isup_grade'].value_counts()))
X_train, X_val = train_test_split(train, test_size=.2, stratify=train['isup_grade'], random_state=SEED)
 lbl_value_counts = X_train['isup_grade'].value_counts()
class_weights = {i: max(lbl_value_counts)/ v for i, v in lbl_value_counts.items() }
print('classes weigths:', class_weights )<train_model>
class Regressor(pl.LightningModule):
def __init__(self, input_size, output_size, params, model_path='models/'):
super(Regressor, self ).__init__()
dim_1 = params['dim_1']
dim_2 = params['dim_2']
dim_3 = params['dim_3']
dim_4 = params['dim_4']
self.dropout_prob = params['dropout']
self.lr = params['lr']
self.activation = params['activation']
self.input_size = input_size
self.output_size = output_size
self.loss = nn.BCEWithLogitsLoss()
self.weight_decay = params['weight_decay']
self.amsgrad = params['amsgrad']
self.label_smoothing = params['label_smoothing']
self.model_path = model_path
self.encoder = nn.Sequential(
 nn.BatchNorm1d(input_size),
nn.Linear(input_size, dim_1, bias=False),
nn.BatchNorm1d(dim_1),
self.activation() ,
nn.Dropout(p=self.dropout_prob),
nn.Linear(dim_1, dim_2, bias=False),
nn.BatchNorm1d(dim_2),
self.activation() ,
nn.Dropout(p=self.dropout_prob),
nn.Linear(dim_2, dim_3, bias=False),
nn.BatchNorm1d(dim_3),
self.activation() ,
nn.Dropout(p=self.dropout_prob),
nn.Linear(dim_3, dim_4, bias=False),
nn.BatchNorm1d(dim_4),
self.activation() ,
nn.Dropout(p=self.dropout_prob),
nn.Linear(dim_4, self.output_size, bias=False)
)
 def forward(self, x):
out = self.encoder(x)
 return out
def training_step(self, batch, batch_idx):
x, y = batch['data'], batch['target']
x = x.view(x.size(1), -1)
 y = y.view(y.size(1), -1)
 logits = self(x)
 loss = self.loss(input=logits, target=y)
 logits = torch.sigmoid(logits)
 auc_metric = roc_auc_score(y_true=y.cpu().numpy() ,
y_score=logits.cpu().detach().numpy())
 self.log('train_auc', auc_metric, on_step=False,
on_epoch=True, prog_bar=True)
 self.log('train_loss', loss, prog_bar=True)
 return {'loss': loss}
def validation_step(self, batch, batch_idx):
x, y = batch['data'], batch['target']
x = x.view(x.size(1), -1)
 y = y.view(y.size(1), -1)
 logits = self(x)
 loss = self.loss(input=logits,
target=y)
 logits = torch.sigmoid(logits)
 auc = roc_auc_score(y_true=y.cpu().numpy() ,
y_score=logits.cpu().detach().numpy())
 return {'loss': loss, 'y': y, 'logits': logits, 'auc': auc}
def validation_epoch_end(self, val_step_outputs):
epoch_loss = torch.tensor([x['loss'] for x in val_step_outputs] ).mean()
epoch_auc = torch.tensor([x['auc'] for x in val_step_outputs] ).mean()
self.log('val_loss', epoch_loss, prog_bar=True)
 self.log('val_auc', epoch_auc, prog_bar=True)
 def test_step(self, batch, batch_idx):
return self.validation_step(batch, batch_idx)
 def test_epoch_end(self, outputs):
epoch_loss = torch.tensor([x['loss'] for x in outputs] ).mean()
epoch_auc = torch.tensor([x['auc'] for x in outputs] ).mean()
self.log('test_loss', epoch_loss)
 self.log('test_auc', epoch_auc)
 def configure_optimizers(self):
optimizer = torch.optim.Adam(self.parameters() , lr=self.lr,
amsgrad=self.amsgrad)
 scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
 optimizer, patience=5, factor=0.1, min_lr=1e-7, eps=1e-08 
)
 return {'optimizer': optimizer, 'lr_scheduler': scheduler, 'monitor': 'val_auc'}
def init_weights(m, func):
if type(m)== nn.Linear:
nn.init.xavier_normal_(m.weight, nn.init.calculate_gain(func))
def cross_val(p)-> object:
data_ = load_data(root_dir='./data/', mode='train')
 data_, target_, features, date = preprocess_data(
 data_, nn=True, action='multi')
 input_size = data_.shape[-1]
output_size = target_.shape[-1]
gts = PurgedGroupTimeSeriesSplit(n_splits=5, group_gap=5)
 models = []
tb_logger = pl_loggers.TensorBoardLogger('logs/multiclass_')
 for i,(train_idx, val_idx)in enumerate(gts.split(data_, groups=date)) :
idx = np.concatenate([train_idx, val_idx])
 data = copy.deepcopy(data_[idx])
 target = copy.deepcopy(target_[idx])
 checkpoint_callback = pl.callbacks.ModelCheckpoint(
 os.path.join('models/', 'multi_class_fold_{}'.format(i)) , monitor='val_auc', save_top_k=1, period=10,
mode='max' 
)
 model = Regressor(input_size, output_size, p)
 if p['activation'] == nn.ReLU:
model.apply(lambda m: init_weights(m, 'relu'))
elif p['activation'] == nn.LeakyReLU:
model.apply(lambda m: init_weights(m, 'leaky_relu'))
train_idx = [i for i in range(0, max(train_idx)+ 1)]
val_idx = [i for i in range(len(train_idx), len(idx)) ]
data[train_idx] = calc_data_mean(
 data[train_idx], './cache', train=True, mode='mean')
 data[val_idx] = calc_data_mean(
 data[val_idx], './cache', train=False, mode='mean')
 dataset = FinData(data=data, target=target, date=date, multi=True)
 dataloaders = create_dataloaders(
 dataset, indexes={'train': train_idx, 'val': val_idx}, batch_size=p['batch_size'])
 es = EarlyStopping(monitor='val_auc', patience=10,
min_delta=0.0005, mode='max')
 trainer = pl.Trainer(logger=tb_logger,
max_epochs=1,
gpus=1,
callbacks=[checkpoint_callback, es],
precision=16)
 trainer.fit(
 model, train_dataloader=dataloaders['train'], val_dataloaders=dataloaders['val'])
 torch.save(model.state_dict() , f'models/fold_{i}_state_dict.pth')
 models.append(model)
 return models, features
aug = albu.Compose(
 [
albu.OneOf([albu.RandomBrightness(limit=.15), albu.RandomContrast(limit=.3), albu.RandomGamma() ], p=.25),
albu.HorizontalFlip(p=.25),
albu.VerticalFlip(p=.25),
albu.ShiftScaleRotate(shift_limit=.1, scale_limit=.1, rotate_limit=20, p=.25)
 ] 
)
 train_datagen = DataGenPanda(
 imgs_path='{}/train_images'.format(DATA_PATH),
df=X_train,
batch_size=BATCH_SIZE,
mode='fit',
shuffle=True,
aug=aug,
seq_len=SEQ_LEN,
img_size=IMG_SIZE,
n_classes=6 
)
 val_datagen = DataGenPanda(
 imgs_path='{}/train_images'.format(DATA_PATH),
df=X_val,
batch_size=BATCH_SIZE,
mode='fit',
shuffle=False,
aug=None,
seq_len=SEQ_LEN,
img_size=IMG_SIZE,
n_classes=6
)<load_pretrained>
def final_train(p, load=False):
data_ = load_data(root_dir='./data/', mode='train',overide='/kaggle/input/jane-street-market-prediction/train.csv')
 data, target, features, date = preprocess_data(data_, nn=True)
 dataset = FinData(data=data, target=target, date=date)
 input_size = data.shape[-1]
output_size = 1
train_idx, val_idx = date[date <= 498].index.values.tolist() , date[date > 498].index.values.tolist()
data[train_idx] = calc_data_mean(data[train_idx], './cache', train=True)
 data[val_idx] = calc_data_mean(data[val_idx], './cache', train=False)
 checkpoint_callback = pl.callbacks.ModelCheckpoint(filepath='./cache/full_train',
monitor="val_auc", mode='max', save_top_k=1, period=10)
 model = Regressor(input_size=input_size,
output_size=output_size, params=p)
 if p['activation'] == nn.ReLU:
model.apply(lambda m: init_weights(m, 'relu'))
elif p['activation'] == nn.LeakyReLU:
model.apply(lambda m: init_weights(m, 'leaky_relu'))
dataset = FinData(data, target, date)
 dataloaders = create_dataloaders(dataset, indexes={'train': train_idx, 'val': val_idx}, batch_size=p['batch_size'])
 es = EarlyStopping(monitor='val_auc', patience=10,
min_delta=0.0005, mode='max')
 trainer = pl.Trainer(max_epochs=25,
gpus=1,
callbacks=[checkpoint_callback, es],
precision=16)
 trainer.fit(model, train_dataloader=dataloaders['train'], val_dataloaders=dataloaders['val'])
 torch.save(model.state_dict() , './cache/final_train.pth')
 return model, features
MODEL_DIR = '.. /input/pandamodels/'
models = {'B0': 'EfficientNetB0.h5',
'B2':'EfficientNetB2.h5',
'B4':'EfficientNetB4.h5',
'B6':'EfficientNetB6.h5',
'Xception':'Xception.h5'}
model = tf.keras.models.load_model(MODEL_DIR + models['B6'], compile=False)
 print('Model loaded:', model )<load_from_csv>
def fillna_npwhere(array, values):
if np.isnan(array.sum()):
array = np.nan_to_num(array)+ np.isnan(array)* values
return array
def test_model(models, features, cache_dir='cache'):
env = janestreet.make_env()
iter_test = env.iter_test()
if type(models)== list:
models = [model.eval() for model in models]
else:
models.eval()
f_mean = np.load(f'{cache_dir}/f_mean.npy')
 for(test_df, sample_prediction_df)in tqdm(iter_test):
if test_df['weight'].item() > 0:
vals = torch.FloatTensor(
 fillna_npwhere(test_df[features].values, f_mean))
if type(models)== list:
preds = [torch.sigmoid(model.forward(vals.view(1, -1)) ).detach().numpy()
for model in models]
pred = np.mean(np.mean(preds, axis=1))
else:
pred = torch.sigmoid(models.forward(vals.view(1, -1)) ).item()
sample_prediction_df.action = np.where(
 pred > 0.502, 1, 0 ).astype(int ).item()
else:
sample_prediction_df.action = 0
env.predict(sample_prediction_df )
test = pd.read_csv('{}/test.csv'.format(DATA_PATH))
preds = [[0] * 6] * len(test)
 if os.path.exists('.. /input/prostate-cancer-grade-assessment/test_images'):
print('Inference on test images...')
 subm_datagen = DataGenPanda(
 imgs_path='{}/test_images'.format(DATA_PATH),
df=test,
batch_size=1,
mode='predict',
shuffle=False,
aug=None,
seq_len=SEQ_LEN,
img_size=IMG_SIZE,
n_classes=6 
)
 preds = model.predict_generator(subm_datagen)
 print('Total preds test:', len(preds))
else:
print('Inference on validation images...')
 preds = model.predict_generator(val_datagen.__getitem__(1)) [:len(preds)]
print('Total preds val:', len(preds))<save_to_csv>
def main(train=True):

 p = {'dim_1': 167,
'dim_2': 454,
'dim_3': 371,
'dim_4': 369,
'dim_5': 155,
'activation': nn.LeakyReLU,
'dropout': 0.21062362698532755,
'lr': 0.0022252024054478523,
'label_smoothing': 0.05564974140461841,
'weight_decay': 0.04106097088288333,
'amsgrad': True,
'batch_size': 10072}
if train:
models, features = final_train(p, load=False)
 else:
data_ = load_data(root_dir='./data/', mode='train',overide='/kaggle/input/jane-street-market-prediction/train.csv')
 data_, target_, features, date = preprocess_data(data_, nn=True)
 model_path = '/kaggle/input/model-states'
f_mean = calc_data_mean(data_, 'cache')
 models = load_model(model_path, data_.shape[-1], 5, p, False)
 test_model(models, features)
 return models