This commit is contained in:
ayush chaurasia
2024-02-21 21:43:23 +05:30
parent 22c196b3e3
commit bd4e8341fe
3 changed files with 58 additions and 29 deletions

View File

@@ -102,7 +102,7 @@ def test_linear_combination(tmp_path):
query = "Our father who art in heaven"
query_vector = table.to_pandas()["vector"][0]
result = (
table.search((query_vector, query))
table.search(vector=query_vector, text=query, query_type="vector")
.limit(30)
.rerank(normalize="score")
.to_arrow()
@@ -116,6 +116,13 @@ def test_linear_combination(tmp_path):
"be descending."
)
result = (
table.search(vector=query_vector, text=query)
.limit(30)
.rerank(normalize="score")
.to_arrow()
)
@pytest.mark.skipif(
os.environ.get("COHERE_API_KEY") is None, reason="COHERE_API_KEY not set"
@@ -139,7 +146,7 @@ def test_cohere_reranker(tmp_path):
query = "Our father who art in heaven"
query_vector = table.to_pandas()["vector"][0]
result = (
table.search((query_vector, query))
table.search(vector=query_vector, text=query)
.limit(30)
.rerank(reranker=CohereReranker())
.to_arrow()
@@ -173,7 +180,7 @@ def test_cross_encoder_reranker(tmp_path):
query = "Our father who art in heaven"
query_vector = table.to_pandas()["vector"][0]
result = (
table.search((query_vector, query), query_type="hybrid")
table.search(vector=query_vector, text=query, query_type="hybrid")
.limit(30)
.rerank(reranker=CrossEncoderReranker())
.to_arrow()
@@ -207,7 +214,7 @@ def test_colbert_reranker(tmp_path):
query = "Our father who art in heaven"
query_vector = table.to_pandas()["vector"][0]
result = (
table.search((query_vector, query))
table.search(vector=query_vector, text=query)
.limit(30)
.rerank(reranker=ColbertReranker())
.to_arrow()
@@ -244,7 +251,7 @@ def test_openai_reranker(tmp_path):
query = "Our father who art in heaven"
query_vector = table.to_pandas()["vector"][0]
result = (
table.search((query_vector, query))
table.search(vector=query_vector, text=query)
.limit(30)
.rerank(reranker=OpenaiReranker())
.to_arrow()