Add option to remove duplicates

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")