-
Notifications
You must be signed in to change notification settings - Fork 5
/
rank.go
87 lines (74 loc) · 1.63 KB
/
rank.go
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
package eval
import (
"fmt"
"github.com/hscells/trecresults"
"math"
"sort"
)
type DCG struct{ K int }
type NDCG struct{ K int }
var (
AP = ap{}
)
type ap struct{}
func (e ap) Score(results *trecresults.ResultList, qrels trecresults.Qrels) float64 {
var sum float64
var numRelSeen float64
nr := NumRel.Score(results, qrels)
for i, res := range *results {
if numRelSeen == nr {
break
}
if _, ok := qrels[res.DocId]; ok {
if qrels[res.DocId].Score > RelevanceGrade {
numRelSeen++
sum += PrecisionAtK{K: i + 1}.Score(results, qrels)
}
}
}
return sum / nr
}
func (e ap) Name() string {
return "AP"
}
func (e DCG) Score(results *trecresults.ResultList, qrels trecresults.Qrels) float64 {
var score float64
for i, item := range *results {
// Compute DCG at a cutoff.
if e.K != 0 && i >= e.K {
break
}
if _, ok := qrels[item.DocId]; ok {
score += float64(qrels[item.DocId].Score) / math.Log2(float64(i)+2)
}
}
return score
}
func (e DCG) Name() string {
return "DCG"
}
func (e NDCG) Score(results *trecresults.ResultList, qrels trecresults.Qrels) float64 {
// Compute ideal discounted cumulative gain.
ideal := make(trecresults.ResultList, len(qrels))
i := 0
for _, rel := range qrels {
ideal[i] = &trecresults.Result{
Topic: rel.Topic,
DocId: rel.DocId,
Score: float64(rel.Score),
}
i++
}
sort.Slice(ideal, func(i, j int) bool {
return ideal[i].Score > ideal[j].Score
})
dcg := DCG{K: e.K}.Score(results, qrels)
idcg := DCG{K: e.K}.Score(&ideal, qrels)
return dcg / idcg
}
func (e NDCG) Name() string {
if e.K > 0 {
return fmt.Sprintf("nDCG@%d", e.K)
}
return "nDCG"
}