expeval/data.py at master · diazf/expeval · 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
from permutation import Permutation

#
# read_qrels
#
# QID GID DID REL
#

def read_qrels(fn, binarize, complete):
    qrels={}
    did2gids = {}
    #
    # read qrels
    #
    fp = open(fn,"r")
    for line in fp:
        fields = line.strip().split()
        if (len(fields) == 3):
            qid = fields[0]
            itr = "-1"
            did = fields[1]
            rel = fields[2]
        else:
            qid = fields[0]
            itr = fields[1]
            did = fields[2]
            rel = fields[3]
        gids = []
        gids = map(lambda x: int(x), itr.split("|"))

        rel=int(rel)
        if (rel>0) and binarize:
            rel = 1
        if complete or (rel>0):
            if not(qid in qrels):
                qrels[qid] = {}
                did2gids[qid] = {}
            qrels[qid][did] = rel
            if not(did in did2gids[qid]):
                did2gids[qid][did] = []
            for gid in gids:
                if not (gid in did2gids[qid][did]):
                    did2gids[qid][did].append(gid)
    fp.close()
    return qrels, did2gids

def read_topfile(fn):
    #
    # get ranked lists
    #
    fp = open(fn,"r")
    sample_ids=set([])
    rls = {} # qid x iteration -> permutation
    for line in fp:
        fields = line.strip().split()
        qid,itr,did,rank,score=fields[:5]
        sample_ids.add(itr)
        score = float(score)
        rank = int(rank)
        if not(qid in rls):
            rls[qid] = {}
        if not(itr in rls[qid]):
            rls[qid][itr] = Permutation()
        rls[qid][itr].add(rank,did)
    fp.close()
    return rls