Combine and return multiple datasets using Tensorflow2

I'm using Tensorflow 2.1.0 & i'm trying to import base from tensorflow.contrib.learn.python.learn.datasets in order to return multiple datasets in a function called read_data_sets:

def read_data_sets():
    flow, labels = create_data_sets()
    validation_size = 8000
    train_flow, test_flow, train_labels, test_labels = train_test_split(flow, 
                                                                    labels, 
                                                                    test_size = 0.2,
                                                                    random_state=0)
    validation_flow = train_flow[:validation_size]
    validation_labels = train_labels[:validation_size]
    train_flow = train_flow[validation_size:]
    train_labels = train_labels[validation_size:]
    train = DataSet(train_flow, train_labels)
    test = DataSet(test_flow, test_labels)
    validation = DataSet(validation_flow, validation_labels)
    return base.Datasets(train=train, validation=validation, test=test)

I got this error: ModuleNotFoundError: No module named 'tensorflow.contrib'.

I'm searching for the equivalent of from tensorflow.contrib.learn.python.learn.datasets import base in Tensorflow2 knowing that the code works using Tensorflow1.

The full code:

import numpy as np
import data_preprocess
from tensorflow.contrib.learn.python.learn.datasets import base
import tensorflow as tf
from sklearn.model_selection import train_test_split

class DataSet(object):
def __init__(self, flow, labels):
    self._index_in_epoch = 0
    self._epochs_completed = 0
    self._flow = flow
    self._labels = labels
    self._num_examples = flow.shape[0]
    pass

@property
def flow(self):
    return self._flow

@property
def labels(self):
    return self._labels

@property
def num_examples(self):
    return self._num_examples

@property
def epochs_completed(self):
    return self._epochs_completed

def next_batch(self, batch_size, shuffle = True):
    """
    Return the next 'batch_size' examples from this dataset
    """
    start = self._index_in_epoch
    # Shuffle for the first epoch
    if start == 0 and self._epochs_completed == 0 and shuffle:
        idx = np.arange(self._num_examples)
        np.random.shuffle(idx)
        self._flow = self.flow[idx]
        self._labels = self.labels[idx]

    # Go to the next epoch
    if start + batch_size > self._num_examples:
        # Finished epoch
        self._epochs_completed += 1
        # Get the rest examples in this epoch
        rest_num_examples = self._num_examples - start
        flow_rest_part = self._flow[start:self._num_examples]
        labels_rest_part = self._labels[start:self._num_examples]

        # Shuffle the data
        if shuffle:
            idx0 = numpy.arrange(self._num_examples)
            numpy.random.shuffle(idx0)
            self._flow = self.flow[idx0]
            self._labels = self.labels[idx0]
        # Start next epoch
        start = 0
        self._index_in_epoch = batch_size - rest_num_examples
        end = self._index_in_epoch
        flow_new_part = self._flow[start:end]
        labels_new_part = self._labels[start:end]
        return numpy.concatenate((flow_rest_part, flow_new_part), axis=0), numpy.concatenate((labels_rest_part, labels_new_part), axis=0)

    else:
        self._index_in_epoch += batch_size
        end = self._index_in_epoch
        return self._flow[start:end], self._labels[start:end]

def create_data_sets():
    samples = data_preprocess.samples
    look_back = 8
    interval = 0
    flow, labels = [], []
    for i in range(len(samples)-look_back-interval):
        flow.append(samples[i:(i+look_back)])
        labels.append(samples[i+look_back+interval])
    return np.asarray(flow), np.asarray(labels)

def read_data_sets():
    flow, labels = create_data_sets()
    validation_size = 8000
    train_flow, test_flow, train_labels, test_labels = train_test_split(flow, 
                                                                    labels, 
                                                                    test_size = 0.2,
                                                                    random_state=0)
    validation_flow = train_flow[:validation_size]
    validation_labels = train_labels[:validation_size]
    train_flow = train_flow[validation_size:]
    train_labels = train_labels[validation_size:]
    train = DataSet(train_flow, train_labels)
    test = DataSet(test_flow, test_labels)
    validation = DataSet(validation_flow, validation_labels)
    return base.Datasets(train=train, validation=validation, test=test)

train, validation, test = read_data_sets()
print(train.flow.shape)

The data_preprocess file :

import pandas as pd
import numpy as np
import glob
from sklearn import preprocessing

FILE_NUMS = 33

def read_files():
    allframes = pd.DataFrame()
    list_ = []
    for i in range(FILE_NUMS):
        path = r'./PEMSs/station' + str(i+1)
        allfiles = glob.glob(path + "/*.txt")
        frame = pd.DataFrame()
        frame_ = []
        for file_ in allfiles:
            table = pd.read_table(file_, usecols=[0,1])
            frame_.append(table)
        frame = pd.concat(frame_)
        list_.append(frame)
    allframes = pd.concat(list_)
    return allframes

def group_by_time():
    frame = read_files()
    frame['5 Minutes'] = pd.to_datetime(frame['5 Minutes'], format='%m/%d/%Y %H:%M')
    values = frame.groupby('5 Minutes')['Flow (Veh/5 Minutes)'].apply(list)
    vehicles = []
    for i in range(len(values)):
       vehicles.append(values[i])
    return vehicles
vehicles = group_by_time()
scaler = preprocessing.MinMaxScaler()
samples = scaler.fit_transform(vehicles)
  
if __name__ == "__main__":
    save = group_by_time()
    print(save) 

You have to go to this github link in order to unzip the PEMSs-.zip file to the main directory. This is the data for the training.