lemonade-sdk 7.0.0__py3-none-any.whl

lemonade/state.py

@@ -0,0 +1,159 @@
+import os
+import sys
+from typing import Dict, Optional, Any
+import yaml
+import lemonade.common.build as build
+import lemonade.common.filesystem as fs
+from lemonade.version import __version__ as lemonade_version
+
+
+def _is_nice_to_write(value):
+    """
+    Checks whether a value is nice to write to YAML.
+    Returns True if the value is a string, int, float, bool, list, dict, or tuple.
+    Returns False otherwise.
+    """
+    if isinstance(value, (str, int, float, bool)):
+        return True
+    elif isinstance(value, list) or isinstance(value, tuple):
+        # Check if all elements in the list are nice to write
+        return all(_is_nice_to_write(item) for item in value)
+    elif isinstance(value, dict):
+        # Check if all values in the dictionary are nice to write
+        return all(_is_nice_to_write(item) for item in value.values())
+    return False
+
+
+def _sanitize_for_yaml(input_dict: Dict) -> Dict:
+    """
+    Creates a new dictionary containing only nice-to-write values
+    from the original dictionary.
+    """
+    result = {}
+    for key, value in input_dict.items():
+        if _is_nice_to_write(value):
+            result[key] = value
+    return result
+
+
+class State:
+    """
+    The State class is meant to carry build state, starting with the user's
+    initial arguments, through each build Tool in the Sequence, and finally
+    to the disk, where it is used to assess cache hits.
+
+    State is initialized with the key members that are shared by every build,
+    and reasonable default values are assigned as appropriate.
+
+    Tool developers can also add any members they wish. To get or set an
+    attribute, reference it as an attribute:
+        1. get: `my_variable = state.attribute_name`
+        2. set: `state.attribute_name = my_variable`
+
+    Build State can be saved and loaded from disk in the form of a state.yaml file
+    via State.save() and load_state(), respectively. Note that while State can
+    contain members of any type, only YAML-safe members (str, int, bool, float,
+    list, dict, tuple) will be saved and loaded.
+    """
+
+    def __init__(
+        self,
+        cache_dir: str,
+        build_name: Optional[str] = None,
+        sequence_info: Dict[str, Dict] = None,
+        **kwargs,
+    ):
+
+        # The default model name is the name of the python file that calls build_model()
+        if build_name is None:
+            build_name = os.path.basename(sys.argv[0])
+
+        # Support "~" in the cache_dir argument
+        parsed_cache_dir = os.path.expanduser(cache_dir)
+
+        # Save settings as State members
+        self.cache_dir = parsed_cache_dir
+        self.build_name = build_name
+        self.sequence_info = sequence_info
+        self.lemonade_version = lemonade_version
+        self.build_status = build.FunctionStatus.NOT_STARTED
+        self.downcast_applied = False
+        self.uid = build.unique_id()
+        self.results = None
+
+        # Store any additional kwargs as members
+        for key, value in kwargs.items():
+            self.__dict__[key] = value
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        """
+        Tool developers can add a new member to State by simply
+        assigning it as an attribute, i.e., `state.new_member = value`.
+        """
+        return super().__setattr__(name, value)
+
+    def save_stat(self, key: str, value):
+        """
+        Save statistics to an yaml file in the build directory
+        """
+
+        stats = fs.Stats(self.cache_dir, self.build_name)
+        stats.save_stat(key, value)
+
+    def save_sub_stat(self, parent_key: str, key: str, value):
+        """
+        Save statistics to an yaml file in the build directory
+        """
+
+        stats = fs.Stats(self.cache_dir, self.build_name)
+        stats.save_sub_stat(parent_key, key, value)
+
+    def save(self):
+        """
+        Save all YAML-friendly members to disk as a state.yaml file.
+
+        Note that `model` and `inputs` will typically not be saved since
+        they are typically in non-YAML-friendly types such as `torch.nn.Module`
+        and `torch.tensor`.
+        """
+
+        state_to_save = _sanitize_for_yaml(vars(self))
+
+        # Create a build directory in the cache
+        fs.make_build_dir(self.cache_dir, self.build_name)
+
+        with open(
+            build.state_file(self.cache_dir, self.build_name),
+            "w",
+            encoding="utf8",
+        ) as outfile:
+            yaml.dump(state_to_save, outfile)
+
+
+def load_state(
+    cache_dir=None,
+    build_name=None,
+    state_path=None,
+) -> State:
+    """
+    Read a state.yaml file corresponding to a specific build in a specific
+    cache, and use its contents to initialize a State instance.
+    """
+
+    if state_path is not None:
+        file_path = state_path
+    elif build_name is not None and cache_dir is not None:
+        file_path = build.state_file(cache_dir, build_name)
+    else:
+        raise ValueError(
+            "This function requires either build_name and cache_dir to be set, "
+            "or state_path to be set, not both or neither"
+        )
+
+    state_dict = build.load_yaml(file_path)
+
+    return State(**state_dict)
+
+
+# This file was originally licensed under Apache 2.0. It has been modified.
+# Modifications Copyright (c) 2025 AMD

