File size: 17,999 Bytes
a54024a
 
 
 
 
 
 
27a1559
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a54024a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
27a1559
a54024a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
27a1559
a54024a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
import unittest

import numpy as np
import torch
from sentence_transformers import SentenceTransformer

from .encoder_models import SBertEncoder, get_encoder, get_sbert_encoder
from .semncg import (
    RankedGains,
    compute_cosine_similarity,
    compute_gain,
    score_ncg,
    compute_ncg,
    _validate_input_format,
    SemNCG
)
from .utils import (
    get_gpu,
    slice_embeddings,
    is_nested_list_of_type,
    flatten_list,
    prep_sentences,
    tokenize_and_prep_document
)


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_prep_sentences(self):
        # Test normal case
        self.assertEqual(prep_sentences(["Hello, world!", " This is a test. ", "!!!"]),
                         ['Hello, world!', 'This is a test.'])

        # Test case with only punctuations
        with self.assertRaises(ValueError):
            prep_sentences(["!!!", "..."])

        # Test case with empty list
        with self.assertRaises(ValueError):
            prep_sentences([])

    def test_tokenize_and_prep_document(self):
        # Test tokenize=True with string input
        self.assertEqual(tokenize_and_prep_document("Hello, world! This is a test.", True),
                         ['Hello, world!', 'This is a test.'])

        # Test tokenize=False with list of strings input
        self.assertEqual(tokenize_and_prep_document(["Hello, world!", "This is a test."], False),
                         ['Hello, world!', 'This is a test.'])

        # Test tokenize=True with empty document
        with self.assertRaises(ValueError):
            tokenize_and_prep_document("!!! ...", True)

    def test_slice_embeddings(self):
        # Case 1
        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))
        )

        # Case 2
        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
        )

        # Case 3
        document_sentences_count = [10, 8, 7]
        reference_sentences_count = [5, 3, 2]
        pred_sentences_count = [2, 2, 1]
        all_embeddings = np.random.rand(
            sum(document_sentences_count + reference_sentences_count + pred_sentences_count), 5,
        )

        embeddings = all_embeddings
        expected_doc_embeddings = [embeddings[:10], embeddings[10:18], embeddings[18:25]]

        embeddings = all_embeddings[25:]
        expected_ref_embeddings = [embeddings[:5], embeddings[5:8], embeddings[8:10]]

        embeddings = all_embeddings[35:]
        expected_pred_embeddings = [embeddings[:2], embeddings[2:4], embeddings[4:5]]

        doc_embeddings = slice_embeddings(all_embeddings, document_sentences_count)
        ref_embeddings = slice_embeddings(all_embeddings[sum(document_sentences_count):], reference_sentences_count)
        pred_embeddings = slice_embeddings(
            all_embeddings[sum(document_sentences_count + reference_sentences_count):], pred_sentences_count
        )

        self.assertTrue(doc_embeddings, expected_doc_embeddings)
        self.assertTrue(ref_embeddings, expected_ref_embeddings)
        self.assertTrue(pred_embeddings, expected_pred_embeddings)

        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])


