lemonade-sdk 7.0.0__py3-none-any.whl

lemonade/tools/ort_genai/oga_bench.py

@@ -0,0 +1,142 @@
+import argparse
+import statistics
+from statistics import StatisticsError
+from lemonade.state import State
+from lemonade.cache import Keys
+from lemonade.tools.adapter import ModelAdapter, TokenizerAdapter
+from lemonade.tools.bench import Bench
+
+
+class OgaBench(Bench):
+    """
+    Benchmark any model that adheres to the ModelAdapter interface.
+
+    Required input state:
+        - MODEL: model instance to benchmark.
+        - TOKENIZER: tokenizer instance used to generate inputs for the model.
+
+    Output state produced: None
+    """
+
+    unique_name = "oga-bench"
+
+    def __init__(self):
+        super().__init__()
+
+        # Additional statistics generated by this bench tool
+        self.status_stats.insert(
+            self.status_stats.index(Keys.TOKEN_GENERATION_TOKENS_PER_SECOND) + 1,
+            Keys.STD_DEV_TOKENS_PER_SECOND,
+        )
+        self.std_dev_token_generation_tokens_per_second_list = []
+
+    @staticmethod
+    def parser(add_help: bool = True) -> argparse.ArgumentParser:
+        parser = __class__.helpful_parser(
+            short_description="Benchmark an LLM in onnxruntime-genai (OGA)",
+            add_help=add_help,
+        )
+
+        parser = Bench.parser(parser)
+
+        return parser
+
+    def get_prompt_str(self, state, token_length):
+        """
+        Returns a string with the prescribed token length.
+        """
+        tokenizer: TokenizerAdapter = state.tokenizer
+        test_prompt = "word " * (token_length - 1)
+        input_ids = tokenizer(test_prompt, return_tensors="pt").input_ids
+        test_token_length = len(input_ids)
+        delta = test_token_length - token_length
+        if delta == 0:
+            return test_prompt
+        return "word " * max(token_length - 1 - delta, 0)
+
+    def run_prompt(
+        self,
+        state: State,
+        report_progress_fn,
+        prompt: str,
+        iterations: int,
+        warmup_iterations: int,
+        output_tokens: int,
+    ) -> State:
+
+        model: ModelAdapter = state.model
+        tokenizer: TokenizerAdapter = state.tokenizer
+
+        input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+        self.input_ids_len_list.append(len(input_ids))
+        per_iteration_time_to_first_token = []
+        per_iteration_tokens_per_second = []
+
+        # Don't capture time for warmup
+        for count in range(warmup_iterations):
+            outputs = model.generate(input_ids, max_new_tokens=output_tokens)
+            self.tokens_out_len_list.append(len(outputs[0]) - len(input_ids))
+            report_progress_fn((count + 1) / (warmup_iterations + iterations))
+
+        for count in range(iterations):
+            outputs = model.generate(
+                input_ids,
+                max_new_tokens=output_tokens,
+                min_new_tokens=output_tokens,
+            )
+            report_progress_fn(
+                (warmup_iterations + count + 1) / (warmup_iterations + iterations)
+            )
+
+            token_len = len(outputs[0]) - len(input_ids)
+            self.tokens_out_len_list.append(token_len)
+
+            # Only count an iteration if it produced enough tokens
+            if token_len >= output_tokens:
+                per_iteration_time_to_first_token.append(model.time_to_first_token)
+                per_iteration_tokens_per_second.append(model.tokens_per_second)
+
+        if not per_iteration_time_to_first_token or not per_iteration_tokens_per_second:
+            raise Bench.not_enough_tokens(output_tokens)
+
+        mean_time_to_first_token = statistics.mean(per_iteration_time_to_first_token)
+        self.mean_time_to_first_token_list.append(mean_time_to_first_token)
+        self.prefill_tokens_per_second_list.append(
+            len(input_ids) / mean_time_to_first_token
+        )
+        self.token_generation_tokens_per_second_list.append(
+            statistics.mean(per_iteration_tokens_per_second)
+        )
+        try:
+            self.std_dev_time_to_first_token_list.append(
+                statistics.stdev(per_iteration_time_to_first_token)
+            )
+        except StatisticsError:
+            # Less than 2 measurements
+            self.std_dev_time_to_first_token_list.append(None)
+        try:
+            self.std_dev_token_generation_tokens_per_second_list.append(
+                statistics.stdev(per_iteration_tokens_per_second)
+            )
+        except StatisticsError:
+            # Less than 2 measurements
+            self.std_dev_token_generation_tokens_per_second_list.append(None)
+
+    def save_stats(self, state):
+        super().save_stats(state)
+
+        # Save additional statistics
+        if not all(
+            element is None
+            for element in self.std_dev_token_generation_tokens_per_second_list
+        ):
+            state.save_stat(
+                Keys.STD_DEV_TOKENS_PER_SECOND,
+                self.get_item_or_list(
+                    self.std_dev_token_generation_tokens_per_second_list
+                ),
+            )
+
+
+# This file was originally licensed under Apache 2.0. It has been modified.
+# Modifications Copyright (c) 2025 AMD