lemonade/tools/__init__.py

@@ -0,0 +1 @@
+from .tool import Tool, FirstTool, NiceHelpFormatter

lemonade/tools/adapter.py

@@ -0,0 +1,104 @@
+import abc
+from transformers import AutoTokenizer
+
+
+class ModelAdapter(abc.ABC):
+    """
+    Base class for adapting an LLM to work with lemonade's standardized tools
+    """
+
+    def __init__(self):
+        """
+        Self-benchmarking ModelAdapters can store their results in the
+        tokens_per_second and time_to_first_token members.
+        """
+        self.tokens_per_second = None
+        self.time_to_first_token = None
+        self.type = "generic"
+
+    @abc.abstractmethod
+    def generate(self, input_ids, max_new_tokens=512):
+        """
+        Generate is the primary method required by lemonade's accuracy tools
+
+        We try to keep the signature here minimal to allow for maximum compatibility
+        with recipe components, which themselves may not support a lot of arguments.
+        """
+
+
+class TokenizerAdapter(abc.ABC):
+    """
+    Base class for adapting an LLM's tokenizer to work with lemonade's standard tools
+    """
+
+    def __init__(self, tokenizer: AutoTokenizer = None):
+        self.auto_tokenizer = tokenizer
+
+    @abc.abstractmethod
+    def __call__(self, prompt: str):
+        """
+        Args:
+            prompt: text that should be encoded and passed to the LLM as input_ids
+
+        Returns: input_ids
+        """
+
+    @abc.abstractmethod
+    def decode(self, response) -> str:
+        """
+        Args:
+            response: tokens from the LLM that should be decoded into text
+
+        Returns: text response of the LLM
+        """
+
+    def apply_chat_template(self, *args, **kwargs):
+        """
+        Convert messages into a single tokenizable string
+        """
+        return self.auto_tokenizer.apply_chat_template(*args, **kwargs)
+
+    @property
+    def chat_template(self):
+        return self.auto_tokenizer.chat_template
+
+    @property
+    def eos_token(self):
+        return self.auto_tokenizer.eos_token
+
+
+class PassthroughTokenizerResult:
+    """
+    Data structure for holding a tokenizer result where the input_ids
+    are packaged in a non-standard way, but we still want to adhere to
+    standard interfaces (e.g., result.input_ids).
+
+    For example: CLI-based tools that have their own internal tokenizer that
+    isn't exposed to the user. In this case we can pass the prompt through as
+    a string.
+    """
+
+    def __init__(self, prompt):
+        self.input_ids = prompt
+
+
+class PassthroughTokenizer(TokenizerAdapter):
+    """
+    Tokenizer adapter that forwards the prompt to input_ids as text,
+    and then forwards a text LLM response through decode() as text.
+
+    Useful for CLI-based tools that have their own internal tokenizer that
+    isn't exposed to the user.
+    """
+
+    # pylint: disable=unused-argument
+    def __call__(self, prompt: str, **kwargs):
+        return PassthroughTokenizerResult(prompt)
+
+    # pylint: disable=unused-argument
+    def decode(self, response: str, **kwargs):
+        return response
+
+
+# This file was originally licensed under Apache 2.0. It has been modified.
+# Modifications Copyright (c) 2025 AMD

lemonade/tools/bench.py

