Update relation_extraction.py

relation_extraction.py

@@ -2,40 +2,32 @@ import evaluate
 import datasets
 import numpy as np
 
-# Add BibTeX citation
+# TODO: Add BibTeX citation
 _CITATION = """\
-@Paper{
-    author = {Bruno Taillé, Vincent Guigue, Geoffrey Scoutheeten, Patrick Gallinari},
-    title = {Let's Stop Incorrect Comparisons in End-to-end Relation Extraction!},
-    year = {2020},
-    link = https://arxiv.org/abs/2009.10684
+@InProceedings{huggingface:module,
+title = {A great new module},
+authors={huggingface, Inc.},
+year={2020}
 }
 """
 
-# Add description of the module here
+# TODO: Add description of the module here
 _DESCRIPTION = """\
-This metric is used for evaluating the quality of relation extraction output. By calculating the Micro and Macro F1 score of every relation extraction outputs to ensure the quality.
+This new module is designed to solve this great ML task and is crafted with a lot of care.
 """
 
 
-# Add description of the arguments of the module here
+# TODO: Add description of the arguments of the module here
 _KWARGS_DESCRIPTION = """
-Calculates how good are predictions given some references, using Micro and Macro F1 scores
+Calculates how good are predictions given some references, using certain scores
 Args:
-    predictions: list of list of dictionary, including relation and its type 
-        dictionary should be key value pair like entity name link to its type
-    references: list of list of dictionary, including relation and its type 
-        dictionary should be entity name pair like entity name link to its type
+    predictions (list of list of dictionary): relation and its type of prediction
+       
+    references (list of list of dictionary): references for each relation and its type
+    
 Returns:
-    evaluation result: 
-    - **sell** (`dictionary`): score of type sell
-  - **tp** : true positive count
-  - **fp** : false positive count
-  - **fn** : false negative count
-  - **p** : precision
-  - **r** : recall
-  - **f1** : micro f1 score
-- **ALL** (`dictionary`): score of all of the type (sell and belongs to)
+    **output** (`dictionary` of `dictionary`s) with multiple key-value pairs
+- **entity type** (`dictionary`): score of all of the type 
   - **tp** : true positive count
   - **fp** : false positive count
   - **fn** : false negative count
@@ -46,25 +38,15 @@ Returns:
   - **Macro_p** : macro precision
   - **Macro_r** : macro recall
 Examples:
-    >>> metric_path = "Ikala-allen/relation_extraction"
-    >>> module = evaluate.load(metric_path)
-    >>> references = [
-    ...    [
-    ...       {"head": "phip igments", "head_type": "brand", "type": "sell", "tail": "國際認證之色乳", "tail_type": "product"},
-    ...        {"head": "tinadaviespigments", "head_type": "brand", "type": "sell", "tail": "國際認證之色乳", "tail_type": "product"},
-    ...    ]
-    ... ]
-    >>> predictions = [
-    ...    [
-    ...        {"head": "phipigments", "head_type": "product", "type": "sell", "tail": "國際認證之色乳", "tail_type": "product"},
-    ...        {"head": "tinadaviespigments", "head_type": "brand", "type": "sell", "tail": "國際認證之色乳", "tail_type": "product"},
-    ...    ]
-    ... ]
-    >>> evaluation_scores = module.compute(predictions=predictions, references=references)
-    >>> print(evaluation_scores)
-    {'sell': {'tp': 1, 'fp': 1, 'fn': 1, 'p': 50.0, 'r': 50.0, 'f1': 50.0}, 'ALL': {'tp': 1, 'fp': 1, 'fn': 1, 'p': 50.0, 'r': 50.0, 'f1': 50.0, 'Macro_f1': 50.0, 'Macro_p': 50.0, 'Macro_r': 50.0}}
+    Examples should be written in doctest format, and should illustrate how
+    to use the function.
+    my_new_module = evaluate.load("my_new_module")
+    results = my_new_module.compute(references=[0, 1], predictions=[0, 1])
+    print(results)
+    {'accuracy': 1.0}
 """
 
+
 def convert_format(data:list):
     """
     Args:
@@ -75,12 +57,13 @@ def convert_format(data:list):
         'head_type': ['product', 'brand'...], 
         'type': ['sell', 'sell'...], 
         'tail': ['國際認證之色乳', '國際認證之色乳'...], 
-        'tail_type': ['product', 'product'...]},  # first element
+        'tail_type': ['product', 'product'...]}, 
+
         {'head': ['SABONTAIWAN', 'SNTAIWAN'...], 
         'head_type': ['brand', 'brand'...], 
         'type': ['sell', 'sell'...], 
         'tail': ['大馬士革玫瑰有機光燦系列', '大馬士革玫瑰有機光燦系列'...], 
-        'tail_type': ['product', 'product'...]}  # second element
+        'tail_type': ['product', 'product'...]}
         ...
         ]
     """
@@ -101,12 +84,7 @@ def convert_format(data:list):
 
 @evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
 class relation_extraction(evaluate.Metric):
-    """
-    evaluation metric of relation extraction
-    inputs:
-        predictions : (`list` of `list`s of `dictionary`s) about relation and its type of prediction
-        references : (`list` of `list`s of `dictionary`s) about references for each relation and its type.
-    """
+    """TODO: Short description of my evaluation module."""
 
     def _info(self):
         # TODO: Specifies the evaluate.EvaluationModuleInfo object
@@ -145,7 +123,7 @@ class relation_extraction(evaluate.Metric):
         # TODO: Download external resources if needed
         pass
 
-    def _compute(self, predictions, references, mode="strict", relation_types=[]):
+    def _compute(self, predictions, references, mode, only_all=True, relation_types=[]):
         """Returns the scores"""
         # TODO: Compute the different scores of the module
         
@@ -229,4 +207,7 @@ class relation_extraction(evaluate.Metric):
         scores["ALL"]["Macro_p"] = np.mean([scores[ent_type]["p"] for ent_type in relation_types])
         scores["ALL"]["Macro_r"] = np.mean([scores[ent_type]["r"] for ent_type in relation_types])
 
+        if only_all:
+            return scores["ALL"]
+        
         return scores