bug fix

Author: voorhs

autointent/modules/embedding/_retrieval.py

@@ -150,5 +150,5 @@ def predict(self, utterances: list[str]) -> list[ListOfLabels]:
         Returns:
             List of labels for each retrieved utterance
         """
-        predictions, _ = self._vector_index.query(utterances, self.k)
-        return predictions
+        _, documents = self._vector_index.query(utterances, self.k)
+        return [[n.label for n in neigs] for neigs in documents]

tests/callback/test_callback.py

@@ -5,7 +5,7 @@
 
 from autointent import Context, Pipeline
 from autointent._callbacks import CallbackHandler, OptimizerCallback
-from autointent.configs import DataConfig, HPOConfig, LoggingConfig
+from autointent.configs import DataConfig, FaissConfig, HPOConfig, LoggingConfig
 from tests.conftest import setup_environment
 
 
@@ -96,6 +96,7 @@ def test_pipeline_callbacks(dataset):
     context.callback_handler = CallbackHandler([DummyCallback])
     context.set_dataset(dataset, DataConfig(scheme="ho"))
     context.configure_hpo(HPOConfig(n_trials=10))
+    context.configure_vector_index(FaissConfig())
 
     pipeline_optimizer._fit(context)
 

tests/context/test_vector_index.py

@@ -1,7 +1,7 @@
 import pytest
 
 from autointent import VectorIndex
-from autointent.configs import EmbedderConfig
+from autointent.configs import EmbedderConfig, FaissConfig
 
 
 @pytest.fixture
@@ -14,7 +14,7 @@ class MockDataHandler:
 
 
 def test_create_collection(data_handler):
-    vector_index = VectorIndex(embedder_config=EmbedderConfig(model_name="bert-base-uncased"))
+    vector_index = VectorIndex(embedder_config=EmbedderConfig(model_name="bert-base-uncased"), config=FaissConfig())
     vector_index.add(
         data_handler.utterances_train,
         data_handler.labels_train,

tests/modules/test_dumper.py

@@ -52,9 +52,10 @@ def check_attributes(self):
         tokenizer_predictions = self.tokenizer(["hello", "world"]).input_ids
         np.testing.assert_array_equal(self._tokenizer_predictions, tokenizer_predictions)
         with torch.no_grad():
-            np.testing.assert_array_equal(
+            np.testing.assert_almost_equal(
                 self._transformer_predictions,
                 self.transformer(input_ids=torch.tensor(tokenizer_predictions)).logits.cpu().numpy(),
+                decimal=4,
             )