Skip to content
Snippets Groups Projects
Commit def1b250 authored by André Pedersen's avatar André Pedersen
Browse files

Add test for verifying that top_k selecting works

parent a2f363f8
No related branches found
No related tags found
No related merge requests found
...@@ -81,11 +81,12 @@ sparse_encoder.fit(["The quick brown fox", "jumps over the lazy dog", "Hello, wo ...@@ -81,11 +81,12 @@ sparse_encoder.fit(["The quick brown fox", "jumps over the lazy dog", "Hello, wo
class TestHybridRouteLayer: class TestHybridRouteLayer:
def test_initialization(self, openai_encoder, routes): def test_initialization(self, openai_encoder, routes):
route_layer = HybridRouteLayer( 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 route_layer.index is not None and route_layer.categories is not None
assert openai_encoder.score_threshold == 0.82 assert openai_encoder.score_threshold == 0.82
assert route_layer.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(route_layer.index) == 5
assert len(set(route_layer.categories)) == 2 assert len(set(route_layer.categories)) == 2
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment