Repository navigation
Expand file tree
/
Copy pathutils.py
More file actions
80 lines (67 loc) · 2.66 KB
/
Copy pathutils.py
File metadata and controls
80 lines (67 loc) · 2.66 KB
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
"""
Utility functions for loading qrels and results files
"""
from collections import defaultdict
from typing import Dict, List, Tuple
def load_qrels(qrels_file: str) -> Dict[str, Dict[str, int]]:
"""
Load qrels file.
Args:
qrels_file: Path to qrels file
Returns:
Dictionary mapping query_id to {passage_id: relevance}
Format (supports both):
query_id \t passage_id \t relevance (3 columns)
query_id \t 0 \t passage_id \t relevance (4 columns, standard TREC)
"""
qrels = defaultdict(dict)
with open(qrels_file, 'r', encoding='utf-8') as f:
for line in f:
parts = line.strip().split('\t')
if len(parts) >= 4:
# Standard TREC format: query_id \t 0 \t passage_id \t relevance
query_id = parts[0]
passage_id = parts[2]
relevance = int(parts[3])
qrels[query_id][passage_id] = relevance
elif len(parts) >= 3:
# Simplified format: query_id \t passage_id \t relevance
query_id = parts[0]
passage_id = parts[1]
relevance = int(parts[2])
qrels[query_id][passage_id] = relevance
return dict(qrels)
def load_results(results_file: str) -> Dict[str, List[Tuple[str, int, float]]]:
"""
Load search results file.
Args:
results_file: Path to results file
Returns:
Dictionary mapping query_id to [(passage_id, rank, score), ...]
Format:
query_id \t passage_id \t rank \t score \t run_name
OR
query_id \t Q0 \t passage_id \t rank \t score \t run_name
"""
results = defaultdict(list)
with open(results_file, 'r', encoding='utf-8') as f:
for line in f:
parts = line.strip().split('\t')
if len(parts) >= 5:
query_id = parts[0]
# Check if second field is "Q0" (TREC format) or passage_id (our format)
if parts[1] == 'Q0':
# Standard TREC format: query_id Q0 passage_id rank score run_name
passage_id = parts[2]
rank = int(parts[3])
score = float(parts[4])
else:
# Our format: query_id passage_id rank score run_name
passage_id = parts[1]
rank = int(parts[2])
score = float(parts[3])
results[query_id].append((passage_id, rank, score))
# Sort by rank to ensure correct order
for query_id in results:
results[query_id].sort(key=lambda x: x[1])
return dict(results)