sklearn fix

model.py

@@ -14,7 +14,6 @@ from torch.utils.data import Dataset, DataLoader
 from transformers import BertTokenizer, BertModel, AdamW, get_linear_schedule_with_warmup
 import pytorch_lightning as pl
 from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
-from sklearn.model_selection import train_test_split
 
 RANDOM_SEED = 42
 np.random.seed(RANDOM_SEED)
@@ -93,7 +92,7 @@ df.drop(['emotion', 'message_id', 'response_id', 'article_id', 'empathy', 'distr
 
 print(df.head())
 
-train_df, val_df = train_test_split(df, test_size=0.05)
+train_df, val_df = sklearn.model_selection.train_test_split(df, test_size=0.05)
 train_df.shape, val_df.shape
 
 LABEL_COLUMNS = ['anger','joy','fear','surprise','sadness', 'neutral']

requirements.txt

@@ -2,4 +2,5 @@ transformers
 pytorch-lightning
 torch
 pandas
-numpy
+numpy
+sklearn==1.0.2