# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Header

# +
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import math as math
import janitor
from sklearn.linear_model import LinearRegression
import os

exec(open("../header.py").read())
# -

# # Import

raw_data = pd.read_csv(processed_root("01-poetryfoundation/poetry_foundation.csv"))

raw_data.head()

raw_data.columns

# # Clean

data_threshold_30 = raw_data.loc[lambda x:x.author_poem_count >= 30]
data_threshold_40 = raw_data.loc[lambda x:x.author_poem_count >= 40]
data_threshold_50 = raw_data.loc[lambda x:x.author_poem_count >= 50]


# ## Train-validation-test split

def train_val_test_split(data, splits = [0.7, 0.2, 0.1]):
    if not np.isclose(sum([0.7,0.2,0.1]), 1):
        raise RuntimeError("Splits must sum to 1.")
    train_pt = splits[0]
    val_pt = splits[0] + splits[1]
    
    train_data = data.loc[lambda x:x.author_poem_pct <= train_pt]
    val_data = data.loc[lambda x:
                                    (x.author_poem_pct > train_pt)&
                                    (x.author_poem_pct <= val_pt)]
    test_data = data.loc[lambda x:x.author_poem_pct > val_pt]
    
    train_data.name = 'train_data.csv'
    val_data.name = 'val_data.csv'
    test_data.name = 'test_data.csv'
    
    return train_data, val_data, test_data


train_30, val_30, test_30 = train_val_test_split(data_threshold_30)
train_40, val_40, test_40 = train_val_test_split(data_threshold_40)
train_50, val_50, test_50 = train_val_test_split(data_threshold_50)


# # Save datasets

def save_datasets(dfs, save_folder):
    for df in dfs:
        try:
            df.to_csv(save_folder + "/" + df.name, index = False)
        except FileNotFoundError:
            os.mkdir(save_folder)
            df.to_csv(save_folder + "/" + df.name, index = False)


save_datasets([train_30, val_30, test_30], 
             save_folder = processed_root("02-train-validation-test-split/threshold-30"))

save_datasets([train_40, val_40, test_40], 
             save_folder = processed_root("02-train-validation-test-split/threshold-40"))

save_datasets([train_50, val_50, test_50], 
             save_folder = processed_root("02-train-validation-test-split/threshold-50"))

StarcoderdataJupyterScriptsDedupFiltered