Spaces:

Leon4gr45
/

agent-0

Running

agent-0 / augmentoolkit /generation_functions /generation_step_class.py

Upload folder using huggingface_hub

42da440 verified 10 days ago

9.08 kB

import re
import os
import traceback
import logging
import yaml
from augmentoolkit.generation_functions.safe_formatter import safe_format


class OutputProcessorError(Exception):
    """Exception raised when an output processor function fails.
    This error is used to provide more context when the output processing
    function encounters an error during execution.
    """


class GenerationStep:
    def __init__(
        self,
        prompt_path="",  # relative to the Inputs directory
        regex=re.compile(r".*", re.DOTALL),  # take whole completion
        sampling_params={
            "temperature": 1,
            "top_p": 1,
            "max_tokens": 3000,
            "stop": [
                "### Response",
                "\n\n\n\n\n",
                "</s>",
                "# Input:",
                "[INST]",
                "### Instruction",
                "### Information",
                "## Information",
                "## Instruction",
                "Name:",
                "<|eot_id|>",
                "<|start_header_id|>",
                "<|end_header_id|>",
            ],
        },
        completion_mode=True,  # Chat vs completion mode
        retries=0,
        engine_wrapper=None,
        logging_level=logging.INFO,  # Default logging level
        output_processor=lambda x: x,  # to ensure that control flow does not need to have decision code handling the outputs of the LLM, you can pass in a function to handle and modify the outputs (post regex) here. By default it's just the identity function and does nothing.
        return_input_too=True,
        default_prompt_folder="prompts",
        prompt_folder="prompts",
        use_stop=True,
        messages=None,
    ):
        self.prompt_path = prompt_path
        self.regex = regex
        self.sampling_params = sampling_params
        if not use_stop:
            if "stop" in self.sampling_params:
                del self.sampling_params["stop"]
        self.completion_mode = completion_mode
        self.retries = retries
        self.logging_level = logging_level
        self.output_processor = output_processor
        self.return_input_too = return_input_too
        if not engine_wrapper:
            raise Exception("Engine wrapper not passed in!")
        self.engine_wrapper = engine_wrapper
        self.prompt_folder = prompt_folder
        self.default_prompt_folder = default_prompt_folder
        logging.basicConfig(
            level=self.logging_level, format="%(asctime)s - %(levelname)s - %(message)s"
        )
        self.messages = messages

    async def generate(self, **kwargs):
        # Current file directory
        current_dir = os.path.dirname(os.path.abspath(__file__))

        if not self.messages:
            # Get the full path of the prompt file
            ideal_path = os.path.join(
                current_dir, "..", "..", self.prompt_folder, self.prompt_path
            )
            if os.path.exists(ideal_path):
                full_prompt_path = ideal_path
            else:
                full_prompt_path = os.path.join(
                    current_dir,
                    "..",
                    "..",
                    self.default_prompt_folder,
                    self.prompt_path,
                )

            with open(full_prompt_path, "r", encoding="utf-8", errors="ignore") as pf:
                prompt = pf.read()

        # Submit generation and return response, retrying as needed
        times_tried = 0
        if self.completion_mode:
            prompt_formatted = safe_format(prompt, **kwargs)
            while times_tried <= self.retries:
                try:
                    response, timeout = await self.engine_wrapper.submit_completion(
                        prompt_formatted, self.sampling_params
                    )
                    filtered_response = re.search(self.regex, response).group(1)
                    try:
                        ret = self.output_processor(filtered_response)
                    except Exception as e:
                        traceback.print_exc()
                        print(f"Output processor failed to execute on output: {e}")
                        raise OutputProcessorError()
                    if self.return_input_too:
                        return (
                            ret,
                            prompt_formatted + filtered_response,
                            filtered_response,
                            prompt_formatted,
                        )
                    return ret, timeout
                except Exception as e:
                    # logging.error(f"Error in Generation Step: {e}")
                    try:
                        if not self.engine_wrapper.mode == "llamacpp":
                            print("Response:")
                            print(response)
                    except Exception as e:
                        print(f"Error {e}")
                        traceback.print_exc()
                        pass
                    traceback.print_exc()
                    times_tried += 1
            raise Exception("Generation step failed -- too many retries!")
        else:
            if not self.messages:
                messages = yaml.safe_load(prompt)
            else:
                messages = self.messages
            new_messages = []
            for message in messages:
                try:
                    new_messages.append(
                        {
                            "role": message["role"],
                            "content": safe_format(message["content"], **kwargs),
                        }
                    )
                except Exception as e:
                    new_messages.append(
                        {"role": message["role"], "content": message["content"]}
                    )
            messages = new_messages

            # messages = [{
            #     "role": message["role"],
            #     "content": safe_format(message["content"],**arguments)
            #     }
            #             for message in messages]
            while times_tried <= self.retries:
                try:

                    # strip whitespace added by yaml load
                    messages = [
                        {
                            "role": message["role"],
                            "content": message["content"].strip(),
                        }
                        for message in messages
                    ]
                    # print("\n\n\nBEGIN DEBUG")
                    # print(messages)
                    # print("END DEBUG\n\n\n")
                    response, timeout = await self.engine_wrapper.submit_chat(
                        messages, self.sampling_params
                    )
                    try:
                        ret = self.output_processor(response)
                    except Exception as e:
                        traceback.print_exc()
                        print(f"Output processor failed to execute on output: {e}")
                        raise OutputProcessorError()
                    if self.return_input_too:
                        return (
                            ret,
                            yaml.dump(
                                messages
                                + [
                                    {
                                        "role": "assistant",
                                        "content": response,
                                        "timeout": timeout,
                                    }
                                ],
                                default_flow_style=False,
                                allow_unicode=True,
                            ),
                            response,
                            messages,
                        )
                    return ret, timeout
                except Exception as e:
                    logging.error(f"Error in Generation Step: {e}")
                    if self.completion_mode:
                        print("Prompt:")
                        print(prompt)
                    else:
                        print("Messages:")
                        print(
                            yaml.dump(
                                messages, default_flow_style=False, allow_unicode=True
                            )
                        )
                    try:
                        print("\n\nResponse:\n-----")
                        print(response)
                    except UnboundLocalError:
                        print("No response to print")
                        pass
                    # if prompt_formatted:
                    #     print(prompt_formatted)
                    logging.error(
                        f"Above prompt resulted in error, probably the model's fault: {e}"
                    )
                    traceback.print_exc()
                    times_tried += 1
            raise Exception("Generation step failed -- too many retries!")

	import re
	import os
	import traceback
	import logging
	import yaml
	from augmentoolkit.generation_functions.safe_formatter import safe_format


	class OutputProcessorError(Exception):
	"""Exception raised when an output processor function fails.

	This error is used to provide more context when the output processing
	function encounters an error during execution.
	"""


	class GenerationStep:
	def __init__(
	self,
	prompt_path="", # relative to the Inputs directory
	regex=re.compile(r".*", re.DOTALL), # take whole completion
	sampling_params={
	"temperature": 1,
	"top_p": 1,
	"max_tokens": 3000,
	"stop": [
	"### Response",
	"\n\n\n\n\n",
	"</s>",
	"# Input:",
	"[INST]",
	"### Instruction",
	"### Information",
	"## Information",
	"## Instruction",
	"Name:",
	"<\|eot_id\|>",
	"<\|start_header_id\|>",
	"<\|end_header_id\|>",
	],
	},
	completion_mode=True, # Chat vs completion mode
	retries=0,
	engine_wrapper=None,
	logging_level=logging.INFO, # Default logging level
	output_processor=lambda x: x, # to ensure that control flow does not need to have decision code handling the outputs of the LLM, you can pass in a function to handle and modify the outputs (post regex) here. By default it's just the identity function and does nothing.
	return_input_too=True,
	default_prompt_folder="prompts",
	prompt_folder="prompts",
	use_stop=True,
	messages=None,
	):
	self.prompt_path = prompt_path
	self.regex = regex
	self.sampling_params = sampling_params
	if not use_stop:
	if "stop" in self.sampling_params:
	del self.sampling_params["stop"]
	self.completion_mode = completion_mode
	self.retries = retries
	self.logging_level = logging_level
	self.output_processor = output_processor
	self.return_input_too = return_input_too
	if not engine_wrapper:
	raise Exception("Engine wrapper not passed in!")
	self.engine_wrapper = engine_wrapper
	self.prompt_folder = prompt_folder
	self.default_prompt_folder = default_prompt_folder
	logging.basicConfig(
	level=self.logging_level, format="%(asctime)s - %(levelname)s - %(message)s"
	)
	self.messages = messages

	async def generate(self, **kwargs):
	# Current file directory
	current_dir = os.path.dirname(os.path.abspath(__file__))

	if not self.messages:
	# Get the full path of the prompt file
	ideal_path = os.path.join(
	current_dir, "..", "..", self.prompt_folder, self.prompt_path
	)
	if os.path.exists(ideal_path):
	full_prompt_path = ideal_path
	else:
	full_prompt_path = os.path.join(
	current_dir,
	"..",
	"..",
	self.default_prompt_folder,
	self.prompt_path,
	)

	with open(full_prompt_path, "r", encoding="utf-8", errors="ignore") as pf:
	prompt = pf.read()

	# Submit generation and return response, retrying as needed
	times_tried = 0
	if self.completion_mode:
	prompt_formatted = safe_format(prompt, **kwargs)
	while times_tried <= self.retries:
	try:
	response, timeout = await self.engine_wrapper.submit_completion(
	prompt_formatted, self.sampling_params
	)
	filtered_response = re.search(self.regex, response).group(1)
	try:
	ret = self.output_processor(filtered_response)
	except Exception as e:
	traceback.print_exc()
	print(f"Output processor failed to execute on output: {e}")
	raise OutputProcessorError()
	if self.return_input_too:
	return (
	ret,
	prompt_formatted + filtered_response,
	filtered_response,
	prompt_formatted,
	)
	return ret, timeout
	except Exception as e:
	# logging.error(f"Error in Generation Step: {e}")
	try:
	if not self.engine_wrapper.mode == "llamacpp":
	print("Response:")
	print(response)
	except Exception as e:
	print(f"Error {e}")
	traceback.print_exc()
	pass
	traceback.print_exc()
	times_tried += 1
	raise Exception("Generation step failed -- too many retries!")
	else:
	if not self.messages:
	messages = yaml.safe_load(prompt)
	else:
	messages = self.messages
	new_messages = []
	for message in messages:
	try:
	new_messages.append(
	{
	"role": message["role"],
	"content": safe_format(message["content"], **kwargs),
	}
	)
	except Exception as e:
	new_messages.append(
	{"role": message["role"], "content": message["content"]}
	)
	messages = new_messages

	# messages = [{
	# "role": message["role"],
	# "content": safe_format(message["content"],**arguments)
	# }
	# for message in messages]
	while times_tried <= self.retries:
	try:

	# strip whitespace added by yaml load
	messages = [
	{
	"role": message["role"],
	"content": message["content"].strip(),
	}
	for message in messages
	]
	# print("\n\n\nBEGIN DEBUG")
	# print(messages)
	# print("END DEBUG\n\n\n")
	response, timeout = await self.engine_wrapper.submit_chat(
	messages, self.sampling_params
	)
	try:
	ret = self.output_processor(response)
	except Exception as e:
	traceback.print_exc()
	print(f"Output processor failed to execute on output: {e}")
	raise OutputProcessorError()
	if self.return_input_too:
	return (
	ret,
	yaml.dump(
	messages
	+ [
	{
	"role": "assistant",
	"content": response,
	"timeout": timeout,
	}
	],
	default_flow_style=False,
	allow_unicode=True,
	),
	response,
	messages,
	)
	return ret, timeout
	except Exception as e:
	logging.error(f"Error in Generation Step: {e}")
	if self.completion_mode:
	print("Prompt:")
	print(prompt)
	else:
	print("Messages:")
	print(
	yaml.dump(
	messages, default_flow_style=False, allow_unicode=True
	)
	)
	try:
	print("\n\nResponse:\n-----")
	print(response)
	except UnboundLocalError:
	print("No response to print")
	pass
	# if prompt_formatted:
	# print(prompt_formatted)
	logging.error(
	f"Above prompt resulted in error, probably the model's fault: {e}"
	)
	traceback.print_exc()
	times_tried += 1
	raise Exception("Generation step failed -- too many retries!")