How to get stratified validation ids?

I have an unbalanced dataset and would like to use stratified validation ids instead of just random val ids. There is a function in sklearn but I was wondering if I can modify get_cv_idxs(n) to get stratified ids instead of just random ids.

Any other suggestions?