Skip to content

Commit

Permalink
Remove values test for inited models
Browse files Browse the repository at this point in the history
  • Loading branch information
NohTow committed Nov 28, 2024
1 parent 37bcf88 commit 20b8061
Showing 1 changed file with 0 additions and 26 deletions.
26 changes: 0 additions & 26 deletions tests/test_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,35 +12,9 @@ def test_model_creation(**kwargs) -> None:
torch.manual_seed(42)
# Creation from a base encoder
model = models.ColBERT(model_name_or_path="bert-base-uncased")
queries_embeddings = model.encode(sentences=query, is_query=True)
documents_embeddings = model.encode(sentences=documents, is_query=False)
reranked_documents = rank.rerank(
documents_ids=[["1", "2"]],
queries_embeddings=queries_embeddings,
documents_embeddings=documents_embeddings,
)
assert math.isclose(
reranked_documents[0][0]["score"], 25.92, rel_tol=0.01, abs_tol=0.01
)
assert math.isclose(
reranked_documents[0][1]["score"], 23.7, rel_tol=0.01, abs_tol=0.01
)

# Creation from a base sentence-transformer
model = models.ColBERT(model_name_or_path="sentence-transformers/all-MiniLM-L6-v2")
queries_embeddings = model.encode(sentences=query, is_query=True)
documents_embeddings = model.encode(sentences=documents, is_query=False)
reranked_documents = rank.rerank(
documents_ids=[["1", "2"]],
queries_embeddings=queries_embeddings,
documents_embeddings=documents_embeddings,
)
assert math.isclose(
reranked_documents[0][0]["score"], 18.77, rel_tol=0.01, abs_tol=0.01
)
assert math.isclose(
reranked_documents[0][1]["score"], 18.63, rel_tol=0.01, abs_tol=0.01
)

# Creation from stanford-nlp (safetensor)
model = models.ColBERT(model_name_or_path="answerdotai/answerai-colbert-small-v1")
Expand Down

0 comments on commit 20b8061

Please sign in to comment.