Function create_group_split

lib/utils/split.py:819–841

split within each stratification defined by the group column for training/validation split Args: train_val (pd.DataFrame): the train+valid dataframe to split on seed (int): the random seed holdout_frac (float): the fraction of validation group_column (str): the name of the gro

(train_val, seed, holdout_frac, group_column)

def create_group_split(train_val, seed, holdout_frac, group_column):
	"""split within each stratification defined by the group column for training/validation split
	
	Args:
	    train_val (pd.DataFrame): the train+valid dataframe to split on
	    seed (int): the random seed
	    holdout_frac (float): the fraction of validation
	    group_column (str): the name of the group column
	
	Returns:
	    dict: a dictionary of splitted dataframes, where keys are train/valid/test and values correspond to each dataframe
	"""
	train_df = pd.DataFrame()
	val_df = pd.DataFrame()

	for i in train_val[group_column].unique():
		train_val_temp = train_val[train_val[group_column] == i]
		np.random.seed(seed)
		msk = np.random.rand(len(train_val_temp)) < (1 - holdout_frac)
		train_df = train_df.append(train_val_temp[msk])
		val_df = val_df.append(train_val_temp[~msk])

	return {'train': train_df.reset_index(drop = True), 'valid': val_df.reset_index(drop = True)}