adjust pipeline to new annotation format

1 files changed, 65 insertions, 97 deletions

diff --git a/notebooks/cc/reference_annotations/data_preprocessing.ipynb b/notebooks/cc/reference_annotations/data_preprocessing.ipynb
index 7a88bb05..76d16326 100644
--- a/notebooks/cc/reference_annotations/data_preprocessing.ipynb
+++ b/notebooks/cc/reference_annotations/data_preprocessing.ipynb
@@ -28,7 +28,6 @@
     "import spacy\n",
     "from sec_certs.utils.parallel_processing import process_parallel\n",
     "import pandas as pd\n",
-    "from tqdm import tqdm\n",
     "import json\n",
     "\n",
     "nlp = spacy.load(\"en_core_web_sm\")\n",
@@ -36,89 +35,100 @@
     "\n",
     "REPO_ROOT = Path(\"../../../\").resolve()\n",
     "\n",
+    "\n",
     "@dataclass\n",
     "class ReferenceRecord:\n",
     "    \"\"\"\n",
     "    Intermediate object to hold references for a given certificate together with sensible attributes to be extracted\n",
     "    for labeling.\n",
     "    \"\"\"\n",
+    "\n",
     "    certificate: CCCertificate | None\n",
-    "    dgst: str\n",
-    "    cert_id: str\n",
-    "    location: str\n",
+    "    referenced_cert_id: str\n",
+    "    source: str\n",
     "    label: str | None = None\n",
     "    sentences: set[str] | None = None\n",
     "\n",
     "    @staticmethod\n",
-    "    def get_reference_sentences(doc, cert_id: str) -> set[str]:\n",
+    "    def get_reference_sentences(doc, referenced_cert_id: str) -> set[str]:\n",
     "        \"\"\"\n",
     "        Return a set of sentences corresponding to the given cert_id for the record\n",
     "        \"\"\"\n",
-    "        return {sent.text for sent in doc.sents if cert_id in sent.text}\n",
+    "        return {sent.text for sent in doc.sents if referenced_cert_id in sent.text}\n",
     "\n",
     "    @staticmethod\n",
-    "    def get_cert_references_with_sentences(record: ReferenceRecord) -> set[tuple[str, str, str]]:\n",
+    "    def get_cert_references_with_sentences(record: ReferenceRecord) -> ReferenceRecord:\n",
     "        pth_to_read = (\n",
     "            record.certificate.state.st_txt_path\n",
-    "            if record.location == \"target\"\n",
+    "            if record.source == \"target\"\n",
     "            else record.certificate.state.report_txt_path\n",
     "        )\n",
     "\n",
     "        with pth_to_read.open(\"r\") as handle:\n",
     "            data = handle.read()\n",
     "\n",
-    "        result = ReferenceRecord.get_reference_sentences(nlp(data), record.cert_id)\n",
+    "        result = ReferenceRecord.get_reference_sentences(nlp(data), record.referenced_cert_id)\n",
     "        record.sentences = result if result else None\n",
     "\n",
     "        return record\n",
     "\n",
     "    def to_pandas_tuple(self) -> tuple[str, str, str, str, set[str] | None]:\n",
-    "        return self.dgst, self.cert_id, self.location, self.label, self.sentences\n",
+    "        return self.certificate.dgst, self.referenced_cert_id, self.source, self.label, self.sentences\n",
+    "\n",
     "\n",
     "def get_df_from_records(records: list[ReferenceRecord]):\n",
     "    \"\"\"\n",
     "    Builds dataframe with [dgst,cert_id,location,reason,sentences] with references from list of ReferenceRecords.\n",
-    "    Reason set to None if not defined. \n",
+    "    Reason set to None if not defined.\n",
     "    \"\"\"\n",
-    "    results =  process_parallel(ReferenceRecord.get_cert_references_with_sentences, records, use_threading=False, progress_bar=True)\n",
-    "    return pd.DataFrame.from_records([x.to_pandas_tuple() for x in results], columns=[\"dgst\", \"cert_id\", \"location\", \"label\", \"sentences\"])\n",
+    "    results = process_parallel(\n",
+    "        ReferenceRecord.get_cert_references_with_sentences, records, use_threading=False, progress_bar=True\n",
+    "    )\n",
+    "    return pd.DataFrame.from_records(\n",
+    "        [x.to_pandas_tuple() for x in results], columns=[\"dgst\", \"referenced_cert_id\", \"source\", \"label\", \"sentences\"]\n",
+    "    )\n",
+    "\n",
     "\n",
     "def preprocess_segment(segment):\n",
     "    segment = segment.replace(\"\\n\", \" \")\n",
     "    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",
+    "def get_split_dict(\n",
+    "    train_path: Path | None = None, valid_path: Path | None = None, test_path: Path | None = None\n",
+    ") -> 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",
+    "\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",
+    "    return {\n",
+    "        **get_single_dct(train_path, \"train\"),\n",
+    "        **get_single_dct(valid_path, \"valid\"),\n",
+    "        **get_single_dct(test_path, \"test\"),\n",
+    "    }\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",
+    "def load_annotated_samples(\n",
+    "    train_path: Path | None = None, valid_path: Path | None = None, test_path: Path | None = None\n",
+    "):\n",
+    "    def load_single_df(pth: Path | None, split_name: str) -> pd.DataFrame:\n",
+    "        if not pth:\n",
+    "            return pd.DataFrame()\n",
+    "        return (\n",
+    "            pd.read_csv(pth)\n",
+    "            .assign(label=lambda df_: df_.label.str.replace(\" \", \"_\").str.upper(), split=split_name)\n",
+    "            .replace(\"NONE\", None)\n",
+    "            .dropna(subset=\"label\")\n",
+    "        )\n",
     "\n",