class TestSBertEncoder(unittest.TestCase):

    def setUp(self) -> None:
        # Set up a test SentenceTransformer model
        self.model_name = "paraphrase-distilroberta-base-v1"
        self.sbert_model = get_sbert_encoder(self.model_name)
        self.device = "cpu"  # For testing on CPU
        self.batch_size = 32
        self.verbose = False
        self.encoder = SBertEncoder(self.sbert_model, self.device, self.batch_size, self.verbose)

    def test_encode_single_sentence(self):
        sentence = "Hello, world!"
        embeddings = self.encoder.encode([sentence])
        self.assertEqual(embeddings.shape, (1, 768))  # Adjust shape based on your model's embedding dimension

    def test_encode_multiple_sentences(self):
        sentences = ["Hello, world!", "This is a test."]
        embeddings = self.encoder.encode(sentences)
        self.assertEqual(embeddings.shape, (2, 768))  # Adjust shape based on your model's embedding dimension

    def test_get_sbert_encoder(self):
        model_name = "paraphrase-distilroberta-base-v1"
        sbert_model = get_sbert_encoder(model_name)
        self.assertIsInstance(sbert_model, SentenceTransformer)

    def test_encode_with_gpu(self):
        if torch.cuda.is_available():
            device = "cuda"
            encoder = get_encoder(self.sbert_model, device, self.batch_size, self.verbose)
            sentences = ["Hello, world!", "This is a test."]
            embeddings = encoder.encode(sentences)
            self.assertEqual(embeddings.shape, (2, 768))  # Adjust shape based on your model's embedding dimension
        else:
            self.skipTest("CUDA not available, skipping GPU test.")

    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"]
            encoder = get_encoder(self.sbert_model, devices, self.batch_size, self.verbose)
            sentences = ["This is a test sentence.", "Here is another sentence.", "This is a test sentence."]
            embeddings = 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):
        sbert_model = get_sbert_encoder(model_name)
        encoder = get_encoder(sbert_model, 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_sbert_encoder(model_name)

    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_sbert_encoder(model_name)


class TestRankedGainsDataclass(unittest.TestCase):
    def test_ranked_gains_dataclass(self):
        # Test initialization and attribute access
        gt_gains = [("doc1", 0.8), ("doc2", 0.6)]
        pred_gains = [("doc2", 0.7), ("doc1", 0.5)]
        k = 2
        ncg = 0.75
        ranked_gains = RankedGains(gt_gains, pred_gains, k, ncg)

        self.assertEqual(ranked_gains.gt_gains, gt_gains)
        self.assertEqual(ranked_gains.pred_gains, pred_gains)
        self.assertEqual(ranked_gains.k, k)
        self.assertEqual(ranked_gains.ncg, ncg)


class TestComputeCosineSimilarity(unittest.TestCase):
    def test_compute_cosine_similarity(self):
        doc_embeds = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]])
        ref_embeds = np.array([[0.2, 0.3, 0.4], [0.5, 0.6, 0.7]])
        # Test compute_cosine_similarity function
        similarity_scores = compute_cosine_similarity(doc_embeds, ref_embeds)
        print(similarity_scores)

        # Example values, change as per actual function output
        expected_scores = [0.980, 0.997]

        self.assertAlmostEqual(similarity_scores[0], expected_scores[0], places=3)
        self.assertAlmostEqual(similarity_scores[1], expected_scores[1], places=3)


class TestComputeGain(unittest.TestCase):
    def test_compute_gain(self):
        # Test compute_gain function
        sim_scores = [0.8, 0.6, 0.7]
        gains = compute_gain(sim_scores)
        print(gains)

        # Example values, change as per actual function output
        expected_gains = [(0, 0.5), (2, 0.3333333333333333), (1, 0.16666666666666666)]

        self.assertEqual(gains, expected_gains)


class TestScoreNcg(unittest.TestCase):
    def test_score_ncg(self):
        # Test score_ncg function
        model_relevance = [0.8, 0.7, 0.6]
        gt_relevance = [1.0, 0.9, 0.8]
        ncg_score = score_ncg(model_relevance, gt_relevance)
        expected_ncg = 0.778  # Example value, change as per actual function output

        self.assertAlmostEqual(ncg_score, expected_ncg, places=3)


class TestComputeNcg(unittest.TestCase):
    def test_compute_ncg(self):
        # Test compute_ncg function
        pred_gains = [(0, 0.8), (2, 0.7), (1, 0.6)]
        gt_gains = [(0, 1.0), (1, 0.9), (2, 0.8)]
        k = 3
        ncg_score = compute_ncg(pred_gains, gt_gains, k)
        expected_ncg = 1.0  # TODO: Confirm this with Dr. Santu

        self.assertAlmostEqual(ncg_score, expected_ncg, places=6)