lemonade/tools/perplexity.py

@@ -0,0 +1,146 @@
+import os
+import argparse
+import pandas as pd
+import torch
+from datasets import load_dataset
+from lemonade.state import State
+from lemonade.tools import Tool
+import lemonade.common.printing as printing
+import lemonade.common.build as build
+
+
+class AccuracyPerplexity(Tool):
+    """
+    Measure perplexity of an LLM using the Wikitext-2 dataset.
+
+    Required input state:
+        - state.model: instance that provides a __call__() method that returns
+        output.logits and supports model.config.max_position_embeddings
+        - state.tokenizer: instance of Hugging Face PretrainedTokenizer
+
+    Output state produced: None
+
+    See docs/lemonade/perplexity.md for more details.
+    """
+
+    unique_name = "accuracy-perplexity"
+
+    def __init__(self):
+        super().__init__(monitor_message="Measuring perplexity")
+
+    @staticmethod
+    def parser(add_help: bool = True) -> argparse.ArgumentParser:
+        parser = __class__.helpful_parser(
+            short_description="Measure perplexity score",
+            add_help=add_help,
+        )
+        return parser
+
+    def run(
+        self,
+        state: State,
+    ) -> State:
+
+        try:
+            printing.log_info("Downloading dataset ...")
+            dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
+        except Exception as e:  # pylint: disable=broad-except
+            printing.log_error(f"Error during dataset load: {e}")
+            raise e
+
+        tokenizer = state.tokenizer
+        model = state.model
+        # Tokenize the entire test dataset text, joining entries with double new lines
+        encodings = tokenizer("\n\n".join(dataset["text"]), return_tensors="pt")
+
+        # Retrieve the maximum input length that the model can handle
+        try:
+            max_length = model.config.max_position_embeddings
+        except AttributeError:
+            # Some LLMs do not have the config.max_position_embeddings attribute
+            # However, most LLMs support at least 2048 context length, so this
+            # try-except will allow a few more LLMs to work
+            max_length = 2048
+        # Set stride to half of the maximum input length for overlapping window processing
+        # Refer to docs/perplexity.md for more information on sliding window
+        stride = max_length // 2
+        # Determine the total sequence length of the tokenized input
+        seq_len = encodings.input_ids.size(1)
+
+        negative_log_likelihoods = []
+        summary_data = []
+        prev_end_location = 0
+
+        model_results_dir = os.path.join(
+            build.output_dir(state.cache_dir, state.build_name), "perplexity"
+        )
+
+        for begin_location in range(0, seq_len, stride):
+            end_location = min(begin_location + max_length, seq_len)
+            target_len = end_location - prev_end_location
+            input_ids = encodings.input_ids[:, begin_location:end_location]
+            target_ids = input_ids.clone()
+            target_ids[:, :-target_len] = -100
+
+            # Forward pass the model to get logits
+            with torch.no_grad():
+                try:
+                    outputs = model(input_ids, labels=target_ids)
+                    logits = outputs.logits
+                except Exception as e:  # pylint: disable=broad-except
+                    printing.log_error(
+                        f"Error during model forward pass execution: {e}"
+                    )
+
+            # Compute loss manually for visualization
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = target_ids[..., 1:].contiguous()
+            effective_token_count = (target_ids != -100).sum().item()
+            negative_log_likelihoods.append(
+                (outputs.loss.item(), effective_token_count)
+            )
+
+            # Decode predicted and actual next words for the last token position
+            predictions = torch.argmax(shift_logits, dim=-1)
+            predicted_tokens = predictions[:, -1]
+            actual_tokens = shift_labels[:, -1]
+
+            predicted_words = tokenizer.batch_decode(
+                predicted_tokens, skip_special_tokens=True
+            )
+            actual_words = tokenizer.batch_decode(
+                actual_tokens, skip_special_tokens=True
+            )
+            context = tokenizer.decode(input_ids[0, :])
+
+            summary_data.append(
+                {
+                    "Context": context[-stride:],
+                    "Predicted next word": predicted_words,
+                    "Actual next word": actual_words,
+                    "Loss for this window": outputs.loss.item(),
+                }
+            )
+            prev_end_location = end_location
+
+        # Total loss calculation considering the number of tokens for each segment
+        total_loss = sum(loss * count for loss, count in negative_log_likelihoods)
+        total_tokens = sum(count for _, count in negative_log_likelihoods)
+
+        # Calculate average negative_log_likelihood and perplexity
+        average_negative_log_likelihood = total_loss / total_tokens
+        perplexity = torch.exp(torch.tensor(average_negative_log_likelihood))
+
+        # Save accuracy results to stats file
+        state.save_stat("perplexity_score", float(perplexity.item()))
+
+        # Save accuracy results to CSV file
+        summary_df = pd.DataFrame(summary_data)
+        summary_df.to_csv(
+            os.path.join(model_results_dir, "summary_results.csv"), index=False
+        )
+        return state
+
+
+# This file was originally licensed under Apache 2.0. It has been modified.
+# Modifications Copyright (c) 2025 AMD

