Fix generator progress bar

prompt2model/dataset_generator/prompt_based.py

@@ -449,16 +449,18 @@ def generate_dataset_split(
                 handle_api_error(e)
 
             # Extract the responses and add new examples to the dataset.
+            prev_length = len(generated_examples)
             self.extract_and_append_responses(responses, all_generated_examples)
             generated_examples = (
                 self.apply_multi_vote_filtering(all_generated_examples)
                 if self.filter_duplicated_examples
                 else all_generated_examples
             )
-            if len(generated_examples) >= num_examples:
-                generated_examples = generated_examples[:num_examples]
 
-            pbar.update(len(generated_examples))
+            pbar.update(len(generated_examples) - prev_length)
+
+        if len(generated_examples) >= num_examples:
+            generated_examples = generated_examples[:num_examples]
 
         return Dataset.from_dict(
             {

prompt2model/dataset_retriever/description_dataset_retriever.py

@@ -77,13 +77,14 @@ def initialize_search_index(self) -> None:
                 "Search index must either be a valid file or not exist yet. "
                 "But {self.search_index_path} is provided."
             )
-        logger.info("Creating dataset descriptions")
-        encode_text(
-            self.encoder_model_name,
-            text_to_encode=[x.description for x in self.dataset_infos],
-            encoding_file=self.search_index_path,
-            device=self.device,
-        )
+        if not os.path.exists(self.search_index_path):
+            logger.info("Creating dataset descriptions")
+            encode_text(
+                self.encoder_model_name,
+                text_to_encode=[x.description for x in self.dataset_infos],
+                encoding_file=self.search_index_path,
+                device=self.device,
+            )
 
     # ---------------------------- Utility Functions ----------------------------
     @staticmethod

tests/dataset_retriever_test.py

@@ -153,7 +153,7 @@ def test_retrieve_dataset_dict_when_search_index_exists(encode_text):
 
         mock_prompt = MockPromptSpec(task_type=TaskType.TEXT_GENERATION)
         retrieved_dataset = retriever.retrieve_dataset_dict(mock_prompt)
-        assert encode_text.call_count == 2
+        assert encode_text.call_count == 1
         for split_name in ["train", "val", "test"]:
             assert split_name in retrieved_dataset
             split = retrieved_dataset[split_name]