ProCIS/metrics.py at main · algoprog/ProCIS · 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
import numpy as np

from typing import List, Dict, Tuple


def calculate_npdcg(retrieved: List[Dict[str, List[Tuple[str, int]]]], cutoffs: List[int]) -> Dict[int, float]:
    npdcg_values = {}
    for cutoff in cutoffs:
        pdcg = calculate_pdcg(retrieved, cutoff)
        ipdcg = calculate_ipdcg(retrieved, cutoff)
        npdcg = pdcg / ipdcg if ipdcg != 0 else 0
        npdcg_values[cutoff] = npdcg
    return npdcg_values

def calculate_pdcg(retrieved: List[Dict[str, List[Tuple[str, int]]]], cutoff: int) -> float:
    pdcg = 0
    Z = 0
    ideal_position_l = {}
    labels = {}
    for i, utterance in enumerate(retrieved):
        for doc, label in utterance["correct_docs"]:
            if doc not in ideal_position_l and label > 0:
                ideal_position_l[doc] = i
                labels[doc] = label
    checked_docs = set()
    for i, utterance in enumerate(retrieved):
        retrieved_docs = utterance["retrieved_docs"][:cutoff]
        if len(retrieved_docs) > 0:
            Z += 1
            dcg = 0
            for j, doc in enumerate(retrieved_docs):
                if doc in ideal_position_l and i >= ideal_position_l[doc]:
                    if doc not in checked_docs:
                        checked_docs.add(doc)
                        dcg += labels[doc] / np.log2(2 + i - ideal_position_l[doc]) / np.log2(j + 2)
            pdcg += dcg
    return pdcg / Z if Z != 0 else 0


def calculate_ipdcg(retrieved: List[Dict[str, List[Tuple[str, int]]]], cutoff: int) -> float:
    ipdcg = 0
    Z = 0
    ideal_position_l = {}
    labels = {}
    for i, utterance in enumerate(retrieved):
        for doc, label in utterance["correct_docs"]:
            if doc not in ideal_position_l and label > 0:
                ideal_position_l[doc] = i
                labels[doc] = label
    checked_docs = set()
    for i, utterance in enumerate(retrieved):
        retrieved_docs = [doc for doc, _ in sorted(utterance["correct_docs"], key=lambda x: x[1], reverse=True)][:cutoff]
        if len(retrieved_docs) > 0:
            Z += 1
            dcg = 0
            for j, doc in enumerate(retrieved_docs):
                if doc not in checked_docs:
                    checked_docs.add(doc)
                    rel = labels[doc]
                    dcg += rel / np.log2(j + 2)
            ipdcg += dcg
    return ipdcg / Z if Z != 0 else 0


# retrieved = [
#     {
#         "retrieved_docs": [1, 2, 30, 4],
#         "correct_docs": [(1, 2), (2, 1), (14, 2)]
#     },
#     {
#         "retrieved_docs": [],
#         "correct_docs": [(9, 1)]
#     },
#     {
#         "retrieved_docs": [3, 7, 6, 9, 10],
#         "correct_docs": [(3, 2)]
#     },
#     {
#         "retrieved_docs": [12, 13, 14, 30],
#         "correct_docs": []
#     }
# ]

# retrieved = [
#     {
#         "retrieved_docs": [1],
#         "correct_docs": [(1, 1), (2, 1)]
#     },
#     {
#         "retrieved_docs": [3, 2, 4],
#         "correct_docs": [(3, 1), (4, 1)]
#     }
# ]

# cutoffs = [100]
# npdcg_values = calculate_npdcg(retrieved, cutoffs)
# print(npdcg_values)