class TestValidateInputFormat(unittest.TestCase):
    def test_validate_input_format(self):
        # Test _validate_input_format function
        tokenize_sentences = True
        predictions = ["Prediction 1", "Prediction 2"]
        references = ["Reference 1", "Reference 2"]
        documents = ["Document 1", "Document 2"]

        # No exception should be raised for valid input
        try:
            _validate_input_format(tokenize_sentences, predictions, references, documents)
        except ValueError as e:
            self.fail(f"_validate_input_format raised ValueError unexpectedly: {str(e)}")

        # Test invalid input format
        predictions_invalid = [["Sentence 1 in prediction 1.", "Sentence 2 in prediction 1."],
                               ["Sentence 1 in prediction 2.", "Sentence 2 in prediction 2."]]
        references_invalid = [["Sentences in reference 1."], ["Sentences in reference 2."]]
        documents_invalid = [["Sentence 1 in document 1.", "Sentence 2 in document 1."],
                             ["Sentence 1 in document 2.", "Sentence 2 in document 2."]]

        with self.assertRaises(ValueError):
            _validate_input_format(tokenize_sentences, predictions_invalid, references, documents)

        with self.assertRaises(ValueError):
            _validate_input_format(tokenize_sentences, predictions, references_invalid, documents)

        with self.assertRaises(ValueError):
            _validate_input_format(tokenize_sentences, predictions, references, documents_invalid)


class TestSemnCG(unittest.TestCase):
    def setUp(self):
        self.model_name = "stsb-distilbert-base"
        self.metric = SemNCG(self.model_name)

    def _basic_assertion(self, result, debug: bool = False):
        self.assertIsInstance(result, tuple)
        self.assertEqual(len(result), 2)
        self.assertIsInstance(result[0], float)
        self.assertTrue(0.0 <= result[0] <= 1.0)
        self.assertIsInstance(result[1], list)
        if debug:
            for ranked_gain in result[1]:
                self.assertTrue(isinstance(ranked_gain, RankedGains))
                self.assertTrue(0.0 <= ranked_gain.ncg <= 1.0)
        else:
            for gain in result[1]:
                self.assertTrue(isinstance(gain, float))
                self.assertTrue(0.0 <= gain <= 1.0)

    def test_compute_basic(self):
        predictions = ["The cat sat on the mat.", "The quick brown fox jumps over the lazy dog."]
        references = ["A cat was sitting on a mat.", "A quick brown fox jumped over a lazy dog."]
        documents = ["There was a cat on a mat.", "The quick brown fox jumped over the lazy dog."]

        result = self.metric.compute(predictions=predictions, references=references, documents=documents)
        self._basic_assertion(result)

    def test_compute_with_tokenization(self):
        predictions = [["The cat sat on the mat."], ["The quick brown fox jumps over the lazy dog."]]
        references = [["A cat was sitting on a mat."], ["A quick brown fox jumped over a lazy dog."]]
        documents = [["There was a cat on a mat."], ["The quick brown fox jumped over the lazy dog."]]

        result = self.metric.compute(
            predictions=predictions, references=references, documents=documents, tokenize_sentences=False
        )
        self._basic_assertion(result)

    def test_compute_with_pre_compute_embeddings(self):
        predictions = ["The cat sat on the mat.", "The quick brown fox jumps over the lazy dog."]
        references = ["A cat was sitting on a mat.", "A quick brown fox jumped over a lazy dog."]
        documents = ["There was a cat on a mat.", "The quick brown fox jumped over the lazy dog."]

        result = self.metric.compute(
            predictions=predictions, references=references, documents=documents, pre_compute_embeddings=True
        )
        self._basic_assertion(result)

    def test_compute_with_debug(self):
        predictions = ["The cat sat on the mat.", "The quick brown fox jumps over the lazy dog."]
        references = ["A cat was sitting on a mat.", "A quick brown fox jumped over a lazy dog."]
        documents = ["There was a cat on a mat.", "The quick brown fox jumped over the lazy dog."]

        result = self.metric.compute(
            predictions=predictions, references=references, documents=documents, debug=True
        )
        self._basic_assertion(result, debug=True)

    def test_compute_invalid_input_format(self):
        predictions = "The cat sat on the mat."
        references = ["A cat was sitting on a mat."]
        documents = ["There was a cat on a mat."]

        with self.assertRaises(ValueError):
            self.metric.compute(predictions=predictions, references=references, documents=documents)


if __name__ == '__main__':
    unittest.main(verbosity=2)