Update test_lora.py

Author: riapush

tests/modules/scoring/test_lora.py

@@ -1,16 +1,65 @@
+import shutil
+import tempfile
+from pathlib import Path
+
 import numpy as np
 import pytest
 
-from autointent._dump_tools import Dumper
 from autointent.context.data_handler import DataHandler
 from autointent.modules import BERTLoRAScorer
 
 
+def test_lora_scorer_dump_load(dataset):
+    """Test that BERTLoRAScorer can be saved and loaded while preserving predictions."""
+    data_handler = DataHandler(dataset)
+
+    # Create and train scorer
+    scorer_original = BERTLoRAScorer(classification_model_config="prajjwal1/bert-tiny", num_train_epochs=1, batch_size=8)
+    scorer_original.fit(data_handler.train_utterances(0), data_handler.train_labels(0))
+
+    # Test data
+    test_data = [
+        "why is there a hold on my account",
+        "why is my bank account frozen",
+    ]
+
+    # Get predictions before saving
+    predictions_before = scorer_original.predict(test_data)
+
+    # Create temp directory and save model
+    temp_dir_path = Path(tempfile.mkdtemp(prefix="lora_scorer_test_"))
+    try:
+        # Save the model
+        scorer_original.dump(str(temp_dir_path))
+
+        # Create a new scorer and load saved model
+        scorer_loaded = BERTLoRAScorer(classification_model_config="prajjwal1/bert-tiny", num_train_epochs=1, batch_size=8)
+        scorer_loaded.load(str(temp_dir_path))
+
+        # Verify model and tokenizer are loaded
+        assert hasattr(scorer_loaded, "_model")
+        assert scorer_loaded._model is not None
+        assert hasattr(scorer_loaded, "_tokenizer")
+        assert scorer_loaded._tokenizer is not None
+
+        # Get predictions after loading
+        predictions_after = scorer_loaded.predict(test_data)
+
+        # Verify predictions match
+        assert predictions_before.shape == predictions_after.shape
+        np.testing.assert_allclose(predictions_before, predictions_after, atol=1e-6)
+
+    finally:
+        # Clean up
+        shutil.rmtree(temp_dir_path, ignore_errors=True)  # workaround for windows permission error
+
+
 def test_lora_prediction(dataset):
-    """Test that the transformer model can fit and make predictions."""
+    """Test that the lora model can fit and make predictions."""
     data_handler = DataHandler(dataset)
 
     scorer = BERTLoRAScorer(classification_model_config="prajjwal1/bert-tiny", num_train_epochs=1, batch_size=8)
+
     scorer.fit(data_handler.train_utterances(0), data_handler.train_labels(0))
 
     test_data = [
@@ -21,49 +70,47 @@ def test_lora_prediction(dataset):
         "can you tell me why is my bank account frozen",
     ]
 
-    initial_predictions = scorer.predict(test_data)
+    predictions = scorer.predict(test_data)
 
-    import tempfile
-    with tempfile.TemporaryDirectory() as tmpdir:
-        import pathlib
-        dump_path = pathlib.Path(tmpdir)
-        Dumper.dump(scorer, dump_path)
+    # Verify prediction shape
+    assert predictions.shape[0] == len(test_data)
+    assert predictions.shape[1] == len(set(data_handler.train_labels(0)))
 
-        new_scorer = BERTLoRAScorer(classification_model_config="prajjwal1/bert-tiny", num_train_epochs=1, batch_size=8)
-        Dumper.load(new_scorer, dump_path)
-
-        loaded_predictions = new_scorer.predict(test_data)
-        np.testing.assert_array_almost_equal(initial_predictions, loaded_predictions, decimal=5)
-
-    assert initial_predictions.shape[0] == len(test_data)
-    assert initial_predictions.shape[1] == len(set(data_handler.train_labels(0)))
-
-    assert 0.0 <= np.min(initial_predictions) <= np.max(initial_predictions) <= 1.0
+    # Verify predictions are probabilities
+    assert 0.0 <= np.min(predictions) <= np.max(predictions) <= 1.0
 
+    # Verify probabilities sum to 1 for multiclass
     if not scorer._multilabel:
-        for pred_row in initial_predictions:
+        for pred_row in predictions:
             np.testing.assert_almost_equal(np.sum(pred_row), 1.0, decimal=5)
 
+    # Test metadata function if available
     if hasattr(scorer, "predict_with_metadata"):
         predictions, metadata = scorer.predict_with_metadata(test_data)
         assert len(predictions) == len(test_data)
         assert metadata is None
 
 
-def test_bert_cache_clearing(dataset):
-    """Test that the transformer model properly handles cache clearing."""
+def test_lora_cache_clearing(dataset):
+    """Test that the lora model properly handles cache clearing."""
     data_handler = DataHandler(dataset)
 
     scorer = BERTLoRAScorer(classification_model_config="prajjwal1/bert-tiny", num_train_epochs=1, batch_size=8)
+
     scorer.fit(data_handler.train_utterances(0), data_handler.train_labels(0))
 
     test_data = ["test text"]
 
+    # Should work before clearing cache
     scorer.predict(test_data)
+
+    # Clear the cache
     scorer.clear_cache()
 
+    # Verify model and tokenizer are removed
     assert not hasattr(scorer, "_model") or scorer._model is None
     assert not hasattr(scorer, "_tokenizer") or scorer._tokenizer is None
 
+    # Should raise exception after clearing cache
     with pytest.raises(RuntimeError):
-        scorer.predict(test_data)
+        scorer.predict(test_data)