Implement parallel diversity calculation

aif_gen/_cli/commands/validate.py

@@ -49,13 +49,20 @@
     default=True,
     help='Perform llm judge validation on the dataset.',
 )
+@click.option(
+    '--num-workers',
+    type=click.IntRange(min=1, max=32, clamp=True),
+    help='Number of sub-process workers to spawn for computing diversity validation.',
+    default=4,
+)
 def validate(
     input_data_file: pathlib.Path,
     output_validation_file: pathlib.Path,
     validate_count: bool,
     validate_entropy: bool,
     validate_diversity: bool,
     validate_llm_judge: bool,
+    num_workers: int,
 ) -> None:
     r"""Validate a ContinualAlignmentDataset.
 
@@ -79,7 +86,7 @@ def validate(
 
     if validate_diversity:
         logging.info('Performing diversity validation')
-        results['diversity_validation'] = diversity_validation(dataset)
+        results['diversity_validation'] = diversity_validation(dataset, num_workers)
         logging.info('Finished diversity validation')
 
     if validate_llm_judge:

aif_gen/dataset/validation/diversity.py

@@ -1,4 +1,5 @@
 import logging
+import multiprocessing as mp
 from typing import Callable, Dict, List, Optional
 
 import nltk
@@ -10,12 +11,13 @@
 
 
 def diversity_validation(
-    dataset: Dataset, ngram: int = 3
+    dataset: Dataset, num_workers: int, ngram: int = 3
 ) -> List[Optional[Dict[str, float]]]:
     r"""Report the inverse Self-BLEU score as a measure of diversity within the generated samples.
 
     Args:
         dataset (Union[ContinualAlignmentDataset, AlignmentDataset]): The dataset to validate.
+        num_workers (int): Number of sub-process workers to spawn for the diversity validation.
         ngram (int): The maximum n-gram order for BLEU calculation. Default of 3 matches the original paper.
 
     Returns:
