lemonade-sdk 7.0.0__py3-none-any.whl

lemonade/tools/management_tools.py

@@ -0,0 +1,273 @@
+import argparse
+import abc
+from typing import List
+import lemonade.common.filesystem as fs
+import lemonade.common.exceptions as exp
+import lemonade.common.printing as printing
+from lemonade.tools.tool import ToolParser
+from lemonade.version import __version__ as lemonade_version
+from lemonade.common.system_info import get_system_info_dict
+from lemonade.common.build import output_dir
+import lemonade.cache as lemonade_cache
+
+
+class ManagementTool(abc.ABC):
+    """
+    Intended for management functions, such as managing the cache
+    or printing the version number.
+    """
+
+    unique_name: str
+
+    @classmethod
+    def helpful_parser(cls, short_description: str, **kwargs):
+        epilog = (
+            f"`{cls.unique_name}` is a Management Tool. It is intended to be invoked by itself "
+            "(i.e., not as part of a sequence), to accomplish a utility function. "
+        )
+
+        return ToolParser(
+            prog=f"lemonade {cls.unique_name}",
+            short_description=short_description,
+            description=cls.__doc__,
+            epilog=epilog,
+            **kwargs,
+        )
+
+    @staticmethod
+    @abc.abstractmethod
+    def parser() -> argparse.ArgumentParser:
+        """
+        Static method that returns an ArgumentParser that defines the command
+        line interface for this Tool.
+        """
+
+    # pylint: disable=unused-argument
+    def parse(self, args, known_only=True) -> argparse.Namespace:
+        """
+        Run the parser and return a Namespace of keyword arguments that the user
+        passed to the Tool via the command line.
+
+        Tools should extend this function only if they require specific parsing
+        logic.
+
+        Args:
+            args: command line arguments passed from the CLI.
+            known_only: this argument allows the CLI framework to
+                incrementally parse complex commands.
+        """
+
+        if known_only:
+            parsed_args = self.__class__.parser().parse_args(args)
+        else:
+            parsed_args, _ = self.__class__.parser().parse_known_args(args)
+
+        return parsed_args
+
+    @abc.abstractmethod
+    def run(self, cache_dir: str):
+        """
+        Execute the functionality of the Tool.
+        """
+
+    def parse_and_run(self, cache_dir: str, args, known_only=True):
+        """
+        Helper function to parse CLI arguments into the args expected
+        by run(), and then forward them into the run() method.
+        """
+
+        parsed_args = self.parse(args, known_only)
+        self.run(cache_dir, **parsed_args.__dict__)
+
+
+class Version(ManagementTool):
+    """
+    Simply prints the version number of the lemonade installation.
+    """
+
+    unique_name = "version"
+
+    @staticmethod
+    def parser(add_help: bool = True) -> argparse.ArgumentParser:
+        parser = __class__.helpful_parser(
+            short_description="Print the lemonade version number",
+            add_help=add_help,
+        )
+
+        return parser
+
+    def run(self, _):
+        print(lemonade_version)
+
+
+class Cache(ManagementTool):
+    # pylint: disable=pointless-statement,f-string-without-interpolation
+    f"""
+    A set of functions for managing the lemonade build cache. The default
+    cache location is {lemonade_cache.DEFAULT_CACHE_DIR}, and can also be 
+    selected with
+    the global --cache-dir option or the LEMONADE_CACHE_DIR environment variable.
+
+    Users must set either "--all" or "--build-names" to let the tool
+    know what builds to operate on.
+
+    Users must also set one of the available actions (e.g., list, stats, etc.).
+
+    That action will be applied to all selected builds.
+    """
+
+    unique_name = "cache"
+
+    @staticmethod
+    def parser(add_help: bool = True) -> argparse.ArgumentParser:
+        # NOTE: `--cache-dir` is set as a global input to the lemonade CLI and
+        # passed directly to the `run()` method
+
+        parser = __class__.helpful_parser(
+            short_description="Manage the build cache " f"",
+            add_help=add_help,
+        )
+
+        build_selection_group = parser.add_mutually_exclusive_group(required=True)
+
+        build_selection_group.add_argument(
+            "-b",
+            "--build-names",
+            nargs="+",
+            help="Name of the specific build(s) to be operated upon, within the cache directory",
+        )
+
+        build_selection_group.add_argument(
+            "-a",
+            "--all",
+            dest="all_builds",
+            help="Operate on all the builds in the cache",
+            action="store_true",
+        )
+
+        action_group = parser.add_mutually_exclusive_group(required=True)
+
+        action_group.add_argument(
+            "-l",
+            "--list",
+            dest="list_builds",
+            action="store_true",
+            help="List all of the builds in the cache",
+        )
+
+        action_group.add_argument(
+            "-s",
+            "--stats",
+            action="store_true",
+            help="Print the collected stats for the selected build(s)",
+        )
+
+        action_group.add_argument(
+            "--delete",
+            action="store_true",
+            help="Permanently delete the selected build(s)",
+        )
+
+        action_group.add_argument(
+            "--clean",
+            action="store_true",
+            help="Remove the build artifacts from the selected build(s)",
+        )
+
+        return parser
+
+    def run(
+        self,
+        cache_dir: str,
+        all_builds: bool = False,
+        build_names: List[str] = None,
+        list_builds: bool = False,
+        stats: bool = False,
+        delete: bool = False,
+        clean: bool = False,
+    ):
+        fs.check_cache_dir(cache_dir)
+
+        if all_builds and build_names:
+            raise ValueError(
+                "all_builds and build_names are mutually exclusive, "
+                "but both are used in this call."
+            )
+        elif all_builds:
+            builds = fs.get_available_builds(cache_dir)
+        elif build_names:
+            builds = build_names
+        else:
+            raise ValueError(
+                "Either all_builds or build_names must be set, "
+                "but this call sets neither."
+            )
+
+        # Print a nice heading
+        printing.log_info(f"Operating on cache directory {cache_dir}")
+
+        if not builds:
+            printing.log_warning("No builds found.")
+
+        for build in builds:
+            build_path = output_dir(cache_dir, build_name=build)
+            if fs.is_build_dir(cache_dir, build):
+                # Run actions on the build
+                # These actions are intended to be mutually exclusive, so we
+                # use an if-elif block in order from least to most destructive
+                if list_builds:
+                    print(build)
+                elif stats:
+                    fs.print_yaml_file(fs.Stats(cache_dir, build).file, "stats")
+                elif clean:
+                    fs.clean_output_dir(cache_dir, build)
+                    printing.log_info(f"Removed the build artifacts from: {build}")
+
+                elif delete:
+                    fs.rmdir(build_path)
+                    printing.log_info(f"Deleted build: {build}")
+            else:
+                raise exp.CacheError(
+                    f"No build found with name: {build}. "
+                    "Try running `lemonade cache --list` to see the builds in your build cache."
+                )
+
+        print()
+
+
+class SystemInfo(ManagementTool):
+    """
+    Prints system information for the lemonade installation.
+    """
+
+    unique_name = "system-info"
+
+    @staticmethod
+    def parser(add_help: bool = True) -> argparse.ArgumentParser:
+        parser = __class__.helpful_parser(
+            short_description="Print system information",
+            add_help=add_help,
+        )
+
+        return parser
+
+    @staticmethod
+    def pretty_print(my_dict: dict, level=0):
+        for k, v in my_dict.items():
+            if isinstance(v, dict):
+                print("    " * level + f"{k}:")
+                SystemInfo.pretty_print(v, level + 1)
+            elif isinstance(v, list):
+                print("    " * level + f"{k}:")
+                for item in v:
+                    print("    " * (level + 1) + f"{item}")
+            else:
+                print("    " * level + f"{k}: {v}")
+
+    def run(self, _):
+        system_info_dict = get_system_info_dict()
+        self.pretty_print(system_info_dict)
+
+
+# This file was originally licensed under Apache 2.0. It has been modified.
+# Modifications Copyright (c) 2025 AMD

