Eu tenho um arquivo CSV razoavelmente grande contendo dados de revisão da amazon que li em um quadro de dados do pandas. Quero dividir os dados 80-20 (teste de trem), mas, ao fazê-lo, quero garantir que os dados divididos estejam representando proporcionalmente os valores de uma coluna (Categorias), ou seja, todas as diferentes categorias de revisões estão presentes no trem e testar dados proporcionalmente.

Os dados são assim:

**ReviewerID**       **ReviewText**        **Categories**       **ProductId**

1212                   good product         Mobile               14444425
1233                   will buy again       drugs                324532
5432                   not recomended       dvd                  789654123 

Estou usando o seguinte código para fazer isso:

import pandas as pd
Meta = pd.read_csv('C:\\Users\\xyz\\Desktop\\WM Project\\Joined.csv')
import numpy as np
from sklearn.cross_validation import train_test_split

train, test = train_test_split(Meta.categories, test_size = 0.2, stratify=y)

dá o seguinte erro

NameError: name 'y' is not defined

Como sou relativamente novo em python, não consigo descobrir o que estou fazendo de errado ou se esse código estratificará com base nas categorias de colunas. Parece funcionar bem quando removo a opção estratificar, bem como a coluna de categorias da divisão de teste de trem.

Qualquer ajuda será apreciada.