KnowledgeSelectionSimulation/dataset_util.py at main · jacklxc/KnowledgeSelectionSimulation · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
import json
import nltk
from rouge import Rouge

class WoWDataset:
    def __init__(self, path, top_k=None):
        with open(path) as f:
            wow_dataset = json.load(f)
        if top_k:
            self.wow_dataset = wow_dataset[:top_k]

        self.examples = self.preprocess(self.wow_dataset)

    def preprocess(self, wow_dataset):
        examples = []
        for conversation in wow_dataset:
            persona = conversation["persona"]
            full_history = []
            for dialog in conversation["dialog"]:
                full_history.append({"speaker": dialog["speaker"], "text": dialog["text"]})
            for ui, utterance in enumerate(conversation["dialog"]):
                if "checked_sentence" in utterance:
                    history = full_history[:ui]

                    retrieved_passages = {}
                    retrieved_passages_reversed = {}
                    for obj in utterance["retrieved_passages"]:
                        for title, passage in obj.items():
                            retrieved_passages[title] = passage
                            for sentence in passage:
                                retrieved_passages_reversed[sentence] = title

                    gold_passages = list(utterance["checked_passage"].values())
                    if gold_passages:
                        gold_passage = gold_passages[0]
                    else:
                        gold_passage = "Unknown title"

                    gold_knowledge = []
                    #gold_sentences = list(utterance["checked_sentence"].values())
                    gold_sentences = []
                    for knowledge in utterance["multi_source_knowledge"]:
                        if knowledge["gold_sentence"]:
                            gold_sentences.append(knowledge["sentence"])
                    for gold_sentence in gold_sentences:
                        gold_topic = retrieved_passages_reversed.get(gold_sentence, gold_passage)
                        gold_knowledge.append((gold_topic, gold_sentence))

                    examples.append({
                        "id": len(examples),
                        "persona": persona,
                        "history": history,
                        "response": utterance["text"],
                        "gold_knowledge": gold_knowledge,
                        "retrieved_passages": retrieved_passages
                    })
        return examples

    def __len__(self):
        return len(self.examples)

    def __getitem__(self, idx):
        return self.examples[idx]


def make_context(selected_knowledge, example):
    retrieved_passages = example["retrieved_passages"]
    context = []
    for title, indices in selected_knowledge.items():
        context.append({
            "title": title,
            "sentences": [retrieved_passages[title][i] for i in indices]
        })
    return context

def evaluate_retrieval(all_retrieved_knowledge, all_gold_knowledge):
    true_positives = 0
    false_positives = 0
    false_negatives = 0

    for retrieved_knowledge, gold_knowledge in zip(all_retrieved_knowledge, all_gold_knowledge):
        retrieved_knowledge_setences = set([])
        for entry in retrieved_knowledge:
            for sentence in entry["sentences"]:
                retrieved_knowledge_setences.add(sentence)

        gold_knowledge_setences = set([sentence for title, sentence in gold_knowledge])

        for gold in gold_knowledge_setences:
            if gold in retrieved_knowledge_setences:
                true_positives += 1
            else:
                false_negatives += 1

        for retrieved in retrieved_knowledge_setences:
            if retrieved not in gold_knowledge_setences:
                false_positives += 1
    precision = true_positives / (true_positives + false_positives) if (true_positives + false_positives) > 0 else 0
    recall = true_positives / (true_positives + false_negatives) if (true_positives + false_negatives) > 0 else 0
    f1 = 2 * (precision * recall) / (precision + recall) if (precision + recall) > 0 else 0

    return {
        "precision": precision,
        "recall": recall,
        "f1": f1
    }

def evaluate_response_f1(predictions, references):
    tp = 0
    fn = 0
    fp = 0
    for ref, pred in zip(references, predictions):
        ref_tokens = set(nltk.word_tokenize(ref.lower()))
        pred_tokens = set(nltk.word_tokenize(pred.lower()))
        for pred_token in pred_tokens:
            if pred_token in ref_tokens:
                tp += 1
            else:
                fp += 1
        for ref_token in ref_tokens:
            if ref_token not in pred_tokens:
                fn += 1
    precision = tp / (tp + fp)
    recall = tp / (tp + fn)
    f1 = 2 * precision * recall / (precision + recall)
    return {"precision": precision, "recall": recall, "f1": f1}

def evaluate_response_rouge(predictions, references):
    rouge = Rouge()
    scores = rouge.get_scores(predictions, references, avg=True)
    return scores