Update classification_report

classification_report.py

@@ -12,26 +12,67 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """ classification_report metric. """
-from typing import Optional
-
-import sklearn
+from sklearn.metrics import classification_report
 import evaluate
 import datasets
 
 
+_DESCRIPTION = """
+Build a text report showing the main classification metrics that are accuracy, precision, recall and F1.
+"""
+
+_KWARGS_DESCRIPTION = """
+Args:
+    predictions (`list` of `int`): Predicted labels.
+    references (`list` of `int`): Ground truth labels.
+    labels (`list` of `int`): Optional list of label indices to include in the report. Defaults to None.
+    target_names: (`list` of `str`): Optional display names matching the labels (same order). Defaults to None.
+    sample_weight (`list` of `float`): Sample weights. Defaults to None.
+    digits (`int`): Number of digits for formatting output floating point values. When output_dict is True, this will be ignored and the returned values will not be rounded. Defaults to 2.
+    zero_division (`warn`, `0` or `1`): Sets the value to return when there is a zero division. If set to “warn”, this acts as 0, but warnings are also raised. Defaults to `warn`.
+Returns:
+    report (`str` or `dict`): Text summary of the precision, recall, F1 score for each class. Dictionary returned if output_dict is True. Dictionary has the following structure:
+        ```
+        {'label 1': {'precision':0.5,
+             'recall':1.0,
+             'f1-score':0.67,
+             'support':1},
+         'label 2': { ... },
+             ...
+        }
+        ```
+        The reported averages include macro average (averaging the unweighted mean per label), weighted average (averaging the support-weighted mean per label), and sample average (only for multilabel classification). Micro average (averaging the total true positives, false negatives and false positives) is only shown for multi-label or multi-class with a subset of classes, because it corresponds to accuracy otherwise and would be the same for all metrics. See also precision_recall_fscore_support for more details on averages.
+        Note that in binary classification, recall of the positive class is also known as “sensitivity”; recall of the negative class is “specificity”.
+Examples:
+    Simple example
+        >>> accuracy_metric = evaluate.load("bstrai/classification_report")
+        >>> results = accuracy_metric.compute(references=[0, 1, 2, 0, 1, 2], predictions=[0, 1, 1, 2, 1, 0])
+        >>> print(results)
+        {'0': {'precision': 0.5, 'recall': 0.5, 'f1-score': 0.5, 'support': 2}, '1': {'precision': 0.6666666666666666, 'recall': 1.0, 'f1-score': 0.8, 'support': 2}, '2': {'precision': 0.0, 'recall': 0.0, 'f1-score': 0.0, 'support': 2}, 'accuracy': 0.5, 'macro avg': {'precision': 0.38888888888888884, 'recall': 0.5, 'f1-score': 0.43333333333333335, 'support': 6}, 'weighted avg': {'precision': 0.38888888888888884, 'recall': 0.5, 'f1-score': 0.43333333333333335, 'support': 6}}
+"""
+
+
+_CITATION = """
+@article{scikit-learn,
+  title={Scikit-learn: Machine Learning in {P}ython},
+  author={Pedregosa, F. and Varoquaux, G. and Gramfort, A. and Michel, V.
+         and Thirion, B. and Grisel, O. and Blondel, M. and Prettenhofer, P.
+         and Weiss, R. and Dubourg, V. and Vanderplas, J. and Passos, A. and
+         Cournapeau, D. and Brucher, M. and Perrot, M. and Duchesnay, E.},
+  journal={Journal of Machine Learning Research},
+  volume={12},
+  pages={2825--2830},
+  year={2011}
+}
+"""
+
+@evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
 class ClassificationReportModule(evaluate.Metric):
-    """
-    Local metric used for classification task based on sklearn classiication_report().
-    a classification report is a simple tool to compute multiple metrics such as:
-        - accuracy
-        - precision/recall/f1-score by class.
-        - mean/weighted average.
-    """
     def _info(self) -> evaluate.MetricInfo:
         return evaluate.MetricInfo(
-            description="Metric based on sklearn classification_report() method.",
-            citation="",
-            inputs_description="",
+            description=_DESCRIPTION,
+            citation=_CITATION,
+            inputs_description=_KWARGS_DESCRIPTION,
             features=datasets.Features(
                 {
                     "predictions": datasets.Sequence(datasets.Value("int32")),
@@ -43,8 +84,8 @@ class ClassificationReportModule(evaluate.Metric):
                     "references": datasets.Value("int32"),
                 }
             ),
-            reference_urls=[""],
+            reference_urls=["https://scikit-learn.org/stable/modules/generated/sklearn.metrics.classification_report.html"],
         )
 
-    def _compute(self, *, predictions=None, references=None, **kwargs) -> Optional[dict]:
-        return sklearn.metrics.classification_report(references, predictions)
+    def _compute(self, predictions, references, labels=None, target_names=None, sample_weight=None, digits=2, zero_division="warn") -> dict:
+        return classification_report(y_true=references, y_pred=predictions, labels=labels, target_names=target_names, sample_weight=sample_weight, digits=digits, output_dict=True, zero_division=zero_division)