diff --git a/split-cls.py b/split-cls.py index c89c428..b3b4415 100644 --- a/split-cls.py +++ b/split-cls.py @@ -3,14 +3,48 @@ import fire from pathlib import Path from sys import stderr from sklearn.model_selection import train_test_split +import re def to_csv(df, path): df.to_csv(path, header=None, index=None) -def split(data_dir): +def remove_rt(df): + return df.assign(text=df.text.str.replace('^RT @anonymized_account ','')) + +def remove_duplicates(df): + exact = df[~df.duplicated('text')] + prefixes = exact.text.map(lambda t: t.endswith('…') and exact.text.str.startswith(t[:-1]).sum()>1) + return exact[~prefixes] + +def cross_remove_duplicates(from_df, other_df): + exact = from_df[~from_df.text.isin(other_df.text)] + other_prefixes = other_df.text[other_df.text.str.endswith('…')].str[:-1] + if len(other_prefixes): + other_prefixes_re = re.compile('^'+'|'.join([f'({re.escape(t)})' for t in other_prefixes])) + else: + other_prefixes_re = re.compile('^$') + prefixes = exact.text.map(lambda t: + (t.endswith('…') and other_df.text.str.startswith(t[:-1]).any()) or + other_prefixes_re.match(t) is not None + ) + return exact[~prefixes] + +def split(data_dir, dedup=False): data_dir = Path(data_dir) - train = pd.read_csv(data_dir / "pl.unsup.csv", header=None) - trn, val = train_test_split(train, test_size=0.1, random_state=12345, stratify=train[0]) + train = pd.read_csv(data_dir / "pl.unsup.csv", header=None, names=["label", "text"]) + val_ratio = 0.1 + train = remove_rt(train) + trn, val = train_test_split(train, test_size=val_ratio, random_state=12345, stratify=train.label) + + if dedup: + trn = remove_duplicates(trn) + val = remove_duplicates(val) + val = cross_remove_duplicates(val, trn) + l1, l2, l3 = len(remove_duplicates(train)), len(trn), len(val) + if l1 != l2 + l3: + print("Warning: some condition believed by me to be invariant is not hold") + print(f"{l1} should be equal to {l2} + {l3} = {l2+l3}") + to_csv(trn, data_dir / "pl.train.csv") to_csv(val, data_dir / "pl.dev.csv")