18 lines
584 B
Python
18 lines
584 B
Python
import torch
|
|
|
|
from data.schema import FACT_TYPE
|
|
from eval.benchmark_extractor import hard_negative_recall_at_one, recall_at_one
|
|
|
|
|
|
def test_retrieval_metrics_distinguish_hard_negative():
|
|
labels = torch.tensor(
|
|
[
|
|
[FACT_TYPE, 0, 0, 4, -1, -1],
|
|
[FACT_TYPE, 0, 0, 5, -1, -1],
|
|
]
|
|
)
|
|
candidates = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
|
|
queries = torch.tensor([[0.9, 0.1], [0.1, 0.9]])
|
|
assert recall_at_one(queries, labels, candidates, labels) == 1.0
|
|
assert hard_negative_recall_at_one(queries, labels, candidates, labels) == 1.0
|