lemonade/tools/prompt.py

@@ -0,0 +1,228 @@
+import argparse
+import os
+import matplotlib.pyplot as plt
+import lemonade.common.build as build
+import lemonade.common.printing as printing
+from lemonade.state import State
+from lemonade.tools import Tool
+from lemonade.tools.adapter import ModelAdapter, TokenizerAdapter
+from lemonade.cache import Keys
+
+DEFAULT_GENERATE_PARAMS = {
+    "do_sample": True,
+    "top_k": 50,
+    "top_p": 0.95,
+    "temperature": 0.7,
+}
+
+DEFAULT_MAX_NEW_TOKENS = 512
+DEFAULT_N_TRIALS = 1
+
+
+def sanitize_string(input_string):
+    return input_string.encode("charmap", "ignore").decode("charmap")
+
+
+def sanitize_text(text):
+    if isinstance(text, str):
+        return sanitize_string(text)
+    elif isinstance(text, list):
+        return [sanitize_string(item) for item in text]
+    else:
+        raise TypeError("Input must be a string or a list of strings.")
+
+
+def positive_int(x):
+    """Conversion function for argparse"""
+    i = int(x)
+    if i < 1:
+        raise ValueError("Non-positive values are not allowed")
+    return i
+
+
+class LLMPrompt(Tool):
+    """
+    Send a prompt to an LLM instance and print the response to the screen.
+
+    Required input state:
+        - state.model: LLM instance that supports the generate() method.
+        - state.tokenizer: LLM tokenizer instance that supports the __call__() (ie, encode)
+            and decode() methods.
+
+    Output state produced:
+        - "response": text response from the LLM.
+    """
+
+    unique_name = "llm-prompt"
+
+    def __init__(self):
+        super().__init__(monitor_message="Prompting LLM")
+
+        self.status_stats = [
+            Keys.PROMPT_TOKENS,
+            Keys.PROMPT,
+            Keys.PROMPT_TEMPLATE,
+            Keys.RESPONSE_TOKENS,
+            Keys.RESPONSE,
+            Keys.RESPONSE_LENGTHS_HISTOGRAM,
+        ]
+
+    @staticmethod
+    def parser(add_help: bool = True) -> argparse.ArgumentParser:
+        parser = __class__.helpful_parser(
+            short_description="Prompt an LLM and print the result",
+            add_help=add_help,
+        )
+
+        parser.add_argument(
+            "--prompt",
+            "-p",
+            help="Input prompt to the LLM. Two formats are supported: "
+            "1) str: use a user-provided prompt string, and "
+            "2) path/to/prompt.txt: load the prompt from a .txt file.",
+            required=True,
+        )
+
+        parser.add_argument(
+            "--template",
+            "-t",
+            action="store_true",
+            help="Insert the prompt into the model's chat template before processing.",
+        )
+
+        parser.add_argument(
+            "--max-new-tokens",
+            "-m",
+            default=DEFAULT_MAX_NEW_TOKENS,
+            type=int,
+            help=f"Maximum number of new tokens in the response "
+            f"(default is {DEFAULT_MAX_NEW_TOKENS})",
+        )
+
+        parser.add_argument(
+            "--n-trials",
+            "-n",
+            default=DEFAULT_N_TRIALS,
+            type=positive_int,
+            help=f"Number of responses the LLM will generate for the prompt "
+            f"(useful for testing, default is {DEFAULT_N_TRIALS})",
+        )
+
+        return parser
+
+    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)
+
+        # Decode prompt arg into a string prompt
+        if parsed_args.prompt.endswith(".txt") and os.path.exists(parsed_args.prompt):
+            with open(parsed_args.prompt, "r", encoding="utf-8") as f:
+                parsed_args.prompt = f.read()
+
+        return parsed_args
+
+    def run(
+        self,
+        state: State,
+        prompt: str = "Hello",
+        max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
+        n_trials: int = DEFAULT_N_TRIALS,
+        template: bool = False,
+    ) -> State:
+
+        model: ModelAdapter = state.model
+        tokenizer: TokenizerAdapter = state.tokenizer
+
+        # If template flag is set, then wrap prompt in template
+        if template:
+            # Embed prompt in model's chat template
+            if tokenizer.chat_template:
+                # Use the model's built-in chat template if available
+                messages_dict = [{"role": "user", "content": prompt}]
+                prompt = tokenizer.apply_chat_template(
+                    messages_dict, tokenize=False, add_generation_prompt=True
+                )
+                state.save_stat(Keys.PROMPT_TEMPLATE, "Model-specific")
+            else:
+                # Fallback to a standardized template
+                printing.log_info("No chat template found. Using default template.")
+                prompt = f"<|user|>\n{prompt} <|end|>\n<|assistant|>"
+                state.save_stat(Keys.PROMPT_TEMPLATE, "Default")
+
+        input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+        if isinstance(input_ids, (list, str)):
+            # OGA models return a list of tokens
+            # Our llama.cpp adapter returns a string
+            len_tokens_in = len(input_ids)
+        else:
+            # HF models return a 2-D tensor
+            len_tokens_in = input_ids.shape[1]
+
+        len_tokens_out = []
+        response_texts = []
+        for trial in range(n_trials):
+            if n_trials > 1:
+                self.set_percent_progress(100.0 * trial / n_trials)
+
+            # Get the response from the LLM, which may include the prompt in it
+            response = model.generate(
+                input_ids, max_new_tokens=max_new_tokens, **DEFAULT_GENERATE_PARAMS
+            )
+
+            # Flatten the input and response
+            input_ids_array = (
+                input_ids if isinstance(input_ids, (list, str)) else input_ids[0]
+            )
+            response_array = response if isinstance(response, str) else response[0]
+
+            # Separate the prompt from the response
+            len_tokens_out.append(len(response_array) - len_tokens_in)
+
+            input_token = 0
+            while (
+                input_token < len_tokens_in
+                and input_ids_array[input_token] == response_array[input_token]
+            ):
+                input_token += 1
+
+            # Only decode the actual response (not the prompt)
+            response_text = tokenizer.decode(
+                response_array[input_token:], skip_special_tokens=True
+            ).strip()
+            response_texts.append(response_text)
+
+        state.response = response_texts
+
+        if n_trials == 1:
+            len_tokens_out = len_tokens_out[0]
+            response_texts = response_texts[0]
+        else:
+            self.set_percent_progress(None)
+
+            # Plot data
+            plt.figure()
+            plt.hist(len_tokens_out, bins=20)
+            plt.xlabel("Response Length (tokens)")
+            plt.ylabel("Frequency")
+            plt.title(f"Histogram of Response Lengths\n{state.build_name}")
+            figure_path = os.path.join(
+                build.output_dir(state.cache_dir, state.build_name),
+                "response_lengths.png",
+            )
+            plt.savefig(figure_path)
+            state.save_stat(Keys.RESPONSE_LENGTHS_HISTOGRAM, figure_path)
+
+        state.save_stat(Keys.PROMPT_TOKENS, len_tokens_in)
+        state.save_stat(Keys.PROMPT, prompt)
+        state.save_stat(Keys.RESPONSE_TOKENS, len_tokens_out)
+        state.save_stat(Keys.RESPONSE, sanitize_text(response_texts))
+
+        return state
+
+
+# This file was originally licensed under Apache 2.0. It has been modified.
+# Modifications Copyright (c) 2025 AMD

