mlops_dvc/src/split_data.py

# raw data split
# save it in data/processed folder

import os
import argparse
import pandas as pd
from sklearn.model_selection import train_test_split
from get_data import read_params

def split_and_saved_data(config_path):
    # split and storing the file in respective path
    config = read_params(config_path)
    test_data_path = config["split_data"]["test_path"]
    train_data_path = config["split_data"]["train_path"]
    raw_data_path = config["load_data"]["raw_dataset_csv"]
    split_ratio = config["split_data"]["test_size"]
    random_state = config["base"]["random_state"]
    df = pd.read_csv(raw_data_path,sep=',')
    train,test = train_test_split(df,test_size=split_ratio,random_state=random_state)

    test.to_csv(test_data_path,sep=",",index=False,encoding = "utf-8")
    train.to_csv(train_data_path,sep=",",index=False,encoding = "utf-8")


# we are using main in every py file just to check that py file execution
if __name__ == "__main__":
    args = argparse.ArgumentParser()
    # we just reading the params.yaml
    args.add_argument("--config",default = "params.yaml")
    parsed_args = args.parse_args()
    split_and_saved_data(config_path = parsed_args.config)