58 lines
2.2 KiB
Python
58 lines
2.2 KiB
Python
import os
|
|
|
|
import pandas as pd
|
|
|
|
from misc import DATA_DIR
|
|
|
|
|
|
def clean(filepath):
|
|
encodings = ['utf-8', 'utf-16', 'latin1']
|
|
for enc in encodings:
|
|
try:
|
|
print(f">> Trying to read {filepath} with encoding: {enc}")
|
|
df = pd.read_csv(filepath, encoding=enc, on_bad_lines='skip')
|
|
|
|
print(">> Remove null bytes and non-breaking spaces from all string columns")
|
|
for col in df.select_dtypes(include=['object']).columns:
|
|
df[col] = df[col].astype(str).str.replace('\x00', ' ', regex=False)
|
|
df[col] = df[col].str.replace('\u00a0', ' ', regex=False)
|
|
df[col] = df[col].str.replace(' +', ' ', regex=True)
|
|
|
|
print(f">> Successfully read with encoding: {enc}")
|
|
df = df.dropna(subset=['name', 'sex', 'region'])
|
|
df.to_csv(filepath, index=False, encoding='utf-8')
|
|
return df
|
|
except Exception:
|
|
continue
|
|
raise UnicodeDecodeError(f"Unable to decode {filepath} with common encodings.")
|
|
|
|
|
|
def main():
|
|
df = clean(os.path.join(DATA_DIR, 'names.csv'))
|
|
|
|
df['name'] = df['name'].str.strip().str.lower()
|
|
df['words'] = df['name'].str.split().apply(len)
|
|
df['length'] = df['name'].str.replace(' ', '', regex=False).str.len()
|
|
df['probable_native'] = df['name'].str.split().apply(lambda x: ' '.join(x[:-1]) if len(x) > 1 else '')
|
|
df['probable_surname'] = df['name'].str.split().apply(lambda x: x[-1] if len(x) > 0 else '')
|
|
|
|
print(f">> Arranging columns")
|
|
cols = [c for c in df.columns if c != 'sex'] + ['sex']
|
|
df = df[cols]
|
|
|
|
print(f">> Saving evaluation dataset")
|
|
df_evaluation = df.sample(frac=0.2, random_state=42)
|
|
df_evaluation.to_csv(os.path.join(DATA_DIR, 'names_evaluation.csv'), index=False)
|
|
|
|
print(f">> Saving featured dataset")
|
|
df_featured = df.drop(df_evaluation.index)
|
|
df_featured.to_csv(os.path.join(DATA_DIR, 'names_featured.csv'), index=False)
|
|
|
|
print(f">> Splitting dataset by sex")
|
|
df[df['sex'].str.lower() == 'm'].to_csv(os.path.join(DATA_DIR, 'names_males.csv'), index=False)
|
|
df[df['sex'].str.lower() == 'f'].to_csv(os.path.join(DATA_DIR, 'names_females.csv'), index=False)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|