make post-processing more efficient

run_prompt_creation.py

@@ -7,6 +7,7 @@ from typing import Any, Dict, List, Optional, Union
 
 import torch
 from accelerate import Accelerator
+from accelerate.logging import get_logger
 from datasets import DatasetDict, load_dataset
 from torch.utils.data import DataLoader
 from tqdm import tqdm
@@ -18,7 +19,7 @@ from transformers import (
 )
 
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__, log_level="INFO")
 
 
 @dataclass
@@ -223,6 +224,7 @@ For example, given the following keywords: 'female', 'slightly roomy sounding',
 For the keywords: '[gender]', '[reverberation]', '[noise]', '[speech_monotony]', '[pitch]', '[speaking_rate]', the corresponding description is:"
 """
 
+
 def main():
     # 1. Parse input arguments
     parser = HfArgumentParser((ModelArguments, DataArguments))
@@ -235,7 +237,6 @@ def main():
 
     # 2. Setup logging
     # Make one log on every process with the configuration for debugging.
-    logger.setLevel(logging.INFO)
     logging.basicConfig(
         format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
         datefmt="%m/%d/%Y %H:%M:%S",
@@ -368,6 +369,12 @@ def main():
         output_ids = accelerator.pad_across_processes(output_ids, dim=1, pad_index=tokenizer.pad_token_id)
         return output_ids
 
+    def postprocess_dataset(sample):
+        prompt_text = tokenizer.decode(sample["input_ids"], skip_special_tokens=True)
+        generated_text = tokenizer.decode(sample["generated_ids"], skip_special_tokens=True)
+        sample["text_description"] = generated_text[len(prompt_text) :]
+        return sample
+
     for split in vectorized_datasets:
         data_loader = DataLoader(
             vectorized_datasets[split],
@@ -382,21 +389,16 @@ def main():
         for batch in tqdm(data_loader, disable=not accelerator.is_local_main_process):
             generated_ids = generate_step(batch)
             generated_ids = accelerator.gather_for_metrics(generated_ids)
-            all_generated_ids.extend(generated_ids.cpu())
+            all_generated_ids.extend(generated_ids.cpu().numpy())
 
-        def postprocess_dataset(sample, idx):
-            prompt_text = tokenizer.decode(sample["input_ids"], skip_special_tokens=True)
-            generated_text = tokenizer.decode(all_generated_ids[idx], skip_special_tokens=True)
-            sample["text_description"] = generated_text[len(prompt_text) :]
-            return sample
+        vectorized_datasets[split] = vectorized_datasets[split].add_column("generated_ids", all_generated_ids)
 
         if accelerator.is_main_process:
             vectorized_datasets[split] = vectorized_datasets[split].map(
                 postprocess_dataset,
                 num_proc=data_args.preprocessing_num_workers,
                 desc="Postprocessing dataset",
-                remove_columns=["input_ids"],
-                with_indices=True,
+                remove_columns=["input_ids", "generated_ids"],
             )
 
     if accelerator.is_main_process: