BridgeAI-Lab/Sem-nCG
1
1import unittest2 3import numpy as np4import torch5from sentence_transformers import SentenceTransformer6 7from .encoder_models import SBertEncoder, get_encoder, get_sbert_encoder8from .semncg import (9 RankedGains,10 compute_cosine_similarity,11 compute_gain,12 score_ncg,13 compute_ncg,14 _validate_input_format,15 SemNCG16)17from .utils import (18 get_gpu,19 slice_embeddings,20 is_nested_list_of_type,21 flatten_list,22 prep_sentences,23 tokenize_and_prep_document24)25 26 27class TestUtils(unittest.TestCase):28 def test_get_gpu(self):29 gpu_count = torch.cuda.device_count()30 gpu_available = torch.cuda.is_available()31 32 # Test single boolean input33 self.assertEqual(get_gpu(True), 0 if gpu_available else "cpu")34 self.assertEqual(get_gpu(False), "cpu")35 36 # Test single string input37 self.assertEqual(get_gpu("cpu"), "cpu")38 self.assertEqual(get_gpu("gpu"), 0 if gpu_available else "cpu")39 self.assertEqual(get_gpu("cuda"), 0 if gpu_available else "cpu")40 41 # Test single integer input42 self.assertEqual(get_gpu(0), 0 if gpu_available else "cpu")43 self.assertEqual(get_gpu(1), 1 if gpu_available else "cpu")44 45 # Test list input with unique elements46 self.assertEqual(get_gpu([True, "cpu", 0]), [0, "cpu"] if gpu_available else ["cpu", "cpu", "cpu"])47 48 # Test list input with duplicate elements49 self.assertEqual(get_gpu([0, 0, "gpu"]), 0 if gpu_available else ["cpu", "cpu", "cpu"])50 51 # Test list input with duplicate elements of different types52 self.assertEqual(get_gpu([True, 0, "gpu"]), 0 if gpu_available else ["cpu", "cpu", "cpu"])53 54 # Test list input but only one element55 self.assertEqual(get_gpu([True]), 0 if gpu_available else "cpu")56 57 # Test list input with all integers58 self.assertEqual(get_gpu(list(range(gpu_count))),59 list(range(gpu_count)) if gpu_available else gpu_count * ["cpu"])60 61 with self.assertRaises(ValueError):62 get_gpu("invalid")63 64 with self.assertRaises(ValueError):65 get_gpu(torch.cuda.device_count())66 67 def test_prep_sentences(self):68 # Test normal case69 self.assertEqual(prep_sentences(["Hello, world!", " This is a test. ", "!!!"]),70 ['Hello, world!', 'This is a test.'])71 72 # Test case with only punctuations73 with self.assertRaises(ValueError):74 prep_sentences(["!!!", "..."])75 76 # Test case with empty list77 with self.assertRaises(ValueError):78 prep_sentences([])79 80 def test_tokenize_and_prep_document(self):81 # Test tokenize=True with string input82 self.assertEqual(tokenize_and_prep_document("Hello, world! This is a test.", True),83 ['Hello, world!', 'This is a test.'])84 85 # Test tokenize=False with list of strings input86 self.assertEqual(tokenize_and_prep_document(["Hello, world!", "This is a test."], False),87 ['Hello, world!', 'This is a test.'])88 89 # Test tokenize=True with empty document90 with self.assertRaises(ValueError):91 tokenize_and_prep_document("!!! ...", True)92 93 def test_slice_embeddings(self):94 # Case 195 embeddings = np.random.rand(10, 5)96 num_sentences = [3, 2, 5]97 expected_output = [embeddings[:3], embeddings[3:5], embeddings[5:]]98 self.assertTrue(99 all(np.array_equal(a, b) for a, b in zip(slice_embeddings(embeddings, num_sentences),100 expected_output))101 )102 103 # Case 2104 num_sentences_nested = [[2, 1], [3, 4]]105 expected_output_nested = [[embeddings[:2], embeddings[2:3]], [embeddings[3:6], embeddings[6:]]]106 self.assertTrue(107 slice_embeddings(embeddings, num_sentences_nested), expected_output_nested108 )109 110 # Case 3111 document_sentences_count = [10, 8, 7]112 reference_sentences_count = [5, 3, 2]113 pred_sentences_count = [2, 2, 1]114 all_embeddings = np.random.rand(115 sum(document_sentences_count + reference_sentences_count + pred_sentences_count), 5,116 )117 118 embeddings = all_embeddings119 expected_doc_embeddings = [embeddings[:10], embeddings[10:18], embeddings[18:25]]120 121 embeddings = all_embeddings[25:]122 expected_ref_embeddings = [embeddings[:5], embeddings[5:8], embeddings[8:10]]123 124 embeddings = all_embeddings[35:]125 expected_pred_embeddings = [embeddings[:2], embeddings[2:4], embeddings[4:5]]126 127 doc_embeddings = slice_embeddings(all_embeddings, document_sentences_count)128 ref_embeddings = slice_embeddings(all_embeddings[sum(document_sentences_count):], reference_sentences_count)129 pred_embeddings = slice_embeddings(130 all_embeddings[sum(document_sentences_count + reference_sentences_count):], pred_sentences_count131 )132 133 self.assertTrue(doc_embeddings, expected_doc_embeddings)134 self.assertTrue(ref_embeddings, expected_ref_embeddings)135 self.assertTrue(pred_embeddings, expected_pred_embeddings)136 137 with self.assertRaises(TypeError):138 slice_embeddings(embeddings, "invalid")139 140 def test_is_nested_list_of_type(self):141 # Test case: Depth 0, single element matching element_type142 self.assertEqual(is_nested_list_of_type("test", str, 0), (True, ""))143 144 # Test case: Depth 0, single element not matching element_type145 is_valid, err_msg = is_nested_list_of_type("test", int, 0)146 self.assertEqual(is_valid, False)147 148 # Test case: Depth 1, list of elements matching element_type149 self.assertEqual(is_nested_list_of_type(["apple", "banana"], str, 1), (True, ""))150 151 # Test case: Depth 1, list of elements not matching element_type152 is_valid, err_msg = is_nested_list_of_type([1, 2, 3], str, 1)153 self.assertEqual(is_valid, False)154 155 # Test case: Depth 0 (Wrong), list of elements matching element_type156 is_valid, err_msg = is_nested_list_of_type([1, 2, 3], str, 0)157 self.assertEqual(is_valid, False)158 159 # Depth 2160 self.assertEqual(is_nested_list_of_type([[1, 2], [3, 4]], int, 2), (True, ""))161 self.assertEqual(is_nested_list_of_type([['1', '2'], ['3', '4']], str, 2), (True, ""))162 is_valid, err_msg = is_nested_list_of_type([[1, 2], ["a", "b"]], int, 2)163 self.assertEqual(is_valid, False)164 165 # Depth 3166 is_valid, err_msg = is_nested_list_of_type([[[1], [2]], [[3], [4]]], list, 3)167 self.assertEqual(is_valid, False)168 self.assertEqual(is_nested_list_of_type([[[1], [2]], [[3], [4]]], int, 3), (True, ""))169 170 # Test case: Depth is negative, expecting ValueError171 with self.assertRaises(ValueError):172 is_nested_list_of_type([1, 2], int, -1)173 174 def test_flatten_list(self):175 self.assertEqual(flatten_list([1, [2, 3], [[4], 5]]), [1, 2, 3, 4, 5])176 self.assertEqual(flatten_list([]), [])177 self.assertEqual(flatten_list([1, 2, 3]), [1, 2, 3])178 self.assertEqual(flatten_list([[[[1]]]]), [1])179 180 181class TestSBertEncoder(unittest.TestCase):182 183 def setUp(self) -> None:184 # Set up a test SentenceTransformer model185 self.model_name = "paraphrase-distilroberta-base-v1"186 self.sbert_model = get_sbert_encoder(self.model_name)187 self.device = "cpu" # For testing on CPU188 self.batch_size = 32189 self.verbose = False190 self.encoder = SBertEncoder(self.sbert_model, self.device, self.batch_size, self.verbose)191 192 def test_encode_single_sentence(self):193 sentence = "Hello, world!"194 embeddings = self.encoder.encode([sentence])195 self.assertEqual(embeddings.shape, (1, 768)) # Adjust shape based on your model's embedding dimension196 197 def test_encode_multiple_sentences(self):198 sentences = ["Hello, world!", "This is a test."]199 embeddings = self.encoder.encode(sentences)200 self.assertEqual(embeddings.shape, (2, 768)) # Adjust shape based on your model's embedding dimension201 202 def test_get_sbert_encoder(self):203 model_name = "paraphrase-distilroberta-base-v1"204 sbert_model = get_sbert_encoder(model_name)205 self.assertIsInstance(sbert_model, SentenceTransformer)206 207 def test_encode_with_gpu(self):208 if torch.cuda.is_available():209 device = "cuda"210 encoder = get_encoder(self.sbert_model, device, self.batch_size, self.verbose)211 sentences = ["Hello, world!", "This is a test."]212 embeddings = encoder.encode(sentences)213 self.assertEqual(embeddings.shape, (2, 768)) # Adjust shape based on your model's embedding dimension214 else:215 self.skipTest("CUDA not available, skipping GPU test.")216 217 def test_encode_multi_device(self):218 if torch.cuda.device_count() < 2:219 self.skipTest("Multi-GPU test requires at least 2 GPUs.")220 else:221 devices = ["cuda:0", "cuda:1"]222 encoder = get_encoder(self.sbert_model, devices, self.batch_size, self.verbose)223 sentences = ["This is a test sentence.", "Here is another sentence.", "This is a test sentence."]224 embeddings = encoder.encode(sentences)225 self.assertIsInstance(embeddings, np.ndarray)226 self.assertEqual(embeddings.shape[0], 3)227 self.assertEqual(embeddings.shape[1], self.encoder.model.get_sentence_embedding_dimension())228 229 230class TestGetEncoder(unittest.TestCase):231 def setUp(self):232 self.device = "cuda" if torch.cuda.is_available() else "cpu"233 self.batch_size = 8234 self.verbose = False235 236 def _base_test(self, model_name):237 sbert_model = get_sbert_encoder(model_name)238 encoder = get_encoder(sbert_model, self.device, self.batch_size, self.verbose)239 240 # Assert241 self.assertIsInstance(encoder, SBertEncoder)242 self.assertEqual(encoder.device, self.device)243 self.assertEqual(encoder.batch_size, self.batch_size)244 self.assertEqual(encoder.verbose, self.verbose)245 246 def test_get_sbert_encoder(self):247 model_name = "stsb-roberta-large"248 self._base_test(model_name)249 250 def test_sbert_model(self):251 model_name = "all-mpnet-base-v2"252 self._base_test(model_name)253 254 def test_huggingface_model(self):255 """Test Huggingface models which work with SBert library"""256 model_name = "roberta-base"257 self._base_test(model_name)258 259 def test_get_encoder_environment_error(self): # This parameter is used when using patch decorator260 model_name = "abc" # Wrong model_name261 with self.assertRaises(EnvironmentError):262 get_sbert_encoder(model_name)263 264 def test_get_encoder_other_exception(self):265 model_name = "apple/OpenELM-270M" # This model is not supported by SentenceTransformer lib266 with self.assertRaises(RuntimeError):267 get_sbert_encoder(model_name)268 269 270class TestRankedGainsDataclass(unittest.TestCase):271 def test_ranked_gains_dataclass(self):272 # Test initialization and attribute access273 gt_gains = [("doc1", 0.8), ("doc2", 0.6)]274 pred_gains = [("doc2", 0.7), ("doc1", 0.5)]275 k = 2276 ncg = 0.75277 ranked_gains = RankedGains(gt_gains, pred_gains, k, ncg)278 279 self.assertEqual(ranked_gains.gt_gains, gt_gains)280 self.assertEqual(ranked_gains.pred_gains, pred_gains)281 self.assertEqual(ranked_gains.k, k)282 self.assertEqual(ranked_gains.ncg, ncg)283 284 285class TestComputeCosineSimilarity(unittest.TestCase):286 def test_compute_cosine_similarity(self):287 doc_embeds = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]])288 ref_embeds = np.array([[0.2, 0.3, 0.4], [0.5, 0.6, 0.7]])289 # Test compute_cosine_similarity function290 similarity_scores = compute_cosine_similarity(doc_embeds, ref_embeds)291 print(similarity_scores)292 293 # Example values, change as per actual function output294 expected_scores = [0.980, 0.997]295 296 self.assertAlmostEqual(similarity_scores[0], expected_scores[0], places=3)297 self.assertAlmostEqual(similarity_scores[1], expected_scores[1], places=3)298 299 300class TestComputeGain(unittest.TestCase):301 def test_compute_gain(self):302 # Test compute_gain function303 sim_scores = [0.8, 0.6, 0.7]304 gains = compute_gain(sim_scores)305 print(gains)306 307 # Example values, change as per actual function output308 expected_gains = [(0, 0.5), (2, 0.3333333333333333), (1, 0.16666666666666666)]309 310 self.assertEqual(gains, expected_gains)311 312 313class TestScoreNcg(unittest.TestCase):314 def test_score_ncg(self):315 # Test score_ncg function316 model_relevance = [0.8, 0.7, 0.6]317 gt_relevance = [1.0, 0.9, 0.8]318 ncg_score = score_ncg(model_relevance, gt_relevance)319 expected_ncg = 0.778 # Example value, change as per actual function output320 321 self.assertAlmostEqual(ncg_score, expected_ncg, places=3)322 323 324class TestComputeNcg(unittest.TestCase):325 def test_compute_ncg(self):326 # Test compute_ncg function327 pred_gains = [(0, 0.8), (2, 0.7), (1, 0.6)]328 gt_gains = [(0, 1.0), (1, 0.9), (2, 0.8)]329 k = 3330 ncg_score = compute_ncg(pred_gains, gt_gains, k)331 expected_ncg = 1.0 # TODO: Confirm this with Dr. Santu332 333 self.assertAlmostEqual(ncg_score, expected_ncg, places=6)334 335 336class TestValidateInputFormat(unittest.TestCase):337 def test_validate_input_format(self):338 # Test _validate_input_format function339 tokenize_sentences = True340 predictions = ["Prediction 1", "Prediction 2"]341 references = ["Reference 1", "Reference 2"]342 documents = ["Document 1", "Document 2"]343 344 # No exception should be raised for valid input345 try:346 _validate_input_format(tokenize_sentences, predictions, references, documents)347 except ValueError as e:348 self.fail(f"_validate_input_format raised ValueError unexpectedly: {str(e)}")349 350 # Test invalid input format351 predictions_invalid = [["Sentence 1 in prediction 1.", "Sentence 2 in prediction 1."],352 ["Sentence 1 in prediction 2.", "Sentence 2 in prediction 2."]]353 references_invalid = [["Sentences in reference 1."], ["Sentences in reference 2."]]354 documents_invalid = [["Sentence 1 in document 1.", "Sentence 2 in document 1."],355 ["Sentence 1 in document 2.", "Sentence 2 in document 2."]]356 357 with self.assertRaises(ValueError):358 _validate_input_format(tokenize_sentences, predictions_invalid, references, documents)359 360 with self.assertRaises(ValueError):361 _validate_input_format(tokenize_sentences, predictions, references_invalid, documents)362 363 with self.assertRaises(ValueError):364 _validate_input_format(tokenize_sentences, predictions, references, documents_invalid)365 366 367class TestSemNCG(unittest.TestCase):368 def setUp(self):369 self.model_name = "stsb-distilbert-base"370 self.metric = SemNCG(self.model_name)371 372 def _basic_assertion(self, result, debug: bool = False):373 self.assertIsInstance(result, tuple)374 self.assertEqual(len(result), 2)375 self.assertIsInstance(result[0], float)376 self.assertTrue(0.0 <= result[0] <= 1.0)377 self.assertIsInstance(result[1], list)378 if debug:379 for ranked_gain in result[1]:380 self.assertTrue(isinstance(ranked_gain, RankedGains))381 self.assertTrue(0.0 <= ranked_gain.ncg <= 1.0)382 else:383 for gain in result[1]:384 self.assertTrue(isinstance(gain, float))385 self.assertTrue(0.0 <= gain <= 1.0)386 387 def test_compute_basic(self):388 predictions = ["The cat sat on the mat.", "The quick brown fox jumps over the lazy dog."]389 references = ["A cat was sitting on a mat.", "A quick brown fox jumped over a lazy dog."]390 documents = ["There was a cat on a mat.", "The quick brown fox jumped over the lazy dog."]391 392 result = self.metric.compute(predictions=predictions, references=references, documents=documents)393 self._basic_assertion(result)394 395 def test_compute_with_tokenization(self):396 predictions = [["The cat sat on the mat."], ["The quick brown fox jumps over the lazy dog."]]397 references = [["A cat was sitting on a mat."], ["A quick brown fox jumped over a lazy dog."]]398 documents = [["There was a cat on a mat."], ["The quick brown fox jumped over the lazy dog."]]399 400 result = self.metric.compute(401 predictions=predictions, references=references, documents=documents, tokenize_sentences=False402 )403 self._basic_assertion(result)404 405 def test_compute_with_pre_compute_embeddings(self):406 predictions = ["The cat sat on the mat.", "The quick brown fox jumps over the lazy dog."]407 references = ["A cat was sitting on a mat.", "A quick brown fox jumped over a lazy dog."]408 documents = ["There was a cat on a mat.", "The quick brown fox jumped over the lazy dog."]409 410 result = self.metric.compute(411 predictions=predictions, references=references, documents=documents, pre_compute_embeddings=True412 )413 self._basic_assertion(result)414 415 def test_compute_with_debug(self):416 predictions = ["The cat sat on the mat.", "The quick brown fox jumps over the lazy dog."]417 references = ["A cat was sitting on a mat.", "A quick brown fox jumped over a lazy dog."]418 documents = ["There was a cat on a mat.", "The quick brown fox jumped over the lazy dog."]419 420 result = self.metric.compute(421 predictions=predictions, references=references, documents=documents, debug=True422 )423 self._basic_assertion(result, debug=True)424 425 def test_compute_invalid_input_format(self):426 predictions = "The cat sat on the mat."427 references = ["A cat was sitting on a mat."]428 documents = ["There was a cat on a mat."]429 430 with self.assertRaises(ValueError):431 self.metric.compute(predictions=predictions, references=references, documents=documents)432 433 def test_bad_inputs(self):434 def _call_metric(preds, refs, docs, tok):435 with self.assertRaises(Exception) as ctx:436 _ = self.metric.compute(437 predictions=preds,438 references=refs,439 documents=docs,440 tokenize_sentences=tok,441 pre_compute_embeddings=True,442 )443 print(f"Raised Exception with message: {ctx.exception}")444 return ""445 446 # None Inputs447 # Case I448 tokenize_sentences = True449 predictions = [None]450 references = ["A cat was sitting on a mat."]451 documents = ["There was a cat on a mat."]452 print(f"Case I\n{_call_metric(predictions, references, documents, tokenize_sentences)}\n")453 454 # Case II455 tokenize_sentences = False456 predictions = [["A cat was sitting on a mat.", None]]457 references = [["A cat was sitting on a mat.", "A cat was sitting on a mat."]]458 documents = [["There was a cat on a mat.", "There was a cat on a mat."]]459 print(f"Case II\n{_call_metric(predictions, references, documents, tokenize_sentences)}\n")460 461 # Empty Input462 tokenize_sentences = True463 predictions = []464 references = ["A cat was sitting on a mat."]465 documents = ["There was a cat on a mat."]466 print(f"Case: Empty Input\n{_call_metric(predictions, references, documents, tokenize_sentences)}\n")467 468 # Empty String Input469 tokenize_sentences = True470 predictions = [""]471 references = ["A cat was sitting on a mat."]472 documents = ["There was a cat on a mat."]473 print(f"Case: Empty String Input\n{_call_metric(predictions, references, documents, tokenize_sentences)}\n")474 475 def _test_check_verbose(self):476 """UNUSED: previously used to manually check the progress bar477 478 This test should not be used since they rely on files that are479 not kept in version control. this is purely just left here for480 historical purposes and has the '_' prepended to the function481 name to avoid being executed.482 """483 import sqlite3484 import string485 486 con = sqlite3.connect('sem_ncg_samples.db')487 cur = con.cursor()488 data = cur.execute(489 'SELECT * FROM sem_ncg_samples').fetchmany(100)490 data = list(filter(491 lambda x: x[0].translate(492 str.maketrans('', '', string.punctuation)493 ).strip() != '',494 data495 ))496 preds, refs, docs = list(zip(*data))497 result = self.metric.compute(498 predictions=preds, references=refs, 499 documents=docs, verbose=True,500 gpu=2501 )502 503 breakpoint()504 505 506if __name__ == '__main__':507 unittest.main(verbosity=2)508