-    "    return duplicate_df"
+    "    return pd.concat(\n",
+    "        [load_single_df(train_path, \"train\"), load_single_df(valid_path, \"valid\"), load_single_df(test_path, \"test\")]\n",
+    "    )[[\"dgst\", \"referenced_cert_id\", \"source\", \"label\", \"comment\"]]"
    ]
   },
   {
@@ -131,60 +141,43 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "100%|█████████▉| 943/944 [01:17<00:00, 12.09it/s]\n",
-      "100%|██████████| 944/944 [01:06<00:00, 14.24it/s]\n",
-      "100%|██████████| 2288/2288 [00:33<00:00, 69.10it/s]\n"
+      "100%|██████████| 944/944 [01:05<00:00, 14.48it/s]\n",
+      "100%|██████████| 2288/2288 [00:32<00:00, 69.50it/s]\n"
      ]
     }
    ],
    "source": [
     "# Load annotated references from CSV\n",
-    "annotations_df = (\n",
-    "    pd.read_csv(REPO_ROOT / \"data/cert_id_eval/random_references.csv\")\n",
-    "    .rename(columns={\"id\": \"dgst\", \"reason\": \"label\"})\n",
-    "    .dropna(subset=\"label\")\n",
-    "    .query(\"label != 'self'\")\n",
-    "    .assign(label=lambda row: row.label.str.replace(\" \", \"_\").str.upper())\n",
-    ")\n",
+    "annotations_df = load_annotated_samples(REPO_ROOT / \"data/reference_annotations/manual_annotations/train.csv\", REPO_ROOT / \"data/reference_annotations/manual_annotations/valid.csv\")\n",
     "\n",
     "# Load dataset\n",
     "# dset = CCDataset.from_web_latest()\n",
     "dset = CCDataset.from_json(REPO_ROOT / \"datasets/cc/cc_dataset.json\")\n",
     "\n",
-    "annotated_records = [\n",
-    "    ReferenceRecord(dset[x.dgst], x.dgst, x.cert_id, x.location, x.label)\n",
-    "    for x in annotations_df.itertuples(index=False)\n",
-    "]\n",
-    "\n",
-    "# Reference records without annotations\n",
     "target_certs = [x for x in dset if x.heuristics.st_references.directly_referencing and x.state.st_txt_path]\n",
     "report_certs = [x for x in dset if x.heuristics.report_references.directly_referencing and x.state.report_txt_path]\n",
     "target_records = [\n",
-    "    ReferenceRecord(x, x.dgst, y, \"target\", None, None)\n",
+    "    ReferenceRecord(x, y, \"target\", None, None)\n",
     "    for x in target_certs\n",
     "    for y in x.heuristics.st_references.directly_referencing\n",
     "]\n",
     "report_records = [\n",
-    "    ReferenceRecord(x, x.dgst, y, \"report\", None, None)\n",
+    "    ReferenceRecord(x, y, \"report\", None, None)\n",
     "    for x in report_certs\n",
     "    for y in x.heuristics.report_references.directly_referencing\n",
     "]\n",
     "\n",
-    "# Filter annotated_records from report_records to avoid duplicities\n",
-    "annotated_keys = {(x.dgst, x.cert_id) for x in annotated_records}\n",
-    "report_records = [x for x in report_records if (x.dgst, x.cert_id) not in annotated_keys]\n",
-    "\n",
-    "df_labeled = get_df_from_records(annotated_records)\n",
+    "# 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",
-    "df = pd.concat([df_labeled, df_targets, df_reports])"
+    "df = pd.concat([df_targets, df_reports])"
    ]
   },
   {
@@ -199,24 +192,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 5,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "checking for label noise: 100%|██████████| 2541/2541 [00:02<00:00, 1154.36it/s]\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "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",
     "split_dct = get_split_dict(\n",
     "    REPO_ROOT / \"data/reference_annotations/split/train.json\",\n",
@@ -224,35 +203,24 @@
     "    REPO_ROOT / \"data/reference_annotations/split/test.json\",\n",
     ")\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",
-    "# 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",
+    "# Creates dictionary `(dgst, cert_id): label`` to populate instances with manually assigned annotations.\n",
+    "annotations_dict = (\n",
+    "    annotations_df[[\"dgst\", \"referenced_cert_id\", \"label\"]].set_index([\"dgst\", \"referenced_cert_id\"]).label.to_dict()\n",
     ")\n",
     "\n",
-    "# TODO: We should investigate the cases when we match no sentence\n",
+    "# TODO: We should investigate the cases when we match no sentence, they may be new-lines and stuff\n",
     "# TODO: Add language detection\n",
-    "\n",
     "# 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",
+    "    df.assign(\n",
+    "        split=df.dgst.map(split_dct),\n",
+    "        label=lambda df_: [annotations_dict.get(x) for x in zip(df_[\"dgst\"], df_[\"referenced_cert_id\"])],\n",
+    "    )\n",
+    "    .loc[lambda df_: (df_[\"sentences\"].notnull()) & (df_[\"split\"] != \"test\")]\n",
+    "    .groupby([\"dgst\", \"referenced_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"
+    "df.to_csv(REPO_ROOT / \"datasets/reference_classification_dataset.csv\", index=False)"
    ]
   }
  ],