Skip to content

Commit

Permalink
chore: remove redundant np.round
Browse files Browse the repository at this point in the history
  • Loading branch information
hh-space-invader committed Oct 30, 2024
1 parent 6f33e60 commit 9851be2
Showing 1 changed file with 4 additions and 4 deletions.
8 changes: 4 additions & 4 deletions tests/test_text_cross_encoder.py
Original file line number Diff line number Diff line change
@@ -28,9 +28,9 @@ def test_rerank():

query = "What is the capital of France?"
documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."]
scores = np.round(np.array(list(model.rerank(query, documents))), decimals=4)
scores = np.array(list(model.rerank(query, documents)))

canonical_scores = np.round(CANONICAL_SCORE_VALUES[model_name], decimals=4)
canonical_scores = CANONICAL_SCORE_VALUES[model_name]
assert np.allclose(
scores, canonical_scores, atol=1e-3
), f"Model: {model_name}, Scores: {scores}, Expected: {canonical_scores}"
@@ -49,9 +49,9 @@ def test_batch_rerank(model_name):

query = "What is the capital of France?"
documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."] * 50
scores = np.round(np.array(list(model.rerank(query, documents, batch_size=10))), decimals=4)
scores = np.array(list(model.rerank(query, documents, batch_size=10)))

canonical_scores = np.round(np.tile(CANONICAL_SCORE_VALUES[model_name], 50), decimals=4)
canonical_scores = np.tile(CANONICAL_SCORE_VALUES[model_name], 50)

assert scores.shape == canonical_scores.shape, f"Unexpected shape for model {model_name}"
assert np.allclose(

0 comments on commit 9851be2

Please sign in to comment.