lemonade/tools/mmlu.py

@@ -0,0 +1,327 @@
+import argparse
+import os
+import tarfile
+from pathlib import Path
+from typing import List, Optional
+import subprocess
+import numpy as np
+import pandas as pd
+import requests
+from lemonade.state import State
+from lemonade.tools import Tool
+import lemonade.common.printing as printing
+import lemonade.common.build as build
+import lemonade.common.filesystem as fs
+
+# Constants
+choices = ["A", "B", "C", "D"]
+dataset_url = "https://people.eecs.berkeley.edu/~hendrycks/data.tar"
+
+
+def min_handle_none(*args: int):
+    """
+    Returns the minimum of the arguments. If one of the arguments is none,
+    it doesn't count towards the min.
+    """
+
+    filter_out_none = (value for value in args if value is not None)
+    return min(filter_out_none)
+
+
+class AccuracyMMLU(Tool):
+    """
+    See docs/lemonade/mmlu_accuracy.md for more details
+    """
+
+    unique_name = "accuracy-mmlu"
+
+    def __init__(self):
+        super().__init__(monitor_message="Measuring accuracy with MMLU")
+        self.status_stats = []
+
+    @staticmethod
+    def parser(add_help: bool = True) -> argparse.ArgumentParser:
+        parser = __class__.helpful_parser(
+            short_description="Measure accuracy with Massive Multitask "
+            "Language Understanding (MMLU)",
+            add_help=add_help,
+        )
+
+        parser.add_argument(
+            "--ntrain",
+            type=int,
+            default=5,
+            help="Number of training examples to use. Default set to 5 for `5 Shot`",
+        )
+        parser.add_argument(
+            "--max-evals",
+            type=int,
+            default=None,
+            help="Maximum evaluations to run per test",
+        )
+        parser.add_argument(
+            "--data-dir",
+            type=str,
+            required=False,
+            help="Directory containing test and dev data (default: lemonade cache).",
+        )
+        parser.add_argument(
+            "--tests",
+            nargs="+",
+            help=(
+                "Specific tests to run. For a single quick test, we suggest 'management'."
+                + "Default: run all tests."
+            ),
+        )
+        return parser
+
+    def run(
+        self,
+        state: State,
+        ntrain: int = 5,
+        max_evals: int = None,
+        data_dir: Optional[str] = None,
+        tests: List[str] = None,
+    ) -> State:
+
+        if data_dir:
+            data_dir_to_use = data_dir
+        else:
+            data_dir_to_use = os.path.join(state.cache_dir, "data", "mmlu")
+
+        # Setup MMLU dataset
+        dataset_dir = download_and_extract_dataset(data_dir_to_use, dataset_url)
+
+        model_results_dir = os.path.join(
+            build.output_dir(state.cache_dir, state.build_name), "mmlu"
+        )
+        os.makedirs(model_results_dir, exist_ok=True)
+
+        tests_to_run = [
+            f.replace("_test.csv", "")
+            for f in sorted(os.listdir(os.path.join(dataset_dir, "test")))
+            if f.endswith("_test.csv")
+        ]
+        if tests is not None:
+            unsupported_tests = set(tests) - set(tests_to_run)
+            if unsupported_tests:
+                raise ValueError(
+                    f"Invalid test names provided: {', '.join(unsupported_tests)}. "
+                    f"Valid tests are: {', '.join(tests_to_run)}"
+                )
+            tests_to_run = [test for test in tests if test in tests_to_run]
+
+        tokenizer = state.tokenizer
+        model = state.model
+
+        # Update Tool progress monitor
+        self.set_percent_progress(0.0)
+        number_of_questions = float(
+            sum(
+                [
+                    min_handle_none(
+                        len(
+                            _safe_read_csv(
+                                os.path.join(dataset_dir, "test", f"{subject}_test.csv")
+                            )
+                        ),
+                        max_evals,
+                    )
+                    for subject in tests_to_run
+                ]
+            )
+        )
+
+        questions_completed = 0
+
+        summary_data = []
+        for subject in tests_to_run:
+            dev_df = _safe_read_csv(
+                os.path.join(dataset_dir, "dev", f"{subject}_dev.csv")
+            )[:ntrain]
+            test_df = _safe_read_csv(
+                os.path.join(dataset_dir, "test", f"{subject}_test.csv")
+            )
+
+            # Evaluate the model on the test data for a given subject
+            detailed_results = []
+
+            for i in range(min_handle_none(test_df.shape[0], max_evals)):
+                prompt = _gen_prompt(dev_df, subject, ntrain) + _format_example(
+                    test_df, i, include_answer=False
+                )
+                input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+
+                response_text = _generate_response(tokenizer, model, input_ids)
+                try:
+                    pred_label = response_text[-1].upper()
+                # Handle models generating empty outputs
+                except IndexError:
+                    pred_label = "-"
+
+                label = test_df.iloc[i, -1].strip().upper()
+                detailed_results.append(
+                    {
+                        "Question": test_df.iloc[i, 0],
+                        "Prompt": prompt,
+                        "Correct Answer": label,
+                        "Generated Answer": pred_label,
+                        "Correct": pred_label == label,
+                    }
+                )
+
+                # Update progress monitor
+                questions_completed = questions_completed + 1
+                percent_completed = questions_completed / number_of_questions * 100
+                self.set_percent_progress(percent_completed)
+
+            acc = np.mean([res["Correct"] for res in detailed_results])
+
+            subject_results_df = pd.DataFrame(detailed_results)
+            subject_csv_path = os.path.join(
+                model_results_dir, f"{subject}_detailed_results.csv"
+            )
+            subject_results_df.to_csv(subject_csv_path, index=False)
+
+            # Update summary_data with total questions and correct answers
+            correct_answers_count = sum(
+                result["Correct"] for result in detailed_results
+            )
+
+            summary_data.append(
+                {
+                    "Subject": subject,
+                    "Accuracy": acc,
+                    "Total Questions": len(test_df),
+                    "Evaluated Questions": (
+                        max_evals
+                        if max_evals is not None and max_evals < len(test_df)
+                        else len(test_df)
+                    ),
+                    "Correct Answers": correct_answers_count,
+                }
+            )
+
+            # Save accuracy results to stats file
+            # And display in the CLI
+            stat_name = f"mmlu_{subject}_accuracy"
+            stat_units_name = f"{stat_name}_units"
+            state.save_stat(stat_name, float(acc) * 100)
+            state.save_stat(stat_units_name, "%")
+            self.status_stats.append(stat_name)
+
+        # Calculate average of mmlu accuracy and display in the CLI
+        acc_avg = np.mean([accuracy_data["Accuracy"] for accuracy_data in summary_data])
+        state.save_stat(fs.Keys.AVERAGE_MMLU_ACCURACY, float(acc_avg) * 100)
+        state.save_stat(f"{fs.Keys.AVERAGE_MMLU_ACCURACY}_units", "%")
+        self.status_stats.append(fs.Keys.AVERAGE_MMLU_ACCURACY)
+
+        # 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
+
+
+def _list_tests(data_dir):
+    """Lists all available tests based on the files in the test data directory."""
+    test_files = [
+        f for f in os.listdir(os.path.join(data_dir, "test")) if f.endswith("_test.csv")
+    ]
+    print(
+        "Available tests:",
+        *[f.replace("_test.csv", "") for f in sorted(test_files)],
+        sep="\n",
+    )
+
+
+def _format_subject(subject):
+    """Formats a subject string by replacing underscores with spaces."""
+    return " ".join(subject.split("_"))
+
+
+def _safe_read_csv(path):
+    """Safely reads a CSV file and returns a DataFrame."""
+    try:
+        return pd.read_csv(path, header=None)
+    except FileNotFoundError:
+        printing.log_error(f"Error: File not found - {path}")
+    except Exception as e:  # pylint: disable=broad-except
+        printing.log_error(f"An error occurred while reading {path}: {e}")
+
+
+def _format_example(df, idx, include_answer=True):
+    """Formats an example from the dataframe into a prompt string."""
+    prompt = df.iloc[idx, 0]
+    for j in range(1, df.shape[1] - 1):
+        prompt += f"\n{choices[j-1]}. {df.iloc[idx, j]}"
+    prompt += "\nAnswer_:"
+    if include_answer:
+        prompt += f" {df.iloc[idx, -1]}\n\n"
+    return prompt
+
+
+def _gen_prompt(train_df, subject, k=-1):
+    """Generates a prompt string from multiple examples."""
+    prompt = (
+        "The following are multiple choice questions (with answers) about "
+        + f"{_format_subject(subject)}.\n\n"
+    )
+    for i in range(min(k, train_df.shape[0]) if k != -1 else train_df.shape[0]):
+        prompt += _format_example(train_df, i)
+    return prompt
+
+
+def _generate_response(tokenizer, model, input_ids):
+    """Generates a model response for the given input IDs."""
+    try:
+        response = model.generate(input_ids, max_new_tokens=1)
+        return tokenizer.decode(response[0], skip_special_tokens=True).strip()
+    except subprocess.CalledProcessError as e:
+        printing.log_warning(
+            f"Subprocess failed with command: {e} and error message: {e.stderr}"
+        )
+    except Exception as e:  # pylint: disable=broad-except
+        printing.log_warning(f"Error during model generation: {e}")
+    return ""  # Return an empty string on failure
+
+
+def download_and_extract_dataset(data_cache_dir: str, dataset_url: str):
+    """
+    Download the dataset from the given URL and extract it into the target directory.
+    """
+
+    # Create the directory if it does not exist
+    Path(data_cache_dir).mkdir(parents=True, exist_ok=True)
+
+    # Check if the data already exists to avoid re-downloading
+    if not os.listdir(data_cache_dir):  # Checks if the directory is empty
+        printing.log_info(f"Downloading dataset to {data_cache_dir}")
+
+        # Download the dataset
+        response = requests.get(dataset_url, stream=True)
+        if response.status_code == 200:
+            tar_path = os.path.join(data_cache_dir, "data.tar")
+            with open(tar_path, "wb") as f:
+                f.write(response.raw.read())
+
+            printing.log_info("Extracting dataset...")
+            # Extract the tar file
+            with tarfile.open(tar_path) as tar:
+                tar.extractall(path=data_cache_dir)
+            os.remove(tar_path)
+            printing.log_info("Dataset ready.")
+        else:
+            printing.log_info("Failed to download the dataset.")
+    else:
+        printing.log_info(
+            f"Dataset already exists in {data_cache_dir}, skipping download."
+        )
+
+    # MMLU data is stored in data.tar/data
+    return os.path.join(data_cache_dir, "data")
+
+
+# This file was originally licensed under Apache 2.0. It has been modified.
+# Modifications Copyright (c) 2025 AMD