tests: change some tests

Author: FelSiq

tests/paths.py

@@ -6,8 +6,8 @@
 class TestPaths(t.NamedTuple):
     """Resource paths for tests."""
 
-    model_bert: str = "2_layer_6000_vocab_size_bert"
-    model_lstm: str = "512_hidden_dim_6000_vocab_size_1_layer_lstm"
+    model_bert: str = "2_layer_6000_vocab_size_bert_v2"
+    model_lstm: str = "256_hidden_dim_6000_vocab_size_1_layer_lstm_v2"
     tokenizer: str = "6000_subword_tokenizer"
     legal_text_long: str = "tests/resources/test_legal_text_long.txt"
     legal_text_short: str = "tests/resources/test_legal_text_short.txt"

tests/test_segmentation.py

@@ -39,7 +39,7 @@ def test_inference_pooling_operation_argument_with_short_text_and_lstm(
         uri_tokenizer=fixture_test_paths.tokenizer,
         inference_pooling_operation=pooling_operation,
         device="cpu",
-        lstm_hidden_layer_size=512,
+        lstm_hidden_layer_size=256,
         lstm_num_layers=1,
         local_files_only=False,
         cache_dir_model=fixture_test_paths.cache_dir_models,
@@ -111,7 +111,7 @@ def test_batch_size_with_long_text(
     batch_size: int,
 ):
     segs = fixture_model_bert_2_layers(fixture_legal_text_long, batch_size=batch_size)
-    assert len(segs) == 63 and no_segmentation_at_middle_subwords(segs)
+    assert len(segs) >= 60 and no_segmentation_at_middle_subwords(segs)
 
 
 @pytest.mark.parametrize("window_shift_size", (1024, 512, 256, 1.0, 0.5, 0.25))
@@ -121,7 +121,7 @@ def test_window_shift_size(
     window_shift_size: int,
 ):
     segs = fixture_model_bert_2_layers(fixture_legal_text_long, window_shift_size=window_shift_size)
-    assert len(segs) >= 59 and no_segmentation_at_middle_subwords(segs)
+    assert len(segs) >= 50 and no_segmentation_at_middle_subwords(segs)
 
 
 @pytest.mark.parametrize("input_type_fn", [tuple, list, pd.Series])