Estava dando uma estudada no sklearn e vi que ele tem um método que realiza a partição dos dados em treino e teste randomicamente, acho que seria mais fácil do que fazer toda essa parte que fizemos na unha.
from sklearn.model_selection import train_test_splitdependendo dos parâmetros o mesmo já retorna o X_test, X_train, Y_text e Y_train, acho que seria uma boa abordagem para treino-teste, seus parâmetros são basicamente os arrays de treino e teste e a porcentágem (de 0 à 1) que deseja particionar os dados