@@ -0,0 +1,284 @@
+from abc import ABC, abstractmethod
+import argparse
+import os
+import platform
+import psutil
+from lemonade.state import State
+from lemonade.tools import Tool
+from lemonade.cache import Keys
+
+default_iterations = 10
+default_warmup_runs = 5
+default_prompt_length = 64
+default_output_tokens = 32
+default_prompt = "Hello, I am conscious and"
+
+
+class Bench(Tool, ABC):
+    """
+    Abstract parent class for tools that benchmark the performance of the generate()
+    method of an LLM.
+    """
+
+    def __init__(self, monitor_message="Benchmarking LLM"):
+        super().__init__(monitor_message)
+
+        # The minimum set of statistics that a benchmark tool will produce
+        # Inherited tools should append any additional statistics they generate to this list
+        self.status_stats = [
+            Keys.SECONDS_TO_FIRST_TOKEN,
+            Keys.STD_DEV_SECONDS_TO_FIRST_TOKEN,
+            Keys.TOKEN_GENERATION_TOKENS_PER_SECOND,
+            Keys.PREFILL_TOKENS_PER_SECOND,
+            Keys.PROMPT_TOKENS,
+            Keys.RESPONSE_TOKENS,
+            Keys.MAX_MEMORY_USED_GBYTE,
+        ]
+
+        # Minimum per measurement statistics
+        # Inherited tools should add additional lists for other per prompt statistics
+        self.input_ids_len_list = []
+        self.tokens_out_len_list = []
+        self.mean_time_to_first_token_list = []
+        self.std_dev_time_to_first_token_list = []
+        self.prefill_tokens_per_second_list = []
+        self.token_generation_tokens_per_second_list = []
+        self.max_memory_used_gb_list = []
+
+        # Max memory used can only be measured on Windows systems
+        self.save_max_memory_used = platform.system() == "Windows"
+
+        # This is set to True only for the duration of the first call to run_prompt
+        self.first_run_prompt = None
+
+    @staticmethod
+    def parser(parser: argparse.ArgumentParser = None, add_help: bool = True):
+        # Allow inherited classes to initialize and pass in a parser, add parameters to it if so
+        if parser is None:
+            parser = __class__.helpful_parser(
+                short_description="Benchmark an LLM", add_help=add_help
+            )
+
+        parser.add_argument(
+            "--iterations",
+            "-i",
+            required=False,
+            type=int,
+            default=default_iterations,
+            help="Number of benchmarking iterations to run (default: "
+            f"{default_iterations})",
+        )
+
+        parser.add_argument(
+            "--warmup-iterations",
+            "-w",
+            required=False,
+            type=int,
+            default=default_warmup_runs,
+            help="Number of benchmarking iterations to use for cache warmup "
+            "(the results of these iterations "
+            f"are not included in the results; default: {default_warmup_runs})",
+        )
+
+        parser.add_argument(
+            "--prompts",
+            "-p",
+            nargs="+",
+            required=False,
+            default=[str(default_prompt_length)],
+            metavar="PROMPT",
+            help="Input one or more prompts to the LLM. Three formats are supported. "
+            "1) integer: use a synthetic prompt with the specified length "
+            "2) str: use a user-provided prompt string "
+            "3) path/to/prompt.txt: load the prompt from a text file. "
+            f"(default: {default_prompt_length}) ",
+        )
+
+        parser.add_argument(
+            "--output-tokens",
+            required=False,
+            type=int,
+            default=default_output_tokens,
+            help="Number of new tokens the LLM should make (default: "
+            f"{default_output_tokens})",
+        )
+
+        return parser
+
+    def get_prompt_str(self, _state, token_length):
+        """
+        Returns a string with approximately the prescribed token length.
+        Note: Actual token length is dependent on the tokenizer.
+        """
+        return "word " * (token_length - 1)
+
+    def parse(self, state: State, args, known_only=True) -> argparse.Namespace:
+        """
+        Helper function to parse CLI arguments into the args expected by run()
+        """
+
+        parsed_args = super().parse(state, args, known_only)
+
+        if parsed_args.prompts is None:
+            parsed_args.prompts = [str(default_prompt_length)]
+
+        # Decode prompt arg into a list of prompt strings
+        prompt_strings = []
+        for prompt_item in parsed_args.prompts:
+            if prompt_item.isdigit():
+                # Generate a prompt with the requested length
+                token_length = int(prompt_item)
+                prompt_strings.append(self.get_prompt_str(state, token_length))
+
+            elif os.path.exists(prompt_item):
+                with open(prompt_item, "r", encoding="utf-8") as f:
+                    prompt_strings.append(f.read())
+
+            else:
+                # No change to the prompt
+                prompt_strings.append(prompt_item)
+        parsed_args.prompts = prompt_strings
+
+        return parsed_args
+
+    def run(
+        self,
+        state: State,
+        prompts: list[str] = None,
+        iterations: int = default_iterations,
+        warmup_iterations: int = default_warmup_runs,
+        output_tokens: int = default_output_tokens,
+        **kwargs,
+    ) -> State:
+        """
+        Args:
+            - prompts: List of input prompts used as starting points for LLM text generation
+            - iterations: number of benchmarking samples to take; results are
+                reported as the median and mean of the samples.
+            - warmup_iterations: subset of the iterations to treat as warmup,
+                and not included in the results.
+            - output_tokens: Number of new tokens LLM to create.
+            - kwargs: Additional parameters used by bench tools
+        """
+
+        if prompts is None:
+            prompts = ["word " * (default_prompt_length - 2)]
+        elif isinstance(prompts, str):
+            prompts = [prompts]
+
+        state.save_stat("prompts", prompts)
+        state.save_stat("iterations", iterations)
+        state.save_stat("warmup_iterations", warmup_iterations)
+        state.save_stat("output_tokens", output_tokens)
+
+        counter = 0
+        report_progress_fn = lambda x: self.set_percent_progress(
+            100 * (counter + x) / len(prompts)
+        )
+        self.first_run_prompt = True
+        for counter, prompt in enumerate(prompts):
+            report_progress_fn(0)
+
+            self.run_prompt(
+                state,
+                report_progress_fn,
+                prompt,
+                iterations,
+                warmup_iterations,
+                output_tokens,
+                **kwargs,
+            )
+            self.first_run_prompt = False
+
+            if self.save_max_memory_used:
+                self.max_memory_used_gb_list.append(
+                    psutil.Process().memory_info().peak_wset / 1024**3
+                )
+
+        self.set_percent_progress(None)
+        self.save_stats(state)
+
+        return state
+
+    @abstractmethod
+    def run_prompt(
+        self,
+        state,
+        report_progress_fn,
+        prompt,
+        iterations,
+        warmup_iterations,
+        output_tokens,
+        **kwargs,
+    ):
+        pass
+
+    @staticmethod
+    def get_item_or_list(lst):
+        """
+        If the list is just a single item then return the item, else return the list
+        """
+        if len(lst) == 1:
+            return lst[0]
+        else:
+            return lst
+
+    def save_stats(self, state):
+        # Save performance data to stats
+        state.save_stat(
+            Keys.PROMPT_TOKENS, self.get_item_or_list(self.input_ids_len_list)
+        )
+        state.save_stat(
+            Keys.RESPONSE_TOKENS, self.get_item_or_list(self.tokens_out_len_list)
+        )
+        state.save_stat(
+            Keys.SECONDS_TO_FIRST_TOKEN,
+            self.get_item_or_list(self.mean_time_to_first_token_list),
+        )
+        if not all(
+            element is None for element in self.std_dev_time_to_first_token_list
+        ):
+            state.save_stat(
+                Keys.STD_DEV_SECONDS_TO_FIRST_TOKEN,
+                self.get_item_or_list(self.std_dev_time_to_first_token_list),
+            )
+        state.save_stat(
+            Keys.PREFILL_TOKENS_PER_SECOND,
+            self.get_item_or_list(self.prefill_tokens_per_second_list),
+        )
+        state.save_stat(
+            Keys.TOKEN_GENERATION_TOKENS_PER_SECOND,
+            self.get_item_or_list(self.token_generation_tokens_per_second_list),
+        )
+        if self.save_max_memory_used:
+            state.save_stat(
+                Keys.MAX_MEMORY_USED_GBYTE,
+                self.get_item_or_list(self.max_memory_used_gb_list),
+            )
+
+    @staticmethod
+    def not_enough_tokens(output_tokens: int):
+        """
+        Raise an exception that explains why a benchmark did not produce any results
+        """
+
+        raise ValueError(
+            "Your model was benchmarked, however none of the benchmarking "
+            "iterations produced the requested amount of output tokens "
+            f"(currently {output_tokens}), so "
+            "the results have been discarded. You have the following options "
+            "to solve this: \n"
+            "1. Use the -p option to change the prompt to something that will "
+            "produce more output tokens. For example, 'The extremely long "
+            "story of my life, told in excruciating details is:' "
+            "is an example of a prompt that will result in a lot of output. \n"
+            "2. Set a lower value for --output-tokens to make it more likely "
+            "that the model will produce enough. \n"
+            "3. Set more verbose hyperparameters. \n"
+            "4. Run more benchmarking iterations, to improve the chance of "
+            "getting at least one with enough output tokens. \n"
+        )
+
+
+# This file was originally licensed under Apache 2.0. It has been modified.
+# Modifications Copyright (c) 2025 AMD