lemonade/tools/quark/quark_load.py

@@ -0,0 +1,172 @@
+import argparse
+import os
+import sys
+
+import torch
+from lemonade.state import State
+from lemonade.tools import Tool
+import lemonade.common.printing as printing
+import lemonade.common.build as build
+from lemonade_install.install import DEFAULT_QUARK_DIR
+
+
+class QuarkLoad(Tool):
+    """
+    Load a model Quantized and exported using Quark.
+    Required Input State:
+        - state.model: Pretrained model instance to be quantized.
+        - state.tokenizer: Tokenizer instance from Hugging Face.
+    Output:
+        - state of the loaded model
+
+    See docs/quark.md for more details.
+    """
+
+    unique_name = "quark-load"
+
+    def __init__(self):
+        super().__init__(monitor_message="Load Quark Quantized model")
+
+    @staticmethod
+    def parser(add_help: bool = True) -> argparse.ArgumentParser:
+        parser = __class__.helpful_parser(
+            short_description="Load a quantized model using Quark",
+            add_help=add_help,
+        )
+
+        parser.add_argument(
+            "--quant-scheme",
+            type=str,
+            required=True,
+            default=None,
+            help="Supported quantization schemes in Quark",
+        )
+
+        parser.add_argument(
+            "--quant-algo",
+            type=str,
+            required=True,
+            default=None,
+            choices=["awq", "gptq", "autosmoothquant", None],
+            help="Supported quantization algorithms in Quark",
+        )
+
+        parser.add_argument(
+            "--torch-compile", action="store_true", help="Model torch compile"
+        )
+
+        parser.add_argument(
+            "--safetensors-model-reload",
+            action="store_true",
+            help="Safetensors model reload",
+        )
+
+        parser.add_argument(
+            "--safetensors-model-dir",
+            default=None,
+            help="Directory of safetensors model",
+        )
+
+        parser.add_argument(
+            "--params-load", action="store_true", help="Model parameters load"
+        )
+
+        parser.add_argument("--json-path", help="Specify the path of saved json file")
+
+        parser.add_argument(
+            "--safetensors-path",
+            default=None,
+            help="Specify the path of saved safetensors file",
+        )
+
+        return parser
+
+    def run(
+        self,
+        state: State,
+        quant_scheme: str,
+        quant_algo: str,
+        torch_compile: bool = False,
+        safetensors_model_reload: bool = False,
+        safetensors_model_dir: str = None,
+        params_load: bool = False,
+        json_path: str = None,
+        safetensors_path: str = None,
+    ) -> State:
+        """
+        Executes the QuarkLoad process.
+        Returns:
+            State: The updated state after loading the model.
+        Raises:
+            Exception: If an error occurs during the QuarkLoad process.
+        """
+
+        try:
+            if os.path.isdir(DEFAULT_QUARK_DIR):
+                quark_llm_path = os.path.join(
+                    DEFAULT_QUARK_DIR, "examples", "torch", "language_modeling"
+                )
+                sys.path.insert(0, quark_llm_path)
+            else:
+                raise FileNotFoundError(
+                    f"The directory {DEFAULT_QUARK_DIR} does not exist. \
+                        Please check your installation."
+                )
+
+            # Default load path specific to recipe
+            # This will NOT work
+            # The default path is now uniquely craeated with timestamp
+            # Default load path will not work. Need to pass explicit load path
+            model_export_path = os.path.join(
+                build.output_dir(state.cache_dir, state.build_name),
+                "exported_model",
+                quant_scheme,
+                quant_algo,
+            )
+
+            # Set default paths only if current values are None
+            if safetensors_model_dir is None:
+                safetensors_model_dir = model_export_path
+            if safetensors_path is None:
+                safetensors_path = os.path.join(model_export_path, "model.safetensors")
+            printing.log_info("Loading model ...")
+            if not params_load and not safetensors_model_reload:
+                raise ValueError(
+                    " Specify load format: 'params_load' or 'safetensors_model_reload'."
+                )
+
+            # Reload quantized model if specified
+            from quark.torch import load_params, import_model_info
+
+            if params_load:
+                printing.log_info(
+                    "Restoring quantized model from JSON/safetensors files"
+                )
+                model = load_params(
+                    model,
+                    json_path=json_path,
+                    safetensors_path=safetensors_path,
+                )
+            elif safetensors_model_reload:
+                printing.log_info(
+                    "Restoring quantized model from quark_safetensors files"
+                )
+                model = import_model_info(model, model_info_dir=safetensors_model_dir)
+
+            if torch_compile:
+                printing.log_info("torch.compile...")
+                model = torch.compile(model)
+
+            state.model = model
+            state.dtype = model.dtype
+
+            printing.log_info("Quark Load process completed.")
+
+        except Exception as e:
+            printing.log_error(f"An error occurred during the QuarkLoad process: {e}")
+            raise
+        return state
+
+
+# This file was originally licensed under Apache 2.0. It has been modified.
+# Modifications Copyright (c) 2025 AMD