How do I split a dataset into test/train/validation according to a particular group?

[bowlofred]

Looks like you're doing that shuffle/split in lines 19-30, most of which don't appear to be dataframe specific. The actual shuffle on line 19 could be replaced with something like random.shuffle() assuming you are starting with a list.