WiP production-level reference annotation

10 files changed, 495 insertions, 137 deletions

diff --git a/src/sec_certs/dataset/cc.py b/src/sec_certs/dataset/cc.py
index c2849772..4aa73594 100644
--- a/src/sec_certs/dataset/cc.py
+++ b/src/sec_certs/dataset/cc.py
@@ -23,9 +23,14 @@ from sec_certs.dataset.cpe import CPEDataset
 from sec_certs.dataset.cve import CVEDataset
 from sec_certs.dataset.dataset import AuxiliaryDatasets, Dataset, logger
 from sec_certs.dataset.protection_profile import ProtectionProfileDataset
-from sec_certs.model.reference_finder import ReferenceFinder
-from sec_certs.model.sar_transformer import SARTransformer
-from sec_certs.model.transitive_vulnerability_finder import TransitiveVulnerabilityFinder
+from sec_certs.model import (
+    ReferenceAnnotator,
+    ReferenceAnnotatorTrainer,
+    ReferenceFinder,
+    SARTransformer,
+    TransitiveVulnerabilityFinder,
+)
+from sec_certs.model.references.segment_extractor import ReferenceSegmentExtractor
 from sec_certs.sample.cc import CCCertificate
 from sec_certs.sample.cc_certificate_id import CertificateId
 from sec_certs.sample.cc_maintenance_update import CCMaintenanceUpdate
@@ -33,6 +38,7 @@ from sec_certs.sample.protection_profile import ProtectionProfile
 from sec_certs.serialization.json import ComplexSerializableType, CustomJSONDecoder, serialize
 from sec_certs.utils import helpers
 from sec_certs.utils import parallel_processing as cert_processing
+from sec_certs.utils.nlp import prec_recall_metric
 from sec_certs.utils.sanitization import sanitize_navigable_string as sns
 
 
@@ -160,6 +166,10 @@ class CCDataset(Dataset[CCCertificate, CCAuxiliaryDatasets], ComplexSerializable
         """
         return self.mu_dataset_dir / "maintenance_updates.json"
 
+    @property
+    def reference_annotator_dir(self) -> Path:
+        return self.root_dir / "reference_annotator"
+
     BASE_URL: ClassVar[str] = "https://www.commoncriteriaportal.org"
 
     HTML_PRODUCTS_URL = {
@@ -801,6 +811,27 @@ class CCDataset(Dataset[CCCertificate, CCAuxiliaryDatasets], ComplexSerializable
 
         return update_dset
 
+    def annotate_references(self, fresh: bool = True):
+        df = ReferenceSegmentExtractor().prepare_df_from_cc_certs(list(self.certs.values()))
+        if fresh:
+            annotator = self._train_reference_annotator(df)
+        else:
+            annotator = ReferenceAnnotator.from_pretrained(self.reference_annotator_dir)
+
+        df = annotator.predict_df(df)
+
+        # TODO: Now iterate over DF, fill-in references
+
+    def _train_reference_annotator(self, df: pd.DataFrame, save_model: bool = True) -> ReferenceAnnotator:
+        trainer = ReferenceAnnotatorTrainer.from_df(df, prec_recall_metric, "transformer", "production")
+        trainer.train()
+        logger.info(trainer.evaluate())
+
+        if save_model:
+            trainer.slf.save_pretrained(self.reference_annotator_dir)
+
+        return trainer.clf
+
 
 class CCDatasetMaintenanceUpdates(CCDataset, ComplexSerializableType):
     """
diff --git a/src/sec_certs/dataset/fips.py b/src/sec_certs/dataset/fips.py
index ce7563a0..84ae8b3f 100644
--- a/src/sec_certs/dataset/fips.py
+++ b/src/sec_certs/dataset/fips.py
@@ -17,7 +17,7 @@ from sec_certs.dataset.cpe import CPEDataset
 from sec_certs.dataset.cve import CVEDataset
 from sec_certs.dataset.dataset import AuxiliaryDatasets, Dataset
 from sec_certs.dataset.fips_algorithm import FIPSAlgorithmDataset
-from sec_certs.model.reference_finder import ReferenceFinder
+from sec_certs.model.references.reference_finder import ReferenceFinder
 from sec_certs.model.transitive_vulnerability_finder import TransitiveVulnerabilityFinder
 from sec_certs.sample.fips import FIPSCertificate
 from sec_certs.serialization.json import ComplexSerializableType, serialize
diff --git a/src/sec_certs/model/__init__.py b/src/sec_certs/model/__init__.py
index fe1024f9..881a8ae3 100644
--- a/src/sec_certs/model/__init__.py
+++ b/src/sec_certs/model/__init__.py
@@ -4,8 +4,19 @@ leveraged by members of Dataset package and are directly applied on members of S
 """
 
 from sec_certs.model.cpe_matching import CPEClassifier
-from sec_certs.model.reference_finder import ReferenceFinder
+from sec_certs.model.references.annotator import ReferenceAnnotator
+from sec_certs.model.references.annotator_trainer import ReferenceAnnotatorTrainer
+from sec_certs.model.references.reference_finder import ReferenceFinder
+from sec_certs.model.references.segment_extractor import ReferenceSegmentExtractor
 from sec_certs.model.sar_transformer import SARTransformer
 from sec_certs.model.transitive_vulnerability_finder import TransitiveVulnerabilityFinder
 
-__all__ = ["CPEClassifier", "ReferenceFinder", "TransitiveVulnerabilityFinder", "SARTransformer"]
+__all__ = [
+    "CPEClassifier",
+    "ReferenceFinder",
+    "TransitiveVulnerabilityFinder",
+    "SARTransformer",
+    "ReferenceAnnotator",
+    "ReferenceAnnotatorTrainer",
+    "ReferenceSegmentExtractor",
+]
diff --git a/src/sec_certs/model/reference_classification.py b/src/sec_certs/model/reference_classification.py
deleted file mode 100644
index ce4fc473..00000000
--- a/src/sec_certs/model/reference_classification.py
+++ /dev/null
@@ -1,131 +0,0 @@
-from __future__ import annotations
-
-from dataclasses import dataclass
-from typing import Any, Callable, Literal
-
-import numpy as np
-import pandas as pd
-from datasets import ClassLabel, Dataset, Features, NamedSplit, Value
-from sentence_transformers.losses import CosineSimilarityLoss
-from setfit import SetFitModel, SetFitTrainer
-
-from sec_certs.utils.nlp import softmax
-
-
-class ReferenceClassifierTrainer:
-    def __init__(
-        self,
-        train_dataset: pd.DataFrame,
-        eval_dataset: pd.DataFrame,
-        metric: Callable,
-        mode: Literal["transformer", "baseline"] = "transformer",
-    ):
-        self._train_dataset = train_dataset
-        self._validation_dataset = eval_dataset
-        self._metric = metric
-        self._model, self._trainer, self.label_mapping = self._init_trainer(mode)
-        self.clf = RerefenceClassifier(self._model, self.label_mapping)
-
-    def _init_trainer(self, mode: Literal["transformer", "baseline"]):
-        return (
-            self._init_transformer_model_and_trainer()
-            if mode == "transformer"
-            else self._init_baseline_model_and_trainer()
-        )
-
-    def _init_transformer_model_and_trainer(self):
-        model = SetFitModel.from_pretrained("all-mpnet-base-v2")
-
-        internal_train_dataset = self._get_hugging_face_datasets_from_df(self._train_dataset, "train")
-        internal_validation_dataset = self._get_hugging_face_datasets_from_df(self._validation_dataset, "validation")
-
-        trainer = SetFitTrainer(
-            model=model,
-            train_dataset=internal_train_dataset,
-            eval_dataset=internal_validation_dataset,
-            loss_class=CosineSimilarityLoss,
-            metric=self._metric,
-            batch_size=16,
-            num_iterations=40,  # The number of text pairs to generate for contrastive learning
-            num_epochs=1,  # The number of epochs to use for contrastive learning
-            column_mapping={
-                "sentence": "text",
-                "label": "label",
-            },  # Map dataset columns to text/label expected by trainer
-        )
-        label_mapping = {index: x for index, x in enumerate(internal_train_dataset.features["label"].names)}
-        return model, trainer, label_mapping
-
-    @staticmethod
-    def _get_hugging_face_datasets_from_df(df: pd.DataFrame, split: NamedSplit) -> Dataset:
-        df_to_use = df.explode("sentences").rename(columns={"sentences": "sentence"})
-        features = Features(
-            {
-                "dgst": Value("string"),
-                "referenced_cert_id": Value("string"),
-                "sentence": Value("string"),
-                "label": ClassLabel(names=list(df_to_use.label.unique())),
-            }
-        )
-        return Dataset.from_pandas(df_to_use, features=features, split=split, preserve_index=False)
-
-    def _init_baseline_model_and_trainer(self):
-        # Process the datasets so that BaselineTrainer can work with them and init the trainer.
-        raise NotImplementedError("Not yet implemented.")
-
-    def train(self):
-        self._trainer.train(show_progress_bar=True)
-
-    def evaluate(self):
-        print("Internal evaluation (of model working on individual sentences)")
-        print(self._evaluate_raw())
-        print("Actual evaluation after ensemble soft voting")
-        print(self._evaluate_stacked())
-
-    def _evaluate_raw(self):
-        return self._trainer.evaluate()
-
-    def _evaluate_stacked(self):
-        y_pred = self.clf.predict(self._validation_dataset.sentences)
-        y_true = self._validation_dataset.label
-        return self._metric(y_pred, y_true)
-
-
-# TODO: Implement me
-class BaselineTrainer:
-    """
-    This is where baseline method shall be implemented. It should accept the classifier and fit it on train_dataset.
-    It should then use eval_dataset to evaluate the classifier.
-    """
-
-    def __init__(self, model, train_dataset, eval_dataset, metric):
-        pass
-
-    def train(self):
-        pass
-
-    def evaluate(self):
-        pass
-
-
-@dataclass
-class RerefenceClassifier:
-    _model: Any
-    _label_mapping: dict[int, str]
-
-    def predict(self, X: list[list[str]]) -> list[str]:
-        return [self._predict_single(x) for x in X]
-
-    def _predict_single(self, sample: list[str]) -> str:
-        return self._label_mapping[int(np.argmax(self._predict_proba_single(sample)))]
-
-    def predict_proba(self, X: list[list[str]]) -> list[list[float]]:
-        return [self._predict_proba_single(x) for x in X]
-
-    def _predict_proba_single(self, sample: list[str]) -> list[float]:
-        """
-        1. Get predictions for each sentence
-        2. Square every prediction to reward confidence
-        3. Sum probabilities for each label
-        """
-        return softmax(np.power(self._model.predict_proba(sample), 2).sum(axis=0))
diff --git a/src/sec_certs/model/references/annotator.py b/src/sec_certs/model/references/annotator.py
new file mode 100644
index 00000000..4b02430c
--- /dev/null
+++ b/src/sec_certs/model/references/annotator.py
@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pandas as pd
+from setfit import SetFitModel
+
+from sec_certs.utils.nlp import softmax
+
+
+@dataclass
+class ReferenceAnnotator:
+    """
+    Class for annotating references. Its instances are supposed to by trained by `ReferenceAnnotatorTrainer`.
+    Can be serialized into a directory / load from a directory.
+    """
+
+    _model: Any
+    _label_mapping: dict[int, str]
+
+    @classmethod
+    def from_pretrained(cls, model_dir: str | Path) -> ReferenceAnnotator:
+        """
+        Loads classifier from directory, assuming that:
+          - the SetFitModel was dumped into that directory with model.save_pretrained(model_dir)
+          - json file label_mapping.json exists in model_dir
+
+        :param str | Path model_dir: path to directory to search for model and label mapping
+        :return RerefenceClassifier: classifier with SetFitModel and label mapping
+        """
+        model = SetFitModel.from_pretrained(str(model_dir))
+        with (Path(model_dir) / "label_mapping.json").open("r") as handle:
+            label_mapping = json.load(handle)
+            label_mapping = {int(k): v for k, v in label_mapping.items()}
+
+        return cls(model, label_mapping)
+
+    def save_pretrained(self, model_dir: str | Path):
+        """
+        Will dump _model and _label_mapping into a directory.
+        """
+        model_dir = Path(model_dir)
+        model_dir.mkdir(exist_ok=True, parents=True)
+
+        with (model_dir / "label_mapping.json").open("w") as handle:
+            json.dump(self._label_mapping, handle, indent=4)
+        self._model.save_pretrained(str(model_dir))
+
+    def train(self, train_dataset: pd.DataFrame):
+        pass
+
+    def predict(self, X: list[list[str]]) -> list[str]:
+        return [self._predict_single(x) for x in X]
+
+    def _predict_single(self, sample: list[str]) -> str:
+        return self._label_mapping[int(np.argmax(self._predict_proba_single(sample)))]
+
+    def predict_proba(self, X: list[list[str]]) -> list[list[float]]:
+        return [self._predict_proba_single(x) for x in X]
+
+    def _predict_proba_single(self, sample: list[str]) -> list[float]:
+        """
+        1. Get predictions for each segment
+        2. Square every prediction to reward confidence
+        3. Sum probabilities for each label
+        """
+        return softmax(np.power(self._model.predict_proba(sample), 2).sum(axis=0))
+
+    def predict_df(self, df: pd.DataFrame) -> pd.DataFrame:
+        """
+        WIll read df.segments and populate the dataframe with predictions.
+        """
+        df_new = df.copy()
+        y_proba = self.predict_proba(df.segments)
+        df_new["y_proba"] = y_proba
+        df_new["y_pred"] = df_new.y_proba.map(lambda x: self._label_mapping[int(np.argmax(x))])
+        df_new["correct"] = df_new.label == df_new.y_pred
+        return df_new
diff --git a/src/sec_certs/model/references/annotator_trainer.py b/src/sec_certs/model/references/annotator_trainer.py
new file mode 100644
index 00000000..4ee712fc
--- /dev/null
+++ b/src/sec_certs/model/references/annotator_trainer.py
@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+from typing import Callable, Literal
+
+import pandas as pd
+from datasets import ClassLabel, Dataset, Features, NamedSplit, Value
+from sentence_transformers.losses import CosineSimilarityLoss
+from setfit import SetFitModel, SetFitTrainer
+
+from sec_certs.model.references.annotator import ReferenceAnnotator
+from sec_certs.utils.nlp import prepare_reference_annotations_df
+
+# from enum import Enum, Unique
+# TODO: Do we have any use for the enum below?
+# @Unique
+# class ReferenceMeaning(Enum):
+#     ON_PLATFORM = 0
+#     COMPONENT_USED = 1
+#     PREVIOUS_VERSION = 2
+#     EVALUATION_REUSED = 3
+#     COMPONENT_SHARED = 4
+#     RECERTIFICATION = 5
+
+"""
+Production workflow:
+
+df = ReferenceClassifierTrainer.prepare_df_from_cc_dataset([cert for cert in cc_dset])
+trainer = ReferenceClassifierTrainer.from_df(precision_recall, df, "transformer", "production")
+trainer.train()
+trainer.evaluate()
+
+# Print how great we are on test set
+
+trainer.clf.save_pretrained(/some/directory)
+
+... 48 hours later
+
+cc_dset.annotate_references(/path/to/model)
+
+where
+
+def annotate_references(self, model_directory):
+    clf = ReferenceClassifier.from_pretrained(model_directory)
+    df = ReferenceClassifierTrainer.prepare_df_from_cc_dataset([x for x in self])
+    df = clf.predict_df(df)
+
+    Now iterate over df, take each pair (dgst, referenced_cert_id) and fill_in dictionary in self[dgst].heuristics.references ...
+
+"""
+
+
+class ReferenceAnnotatorTrainer:
+    def __init__(
+        self,
+        train_dataset: pd.DataFrame,
+        eval_dataset: pd.DataFrame,
+        metric: Callable,
+        method: Literal["transformer", "baseline"] = "transformer",
+    ):
+        self._train_dataset = train_dataset
+        self._eval_dataset = eval_dataset
+        self._metric = metric
+        self._model, self._trainer, self.label_mapping = self._init_trainer(method)
+        self.clf = ReferenceAnnotator(self._model, self.label_mapping)
+
+    @classmethod
+    def from_df(
+        cls,
+        df: pd.DataFrame,
+        metric: Callable,
+        method: Literal["transformer", "baseline"] = "transformer",
+        mode: Literal["training", "production"] = "training",
+    ):
+        df = prepare_reference_annotations_df(df)
+        processing_method = {
+            "training": ReferenceAnnotatorTrainer.split_df_for_training,
+            "production": ReferenceAnnotatorTrainer.split_df_for_production,
+        }
+
+        train_dataset, eval_dataset = processing_method[mode](df)
+        return cls(train_dataset, eval_dataset, metric, method)
+
+    @staticmethod
+    def split_df_for_training(df: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]:
+        return df.loc[df.split == "train"].drop(columns="split"), df.loc[df.split == "valid"].drop(columns="split")
+
+    @staticmethod
+    def split_df_for_production(df: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]:
+        df.split = df.split.map({"test": "test", "train": "train", "valid": "train"})
+        return df.loc[df.split == "train"].drop(columns="split"), df.loc[df.split == "test"].drop(columns="split")
+
+    def _init_trainer(self, method: Literal["transformer", "baseline"]):
+        return (
+            self._init_transformer_model_and_trainer()
+            if method == "transformer"
+            else self._init_baseline_model_and_trainer()
+        )
+
+    def _init_transformer_model_and_trainer(self):
+        model = SetFitModel.from_pretrained("paraphrase-multilingual-mpnet-base-v2")
+        # model = SetFitModel.from_pretrained("all-mpnet-base-v2")
+
+        internal_train_dataset = self._get_hugging_face_datasets_from_df(self._train_dataset, "train")
+        internal_validation_dataset = self._get_hugging_face_datasets_from_df(self._eval_dataset, "validation")
+
+        trainer = SetFitTrainer(
+            model=model,
+            train_dataset=internal_train_dataset,
+            eval_dataset=internal_validation_dataset,
+            loss_class=CosineSimilarityLoss,
+            metric=self._metric,
+            batch_size=16,
+            num_iterations=40,  # The number of text pairs to generate for contrastive learning
+            num_epochs=1,  # The number of epochs to use for contrastive learning
+            column_mapping={
+                "segment": "text",
+                "label": "label",
+            },  # Map dataset columns to text/label expected by trainer
+        )
+        label_mapping = {index: x for index, x in enumerate(internal_train_dataset.features["label"].names)}
+        return model, trainer, label_mapping
+
+    @staticmethod
+    def _get_hugging_face_datasets_from_df(df: pd.DataFrame, split: NamedSplit) -> Dataset:
+        df_to_use = df.explode("segments").rename(columns={"segments": "segment"})
+        features = Features(
+            {
+                "dgst": Value("string"),
+                "referenced_cert_id": Value("string"),
+                "segment": Value("string"),
+                "label": ClassLabel(names=list(df_to_use.label.unique())),
+            }
+        )
+        return Dataset.from_pandas(df_to_use, features=features, split=split, preserve_index=False)
+
+    def _init_baseline_model_and_trainer(self):
+        # Process the datasets so that BaselineTrainer can work with them and init the trainer.
+        raise NotImplementedError("Not yet implemented.")
+
+    def train(self):
+        self._trainer.train(show_progress_bar=True)
+
+    def evaluate(self):
+        print("Internal evaluation (of model working on individual segments)")
+        print(self._evaluate_raw())
+        print("Actual evaluation after ensemble soft voting")
+        print(self._evaluate_stacked())
+
+    def _evaluate_raw(self):
+        return self._trainer.evaluate()
+
+    def _evaluate_stacked(self):
+        y_pred = self.clf.predict(self._eval_dataset.segments)
+        y_true = self._eval_dataset.label
+        return self._metric(y_pred, y_true)
+
+
+# TODO: Implement me
+class BaselineTrainer:
+    """
+    This is where baseline method shall be implemented. It should accept the classifier and fit it on train_dataset.
+    It should then use eval_dataset to evaluate the classifier.
+    """
+
+    def __init__(self, model, train_dataset, eval_dataset, metric):
+        pass
+
+    def train(self):
+        pass
+
+    def evaluate(self):
+        pass
diff --git a/src/sec_certs/model/reference_finder.py b/src/sec_certs/model/references/reference_finder.py
index 94a3b29f..94a3b29f 100644
--- a/src/sec_certs/model/reference_finder.py
+++ b/src/sec_certs/model/references/reference_finder.py
diff --git a/src/sec_certs/model/references/segment_extractor.py b/src/sec_certs/model/references/segment_extractor.py
new file mode 100644
index 00000000..f2496447
--- /dev/null
+++ b/src/sec_certs/model/references/segment_extractor.py
@@ -0,0 +1,167 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+import langdetect
+import pandas as pd
+import spacy
+from importlib_resources import files
+
+from sec_certs.sample.cc import CCCertificate
+from sec_certs.utils import parallel_processing
+
+nlp = spacy.load("en_core_web_sm")
+
+
+@dataclass
+class ReferenceRecord:
+    """
+    Data structure to hold objects when extracting text segments from txt files relevant for reference annotations.
+    """
+
+    certificate: CCCertificate
+    referenced_cert_id: str
+    source: str
+    segments: set[str] | None = None
+
+    @staticmethod
+    def fill_reference_segments(record: ReferenceRecord) -> ReferenceRecord:
+        """
+        Open file, read text and extract sentences with `referenced_cert_id` match.
+        Static method to allow for parallelization
+        """
+        pth_to_read = (
+            record.certificate.state.st_txt_path
+            if record.source == "target"
+            else record.certificate.state.report_txt_path
+        )
+
+        with pth_to_read.open("r") as handle:
+            data = handle.read()
+
+        record.segments = {sent.text for sent in nlp(data).sents if record.referenced_cert_id in sent.text}
+        if not record.segments:
+            record.segments = None
+        return record
+
+    def to_pandas_tuple(self) -> tuple[str, str, str, set[str] | None]:
+        return self.certificate.dgst, self.referenced_cert_id, self.source, self.segments
+
+
+class ReferenceSegmentExtractor:
+    """
+    Class to process list of certificates into a dataframe that holds reference segments. Exploses single method
+    Should be only called with ReferenceSegmentExtractor().prepare_df_from_cc_certs(list_of_certificates)
+    """
+
+    def __init__(self):
+        pass
+
+    def prepare_df_from_cc_certs(self, certs: list[CCCertificate]) -> pd.DataFrame:
+        """
+        Prepares processed DataFrame for reference annotator training from a list of certificates. This method:
+        - Extracts text segments relevant for each reference out of the certificates, forms dataframe from those
+        - Loads data splits into train/valid/test
+        - Loads manually annotated samples
+        - Combines all of that into single dataframe
+        """
+        df_targets = self._build_df(
+            [x for x in certs if x.heuristics.st_references.directly_referencing and x.state.st_txt_path], "target"
+        )
+        df_reports = self._build_df(
+            [x for x in certs if x.heuristics.report_references.directly_referencing and x.state.report_txt_path],
+            "report",
+        )
+        df = pd.concat([df_targets, df_reports])
+        return self._process_df(df)
+
+    def _build_df(self, certs: list[CCCertificate], source: Literal["target", "report"]) -> pd.DataFrame:
+        """ """
+        attribute_mapping = {"target": "st_references", "report": "report_references"}
+        records = [
+            ReferenceRecord(x, y, source)
+            for x in certs
+            for y in getattr(x.heuristics, attribute_mapping[source]).directly_referencing
+        ]
+
+        # results = [ReferenceRecord.fill_reference_segments(x) for x in tqdm.tqdm(records)]
+        results = parallel_processing.process_parallel(
+            ReferenceRecord.fill_reference_segments,
+            records,
+            use_threading=False,
+            progress_bar=True,
+            progress_bar_desc=f"Recovering reference segments for {source}s",
+        )
+
+        return pd.DataFrame.from_records(
+            [x.to_pandas_tuple() for x in results],
+            columns=["dgst", "referenced_cert_id", "source", "segments"],
+        )
+
+    def _get_split_dict(self) -> dict[str, str]:
+        """
+        Returns dictionary that maps dgst: split, where split in `train`, `valid`, `test`
+        """
+
+        def get_single_dct(pth: Path, split_name: str) -> dict[str, str]:
+            with pth.open("r") as handle:
+                return dict.fromkeys(json.load(handle), split_name)
+
+        split_directory = files("sec_certs.data") / "reference_annotations/split/"
+        return {
+            **get_single_dct(split_directory / "train.json", "train"),
+            **get_single_dct(split_directory / "valid.json", "valid"),
+            **get_single_dct(split_directory / "test.json", "test"),
+        }
+
+    def _get_annotations_dict(self) -> dict[tuple[str, str], str]:
+        """
+        Returns dictionary mapping tuples `(dgst, referenced_cert_id) -> label`
+        """
+
+        def load_single_df(pth: Path, split_name: str) -> pd.DataFrame:
+            return (
+                pd.read_csv(pth)
+                .assign(label=lambda df_: df_.label.str.replace(" ", "_").str.upper(), split=split_name)
+                .replace("NONE", None)
+                .dropna(subset="label")
+            )
+
+        annotations_directory = files("sec_certs.data") / "reference_annotations/manual_annotations/"
+        df_annot = pd.concat(
+            [
+                load_single_df(annotations_directory / "train.csv", "train"),
+                load_single_df(annotations_directory / "valid.csv", "valid"),
+                load_single_df(annotations_directory / "test.csv", "test"),
+            ]
+        )[["dgst", "referenced_cert_id", "source", "label", "comment"]]
+
+        return (
+            df_annot[["dgst", "referenced_cert_id", "label"]].set_index(["dgst", "referenced_cert_id"]).label.to_dict()
+        )
+
+    def _process_df(self, df: pd.DataFrame) -> pd.DataFrame:
+        """
+        Fully processes the dataframe.
+        """
+        annotations_dict = self._get_annotations_dict()
+        split_dct = self._get_split_dict()
+
+        return (
+            df.loc[df.segments.notnull()]
+            .explode("segments")
+            .assign(lang=lambda df_: df_.segments.map(langdetect.detect))
+            .loc[lambda df_: df_.lang.isin({"en", "fr", "de"})]
+            .groupby(["dgst", "referenced_cert_id", "source"], as_index=False, dropna=False)
+            .agg({"segments": list, "lang": list})
+            .assign(
+                split=lambda df_: df_.dgst.map(split_dct),
+                label=lambda df_: [annotations_dict.get(x) for x in zip(df_["dgst"], df_["referenced_cert_id"])],
+            )
+            .loc[lambda df_: df_["split"] != "test"]
+            .groupby(["dgst", "referenced_cert_id", "label", "split"], as_index=False, dropna=False)
+            .agg({"segments": sum, "lang": sum})
+        )
diff --git a/src/sec_certs/sample/certificate.py b/src/sec_certs/sample/certificate.py
index 5dab9cef..6b87e552 100644
--- a/src/sec_certs/sample/certificate.py
+++ b/src/sec_certs/sample/certificate.py
@@ -20,6 +20,8 @@ P = TypeVar("P", bound="PdfData")
 
 @dataclass
 class References(ComplexSerializableType):
+    # TODO: Refactor me to contain reference labels.
+    # Only direct references will have labels.
     directly_referenced_by: set[str] | None = field(default=None)
     indirectly_referenced_by: set[str] | None = field(default=None)
     directly_referencing: set[str] | None = field(default=None)
diff --git a/src/sec_certs/utils/nlp.py b/src/sec_certs/utils/nlp.py
index 496b71b8..45ad9639 100644
--- a/src/sec_certs/utils/nlp.py
+++ b/src/sec_certs/utils/nlp.py
@@ -1,4 +1,9 @@
+from __future__ import annotations
+
+from ast import literal_eval
+
 import numpy as np
+import pandas as pd
 from sklearn.metrics import precision_score, recall_score
 
 
@@ -11,3 +16,22 @@ def prec_recall_metric(y_pred, y_true):
 
 def softmax(x):
     return np.exp(x - np.max(x)) / np.exp(x - np.max(x)).sum()
+
+
+def eval_strings(series):
+    return [list(literal_eval(x)) for x in series]
+
+
+def filter_short_sentences(sentences, cert_id):
+    return [x for x in sentences if len(x) > len(cert_id) + 20]
+
+
+def prepare_reference_annotations_df(df: pd.DataFrame):
+    df = (
+        df.loc[lambda df_: (df_.label != "SELF") & (df_.label.notnull())]
+        .assign(segments=lambda df_: eval_strings(df_.segments))
+        .drop(columns="lang")
+    )
+    df.segments = df.apply(lambda row: filter_short_sentences(row["segments"], row["referenced_cert_id"]), axis=1)
+    df = df.loc[lambda df_: df_.segments.map(len) > 0]
+    return df