tweet-sentiment-analysis-data/split-train-test.py

import pandas as pd
from sklearn.model_selection import train_test_split
import os

# Const
DATASET_COLUMNS = ["sentiment", "ids", "date", "flag", "user", "text"]
DATASET_ENCODING = "ISO-8859-1"
TARGET_COL = 'sentiment'
CSV_PATH = 'raw-data/twitter-sentiment-analysis-sentiment140dataset.csv'
NEW_DIR = 'split-data'
X_TRAIN_PATH = 'split-data/X_train.csv'
X_TEST_PATH = 'split-data/y_train.csv'
Y_TRAIN_PATH = 'split-data/X_test.csv'
Y_TEST_PATH = '../y_test/y_test.csv'
TEST_SIZE = 0.03

print("Read raw data")
df = pd.read_csv(CSV_PATH, encoding=DATASET_ENCODING, names=DATASET_COLUMNS)
print(f'data set shape {df.shape}')

print("Replace target col values")
df[TARGET_COL] = df[TARGET_COL].replace(0, 1)  # Negative
df[TARGET_COL] = df[TARGET_COL].replace(4, 0)  # Positive

os.makedirs(NEW_DIR, exist_ok=True)

print("Split dataset to train and test")
X_train, X_test, y_train, y_test = train_test_split(df.drop(TARGET_COL, axis=1), df[TARGET_COL],
                                                    test_size=TEST_SIZE, random_state=42,
                                                    stratify=df[TARGET_COL])

_, X_train, _, y_train = train_test_split(X_train, y_train,
                                                    test_size=0.1, random_state=42,
                                                    stratify=y_train)

print("Save data sets to csv")
X_train.to_csv(X_TRAIN_PATH, index=False)
y_train.to_csv(X_TEST_PATH, index=False)
X_test.to_csv(Y_TRAIN_PATH, index=False)

# y_test will be saved outside of the repo - to prevent cheating.
y_test.to_csv(Y_TEST_PATH, index=False)