adjust pipeline to new annotation format

3 files changed, 169 insertions, 297 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)"
    ]
   }
  ],
diff --git a/notebooks/cc/reference_annotations/prediction.ipynb b/notebooks/cc/reference_annotations/prediction.ipynb
index 29b63523..b88c4a60 100644
--- a/notebooks/cc/reference_annotations/prediction.ipynb
+++ b/notebooks/cc/reference_annotations/prediction.ipynb
@@ -14,7 +14,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [
     {
@@ -53,130 +53,36 @@
     "    df_new[\"y_proba\"] = y_proba\n",
     "    df_new[\"y_pred\"] = df_new.y_proba.map(lambda x: label_mapping[np.argmax(x)])\n",
     "    df_new[\"correct\"] = df_new.label == df_new.y_pred\n",
-    "    return df_new"
+    "    return df_new\n",
+    "\n",
+    "def eval_strings(series):\n",
+    "    return [list(literal_eval(x)) for x in series]"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 14,
    "metadata": {},
    "outputs": [],
    "source": [
     "# Prepare dataset\n",
     "\n",
-    "df = pd.read_csv(REPO_ROOT / \"datasets/reference_classification_dataset_merged.csv\", sep=\";\")\n",
-    "df = df.loc[(df.label.notnull())]\n",
-    "df = df.loc[df.label.isin({\"COMPONENT_USED\", \"BASIS_OF_RECERTIFICATION\", \"BASIS_FOR\"})]  # only the most popular labels\n",
-    "df.sentences = df.sentences.map(lambda x: list(literal_eval(x)))\n",
+    "df = pd.read_csv(REPO_ROOT / \"datasets/reference_classification_dataset.csv\").loc[\n",
+    "    lambda df_: (df_.label.notnull()) & (df_.label.isin({\"COMPONENT_USED\", \"RECERTIFICATION\", \"ON_PLATFORM\"}))\n",
+    "].assign(sentences=lambda df_: eval_strings(df_.sentences))\n",
+    "df.label = df.label.map(lambda x: x if x != \"ON_PLATFORM\" else \"COMPONENT_USED\")\n",
     "\n",
-    "# # Split into train/valid\n",
+    "# Split into train/valid\n",
     "df_train = df.loc[df.split == \"train\"].drop(columns=\"split\")\n",
     "df_valid = df.loc[df.split == \"valid\"].drop(columns=\"split\")\n",
     "\n",
     "# Use just few examples for learning\n",
-    "df_train = df_train.sample(n=10)"
+    "df_train = df_train.sample(n=30)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/html": [
-       "<div>\n",
-       "<style scoped>\n",
-       "    .dataframe tbody tr th:only-of-type {\n",
-       "        vertical-align: middle;\n",
-       "    }\n",
-       "\n",
-       "    .dataframe tbody tr th {\n",
-       "        vertical-align: top;\n",
-       "    }\n",
-       "\n",
-       "    .dataframe thead th {\n",
-       "        text-align: right;\n",
-       "    }\n",
-       "</style>\n",
-       "<table border=\"1\" class=\"dataframe\">\n",
-       "  <thead>\n",
-       "    <tr style=\"text-align: right;\">\n",
-       "      <th></th>\n",
-       "      <th>dgst</th>\n",
-       "      <th>cert_id</th>\n",
-       "      <th>label</th>\n",
-       "      <th>sentences</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>12</th>\n",
-       "      <td>99223aca5d9eb3b3</td>\n",
-       "      <td>DCSSI-2009/11</td>\n",
-       "      <td>COMPONENT_USED</td>\n",
-       "      <td>[Toolbox Certificate DCSSI-2009/11\\nTable 1:]</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>4</th>\n",
-       "      <td>0f3900cdcd0c7f3e</td>\n",
-       "      <td>BSI-DSZ-CC-1072-V4-2021-MA-01</td>\n",
-       "      <td>COMPONENT_USED</td>\n",
-       "      <td>[Certification Report NXP Secure Smart Card Co...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>9</th>\n",
-       "      <td>6d6ade44dcc497dd</td>\n",
-       "      <td>BSI-DSZ-CC-0227-2004</td>\n",
-       "      <td>BASIS_OF_RECERTIFICATION</td>\n",
-       "      <td>[This is a\\nre-certification based on BSI-DSZ-...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>5</th>\n",
-       "      <td>0f3900cdcd0c7f3e</td>\n",
-       "      <td>NSCIB-CC-66030-CR5</td>\n",
-       "      <td>COMPONENT_USED</td>\n",
-       "      <td>[certificate identification NSCIB-CC-66030-CR5...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6</th>\n",
-       "      <td>1fb1564dfb0f0b04</td>\n",
-       "      <td>ANSSI-CC-2020/34</td>\n",
-       "      <td>COMPONENT_USED</td>\n",
-       "      <td>[[CER_IC] Rapport de certification ANSSI-CC-20...</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "                dgst                        cert_id                     label  \\\n",
-       "12  99223aca5d9eb3b3                  DCSSI-2009/11            COMPONENT_USED   \n",
-       "4   0f3900cdcd0c7f3e  BSI-DSZ-CC-1072-V4-2021-MA-01            COMPONENT_USED   \n",
-       "9   6d6ade44dcc497dd           BSI-DSZ-CC-0227-2004  BASIS_OF_RECERTIFICATION   \n",
-       "5   0f3900cdcd0c7f3e             NSCIB-CC-66030-CR5            COMPONENT_USED   \n",
-       "6   1fb1564dfb0f0b04               ANSSI-CC-2020/34            COMPONENT_USED   \n",
-       "\n",
-       "                                            sentences  \n",
-       "12      [Toolbox Certificate DCSSI-2009/11\\nTable 1:]  \n",
-       "4   [Certification Report NXP Secure Smart Card Co...  \n",
-       "9   [This is a\\nre-certification based on BSI-DSZ-...  \n",
-       "5   [certificate identification NSCIB-CC-66030-CR5...  \n",
-       "6   [[CER_IC] Rapport de certification ANSSI-CC-20...  "
-      ]
-     },
-     "execution_count": 10,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "df_train.head()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 15,
    "metadata": {},
    "outputs": [
     {
@@ -187,16 +93,16 @@
       "model_head.pkl not found on HuggingFace Hub, initialising classification head with random weights. You should TRAIN this model on a downstream task to use it for predictions and inference.\n",
       "Applying column mapping to training dataset\n",
       "***** Running training *****\n",
-      "  Num examples = 1760\n",
+      "  Num examples = 7200\n",
       "  Num epochs = 1\n",
-      "  Total optimization steps = 110\n",
+      "  Total optimization steps = 450\n",
       "  Total train batch size = 16\n"
      ]
     },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "96e469ad1f984bf6ba2c819884a1c231",
+       "model_id": "26e2525767364e3eba9c9881cd1982be",
        "version_major": 2,
        "version_minor": 0
       },
@@ -210,12 +116,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a9ef28c8c0314e7f831e6e35c2af75db",
+       "model_id": "d62be3272ebd4654b2adcfedee9b9b6a",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "Iteration:   0%|          | 0/110 [00:00<?, ?it/s]"
+       "Iteration:   0%|          | 0/450 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -234,9 +140,9 @@
      "output_type": "stream",
      "text": [
       "Internal evaluation (of model working on individual sentences)\n",
-      "{'precision': 0.45454545454545453, 'recall': 0.45454545454545453}\n",
+      "{'precision': 0.08035714285714286, 'recall': 0.08035714285714286}\n",
       "Actual evaluation after ensemble soft voting\n",
-      "{'precision': 0.2857142857142857, 'recall': 0.2857142857142857}\n"
+      "{'precision': 0.9014084507042254, 'recall': 0.9014084507042254}\n"
      ]
     }
    ],
