Commit 41506b91 authored by Benjamin Murauer's avatar Benjamin Murauer
Browse files

Added Train/Test wrapper

parent f4b1e746
......@@ -8,6 +8,7 @@ from sklearn.model_selection import ParameterGrid
from sklearn.preprocessing import LabelEncoder
from dbispipeline.base import Loader
from dbispipeline.base import TrainTestLoader
class RepeatingLoader(Loader):
......@@ -259,3 +260,30 @@ def _sample(values, strategy, sample_limit):
if strategy == 'random':
return random.sample(values, sample_limit)
class TrainTestWrapper(TrainTestLoader):
"""Wrap two separate loaders to use for train and test loaders."""
def __init__(self, train_loader, test_loader):
"""Initialize class."""
self.train_loader = train_loader
self.test_loader = test_loader
def load_train(self):
"""Load training data from first loader."""
return self.train_loader.load()
def load_test(self):
"""Load testing data from second loader."""
return self.test_loader.load()
def configuration(self):
"""Get db-suitable configuration of this loader."""
return {
'train_loader': self.train_loader.__class__.__name__,
'train_loader_configuration': self.train_loader.configuration,
'test_loader': self.test_loader.__class__.__name__,
'test_loader_configuration': self.test_loader.configuration,
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment