tweak references pre-processing pipeline

1 files changed, 83 insertions, 60 deletions

diff --git a/notebooks/cc/reference_annotations/data_preprocessing.ipynb b/notebooks/cc/reference_annotations/data_preprocessing.ipynb
index 5ee7dce5..0198fc6a 100644
--- a/notebooks/cc/reference_annotations/data_preprocessing.ipynb
+++ b/notebooks/cc/reference_annotations/data_preprocessing.ipynb
@@ -16,7 +16,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -85,7 +85,40 @@
     "\n",
     "def preprocess_segment(segment):\n",
     "    segment = segment.replace(\"\\n\", \" \")\n",
-    "    return segment"
+    "    return segment\n",
+    "\n",
+    "def get_split_dict(train_path: Path | None = None, valid_path: Path | None = None, test_path: Path | None = None) -> dict[str, str]:\n",
+    "    \"\"\"\n",
+    "    Returns dictionary that maps dgst: split, where split in `train`, `valid`, `test`. Expects path to list of dgsts for each split.\n",
+    "    \"\"\"\n",
+    "    def get_single_dct(pth: Path | None, split_name: str) -> dict[str, str]:\n",
+    "        if not pth:\n",
+    "            return dict()\n",
+    "        with pth.open(\"r\") as handle:\n",
+    "            return dict.fromkeys(json.load(handle), split_name)\n",
+    "\n",
+    "    return {**get_single_dct(train_path, \"train\"), **get_single_dct(valid_path, \"valid\"), **get_single_dct(test_path, \"test\")}\n",
+    "\n",
+    "def check_for_label_noise(df: pd.DataFrame) -> pd.DataFrame:\n",
+    "    \"\"\"\n",
+    "    Fills-in a dataframe with samples, such that duplicated labels (for label != None) appear for (dgst, cert_id) tuples.\n",
+    "    \"\"\"\n",
+    "    dgst_cert_id_tuples = (\n",
+    "        df.drop_duplicates(subset=[\"dgst\", \"cert_id\"])\n",
+    "        .loc[:, [\"dgst\", \"cert_id\"]]\n",
+    "        .set_index([\"dgst\", \"cert_id\"])\n",
+    "        .index.tolist()\n",
+    "    )\n",
+    "    duplicate_df = pd.DataFrame()\n",
+    "    for dgst, cert_id in tqdm(dgst_cert_id_tuples, desc=\"checking for label noise\"):\n",
+    "        possible_duplicates = df.loc[(df.dgst == dgst) & (df.cert_id == cert_id) & (df.label.notnull())]\n",
+    "        if (\n",
+    "            possible_duplicates.shape[0] > 1\n",
+    "            and not possible_duplicates.drop_duplicates(subset=[\"dgst\", \"cert_id\", \"label\"], keep=False).empty\n",
+    "        ):\n",
+    "            duplicate_df = pd.concat([duplicate_df, possible_duplicates])\n",
+    "\n",
+    "    return duplicate_df"
    ]
   },
   {
@@ -106,9 +139,8 @@
      "output_type": "stream",
      "text": [
       "100%|██████████| 58/58 [00:07<00:00,  8.27it/s]\n",
-      "100%|██████████| 944/944 [01:07<00:00, 14.02it/s]\n",
-      "100%|██████████| 2259/2259 [00:33<00:00, 68.05it/s]\n",
-      "100%|██████████| 2551/2551 [00:02<00:00, 1053.41it/s]\n"
+      "100%|██████████| 944/944 [01:06<00:00, 14.26it/s]\n",
+      "100%|██████████| 2259/2259 [00:33<00:00, 67.67it/s]\n"
      ]
     }
    ],
@@ -149,46 +181,7 @@
     "df_labeled = get_df_from_records(annotated_records)\n",
     "df_targets = get_df_from_records(target_records)\n",
     "df_reports = get_df_from_records(report_records)\n",
-    "\n",
-    "# Creates dictionary (dgst, cert_id): label to populate instances that have NULL label (on location==target) but were annotated in location==report and could adopt that label\n",
-    "# This helps to avoid duplicities and extends the number of annotated sentences.\n",
-    "dgst_cert_id_to_label_mapping = (\n",
-    "    df_labeled.loc[df_labeled.label.notnull(), [\"dgst\", \"cert_id\", \"label\"]]\n",
-    "    .drop_duplicates(subset=[\"dgst\", \"cert_id\"])\n",
-    "    .set_index([\"dgst\", \"cert_id\"])\n",
-    "    .label.to_dict()\n",
-    ")\n",
-    "\n",
-    "df = pd.concat([df_labeled, df_targets, df_reports])\n",
-    "\n",
-    "# Check for label noise\n",
-    "dgst_cert_id_tuples = (\n",
-    "    df.drop_duplicates(subset=[\"dgst\", \"cert_id\"])\n",
-    "    .loc[:, [\"dgst\", \"cert_id\"]]\n",
-    "    .set_index([\"dgst\", \"cert_id\"])\n",
-    "    .index.tolist()\n",
-    ")\n",
-    "duplicate_df = pd.DataFrame()\n",
-    "for dgst, cert_id in tqdm(dgst_cert_id_tuples):\n",
-    "    possible_duplicates = df.loc[(df.dgst == dgst) & (df.cert_id == cert_id) & (df.label.notnull())]\n",
-    "    if (\n",
-    "        possible_duplicates.shape[0] > 1\n",
-    "        and not possible_duplicates.drop_duplicates(subset=[\"dgst\", \"cert_id\", \"label\"], keep=False).empty\n",
-    "    ):\n",
-    "        duplicate_df = pd.concat([duplicate_df, possible_duplicates])\n",
-    "\n",
-    "if not duplicate_df.empty:\n",
-    "    print(\n",
-    "        \"Warning, label noise detected, see `duplicate_df` for instances that have inconsistent label for `(dgst, cert_id)` key.\"\n",
-    "    )\n",
-    "\n",
-    "# With no label noise, we should be safe to fill in labels for sentences found in targets such that the corresponding report was annotated\n",
-    "df.label = df.apply(\n",
-    "    lambda row: dgst_cert_id_to_label_mapping.get((row[\"dgst\"], row[\"cert_id\"]))\n",
-    "    if pd.isnull(row[\"label\"])\n",
-    "    else row[\"label\"],\n",
-    "    axis=1,\n",
-    ")\n"
+    "df = pd.concat([df_labeled, df_targets, df_reports])"
    ]
   },
   {
@@ -203,30 +196,60 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 13,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "checking for label noise: 100%|██████████| 1165/1165 [00:00<00:00, 1480.46it/s]\n"
+     ]
+    }
+   ],
    "source": [
+    "# Check for label noise\n",
+    "if not (duplicate_df := check_for_label_noise(df)).empty:\n",
+    "    print(\n",
+    "        \"Warning, label noise detected, see `duplicate_df` for instances that have inconsistent label for `(dgst, cert_id)` key.\"\n",
+    "    )\n",
+    "\n",
     "# Load split labels\n",
-    "with (REPO_ROOT / \"data/reference_annotations_split/train.json\").open(\"r\") as handle:\n",
-    "    train_digests = json.load(handle)\n",
+    "split_dct = get_split_dict(\n",
+    "    REPO_ROOT / \"data/reference_annotations_split/train.json\",\n",
+    "    REPO_ROOT / \"data/reference_annotations_split/valid.json\",\n",
+    "    REPO_ROOT / \"data/reference_annotations_split/test.json\",\n",
+    ")\n",
     "\n",
-    "with (REPO_ROOT / \"data/reference_annotations_split/valid.json\").open(\"r\") as handle:\n",
-    "    valid_digests = json.load(handle)\n",
+    "# Creates dictionary (dgst, cert_id): label to populate instances that have NULL label (on location==target) but were annotated in location==report and could adopt that label\n",
+    "# This helps to avoid duplicities and extends the number of annotated sentences.\n",
+    "dgst_cert_id_to_label_mapping = (\n",
+    "    df_labeled.loc[df_labeled.label.notnull(), [\"dgst\", \"cert_id\", \"label\"]]\n",
+    "    .drop_duplicates(subset=[\"dgst\", \"cert_id\"])\n",
+    "    .set_index([\"dgst\", \"cert_id\"])\n",
+    "    .label.to_dict()\n",
+    ")\n",
     "\n",
-    "split_dct = {**dict.fromkeys(train_digests, \"train\"), **dict.fromkeys(valid_digests, \"valid\")}\n",
+    "# With no label noise, we should be safe to fill in labels for sentences found in targets such that the corresponding report was annotated\n",
+    "df.label = df.apply(\n",
+    "    lambda row: dgst_cert_id_to_label_mapping.get((row[\"dgst\"], row[\"cert_id\"]))\n",
+    "    if pd.isnull(row[\"label\"])\n",
+    "    else row[\"label\"],\n",
+    "    axis=1,\n",
+    ")\n",
     "\n",
-    "# Apply filtering\n",
     "# TODO: We should investigate the cases when we match no sentence\n",
-    "df = df.loc[df.sentences.notnull()] \n",
-    "df[\"split\"] = df.dgst.map(split_dct)\n",
-    "df = df.loc[df.split.notnull()]  # Discard test samples\n",
-    "\n",
     "# TODO: Add language detection\n",
     "\n",
-    "# # Aggregate sentences from different sources (target, report) into one row\n",
-    "df = df.groupby([\"dgst\", \"cert_id\", \"label\", \"split\"], as_index=False, dropna=False)[\"sentences\"].agg({\"sentences\": lambda x: set.union(*x)})\n",
-    "df.to_csv(REPO_ROOT / \"datasets/reference_classification_dataset.csv\", sep=';', index=False)"
+    "# Process\n",
+    "df = (\n",
+    "    df.assign(split=df.dgst.map(split_dct))\n",
+    "    .loc[(df.sentences.notnull()) & (df.split != \"test\")]\n",
+    "    .groupby([\"dgst\", \"cert_id\", \"label\", \"split\"], as_index=False, dropna=False)[\"sentences\"]\n",
+    "    .agg({\"sentences\": lambda x: set.union(*x)})\n",
+    ")\n",
+    "\n",
+    "df.to_csv(REPO_ROOT / \"datasets/reference_classification_dataset.csv\", sep=\";\", index=False)\n"
    ]
   }
  ],