@@ -248,7 +154,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 16,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -259,7 +165,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 17,
    "metadata": {},
    "outputs": [
     {
@@ -284,7 +190,7 @@
        "    <tr style=\"text-align: right;\">\n",
        "      <th></th>\n",
        "      <th>dgst</th>\n",
-       "      <th>cert_id</th>\n",
+       "      <th>referenced_cert_id</th>\n",
        "      <th>label</th>\n",
        "      <th>sentences</th>\n",
        "      <th>y_proba</th>\n",
@@ -293,45 +199,17 @@
        "    </tr>\n",
        "  </thead>\n",
        "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>9</th>\n",
-       "      <td>6d6ade44dcc497dd</td>\n",
-       "      <td>BSI-DSZ-CC-0227-2004</td>\n",
-       "      <td>BASIS_OF_RECERTIFICATION</td>\n",
-       "      <td>[This is a\\nre-certification based on BSI-DSZ-...</td>\n",
-       "      <td>[0.5461188093773812, 0.45388119062261884]</td>\n",
-       "      <td>COMPONENT_USED</td>\n",
-       "      <td>False</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>19</th>\n",
-       "      <td>ca5da2fe138af656</td>\n",
-       "      <td>BSI-DSZ-CC-0413-2007</td>\n",
-       "      <td>BASIS_OF_RECERTIFICATION</td>\n",
-       "      <td>[This is a re-certification based on\\nBSI-DSZ-...</td>\n",
-       "      <td>[0.5465589745598575, 0.4534410254401425]</td>\n",
-       "      <td>COMPONENT_USED</td>\n",
-       "      <td>False</td>\n",
-       "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
        "</div>"
       ],
       "text/plain": [
-       "                dgst               cert_id                     label  \\\n",
-       "9   6d6ade44dcc497dd  BSI-DSZ-CC-0227-2004  BASIS_OF_RECERTIFICATION   \n",
-       "19  ca5da2fe138af656  BSI-DSZ-CC-0413-2007  BASIS_OF_RECERTIFICATION   \n",
-       "\n",
-       "                                            sentences  \\\n",
-       "9   [This is a\\nre-certification based on BSI-DSZ-...   \n",
-       "19  [This is a re-certification based on\\nBSI-DSZ-...   \n",
-       "\n",
-       "                                      y_proba          y_pred  correct  \n",
-       "9   [0.5461188093773812, 0.45388119062261884]  COMPONENT_USED    False  \n",
-       "19   [0.5465589745598575, 0.4534410254401425]  COMPONENT_USED    False  "
+       "Empty DataFrame\n",
+       "Columns: [dgst, referenced_cert_id, label, sentences, y_proba, y_pred, correct]\n",
+       "Index: []"
       ]
      },
