update

docred.py

@@ -13,9 +13,10 @@
 # limitations under the License.
 """TODO: Add a description here."""
 
-import evaluate
-import datasets
+import os
 
+import datasets
+import evaluate
 
 # TODO: Add BibTeX citation
 _CITATION = """\
@@ -61,7 +62,30 @@ BAD_WORDS_URL = "http://url/to/external/resource/bad_words.txt"
 class docred(evaluate.Metric):
     """TODO: Short description of my evaluation module."""
 
+    dataset_feat = {
+        "title": datasets.Value("string"),
+        "sents": datasets.Sequence(datasets.Sequence(datasets.Value("string"))),
+        "vertexSet": datasets.Sequence(
+            datasets.Sequence(
+                {
+                    "name": datasets.Value("string"),
+                    "sent_id": datasets.Value("int32"),
+                    "pos": datasets.Sequence(datasets.Value("int32"), length=2),
+                    "type": datasets.Value("string"),
+                }
+            )
+        ),
+        "labels": {
+            "head": datasets.Sequence(datasets.Value("int32")),
+            "tail": datasets.Sequence(datasets.Value("int32")),
+            "relation_id": datasets.Sequence(datasets.Value("string")),
+            "relation_text": datasets.Sequence(datasets.Value("string")),
+            "evidence": datasets.Sequence(datasets.Sequence(datasets.Value("int32"))),
+        },
+    }
+
     def _info(self):
+
         # TODO: Specifies the evaluate.EvaluationModuleInfo object
         return evaluate.MetricInfo(
             # This is the description that will appear on the modules page.
@@ -70,15 +94,12 @@ class docred(evaluate.Metric):
             citation=_CITATION,
             inputs_description=_KWARGS_DESCRIPTION,
             # This defines the format of each prediction and reference
-            features=datasets.Features({
-                'predictions': datasets.Value('int64'),
-                'references': datasets.Value('int64'),
-            }),
+            features=datasets.Features({"predictions": self.dataset_feat, "references": self.dataset_feat}),
             # Homepage of the module for documentation
             homepage="http://module.homepage",
             # Additional links to the codebase or references
             codebase_urls=["http://github.com/path/to/codebase/of/new_module"],
-            reference_urls=["http://path.to.reference.url/new_module"]
+            reference_urls=["http://path.to.reference.url/new_module"],
         )
 
     def _download_and_prepare(self, dl_manager):
@@ -86,10 +107,131 @@ class docred(evaluate.Metric):
         # TODO: Download external resources if needed
         pass
 
-    def _compute(self, predictions, references):
+    def _generate_fact(self, dataset):
+        if dataset is None:
+            return set()
+        facts = set()
+        for data in dataset:
+            vertexSet = data["vertexSet"]
+            labels = self._convert_labels_to_list(data["labels"])
+            for label in labels:
+                rel = label["relation_id"]
+                for n1 in vertexSet[label["head"]]["name"]:
+                    for n2 in vertexSet[label["tail"]]["name"]:
+                        facts.add((n1, n2, rel))
+        return facts
+
+    def _convert_to_relation_set(self, data):
+        relation_set = set()
+        for d in data:
+            labels = d["labels"]
+            labels = self._convert_labels_to_list(labels)
+            for label in labels:
+                relation_set.add((d["title"], label["head"], label["tail"], label["relation_id"]))
+        return relation_set
+
+    def _convert_labels_to_list(self, labels):
+        keys = list(labels.keys())
+        labels = [{key: labels[key][i] for key in keys} for i in range(len(labels[keys[0]]))]
+        return labels
+
+    def _compute(self, predictions, references, train_data=None):
         """Returns the scores"""
-        # TODO: Compute the different scores of the module
-        accuracy = sum(i == j for i, j in zip(predictions, references)) / len(predictions)
-        return {
-            "accuracy": accuracy,
-        }
+
+        fact_in_train_annotated = self._generate_fact(train_data)
+
+        std = {}
+        tot_evidences = 0
+        ref_titleset = set([])
+
+        title2vectexSet = {}
+
+        for x in references:
+            title = x["title"]
+            ref_titleset.add(title)
+
+            vertexSet = x["vertexSet"]
+            title2vectexSet[title] = vertexSet
+            labels = self._convert_labels_to_list(x["labels"])
+            for label in labels:
+                r = label["relation_id"]
+                h_idx = label["head"]
+                t_idx = label["tail"]
+                std[(title, r, h_idx, t_idx)] = set(label["evidence"])
+                tot_evidences += len(label["evidence"])
+
+        tot_relations = len(std)
+        pred_rel = self._convert_to_relation_set(predictions)
+        submission_answer = sorted(pred_rel, key=lambda x: (x[0], x[1], x[2], x[3]))
+
+        correct_re = 0
+        correct_evidence = 0
+        pred_evi = 0
+
+        correct_in_train_annotated = 0
+        titleset2 = set([])
+        for x in submission_answer:
+            title, h_idx, t_idx, r = x
+            titleset2.add(title)
+            if title not in title2vectexSet:
+                continue
+            vertexSet = title2vectexSet[title]
+
+            if "evidence" in x:
+                evi = set(x["evidence"])
+            else:
+                evi = set([])
+            pred_evi += len(evi)
+
+            if (title, r, h_idx, t_idx) in std:
+                correct_re += 1
+                stdevi = std[(title, r, h_idx, t_idx)]
+                correct_evidence += len(stdevi & evi)
+                in_train_annotated = in_train_distant = False
+                for n1 in vertexSet[h_idx]["name"]:
+                    for n2 in vertexSet[t_idx]["name"]:
+                        if (n1, n2, r) in fact_in_train_annotated:
+                            in_train_annotated = True
+
+                if in_train_annotated:
+                    correct_in_train_annotated += 1
+                # if in_train_distant:
+                #     correct_in_train_distant += 1
+
+        re_p = 1.0 * correct_re / (len(submission_answer) + 1e-5)
+        re_r = 1.0 * correct_re / (tot_relations + 1e-5)
+        if re_p + re_r == 0:
+            re_f1 = 0
+        else:
+            re_f1 = 2.0 * re_p * re_r / (re_p + re_r)
+
+        evi_p = 1.0 * correct_evidence / pred_evi if pred_evi > 0 else 0
+        evi_r = 1.0 * correct_evidence / tot_evidences
+        if evi_p + evi_r == 0:
+            evi_f1 = 0
+        else:
+            evi_f1 = 2.0 * evi_p * evi_r / (evi_p + evi_r)
+
+        re_p_ignore_train_annotated = (
+            1.0
+            * (correct_re - correct_in_train_annotated)
+            / (len(submission_answer) - correct_in_train_annotated + 1e-5)
+        )
+        # re_p_ignore_train = (
+        #     1.0 * (correct_re - correct_in_train_distant) / (len(submission_answer) - correct_in_train_distant + 1e-5)
+        # )
+
+        if re_p_ignore_train_annotated + re_r == 0:
+            re_f1_ignore_train_annotated = 0
+        else:
+            re_f1_ignore_train_annotated = (
+                2.0 * re_p_ignore_train_annotated * re_r / (re_p_ignore_train_annotated + re_r)
+            )
+
+        # if re_p_ignore_train + re_r == 0:
+        #     re_f1_ignore_train = 0
+        # else:
+        #     re_f1_ignore_train = 2.0 * re_p_ignore_train * re_r / (re_p_ignore_train + re_r)
+
+        # return re_f1, evi_f1, re_f1_ignore_train_annotated, re_f1_ignore_train, re_p, re_r
+        return {"f1": re_f1, "precision": re_p, "recall": re_r, "ign_f1": re_f1_ignore_train_annotated}

official.py

@@ -0,0 +1,171 @@
+#!/usr/bin/env python
+import json
+import os
+import os.path
+import sys
+
+
+def gen_train_facts(data_file_name, truth_dir):
+    fact_file_name = data_file_name[data_file_name.find("train_") :]
+    fact_file_name = os.path.join(truth_dir, fact_file_name.replace(".json", ".fact"))
+
+    if os.path.exists(fact_file_name):
+        fact_in_train = set([])
+        triples = json.load(open(fact_file_name))
+        for x in triples:
+            fact_in_train.add(tuple(x))
+        return fact_in_train
+
+    fact_in_train = set([])
+    ori_data = json.load(open(data_file_name))
+    for data in ori_data:
+        vertexSet = data["vertexSet"]
+        for label in data["labels"]:
+            rel = label["r"]
+            for n1 in vertexSet[label["h"]]:
+                for n2 in vertexSet[label["t"]]:
+                    fact_in_train.add((n1["name"], n2["name"], rel))
+
+    json.dump(list(fact_in_train), open(fact_file_name, "w"))
+
+    return fact_in_train
+
+
+input_dir = sys.argv[1]
+output_dir = sys.argv[2]
+
+submit_dir = os.path.join(input_dir, "res")
+truth_dir = os.path.join(input_dir, "ref")
+
+if not os.path.isdir(submit_dir):
+    print("%s doesn't exist" % submit_dir)
+
+if os.path.isdir(submit_dir) and os.path.isdir(truth_dir):
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+
+    fact_in_train_annotated = gen_train_facts("../data/train_annotated.json", truth_dir)
+    fact_in_train_distant = gen_train_facts("../data/train_distant.json", truth_dir)
+
+    output_filename = os.path.join(output_dir, "scores.txt")
+    output_file = open(output_filename, "w")
+
+    truth_file = os.path.join(truth_dir, "dev_test.json")
+    truth = json.load(open(truth_file))
+
+    std = {}
+    tot_evidences = 0
+    titleset = set([])
+
+    title2vectexSet = {}
+
+    for x in truth:
+        title = x["title"]
+        titleset.add(title)
+
+        vertexSet = x["vertexSet"]
+        title2vectexSet[title] = vertexSet
+
+        for label in x["labels"]:
+            r = label["r"]
+
+            h_idx = label["h"]
+            t_idx = label["t"]
+            std[(title, r, h_idx, t_idx)] = set(label["evidence"])
+            tot_evidences += len(label["evidence"])
+
+    tot_relations = len(std)
+
+    submission_answer_file = os.path.join(submit_dir, "result.json")
+    tmp = json.load(open(submission_answer_file))
+    tmp.sort(key=lambda x: (x["title"], x["h_idx"], x["t_idx"], x["r"]))
+    submission_answer = [tmp[0]]
+    for i in range(1, len(tmp)):
+        x = tmp[i]
+        y = tmp[i - 1]
+        if (x["title"], x["h_idx"], x["t_idx"], x["r"]) != (y["title"], y["h_idx"], y["t_idx"], y["r"]):
+            submission_answer.append(tmp[i])
+
+    correct_re = 0
+    correct_evidence = 0
+    pred_evi = 0
+
+    correct_in_train_annotated = 0
+    correct_in_train_distant = 0
+    titleset2 = set([])
+    for x in submission_answer:
+        title = x["title"]
+        h_idx = x["h_idx"]
+        t_idx = x["t_idx"]
+        r = x["r"]
+        titleset2.add(title)
+        if title not in title2vectexSet:
+            continue
+        vertexSet = title2vectexSet[title]
+
+        if "evidence" in x:
+            evi = set(x["evidence"])
+        else:
+            evi = set([])
+        pred_evi += len(evi)
+
+        if (title, r, h_idx, t_idx) in std:
+            correct_re += 1
+            stdevi = std[(title, r, h_idx, t_idx)]
+            correct_evidence += len(stdevi & evi)
+            in_train_annotated = in_train_distant = False
+            for n1 in vertexSet[h_idx]:
+                for n2 in vertexSet[t_idx]:
+                    if (n1["name"], n2["name"], r) in fact_in_train_annotated:
+                        in_train_annotated = True
+                    if (n1["name"], n2["name"], r) in fact_in_train_distant:
+                        in_train_distant = True
+
+            if in_train_annotated:
+                correct_in_train_annotated += 1
+            if in_train_distant:
+                correct_in_train_distant += 1
+
+    re_p = 1.0 * correct_re / len(submission_answer)
+    re_r = 1.0 * correct_re / tot_relations
+    if re_p + re_r == 0:
+        re_f1 = 0
+    else:
+        re_f1 = 2.0 * re_p * re_r / (re_p + re_r)
+
+    evi_p = 1.0 * correct_evidence / pred_evi if pred_evi > 0 else 0
+    evi_r = 1.0 * correct_evidence / tot_evidences
+    if evi_p + evi_r == 0:
+        evi_f1 = 0
+    else:
+        evi_f1 = 2.0 * evi_p * evi_r / (evi_p + evi_r)
+
+    re_p_ignore_train_annotated = (
+        1.0 * (correct_re - correct_in_train_annotated) / (len(submission_answer) - correct_in_train_annotated)
+    )
+    re_p_ignore_train = (
+        1.0 * (correct_re - correct_in_train_distant) / (len(submission_answer) - correct_in_train_distant)
+    )
+
+    if re_p_ignore_train_annotated + re_r == 0:
+        re_f1_ignore_train_annotated = 0
+    else:
+        re_f1_ignore_train_annotated = 2.0 * re_p_ignore_train_annotated * re_r / (re_p_ignore_train_annotated + re_r)
+
+    if re_p_ignore_train + re_r == 0:
+        re_f1_ignore_train = 0
+    else:
+        re_f1_ignore_train = 2.0 * re_p_ignore_train * re_r / (re_p_ignore_train + re_r)
+
+    print("RE_F1:", re_f1)
+    print("Evi_F1:", evi_f1)
+    print("RE_ignore_annotated_F1:", re_f1_ignore_train_annotated)
+    print("RE_ignore_distant_F1:", re_f1_ignore_train)
+
+    output_file.write("RE_F1: %f\n" % re_f1)
+    output_file.write("Evi_F1: %f\n" % evi_f1)
+
+    output_file.write("RE_ignore_annotated_F1: %f\n" % re_f1_ignore_train_annotated)
+    output_file.write("RE_ignore_distant_F1: %f\n" % re_f1_ignore_train)
+
+    output_file.close()

sample.py

@@ -0,0 +1,10 @@
+import datasets
+import evaluate
+
+from docred import docred
+
+train_data = datasets.load_dataset("docred", split="train_annotated[:10]")
+data = datasets.load_dataset("docred", split="validation[:10]")
+metric = docred()
+
+print(metric.compute(predictions=data.to_list(), references=data.to_list()))