Spaces:

Leon4gr45
/

agent-0

Running

agent-0 / augmentoolkit /generation_functions /random_variation_step_class.py

Upload folder using huggingface_hub

42da440 verified 10 days ago

13.4 kB

import json
import os
import random
import re
import traceback
import sys
import asyncio
from augmentoolkit.generation_functions.generation_step_class import GenerationStep
from augmentoolkit.generation_functions.pipeline_step_class import (
    PipelineStep,
    filter_out_nonpresent_keys,
)
from augmentoolkit.utils.make_id import make_id
from augmentoolkit.utils.write_output_to_file import write_output_to_file
from tqdm import asyncio as tqdmasyncio


class RandomVariationStep(
    PipelineStep
):  # pipeline steps store the settings and the prompt, and prevent us from having to repeat the "read previous output" code among other things
    def __init__(
        self,
        sampling_params=None,
        output_file=None,
        intermediate_output_path=None,
        save_path=None,
        use_stop=None,
        completion_mode=None,
        validation_function=lambda x, y: {"result": True, "message": "default message"},
        max_retries=3,
        result_key=None,
        prompt_path=None,
        regex=re.compile(r".*", re.DOTALL),
        variation_generator_count=None,
        output_processor=lambda x: x,
        log_full_outputs=False,
        method_overrides={},
        details_key=None,
    ):
        self.variation_generator_count = variation_generator_count
        super().__init__(
            sampling_params=sampling_params,
            output_file=output_file,  # Add output_file parameter for consolidated storage
            intermediate_output_path=intermediate_output_path,
            save_path=save_path,
            use_stop=use_stop,
            completion_mode=completion_mode,
            validation_function=validation_function,
            max_retries=max_retries,
            result_key=result_key,
            prompt_path=prompt_path,
            output_processor=output_processor,
            log_full_outputs=log_full_outputs,
            method_overrides=method_overrides,
            details_key=details_key,
        )

    async def read_previous_output(self, key, input_dict):
        entry = input_dict.get(str(key), {})
        # if not entry:
        #      print(f"  Key '{str(key)}' not found in input_dict.")
        # else:
        #      print(f"  Found entry. Keys: {list(entry.keys())}")

        current_variations = entry.get(self.result_key, [])

        should_skip = False
        if len(current_variations) >= self.variation_generator_count:
            print(
                f"  Found {len(current_variations)} variations, which meets or exceeds target {self.variation_generator_count}. Will skip generation for key {key}."
            )
            should_skip = True
        # else:

        return should_skip

    async def generate_data(
        self,
        processed_data,
        engine_wrapper,
        prompt_path,
        prompt_folder,
        default_prompt_folder,
        completion_mode,
        use_stop,
        **kwargs,
    ):

        # randomly select a yaml file from the prompt folder
        # prompt_path is the name of the subfolder; prompt_folder is, per usual, the path to the overall PROMPTS folder.
        search_dir = os.path.join(prompt_folder, prompt_path)
        if not os.path.exists(search_dir):
            search_dir = os.path.join(default_prompt_folder, prompt_path)

        if not completion_mode:
            prompt_files = [f for f in os.listdir(search_dir) if f.endswith(".yaml")]
        else:
            prompt_files = [
                f for f in os.listdir(search_dir) if f.endswith(".yaml")
            ]  # NOTE: This looks like a potential bug, it lists .yaml for completion mode too.

        if not prompt_files:
            error_msg = f"  ERROR: No suitable prompt files found in {search_dir} for completion_mode={completion_mode}"
            print(error_msg)
            raise FileNotFoundError(error_msg)

        random_prompt_file = random.choice(prompt_files)
        prompt_filepath = os.path.join(search_dir, random_prompt_file)

        try:
            generator = GenerationStep(
                prompt_path=prompt_filepath,
                default_prompt_folder=default_prompt_folder,
                sampling_params=self.sampling_params,
                completion_mode=completion_mode,
                engine_wrapper=engine_wrapper,
                output_processor=self.output_processor,
                retries=1,
                logging_level=self.logging_level,
                use_stop=use_stop,
                prompt_folder=prompt_folder,
                regex=self.regex,
            )

            # Note: We don't log the actual return values here as they could be large
            result, full_output, full_response, full_input = await generator.generate(
                **processed_data, **self.static_arguments, **kwargs
            )
            return result, full_output, full_response, full_input
        except Exception as e:
            print(f"  ERROR during GenerationStep initialization or execution: {e}")
            traceback.print_exc()
            # Re-raise the exception to be handled by the caller (run method)
            raise

    def save(
        self,
        result=None,
        key=None,
        input_dict=None,
        input_data=None,
        full_response=None,
        full_input=None,
        completion_mode=False,
        include_details=False,
    ):

        # Get existing entry or create a new one. Avoid copying input_data unless necessary.
        entry = input_dict.get(str(key))
        if entry is None:
            # If the key truly doesn't exist, we MUST base it on input_data
            # This case should ideally not happen often if execute_pipeline populates correctly
            entry = input_data.copy() if input_data else {}
        # else:

        # Ensure the result key exists and is a list
        if self.result_key not in entry or not isinstance(
            entry.get(self.result_key), list
        ):  # Use .get for safety
            entry[self.result_key] = []

        try:
            entry[self.result_key].append(result)
        except Exception as e:
            print(f"  ERROR appending result: {e}")
            traceback.print_exc()
            raise

        if include_details:
            # Ensure the details key exists and is a list
            if self.details_key not in entry or not isinstance(
                entry.get(self.details_key), list
            ):  # Use .get for safety
                entry[self.details_key] = []

            detail_entry = {
                "full_response": full_response,
                "full_input": full_input,
                "completion_mode": completion_mode,
            }
            try:
                entry[self.details_key].append(detail_entry)
            except Exception as e:
                print(f"  ERROR appending detail entry: {e}")
                traceback.print_exc()
                raise
        # else:

        try:
            input_dict[str(key)] = entry
        except Exception as e:
            print(f"  ERROR assigning entry to input_dict: {e}")
            traceback.print_exc()
            raise

        return entry

    async def run(
        self,
        key=None,
        input_data=None,
        engine_wrapper=None,
        input_dict=None,
        default_prompt_folder=None,
        prompt_folder=None,
        output_dir=None,
        completion_mode=False,
        use_stop=True,
        include_details=False,
        **kwargs,
    ):

        should_skip = await self.read_previous_output(key, input_dict)
        if should_skip:
            return
        # else:

        processed_data, additional_kwargs = self.process_input_data(input_data)
        # existing_variations = [] # This variable seems unused after the FileLock block was removed

        # Determine current count and number needed based *only* on in-memory input_dict
        current_variations_in_memory = input_dict.get(str(key), {}).get(
            self.result_key, []
        )
        num_existing_in_memory = len(current_variations_in_memory)
        num_variations_needed = self.variation_generator_count - num_existing_in_memory

        if num_variations_needed <= 0:
            # This case should technically be caught by read_previous_output, but added for robustness
            return

        # Generate remaining variations
        for i in range(num_variations_needed):
            error_message = ""
            complete = False
            max_retries = self.max_retries
            while not complete and max_retries > 0:
                try:
                    result, full_output, full_response, full_input = (
                        await self.generate_data(
                            processed_data,
                            engine_wrapper,
                            self.prompt_path,
                            prompt_folder,
                            default_prompt_folder,
                            completion_mode,
                            use_stop,
                            error_message=error_message,
                            **kwargs,
                            **additional_kwargs,
                        )
                    )

                    validation_passed = self.validation_function(result, input_data)
                    if validation_passed["result"]:
                        complete = True
                    else:
                        error_message = validation_passed["message"]
                except Exception as e:
                    print(f"      ERROR during generate_data or validation: {e}")
                    error_message = str(e)
                    traceback.print_exc()
                max_retries -= 1
                if not complete:
                    print(
                        f"      Attempt failed or validation failed. Retries left: {max_retries}"
                    )
            # End of while loop

            if not complete:
                print(
                    f"    ERROR: Failed to generate a valid variation after {self.max_retries} attempts for key {key}, iteration {i+1}."
                )
                # We might want to remove the key from input_dict here or handle failure differently
                return
            # else:

            self.save(
                result=result,
                key=key,
                input_dict=input_dict,
                input_data=input_dict.get(
                    str(key)
                ),  # Pass the current state from input_dict
                full_response=full_response,
                full_input=full_input,
                completion_mode=completion_mode,
                include_details=include_details,
            )
        # End of for loop

    async def execute_pipeline(
        self,
        input_dict={},
        engine_wrapper=None,
        rtwl=None,
        default_prompt_folder=None,
        prompt_folder=None,
        output_dir=None,
        completion_mode=None,
        use_stop=None,
        include_details=False,
        **kwargs,
    ):
        self.load_dataset(
            input_dict=input_dict, output_dir=output_dir
        )  # if it is at the start and end of every pipeline execution, we get the same behavior as before, past thing will load and save just fine, without constant reads/writes.
        try:
            data_generations_tasks = [
                self.run(
                    key=key,
                    input_data=value,
                    engine_wrapper=engine_wrapper,
                    input_dict=input_dict,
                    default_prompt_folder=default_prompt_folder,
                    prompt_folder=prompt_folder,
                    output_dir=output_dir,
                    completion_mode=completion_mode,
                    use_stop=use_stop,
                    include_details=include_details,
                    **kwargs,
                )
                for key, value in input_dict.items()
            ]
            coroutines = [rtwl(task) for task in data_generations_tasks]
            TASK_TIMEOUT_SECONDS = 600  # 10 minutes timeout
            processed_count = 0
            total_tasks = len(coroutines)
            for future in tqdmasyncio.tqdm.as_completed(coroutines):
                processed_count += 1
                try:
                    await asyncio.wait_for(future, timeout=TASK_TIMEOUT_SECONDS)
                except asyncio.TimeoutError:
                    print(
                        f"\nWARNING: Task {processed_count}/{total_tasks} for step '{self.output_file}' timed out after {TASK_TIMEOUT_SECONDS} seconds.",
                        file=sys.stderr,
                    )
                    # Task is cancelled by wait_for, loop continues
                except Exception as e_inner:  # Catch exceptions from the future itself
                    print(
                        f"\nError processing task {processed_count}/{total_tasks}: {e_inner}"
                    )
                    traceback.print_exc()  # Optionally log traceback

            filter_out_nonpresent_keys(input_dict, key_to_check=self.result_key)
        except Exception as e:
            print(f"Exception occurred during task execution: {e}")
            # traceback.print_exc()
            raise e
        finally:
            # Robust save with interrupt handling
            self.save_dataset(input_dict=input_dict, output_dir=output_dir)

	import json
	import os
	import random
	import re
	import traceback
	import sys
	import asyncio
	from augmentoolkit.generation_functions.generation_step_class import GenerationStep
	from augmentoolkit.generation_functions.pipeline_step_class import (
	PipelineStep,
	filter_out_nonpresent_keys,
	)
	from augmentoolkit.utils.make_id import make_id
	from augmentoolkit.utils.write_output_to_file import write_output_to_file
	from tqdm import asyncio as tqdmasyncio


	class RandomVariationStep(
	PipelineStep
	): # pipeline steps store the settings and the prompt, and prevent us from having to repeat the "read previous output" code among other things
	def __init__(
	self,
	sampling_params=None,
	output_file=None,
	intermediate_output_path=None,
	save_path=None,
	use_stop=None,
	completion_mode=None,
	validation_function=lambda x, y: {"result": True, "message": "default message"},
	max_retries=3,
	result_key=None,
	prompt_path=None,
	regex=re.compile(r".*", re.DOTALL),
	variation_generator_count=None,
	output_processor=lambda x: x,
	log_full_outputs=False,
	method_overrides={},
	details_key=None,
	):
	self.variation_generator_count = variation_generator_count
	super().__init__(
	sampling_params=sampling_params,
	output_file=output_file, # Add output_file parameter for consolidated storage
	intermediate_output_path=intermediate_output_path,
	save_path=save_path,
	use_stop=use_stop,
	completion_mode=completion_mode,
	validation_function=validation_function,
	max_retries=max_retries,
	result_key=result_key,
	prompt_path=prompt_path,
	output_processor=output_processor,
	log_full_outputs=log_full_outputs,
	method_overrides=method_overrides,
	details_key=details_key,
	)

	async def read_previous_output(self, key, input_dict):
	entry = input_dict.get(str(key), {})
	# if not entry:
	# print(f" Key '{str(key)}' not found in input_dict.")
	# else:
	# print(f" Found entry. Keys: {list(entry.keys())}")

	current_variations = entry.get(self.result_key, [])

	should_skip = False
	if len(current_variations) >= self.variation_generator_count:
	print(
	f" Found {len(current_variations)} variations, which meets or exceeds target {self.variation_generator_count}. Will skip generation for key {key}."
	)
	should_skip = True
	# else:

	return should_skip

	async def generate_data(
	self,
	processed_data,
	engine_wrapper,
	prompt_path,
	prompt_folder,
	default_prompt_folder,
	completion_mode,
	use_stop,
	**kwargs,
	):

	# randomly select a yaml file from the prompt folder
	# prompt_path is the name of the subfolder; prompt_folder is, per usual, the path to the overall PROMPTS folder.
	search_dir = os.path.join(prompt_folder, prompt_path)
	if not os.path.exists(search_dir):
	search_dir = os.path.join(default_prompt_folder, prompt_path)

	if not completion_mode:
	prompt_files = [f for f in os.listdir(search_dir) if f.endswith(".yaml")]
	else:
	prompt_files = [
	f for f in os.listdir(search_dir) if f.endswith(".yaml")
	] # NOTE: This looks like a potential bug, it lists .yaml for completion mode too.

	if not prompt_files:
	error_msg = f" ERROR: No suitable prompt files found in {search_dir} for completion_mode={completion_mode}"
	print(error_msg)
	raise FileNotFoundError(error_msg)

	random_prompt_file = random.choice(prompt_files)
	prompt_filepath = os.path.join(search_dir, random_prompt_file)

	try:
	generator = GenerationStep(
	prompt_path=prompt_filepath,
	default_prompt_folder=default_prompt_folder,
	sampling_params=self.sampling_params,
	completion_mode=completion_mode,
	engine_wrapper=engine_wrapper,
	output_processor=self.output_processor,
	retries=1,
	logging_level=self.logging_level,
	use_stop=use_stop,
	prompt_folder=prompt_folder,
	regex=self.regex,
	)

	# Note: We don't log the actual return values here as they could be large
	result, full_output, full_response, full_input = await generator.generate(
	processed_data, self.static_arguments, **kwargs
	)
	return result, full_output, full_response, full_input
	except Exception as e:
	print(f" ERROR during GenerationStep initialization or execution: {e}")
	traceback.print_exc()
	# Re-raise the exception to be handled by the caller (run method)
	raise

	def save(
	self,
	result=None,
	key=None,
	input_dict=None,
	input_data=None,
	full_response=None,
	full_input=None,
	completion_mode=False,
	include_details=False,
	):

	# Get existing entry or create a new one. Avoid copying input_data unless necessary.
	entry = input_dict.get(str(key))
	if entry is None:
	# If the key truly doesn't exist, we MUST base it on input_data
	# This case should ideally not happen often if execute_pipeline populates correctly
	entry = input_data.copy() if input_data else {}
	# else:

	# Ensure the result key exists and is a list
	if self.result_key not in entry or not isinstance(
	entry.get(self.result_key), list
	): # Use .get for safety
	entry[self.result_key] = []

	try:
	entry[self.result_key].append(result)
	except Exception as e:
	print(f" ERROR appending result: {e}")
	traceback.print_exc()
	raise

	if include_details:
	# Ensure the details key exists and is a list
	if self.details_key not in entry or not isinstance(
	entry.get(self.details_key), list
	): # Use .get for safety
	entry[self.details_key] = []

	detail_entry = {
	"full_response": full_response,
	"full_input": full_input,
	"completion_mode": completion_mode,
	}
	try:
	entry[self.details_key].append(detail_entry)
	except Exception as e:
	print(f" ERROR appending detail entry: {e}")
	traceback.print_exc()
	raise
	# else:

	try:
	input_dict[str(key)] = entry
	except Exception as e:
	print(f" ERROR assigning entry to input_dict: {e}")
	traceback.print_exc()
	raise

	return entry

	async def run(
	self,
	key=None,
	input_data=None,
	engine_wrapper=None,
	input_dict=None,
	default_prompt_folder=None,
	prompt_folder=None,
	output_dir=None,
	completion_mode=False,
	use_stop=True,
	include_details=False,
	**kwargs,
	):

	should_skip = await self.read_previous_output(key, input_dict)
	if should_skip:
	return
	# else:

	processed_data, additional_kwargs = self.process_input_data(input_data)
	# existing_variations = [] # This variable seems unused after the FileLock block was removed

	# Determine current count and number needed based only on in-memory input_dict
	current_variations_in_memory = input_dict.get(str(key), {}).get(
	self.result_key, []
	)
	num_existing_in_memory = len(current_variations_in_memory)
	num_variations_needed = self.variation_generator_count - num_existing_in_memory

	if num_variations_needed <= 0:
	# This case should technically be caught by read_previous_output, but added for robustness
	return

	# Generate remaining variations
	for i in range(num_variations_needed):
	error_message = ""
	complete = False
	max_retries = self.max_retries
	while not complete and max_retries > 0:
	try:
	result, full_output, full_response, full_input = (
	await self.generate_data(
	processed_data,
	engine_wrapper,
	self.prompt_path,
	prompt_folder,
	default_prompt_folder,
	completion_mode,
	use_stop,
	error_message=error_message,
	**kwargs,
	**additional_kwargs,
	)
	)

	validation_passed = self.validation_function(result, input_data)
	if validation_passed["result"]:
	complete = True
	else:
	error_message = validation_passed["message"]
	except Exception as e:
	print(f" ERROR during generate_data or validation: {e}")
	error_message = str(e)
	traceback.print_exc()
	max_retries -= 1
	if not complete:
	print(
	f" Attempt failed or validation failed. Retries left: {max_retries}"
	)
	# End of while loop

	if not complete:
	print(
	f" ERROR: Failed to generate a valid variation after {self.max_retries} attempts for key {key}, iteration {i+1}."
	)
	# We might want to remove the key from input_dict here or handle failure differently
	return
	# else:

	self.save(
	result=result,
	key=key,
	input_dict=input_dict,
	input_data=input_dict.get(
	str(key)
	), # Pass the current state from input_dict
	full_response=full_response,
	full_input=full_input,
	completion_mode=completion_mode,
	include_details=include_details,
	)
	# End of for loop

	async def execute_pipeline(
	self,
	input_dict={},
	engine_wrapper=None,
	rtwl=None,
	default_prompt_folder=None,
	prompt_folder=None,
	output_dir=None,
	completion_mode=None,
	use_stop=None,
	include_details=False,
	**kwargs,
	):
	self.load_dataset(
	input_dict=input_dict, output_dir=output_dir
	) # if it is at the start and end of every pipeline execution, we get the same behavior as before, past thing will load and save just fine, without constant reads/writes.
	try:
	data_generations_tasks = [
	self.run(
	key=key,
	input_data=value,
	engine_wrapper=engine_wrapper,
	input_dict=input_dict,
	default_prompt_folder=default_prompt_folder,
	prompt_folder=prompt_folder,
	output_dir=output_dir,
	completion_mode=completion_mode,
	use_stop=use_stop,
	include_details=include_details,
	**kwargs,
	)
	for key, value in input_dict.items()
	]
	coroutines = [rtwl(task) for task in data_generations_tasks]
	TASK_TIMEOUT_SECONDS = 600 # 10 minutes timeout
	processed_count = 0
	total_tasks = len(coroutines)
	for future in tqdmasyncio.tqdm.as_completed(coroutines):
	processed_count += 1
	try:
	await asyncio.wait_for(future, timeout=TASK_TIMEOUT_SECONDS)
	except asyncio.TimeoutError:
	print(
	f"\nWARNING: Task {processed_count}/{total_tasks} for step '{self.output_file}' timed out after {TASK_TIMEOUT_SECONDS} seconds.",
	file=sys.stderr,
	)
	# Task is cancelled by wait_for, loop continues
	except Exception as e_inner: # Catch exceptions from the future itself
	print(
	f"\nError processing task {processed_count}/{total_tasks}: {e_inner}"
	)
	traceback.print_exc() # Optionally log traceback

	filter_out_nonpresent_keys(input_dict, key_to_check=self.result_key)
	except Exception as e:
	print(f"Exception occurred during task execution: {e}")
	# traceback.print_exc()
	raise e
	finally:
	# Robust save with interrupt handling
	self.save_dataset(input_dict=input_dict, output_dir=output_dir)