Repository: Toloka/crowd-kit

Issue:
[FEATURE] Make DS algorithm takes the information from gold standard questions
### Problem description

In the real world, when doing truth inference, we have a small set of gold standard questions mixed in with other unlabeled questions. I thought it would help if we could use the true labels of these gold standard questions to adjust the estimation of the user confusion matrix and the prior distribution of the options.

### Feature description

Make ds algorithm accept an optional argument, `gt` for example, which can be used to adjust the estimation of user's confusion matrix and the priori distribution of options.

### Potential alternatives

_No response_

### Papers connected with feature

_No response_

### Additional information

This feature may make the EM process in DS more complicated if we add the argument and make some changes directly on the original implementation. Maybe there are better ways to do this?

Discussion:
shenxiangzhuang: Hi, this is a simple implementation that support #107 . I'm wondering if the changes are acceptable for crowd-kit or not? If it's ok, I would like to add more documentation and tests for this.
dustalov: Crowd-Kit is currently in a relatively mature development stage, yet occasionally, we accept the new functionality. The support for the true labels is a reasonable extension that I would happily accept.

Could you please make a separate PR with additional docs and tests?

Patch:
diff --git a/crowdkit/aggregation/classification/dawid_skene.py b/crowdkit/aggregation/classification/dawid_skene.py
index 50b6c2c..2382545 100644
--- a/crowdkit/aggregation/classification/dawid_skene.py
+++ b/crowdkit/aggregation/classification/dawid_skene.py
@@ -146,11 +146,15 @@ class DawidSkene(BaseClassificationAggregator):
         entropy = -(np.log(probas) * probas).sum().sum()
         return float(joint_expectation + entropy)
 
-    def fit(self, data: pd.DataFrame) -> "DawidSkene":
+    def fit(
+        self, data: pd.DataFrame, true_labels: Optional["pd.Series[Any]"] = None
+    ) -> "DawidSkene":
         """Fits the model to the training data with the EM algorithm.
         Args:
             data (DataFrame): The training dataset of workers' labeling results
                 which is represented as the `pandas.DataFrame` data containing `task`, `worker`, and `label` columns.
+            true_labels (Series): The ground truth labels of tasks. The `pandas.Series` data is indexed by `task`
+                        so that `labels.loc[task]` is the task ground truth label.
         Returns:
             DawidSkene: self.
         """
@@ -167,6 +171,9 @@ class DawidSkene(BaseClassificationAggregator):
 
         # Initialization
         probas = MajorityVote().fit_predict_proba(data)
+        # correct the probas by true_labels
+        if true_labels is not None:
+            probas = self._correct_probas_with_golden(probas, true_labels)
         priors = probas.mean()
         errors = self._m_step(data, probas)
         loss = -np.inf
@@ -175,6 +182,9 @@ class DawidSkene(BaseClassificationAggregator):
         # Updating proba and errors n_iter times
         for _ in range(self.n_iter):
             probas = self._e_step(data, priors, errors)
+            # correct the probas by true_labels
+            if true_labels is not None:
+                probas = self._correct_probas_with_golden(probas, true_labels)
             priors = probas.mean()
             errors = self._m_step(data, probas)
             new_loss = self._evidence_lower_bound(data, probas, priors, errors) / len(
@@ -225,6 +235,29 @@ class DawidSkene(BaseClassificationAggregator):
         assert self.labels_ is not None, "no labels_"
         return self.labels_
 
+    @staticmethod
+    def _correct_probas_with_golden(
+        probas: pd.DataFrame, true_labels: "pd.Series[Any]"
+    ) -> pd.DataFrame:
+        """
+        Correct the probas by `true_labels`
+        """
+        corrected_probas = probas
+
+        # Iterate over the unique labels present in true_labels
+        for label in true_labels.unique():
+            # Find the indices in both probas and true_labels where the true label is the current label
+            indices = true_labels[true_labels == label].index.intersection(probas.index)
+            # Set the corresponding probabilities to 1 for the correct label and 0 for others
+            corrected_probas.loc[indices] = (
+                0  # Set all columns to 0 for the given indices
+            )
+            corrected_probas.loc[indices, label] = (
+                1  # Set the correct label column to 1
+            )
+
+        return corrected_probas
+
 
 @attr.s
 class OneCoinDawidSkene(DawidSkene):
@@ -308,7 +341,7 @@ class OneCoinDawidSkene(DawidSkene):
         skills = skilled_data.groupby(["worker"], sort=False)["skill"].mean()
         return skills
 
-    def fit(self, data: pd.DataFrame) -> "OneCoinDawidSkene":
+    def fit(self, data: pd.DataFrame) -> "OneCoinDawidSkene":  # type: ignore[override]
         """Fits the model to the training data with the EM algorithm.
         Args:
             data (DataFrame): The training dataset of workers' labeling results


Test patch:
diff --git a/tests/aggregation/test_ds_aggregation.py b/tests/aggregation/test_ds_aggregation.py
index ba78ae0..51101fe 100644
--- a/tests/aggregation/test_ds_aggregation.py
+++ b/tests/aggregation/test_ds_aggregation.py
@@ -13,6 +13,21 @@ from pandas.testing import assert_frame_equal, assert_series_equal
 from crowdkit.aggregation import DawidSkene, OneCoinDawidSkene
 
 
+@pytest.mark.parametrize("n_iter, tol", [(10, 0), (100500, 1e-5)])
+def test_aggregate_ds_gold_on_toy_ysda(
+    n_iter: int,
+    tol: float,
+    toy_answers_df: pd.DataFrame,
+    toy_ground_truth_df: "pd.Series[Any]",
+    toy_gold_df: "pd.Series[Any]",
+) -> None:
+    np.random.seed(42)
+    assert_series_equal(
+        DawidSkene(n_iter=n_iter, tol=tol).fit(toy_answers_df, toy_gold_df).labels_.sort_index(),  # type: ignore
+        toy_ground_truth_df.sort_index(),
+    )
+
+
 @pytest.mark.parametrize("n_iter, tol", [(10, 0), (100500, 1e-5)])
 def test_aggregate_ds_on_toy_ysda(
     n_iter: int,