Spaces:
Running
Running
| import statistics | |
| import unittest | |
| from unittest.mock import patch, MagicMock | |
| import numpy as np | |
| import torch | |
| from numpy.testing import assert_almost_equal | |
| from sentence_transformers import SentenceTransformer | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| from encoder_models import SBertEncoder, get_encoder | |
| from semf1 import SemF1, _compute_cosine_similarity, _validate_input_format | |
| from utils import get_gpu, slice_embeddings, is_nested_list_of_type, flatten_list, compute_f1, Scores | |
| class TestUtils(unittest.TestCase): | |
| def test_get_gpu(self): | |
| gpu_count = torch.cuda.device_count() | |
| gpu_available = torch.cuda.is_available() | |
| # Test single boolean input | |
| self.assertEqual(get_gpu(True), 0 if gpu_available else "cpu") | |
| self.assertEqual(get_gpu(False), "cpu") | |
| # Test single string input | |
| self.assertEqual(get_gpu("cpu"), "cpu") | |
| self.assertEqual(get_gpu("gpu"), 0 if gpu_available else "cpu") | |
| self.assertEqual(get_gpu("cuda"), 0 if gpu_available else "cpu") | |
| # Test single integer input | |
| self.assertEqual(get_gpu(0), 0 if gpu_available else "cpu") | |
| self.assertEqual(get_gpu(1), 1 if gpu_available else "cpu") | |
| # Test list input with unique elements | |
| self.assertEqual(get_gpu([True, "cpu", 0]), [0, "cpu"] if gpu_available else ["cpu", "cpu", "cpu"]) | |
| # Test list input with duplicate elements | |
| self.assertEqual(get_gpu([0, 0, "gpu"]), 0 if gpu_available else ["cpu", "cpu", "cpu"]) | |
| # Test list input with duplicate elements of different types | |
| self.assertEqual(get_gpu([True, 0, "gpu"]), 0 if gpu_available else ["cpu", "cpu", "cpu"]) | |
| # Test list input but only one element | |
| self.assertEqual(get_gpu([True]), 0 if gpu_available else "cpu") | |
| # Test list input with all integers | |
| self.assertEqual(get_gpu(list(range(gpu_count))), | |
| list(range(gpu_count)) if gpu_available else gpu_count * ["cpu"]) | |
| with self.assertRaises(ValueError): | |
| get_gpu("invalid") | |
| with self.assertRaises(ValueError): | |
| get_gpu(torch.cuda.device_count()) | |
| def test_slice_embeddings(self): | |
| embeddings = np.random.rand(10, 5) | |
| num_sentences = [3, 2, 5] | |
| expected_output = [embeddings[:3], embeddings[3:5], embeddings[5:]] | |
| self.assertTrue( | |
| all(np.array_equal(a, b) for a, b in zip(slice_embeddings(embeddings, num_sentences), | |
| expected_output)) | |
| ) | |
| num_sentences_nested = [[2, 1], [3, 4]] | |
| expected_output_nested = [[embeddings[:2], embeddings[2:3]], [embeddings[3:6], embeddings[6:]]] | |
| self.assertTrue( | |
| slice_embeddings(embeddings, num_sentences_nested), expected_output_nested | |
| ) | |
| with self.assertRaises(TypeError): | |
| slice_embeddings(embeddings, "invalid") | |
| def test_is_nested_list_of_type(self): | |
| # Test case: Depth 0, single element matching element_type | |
| self.assertTrue(is_nested_list_of_type("test", str, 0)) | |
| # Test case: Depth 0, single element not matching element_type | |
| self.assertFalse(is_nested_list_of_type("test", int, 0)) | |
| # Test case: Depth 1, list of elements matching element_type | |
| self.assertTrue(is_nested_list_of_type(["apple", "banana"], str, 1)) | |
| # Test case: Depth 1, list of elements not matching element_type | |
| self.assertFalse(is_nested_list_of_type([1, 2, 3], str, 1)) | |
| # Test case: Depth 0 (Wrong), list of elements matching element_type | |
| self.assertFalse(is_nested_list_of_type([1, 2, 3], str, 0)) | |
| # Depth 2 | |
| self.assertTrue(is_nested_list_of_type([[1, 2], [3, 4]], int, 2)) | |
| self.assertTrue(is_nested_list_of_type([['1', '2'], ['3', '4']], str, 2)) | |
| self.assertFalse(is_nested_list_of_type([[1, 2], ["a", "b"]], int, 2)) | |
| # Depth 3 | |
| self.assertFalse(is_nested_list_of_type([[[1], [2]], [[3], [4]]], list, 3)) | |
| self.assertTrue(is_nested_list_of_type([[[1], [2]], [[3], [4]]], int, 3)) | |
| with self.assertRaises(ValueError): | |
| is_nested_list_of_type([1, 2], int, -1) | |
| def test_flatten_list(self): | |
| self.assertEqual(flatten_list([1, [2, 3], [[4], 5]]), [1, 2, 3, 4, 5]) | |
| self.assertEqual(flatten_list([]), []) | |
| self.assertEqual(flatten_list([1, 2, 3]), [1, 2, 3]) | |
| self.assertEqual(flatten_list([[[[1]]]]), [1]) | |
| def test_compute_f1(self): | |
| self.assertAlmostEqual(compute_f1(0.5, 0.5), 0.5) | |
| self.assertAlmostEqual(compute_f1(1, 0), 0.0) | |
| self.assertAlmostEqual(compute_f1(0, 1), 0.0) | |
| self.assertAlmostEqual(compute_f1(1, 1), 1.0) | |
| def test_scores(self): | |
| scores = Scores(precision=0.8, recall=[0.7, 0.9]) | |
| self.assertAlmostEqual(scores.f1, compute_f1(0.8, statistics.fmean([0.7, 0.9]))) | |
| class TestSBertEncoder(unittest.TestCase): | |
| def setUp(self, device=None): | |
| if device is None: | |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| else: | |
| self.device = device | |
| self.model_name = "stsb-roberta-large" | |
| self.batch_size = 8 | |
| self.verbose = False | |
| self.encoder = SBertEncoder(self.model_name, self.device, self.batch_size, self.verbose) | |
| def test_initialization(self): | |
| self.assertIsInstance(self.encoder.model, SentenceTransformer) | |
| self.assertEqual(self.encoder.device, self.device) | |
| self.assertEqual(self.encoder.batch_size, self.batch_size) | |
| self.assertEqual(self.encoder.verbose, self.verbose) | |
| def test_encode_single_device(self): | |
| sentences = ["This is a test sentence.", "Here is another sentence."] | |
| embeddings = self.encoder.encode(sentences) | |
| self.assertIsInstance(embeddings, np.ndarray) | |
| self.assertEqual(embeddings.shape[0], len(sentences)) | |
| self.assertEqual(embeddings.shape[1], self.encoder.model.get_sentence_embedding_dimension()) | |
| def test_encode_multi_device(self): | |
| if torch.cuda.device_count() < 2: | |
| self.skipTest("Multi-GPU test requires at least 2 GPUs.") | |
| else: | |
| devices = ["cuda:0", "cuda:1"] | |
| self.setUp(devices) | |
| sentences = ["This is a test sentence.", "Here is another sentence.", "This is a test sentence."] | |
| embeddings = self.encoder.encode(sentences) | |
| self.assertIsInstance(embeddings, np.ndarray) | |
| self.assertEqual(embeddings.shape[0], 3) | |
| self.assertEqual(embeddings.shape[1], self.encoder.model.get_sentence_embedding_dimension()) | |
| class TestGetEncoder(unittest.TestCase): | |
| def setUp(self): | |
| self.device = "cuda" if torch.cuda.is_available() else "cpu" | |
| self.batch_size = 8 | |
| self.verbose = False | |
| def _base_test(self, model_name): | |
| encoder = get_encoder(model_name, self.device, self.batch_size, self.verbose) | |
| # Assert | |
| self.assertIsInstance(encoder, SBertEncoder) | |
| self.assertEqual(encoder.device, self.device) | |
| self.assertEqual(encoder.batch_size, self.batch_size) | |
| self.assertEqual(encoder.verbose, self.verbose) | |
| def test_get_sbert_encoder(self): | |
| model_name = "stsb-roberta-large" | |
| self._base_test(model_name) | |
| def test_sbert_model(self): | |
| model_name = "all-mpnet-base-v2" | |
| self._base_test(model_name) | |
| def test_huggingface_model(self): | |
| """Test Huggingface models which work with SBert library""" | |
| model_name = "roberta-base" | |
| self._base_test(model_name) | |
| def test_get_encoder_environment_error(self): # This parameter is used when using patch decorator | |
| model_name = "abc" # Wrong model_name | |
| with self.assertRaises(EnvironmentError): | |
| get_encoder(model_name, self.device, self.batch_size, self.verbose) | |
| def test_get_encoder_other_exception(self): | |
| model_name = "apple/OpenELM-270M" # This model is not supported by SentenceTransformer lib | |
| with self.assertRaises(RuntimeError): | |
| get_encoder(model_name, self.device, self.batch_size, self.verbose) | |
| class TestSemF1(unittest.TestCase): | |
| def setUp(self): | |
| self.semf1_metric = SemF1() # semf1_metric | |
| # Example cases, #Samples = 1 | |
| self.untokenized_single_reference_predictions = [ | |
| "This is a prediction sentence 1. This is a prediction sentence 2."] | |
| self.untokenized_single_reference_references = [ | |
| "This is a reference sentence 1. This is a reference sentence 2."] | |
| self.tokenized_single_reference_predictions = [ | |
| ["This is a prediction sentence 1.", "This is a prediction sentence 2."], | |
| ] | |
| self.tokenized_single_reference_references = [ | |
| ["This is a reference sentence 1.", "This is a reference sentence 2."], | |
| ] | |
| self.untokenized_multi_reference_predictions = [ | |
| "Prediction sentence 1. Prediction sentence 2." | |
| ] | |
| self.untokenized_multi_reference_references = [ | |
| ["Reference sentence 1. Reference sentence 2.", "Alternative reference 1. Alternative reference 2."], | |
| ] | |
| self.tokenized_multi_reference_predictions = [ | |
| ["Prediction sentence 1.", "Prediction sentence 2."], | |
| ] | |
| self.tokenized_multi_reference_references = [ | |
| [ | |
| ["Reference sentence 1.", "Reference sentence 2."], | |
| ["Alternative reference 1.", "Alternative reference 2."] | |
| ], | |
| ] | |
| def test_untokenized_single_reference(self): | |
| scores = self.semf1_metric.compute( | |
| predictions=self.untokenized_single_reference_predictions, | |
| references=self.untokenized_single_reference_references, | |
| tokenize_sentences=True, | |
| multi_references=False, | |
| gpu=False, | |
| batch_size=32, | |
| verbose=False | |
| ) | |
| self.assertIsInstance(scores, list) | |
| self.assertEqual(len(scores), len(self.untokenized_single_reference_predictions)) | |
| def test_tokenized_single_reference(self): | |
| scores = self.semf1_metric.compute( | |
| predictions=self.tokenized_single_reference_predictions, | |
| references=self.tokenized_single_reference_references, | |
| tokenize_sentences=False, | |
| multi_references=False, | |
| gpu=False, | |
| batch_size=32, | |
| verbose=False | |
| ) | |
| self.assertIsInstance(scores, list) | |
| self.assertEqual(len(scores), len(self.tokenized_single_reference_predictions)) | |
| for score in scores: | |
| self.assertIsInstance(score, Scores) | |
| self.assertTrue(0.0 <= score.precision <= 1.0) | |
| self.assertTrue(all(0.0 <= recall <= 1.0 for recall in score.recall)) | |
| def test_untokenized_multi_reference(self): | |
| scores = self.semf1_metric.compute( | |
| predictions=self.untokenized_multi_reference_predictions, | |
| references=self.untokenized_multi_reference_references, | |
| tokenize_sentences=True, | |
| multi_references=True, | |
| gpu=False, | |
| batch_size=32, | |
| verbose=False | |
| ) | |
| self.assertIsInstance(scores, list) | |
| self.assertEqual(len(scores), len(self.untokenized_multi_reference_predictions)) | |
| def test_tokenized_multi_reference(self): | |
| scores = self.semf1_metric.compute( | |
| predictions=self.tokenized_multi_reference_predictions, | |
| references=self.tokenized_multi_reference_references, | |
| tokenize_sentences=False, | |
| multi_references=True, | |
| gpu=False, | |
| batch_size=32, | |
| verbose=False | |
| ) | |
| self.assertIsInstance(scores, list) | |
| self.assertEqual(len(scores), len(self.tokenized_multi_reference_predictions)) | |
| for score in scores: | |
| self.assertIsInstance(score, Scores) | |
| self.assertTrue(0.0 <= score.precision <= 1.0) | |
| self.assertTrue(all(0.0 <= recall <= 1.0 for recall in score.recall)) | |
| def test_same_predictions_and_references(self): | |
| scores = self.semf1_metric.compute( | |
| predictions=self.tokenized_single_reference_predictions, | |
| references=self.tokenized_single_reference_predictions, | |
| tokenize_sentences=False, | |
| multi_references=False, | |
| gpu=False, | |
| batch_size=32, | |
| verbose=False | |
| ) | |
| self.assertIsInstance(scores, list) | |
| self.assertEqual(len(scores), len(self.tokenized_single_reference_predictions)) | |
| for score in scores: | |
| self.assertIsInstance(score, Scores) | |
| self.assertAlmostEqual(score.precision, 1.0, places=6) | |
| assert_almost_equal(score.recall, 1, decimal=5, err_msg="Not all values are almost equal to 1") | |
| def test_exact_output_scores(self): | |
| predictions = [ | |
| ["I go to School.", "You are stupid."], | |
| ["I love adventure sports."], | |
| ] | |
| references = [ | |
| ["I go to playground.", "You are genius.", "You need to be admired."], | |
| ["I love adventure sports."], | |
| ] | |
| scores = self.semf1_metric.compute( | |
| predictions=predictions, | |
| references=references, | |
| tokenize_sentences=False, | |
| multi_references=False, | |
| gpu=False, | |
| batch_size=32, | |
| verbose=False, | |
| model_type="use", | |
| ) | |
| self.assertIsInstance(scores, list) | |
| self.assertEqual(len(scores), len(predictions)) | |
| score = scores[0] | |
| self.assertIsInstance(score, Scores) | |
| self.assertAlmostEqual(score.precision, 0.73, places=2) | |
| self.assertAlmostEqual(score.recall[0], 0.63, places=2) | |
| class TestCosineSimilarity(unittest.TestCase): | |
| def setUp(self): | |
| # Sample embeddings for testing | |
| self.pred_embeds = np.array([ | |
| [1, 0, 0], | |
| [0, 1, 0], | |
| [0, 0, 1] | |
| ]) | |
| self.ref_embeds = np.array([ | |
| [1, 0, 0], | |
| [0, 1, 0], | |
| [0, 0, 1] | |
| ]) | |
| self.pred_embeds_random = np.random.rand(3, 3) | |
| self.ref_embeds_random = np.random.rand(3, 3) | |
| def test_cosine_similarity_perfect_match(self): | |
| precision, recall = _compute_cosine_similarity(self.pred_embeds, self.ref_embeds) | |
| # Expected values are 1.0 for both precision and recall since embeddings are identical | |
| self.assertAlmostEqual(precision, 1.0, places=5) | |
| self.assertAlmostEqual(recall, 1.0, places=5) | |
| def _test_cosine_similarity_base(self, pred_embeds, ref_embeds): | |
| precision, recall = _compute_cosine_similarity(pred_embeds, ref_embeds) | |
| # Calculate expected precision and recall using sklearn's cosine similarity function | |
| cosine_scores = cosine_similarity(pred_embeds, ref_embeds) | |
| expected_precision = np.mean(np.max(cosine_scores, axis=-1)).item() | |
| expected_recall = np.mean(np.max(cosine_scores, axis=0)).item() | |
| self.assertAlmostEqual(precision, expected_precision, places=5) | |
| self.assertAlmostEqual(recall, expected_recall, places=5) | |
| def test_cosine_similarity_random(self): | |
| self._test_cosine_similarity_base(self.pred_embeds_random, self.ref_embeds_random) | |
| def test_cosine_similarity_different_shapes(self): | |
| pred_embeds_diff = np.random.rand(5, 3) | |
| ref_embeds_diff = np.random.rand(3, 3) | |
| self._test_cosine_similarity_base(pred_embeds_diff, ref_embeds_diff) | |
| class TestValidateInputFormat(unittest.TestCase): | |
| def setUp(self): | |
| # Sample predictions and references for different scenarios where number of samples = 1 | |
| # Note: Naming Convention: # When tokenize_sentences = True (i.e. input is untokenized) and vice-versa | |
| # When tokenize_sentences = True (untokenized input) and multi_references = False | |
| self.untokenized_single_reference_predictions = [ | |
| "This is a prediction sentence 1. This is a prediction sentence 2." | |
| ] | |
| self.untokenized_single_reference_references = [ | |
| "This is a reference sentence 1. This is a reference sentence 2." | |
| ] | |
| # When tokenize_sentences = False (tokenized input) and multi_references = False | |
| self.tokenized_single_reference_predictions = [ | |
| ["This is a prediction sentence 1.", "This is a prediction sentence 2."] | |
| ] | |
| self.tokenized_single_reference_references = [ | |
| ["This is a reference sentence 1.", "This is a reference sentence 2."] | |
| ] | |
| # When tokenize_sentences = True (untokenized input) and multi_references = True | |
| self.untokenized_multi_reference_predictions = [ | |
| "This is a prediction sentence 1. This is a prediction sentence 2." | |
| ] | |
| self.untokenized_multi_reference_references = [ | |
| [ | |
| "This is a reference sentence 1. This is a reference sentence 2.", | |
| "Another reference sentence." | |
| ] | |
| ] | |
| # When tokenize_sentences = False (tokenized input) and multi_references = True | |
| self.tokenized_multi_reference_predictions = [ | |
| ["This is a prediction sentence 1.", "This is a prediction sentence 2."] | |
| ] | |
| self.tokenized_multi_reference_references = [ | |
| [ | |
| ["This is a reference sentence 1.", "This is a reference sentence 2."], | |
| ["Another reference sentence."] | |
| ] | |
| ] | |
| def test_tokenized_sentences_true_multi_references_true(self): | |
| # Invalid format should raise an error | |
| with self.assertRaises(ValueError): | |
| _validate_input_format( | |
| True, | |
| True, | |
| self.tokenized_single_reference_predictions, | |
| self.tokenized_single_reference_references, | |
| ) | |
| # Valid format should pass without error | |
| _validate_input_format( | |
| True, | |
| True, | |
| self.untokenized_multi_reference_predictions, | |
| self.untokenized_multi_reference_references, | |
| ) | |
| def test_tokenized_sentences_false_multi_references_true(self): | |
| # Invalid format should raise an error | |
| with self.assertRaises(ValueError): | |
| _validate_input_format( | |
| False, | |
| True, | |
| self.untokenized_single_reference_predictions, | |
| self.untokenized_multi_reference_references, | |
| ) | |
| # Valid format should pass without error | |
| _validate_input_format( | |
| False, | |
| True, | |
| self.tokenized_multi_reference_predictions, | |
| self.tokenized_multi_reference_references, | |
| ) | |
| def test_tokenized_sentences_true_multi_references_false(self): | |
| # Invalid format should raise an error | |
| with self.assertRaises(ValueError): | |
| _validate_input_format( | |
| True, | |
| False, | |
| self.tokenized_single_reference_predictions, | |
| self.tokenized_single_reference_references, | |
| ) | |
| # Valid format should pass without error | |
| _validate_input_format( | |
| True, | |
| False, | |
| self.untokenized_single_reference_predictions, | |
| self.untokenized_single_reference_references, | |
| ) | |
| def test_tokenized_sentences_false_multi_references_false(self): | |
| # Invalid format should raise an error | |
| with self.assertRaises(ValueError): | |
| _validate_input_format( | |
| False, | |
| False, | |
| self.untokenized_single_reference_predictions, | |
| self.untokenized_single_reference_references, | |
| ) | |
| # Valid format should pass without error | |
| _validate_input_format( | |
| False, | |
| False, | |
| self.tokenized_single_reference_predictions, | |
| self.tokenized_single_reference_references, | |
| ) | |
| def test_mismatched_lengths(self): | |
| # Length mismatch should raise an error | |
| with self.assertRaises(ValueError): | |
| _validate_input_format( | |
| True, | |
| True, | |
| self.untokenized_single_reference_predictions, | |
| [self.untokenized_single_reference_predictions[0], self.untokenized_single_reference_predictions[0]], | |
| ) | |
| if __name__ == '__main__': | |
| unittest.main(verbosity=2) | |
| # unittest.main() | |