def test_ndcg_all_relevant():
torch.manual_seed(42)
scores = torch.FloatTensor([[10.0, 5.0, 2.0, 3.0, 4.0],
[5.0, 6.0, 4.0, 2.0, 5.5]])
ys = torch.LongTensor([[1, 1, 1, 1, 1], [1, 1, 1, 1, 1]])
n = torch.LongTensor([5, 4])
out = ndcg(scores, ys, n, k=5, exp=False)
expected = torch.ones(2)
assert out.numpy() == approx(expected.numpy())
def test_ndcg5_nonexp():
torch.manual_seed(42)
scores, ys, n = _generate_data()
out = ndcg(scores, ys, n, k=5, exp=False)
expected = torch.FloatTensor([
1.5177825608059992 / 2.1309297535714578,
3.3234658187877653 / 4.130929753571458
])
assert out.numpy() == approx(expected.numpy())
def test_ndcg5_exp():
torch.manual_seed(42)
scores, ys, n = _generate_data()
out = ndcg(scores, ys, n, k=5, exp=True)
expected = torch.FloatTensor([
1.5177825608059992 / 2.1309297535714578,
5.847184833073595 / 8.130929753571458
])
assert out.numpy() == approx(expected.numpy())
def test_ndcg3():
torch.manual_seed(42)
scores, ys, n = _generate_data()
out = ndcg(scores, ys, n, k=3, exp=True)
expected = torch.FloatTensor([
1.1309297535714575 / 2.1309297535714578,
5.4165082750002025 / 8.130929753571458
])
assert out.numpy() == approx(expected.numpy())
def evaluate():
model.eval()
loader = torch.utils.data.DataLoader(test,
batch_size=2,
shuffle=True,
collate_fn=test.collate_fn())
ndcg_score = 0.0
for batch in loader:
xs, ys, n = batch.features, batch.relevance, batch.n
ndcg_score += float(torch.sum(ndcg(model(xs), ys, n, k=10)))
ndcg_score /= len(test)
model.train()
return ndcg_score