From def1b25056557cd482c1e7b3d82be2b226ea6dd9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9=20Pedersen?= <andrped94@gmail.com>
Date: Wed, 13 Mar 2024 18:56:21 +0100
Subject: [PATCH] Add test for verifying that top_k selecting works
---
tests/unit/test_hybrid_layer.py | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/tests/unit/test_hybrid_layer.py b/tests/unit/test_hybrid_layer.py
index 18bdbddf..52eb6b3e 100644
--- a/tests/unit/test_hybrid_layer.py
+++ b/tests/unit/test_hybrid_layer.py
@@ -81,11 +81,12 @@ sparse_encoder.fit(["The quick brown fox", "jumps over the lazy dog", "Hello, wo
class TestHybridRouteLayer:
def test_initialization(self, openai_encoder, routes):
route_layer = HybridRouteLayer(
- encoder=openai_encoder, sparse_encoder=sparse_encoder, routes=routes
+ encoder=openai_encoder, sparse_encoder=sparse_encoder, routes=routes, top_k=10,
)
assert route_layer.index is not None and route_layer.categories is not None
assert openai_encoder.score_threshold == 0.82
assert route_layer.score_threshold == 0.82
+ assert route_layer.top_k == 10
assert len(route_layer.index) == 5
assert len(set(route_layer.categories)) == 2
--
GitLab