-     "execution_count": 13,
+     "execution_count": 17,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -342,7 +220,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 20,
    "metadata": {},
    "outputs": [
     {
@@ -367,7 +245,7 @@
        "    <tr style=\"text-align: right;\">\n",
        "      <th></th>\n",
        "      <th>dgst</th>\n",
-       "      <th>cert_id</th>\n",
+       "      <th>referenced_cert_id</th>\n",
        "      <th>label</th>\n",
        "      <th>sentences</th>\n",
        "      <th>y_proba</th>\n",
@@ -377,52 +255,72 @@
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
-       "      <th>1</th>\n",
-       "      <td>0c7ef6c32cbdee47</td>\n",
-       "      <td>BSI-DSZ-CC-1074-2019</td>\n",
-       "      <td>BASIS_FOR</td>\n",
-       "      <td>[The BAC+PACE configuration is subject of the ...</td>\n",
-       "      <td>[0.9330686268852108, 0.06693137311478929]</td>\n",
+       "      <th>251</th>\n",
+       "      <td>2c2244c35d126bfb</td>\n",
+       "      <td>BSI-DSZ-CC-0794-2011</td>\n",
+       "      <td>RECERTIFICATION</td>\n",
+       "      <td>[This is a re-certification based on BSI-DSZ-C...</td>\n",
+       "      <td>[0.47754689036832726, 0.5224531096316728]</td>\n",
+       "      <td>COMPONENT_USED</td>\n",
+       "      <td>False</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>364</th>\n",
+       "      <td>4489bfc781a82281</td>\n",
+       "      <td>BSI-DSZ-CC-0817-2013</td>\n",
+       "      <td>RECERTIFICATION</td>\n",
+       "      <td>[This is a re-certification based on BSI-DSZ-C...</td>\n",
+       "      <td>[0.11195564561722356, 0.8880443543827765]</td>\n",
+       "      <td>COMPONENT_USED</td>\n",
+       "      <td>False</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>505</th>\n",
+       "      <td>647d17d44745a532</td>\n",
+       "      <td>BSI-DSZ-CC-0904-2015</td>\n",
+       "      <td>RECERTIFICATION</td>\n",
+       "      <td>[and BSI-DSZ-CC-0904-2015-\\n, This is a re-cer...</td>\n",
+       "      <td>[0.16312949852765818, 0.8368705014723418]</td>\n",
        "      <td>COMPONENT_USED</td>\n",
        "      <td>False</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>2</th>\n",
-       "      <td>0e22fe4e4e58faf4</td>\n",
-       "      <td>BSI-DSZ-CC-1052-V4-2021</td>\n",
-       "      <td>BASIS_OF_RECERTIFICATION</td>\n",
-       "      <td>[basierend auf BSI-DSZ-CC-1052-V4-2021.]</td>\n",
-       "      <td>[0.7070543956916182, 0.2929456043083818]</td>\n",
+       "      <th>616</th>\n",
+       "      <td>7e58bfc14edf68e4</td>\n",
+       "      <td>OCSI/CERT/TEC/01/2013/RC</td>\n",
+       "      <td>RECERTIFICATION</td>\n",
+       "      <td>[OCSI/CERT/TEC/01/2013/RC, versione 1.0,]</td>\n",
+       "      <td>[0.2810731555375018, 0.7189268444624982]</td>\n",
        "      <td>COMPONENT_USED</td>\n",
        "      <td>False</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>7</th>\n",
-       "      <td>238f8edc5eda1358</td>\n",
-       "      <td>BSI-DSZ-CC-0222-2003</td>\n",
-       "      <td>BASIS_OF_RECERTIFICATION</td>\n",
-       "      <td>[This certification is a re-certification of B...</td>\n",
-       "      <td>[0.5998535578550636, 0.4001464421449364]</td>\n",
+       "      <th>628</th>\n",
+       "      <td>81273108dd167b98</td>\n",
+       "      <td>BSI-DSZ-CC-0523-2008</td>\n",
+       "      <td>RECERTIFICATION</td>\n",
+       "      <td>[Specific results from the\\nevaluation process...</td>\n",
+       "      <td>[0.4308279930268643, 0.5691720069731356]</td>\n",
        "      <td>COMPONENT_USED</td>\n",
        "      <td>False</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>8</th>\n",
-       "      <td>29964f32c68b0ce8</td>\n",
-       "      <td>BSI-DSZ-CC-0519-V3-2021</td>\n",
-       "      <td>BASIS_OF_RECERTIFICATION</td>\n",
+       "      <th>709</th>\n",
+       "      <td>9664c0f0ec6401b9</td>\n",
+       "      <td>BSI-DSZ-CC-0891-V2-2016</td>\n",
+       "      <td>RECERTIFICATION</td>\n",
        "      <td>[This is a re-certification based on BSI-DSZ-C...</td>\n",
-       "      <td>[0.8727371952470533, 0.12726280475294682]</td>\n",
+       "      <td>[0.4303467603163502, 0.5696532396836499]</td>\n",
        "      <td>COMPONENT_USED</td>\n",
        "      <td>False</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>13</th>\n",
-       "      <td>a6fac58198296194</td>\n",
-       "      <td>BSI-DSZ-CC-0555-2009</td>\n",
-       "      <td>BASIS_OF_RECERTIFICATION</td>\n",
-       "      <td>[Specific results from the evaluation process\\...</td>\n",
-       "      <td>[0.8670438280210987, 0.13295617197890133]</td>\n",
+       "      <th>719</th>\n",
+       "      <td>983d16512ae92d46</td>\n",
+       "      <td>BSI-DSZ-CC-0891-V3-2018</td>\n",
+       "      <td>RECERTIFICATION</td>\n",
+       "      <td>[The updated documents in compare to the forer...</td>\n",
+       "      <td>[0.48619060832541505, 0.513809391674585]</td>\n",
        "      <td>COMPONENT_USED</td>\n",
        "      <td>False</td>\n",
        "    </tr>\n",
@@ -431,29 +329,35 @@
        "</div>"
       ],
       "text/plain": [
-       "                dgst                  cert_id                     label  \\\n",
-       "1   0c7ef6c32cbdee47     BSI-DSZ-CC-1074-2019                 BASIS_FOR   \n",
-       "2   0e22fe4e4e58faf4  BSI-DSZ-CC-1052-V4-2021  BASIS_OF_RECERTIFICATION   \n",
-       "7   238f8edc5eda1358     BSI-DSZ-CC-0222-2003  BASIS_OF_RECERTIFICATION   \n",
-       "8   29964f32c68b0ce8  BSI-DSZ-CC-0519-V3-2021  BASIS_OF_RECERTIFICATION   \n",
-       "13  a6fac58198296194     BSI-DSZ-CC-0555-2009  BASIS_OF_RECERTIFICATION   \n",
+       "                 dgst        referenced_cert_id            label  \\\n",
+       "251  2c2244c35d126bfb      BSI-DSZ-CC-0794-2011  RECERTIFICATION   \n",
+       "364  4489bfc781a82281      BSI-DSZ-CC-0817-2013  RECERTIFICATION   \n",
+       "505  647d17d44745a532      BSI-DSZ-CC-0904-2015  RECERTIFICATION   \n",
+       "616  7e58bfc14edf68e4  OCSI/CERT/TEC/01/2013/RC  RECERTIFICATION   \n",
+       "628  81273108dd167b98      BSI-DSZ-CC-0523-2008  RECERTIFICATION   \n",
+       "709  9664c0f0ec6401b9   BSI-DSZ-CC-0891-V2-2016  RECERTIFICATION   \n",
+       "719  983d16512ae92d46   BSI-DSZ-CC-0891-V3-2018  RECERTIFICATION   \n",
        "\n",
-       "                                            sentences  \\\n",
-       "1   [The BAC+PACE configuration is subject of the ...   \n",
-       "2            [basierend auf BSI-DSZ-CC-1052-V4-2021.]   \n",
-       "7   [This certification is a re-certification of B...   \n",
-       "8   [This is a re-certification based on BSI-DSZ-C...   \n",
-       "13  [Specific results from the evaluation process\\...   \n",
+       "                                             sentences  \\\n",
+       "251  [This is a re-certification based on BSI-DSZ-C...   \n",
+       "364  [This is a re-certification based on BSI-DSZ-C...   \n",
+       "505  [and BSI-DSZ-CC-0904-2015-\\n, This is a re-cer...   \n",
+       "616          [OCSI/CERT/TEC/01/2013/RC, versione 1.0,]   \n",
+       "628  [Specific results from the\\nevaluation process...   \n",
+       "709  [This is a re-certification based on BSI-DSZ-C...   \n",
+       "719  [The updated documents in compare to the forer...   \n",
        "\n",
-       "                                      y_proba          y_pred  correct  \n",
-       "1   [0.9330686268852108, 0.06693137311478929]  COMPONENT_USED    False  \n",
-       "2    [0.7070543956916182, 0.2929456043083818]  COMPONENT_USED    False  \n",
-       "7    [0.5998535578550636, 0.4001464421449364]  COMPONENT_USED    False  \n",
-       "8   [0.8727371952470533, 0.12726280475294682]  COMPONENT_USED    False  \n",
-       "13  [0.8670438280210987, 0.13295617197890133]  COMPONENT_USED    False  "
+       "                                       y_proba          y_pred  correct  \n",
+       "251  [0.47754689036832726, 0.5224531096316728]  COMPONENT_USED    False  \n",
+       "364  [0.11195564561722356, 0.8880443543827765]  COMPONENT_USED    False  \n",
+       "505  [0.16312949852765818, 0.8368705014723418]  COMPONENT_USED    False  \n",
+       "616   [0.2810731555375018, 0.7189268444624982]  COMPONENT_USED    False  \n",
+       "628   [0.4308279930268643, 0.5691720069731356]  COMPONENT_USED    False  \n",
+       "709   [0.4303467603163502, 0.5696532396836499]  COMPONENT_USED    False  \n",
+       "719   [0.48619060832541505, 0.513809391674585]  COMPONENT_USED    False  "
       ]
      },
-     "execution_count": 14,
+     "execution_count": 20,
      "metadata": {},
      "output_type": "execute_result"
     }
diff --git a/src/sec_certs/model/reference_classification.py b/src/sec_certs/model/reference_classification.py
index 9f1d4ff1..ce4fc473 100644
--- a/src/sec_certs/model/reference_classification.py
+++ b/src/sec_certs/model/reference_classification.py
@@ -62,7 +62,7 @@ class ReferenceClassifierTrainer:
         features = Features(
             {
                 "dgst": Value("string"),
-                "cert_id": Value("string"),
+                "referenced_cert_id": Value("string"),
                 "sentence": Value("string"),
                 "label": ClassLabel(names=list(df_to_use.label.unique())),
             }