@@ -47,44 +49,60 @@ def diversity_validation(
     results: List[Optional[Dict[str, float]]] = []
     for dataset in datasets:
         if len(dataset):
-            result = _diversity_validation(dataset, ngram)
+            result = _diversity_validation(dataset, num_workers, ngram)
         else:
             logging.warning(f'Skipping diversity on empty dataset: {dataset}')
             result = None
         results.append(result)
     return results
 
 
-def _diversity_validation(dataset: AlignmentDataset, ngram: int) -> Dict[str, float]:
+def _diversity_validation(
+    dataset: AlignmentDataset, num_workers: int, ngram: int
+) -> Dict[str, float]:
     weight = [1.0 / ngram for _ in range(ngram)]
     prompts = [sample.prompt for sample in dataset.samples]
     chosens = [sample.chosen for sample in dataset.samples]
     rejected = [sample.rejected for sample in dataset.samples]
 
     results: Dict[str, List[float]] = {}
-    results['prompt_diversity'] = _compute_diversity(prompts, weight)
-    results['chosen_diversity'] = _compute_diversity(chosens, weight)
-    results['rejected_diversity'] = _compute_diversity(rejected, weight)
+    results['prompt_diversity'] = _compute_diversity(prompts, weight, num_workers)
+    results['chosen_diversity'] = _compute_diversity(chosens, weight, num_workers)
+    results['rejected_diversity'] = _compute_diversity(rejected, weight, num_workers)
     return _compute_statistics(results)
 
 
-def _compute_diversity(response_set: List[str], weight: List[float]) -> List[float]:
+def _compute_diversity(
+    response_set: List[str], weight: List[float], num_workers: int
+) -> List[float]:
     if 0 <= len(response_set) < 2:
         return len(response_set) * [0.0]
 
     tokenizer = _get_tokenizer()
     tokenized_responses = [tokenizer(sentence) for sentence in response_set]
-    diversity = []
-    for i, hypothesis in enumerate(tokenized_responses):
-        other_responses = tokenized_responses[:i] + tokenized_responses[i + 1 :]
-        score = sentence_bleu(
-            other_responses,
-            hypothesis,
-            weight,
-            smoothing_function=SmoothingFunction().method1,
+
+    with mp.Pool(num_workers) as pool:
+        return pool.starmap(
+            _diversity_score,
+            [
+                (
+                    tokenized_responses[:i] + tokenized_responses[i + 1 :],
+                    hypothesis,
+                    weight,
+                )
+                for i, hypothesis in enumerate(tokenized_responses)
+            ],
         )
-        diversity.append(1 - score)
-    return diversity
+
+
+def _diversity_score(responses: List[str], hypothesis: str, weight: List[str]) -> float:
+    score = sentence_bleu(
+        responses,
+        hypothesis,
+        weight,
+        smoothing_function=SmoothingFunction().method1,
+    )
+    return 1 - score
 
 
 def _compute_statistics(results: Dict[str, List[float]]) -> Dict[str, float]:
@@ -103,5 +121,8 @@ def _get_tokenizer() -> Callable[[str], List[str]]:
 
 def _download_nltk_resources() -> None:
     logging.info('Downloading NLTK resources.')
-    nltk.download('averaged_perceptron_tagger', quiet=True)
-    nltk.download('punkt', quiet=True)
+    required_resources = ['punkt_tab']
+    for resource in required_resources:
+        logging.info(f'Downloading NLTK: {resource}')
+        nltk.download('punkt_tab')
+    logging.info('Downloaded NLTK resources.')

test/test_validation/test_diversity_validation.py

@@ -21,7 +21,8 @@ def patch_nltk(mocker):
     )
 
 
-def test_diversity_validation():
+@pytest.mark.parametrize('num_workers', [1, 4])
+def test_diversity_validation(num_workers):
     samples = [
         AlignmentDatasetSample(
             'Mock prompt A 1', 'Winning Response A 1', 'Losing Response A 1'
@@ -35,7 +36,7 @@ def test_diversity_validation():
     ]
     dataset = AlignmentDataset(task=None, samples=samples)
 
-    result = diversity_validation(dataset)
+    result = diversity_validation(dataset, num_workers)
     exp_result = [
         {
             'chosen_diversity_max': np.float64(0.767920558319361),
@@ -55,7 +56,8 @@ def test_diversity_validation():
     assert result == exp_result
 
 
-def test_diversity_validation_identical_prompts():
+@pytest.mark.parametrize('num_workers', [1, 4])
+def test_diversity_validation_identical_prompts(num_workers):
     samples = [
         AlignmentDatasetSample(
             'Mock prompt', 'Winning Response A 1', 'Losing Response A 1'
@@ -69,7 +71,7 @@ def test_diversity_validation_identical_prompts():
     ]
     dataset = AlignmentDataset(task=None, samples=samples)
 
-    result = diversity_validation(dataset)
+    result = diversity_validation(dataset, num_workers)
     exp_result = [
         {
             'chosen_diversity_max': np.float64(0.767920558319361),
@@ -89,16 +91,17 @@ def test_diversity_validation_identical_prompts():
     assert result == exp_result
 
 
+@pytest.mark.parametrize('num_workers', [1, 4])
 @pytest.mark.parametrize('ngram', [1, 2, 3])
-def test_diversity_validation_single_sample_dataset(ngram):
+def test_diversity_validation_single_sample_dataset(num_workers, ngram):
     samples = [
         AlignmentDatasetSample(
             'Mock prompt A', 'Winning Response A 1', 'Losing Response A 1'
         ),
     ]
     dataset = AlignmentDataset(task=None, samples=samples)
 
-    result = diversity_validation(dataset, ngram)
+    result = diversity_validation(dataset, num_workers, ngram=ngram)
 
     exp_result = [
         {
@@ -124,7 +127,7 @@ def test_diversity_validation_empty_dataset(ngram):
     samples = []
     dataset = AlignmentDataset(task=None, samples=samples)
 
-    result = diversity_validation(dataset, ngram=ngram)
+    result = diversity_validation(dataset, num_workers=1, ngram=ngram)
     assert result == [None]
 
 
@@ -133,10 +136,11 @@ def test_diversity_validation_bad_ngram(bad_ngram):
     dataset = AlignmentDataset(task=None, samples=[])
 
     with pytest.raises(ValueError):
-        diversity_validation(dataset, ngram=bad_ngram)
+        diversity_validation(dataset, num_workers=1, ngram=bad_ngram)
 
 
-def test_diversity_validation_continual_dataset():
+@pytest.mark.parametrize('num_workers', [1, 4])
+def test_diversity_validation_continual_dataset(num_workers):
     samples = [
         AlignmentDatasetSample(
             'Mock prompt', 'Winning Response A 1', 'Losing Response A 1'
@@ -152,7 +156,7 @@ def test_diversity_validation_continual_dataset():
     dataset2 = AlignmentDataset(task=None, samples=[])
     dataset = ContinualAlignmentDataset(datasets=[dataset1, dataset2])
 
-    result = diversity_validation(dataset)
+    result = diversity_validation(dataset, num_workers)
     exp_result = [
         {
             'chosen_diversity_max': np.float64(0.767920558319361),