lemonade-sdk 9.1.1__py3-none-any.whl

lemonade/tools/huggingface/utils.py

@@ -0,0 +1,359 @@
+from typing import Dict, List, Tuple
+import time
+from contextlib import nullcontext
+import transformers
+import torch
+from lemonade.state import State
+from lemonade.tools.adapter import TokenizerAdapter
+from lemonade.tools.adapter import ModelAdapter
+from lemonade.tools.bench import Bench
+
+# Command line interfaces for tools will use string inputs for data
+# types, however the internal tool logic will need to know the actual
+# torch type
+str_to_dtype = {
+    "float32": torch.float32,
+    "float16": torch.float16,
+    "bfloat16": torch.bfloat16,
+    "int8_static": torch.int8,
+    "int8_dynamic": torch.int8,
+}
+
+
+def make_example_inputs(state: State) -> Dict:
+    """
+    Create a dictionary of LLM inputs that can be passed as an argument
+    into quantization, ONNX export, etc.
+    """
+
+    tokenizer = state.tokenizer
+    inputs_ids = tokenizer("Hello there", return_tensors="pt").input_ids
+    return {"input_ids": inputs_ids}
+
+
+class HuggingfaceTokenizerAdapter(TokenizerAdapter):
+    def __init__(self, tokenizer: transformers.AutoTokenizer, device: str):
+        super().__init__(tokenizer)
+        self.tokenizer = tokenizer
+        self.device = device
+
+    def __call__(self, prompt, **kwargs):
+        tokens = self.tokenizer(prompt, **kwargs)
+        if self.device:
+            return tokens.to(self.device)
+        else:
+            return tokens
+
+    def decode(self, response, **kwargs):
+        return self.tokenizer.decode(response, **kwargs)
+
+    def batch_decode(self, tokens, **kwargs):
+        return self.tokenizer.batch_decode(tokens, **kwargs)
+
+    @property
+    def eos_token_id(self):
+        return self.tokenizer.eos_token_id
+
+    def save_pretrained(self, model_dir, **kwargs):
+        return self.tokenizer.save_pretrained(model_dir, **kwargs)
+
+
+class HuggingfaceAdapter(ModelAdapter):
+    """
+    Wrapper class for Huggingface LLMs that handle generation arguments
+    from callers to match HF specification.
+
+        repetition_penalty: helps the LLM avoid repeating the same short
+            phrase in the response over and over.
+        temperature: helps the LLM stay focused on the prompt.
+        do_sample: apply the temperature.
+    """
+
+    def __init__(self, model, dtype=torch.float32, device="cpu", tokenizer=None):
+        super().__init__()
+        self.model = model
+        self.dtype = dtype
+        self.device = device
+        self.tokenizer = tokenizer
+
+    def generate(
+        self,
+        input_ids,
+        random_seed=1,
+        **kwargs,
+    ):
+
+        # Move input_ids to the same device as the model
+        input_ids = input_ids.to(self.device)
+
+        # Fix temperature handling to avoid errors:
+        # If temperature is 0.0, force do_sample=False (greedy decoding)
+        if kwargs.get("temperature") == 0.0:
+            kwargs["do_sample"] = False
+
+        # If do_sample is False and temperature is 0.0, remove temperature
+        # to avoid the warning from HuggingFace.
+        # Note: This is the same approach taken by LM Eval Harness for handling temperature.
+        generation_kwargs = {
+            "max_new_tokens": kwargs.get("max_new_tokens", 512),
+            "do_sample": kwargs.get("do_sample", True),
+            **kwargs,
+        }
+
+        if random_seed is None:
+            torch.random.seed()
+        else:
+            torch.random.manual_seed(random_seed)
+
+        with torch.no_grad(), torch.inference_mode():
+            outputs = self.model.generate(input_ids=input_ids, **generation_kwargs)
+
+        self.prompt_tokens = input_ids.shape[1]
+        self.response_tokens = len(outputs[0]) - self.prompt_tokens
+        return outputs
+
+    def _model_call(self, input_tensor):
+        """Forward pass through the model to get logits
+
+        This method directly calls the model forward pass rather than using model.generate() for
+        several important reasons:
+        1. Purpose: We need raw logits from a single forward pass, while generate() is for producing
+           multiple tokens through iterative inference
+        2. Efficiency: Direct calls are more efficient for logprob calculations with no sampling
+           overhead
+        3. Precision: Logprob calculations require exact control over input-to-output mapping
+        4. Consistency: Similar approach used in both HF and OGA implementations
+
+        Args:
+            input_tensor: Input token IDs tensor
+
+        Returns:
+            Logits tensor from model forward pass
+        """
+        with torch.no_grad(), torch.inference_mode():
+            outputs = self.model(input_tensor)
+            return outputs.logits
+
+    def _select_cont_toks(self, logits, context_len, cont_toks):
+        """
+        Select logits corresponding to continuation tokens and gather their probabilities
+
+        Args:
+            logits: Model output logits
+            context_len: Length of input context
+            cont_toks: List of continuation token IDs
+
+        Returns:
+            Tensor of log probabilities for continuation tokens
+        """
+        # Get the continuation logits (discard context logits)
+        cont_logits = logits[context_len - 1 : context_len - 1 + len(cont_toks)]
+
+        # Convert cont_toks to tensor if needed
+        if not isinstance(cont_toks, torch.Tensor):
+            cont_toks = torch.tensor(cont_toks, dtype=torch.long, device=logits.device)
+
+        # Gather log probs at the corresponding token indices
+        log_probs = torch.log_softmax(cont_logits, dim=-1)
+        token_log_probs = torch.gather(log_probs, 1, cont_toks.unsqueeze(-1)).squeeze(
+            -1
+        )
+
+        return token_log_probs
+
+    def compute_logprobs(
+        self, text, tokenizer, prompt_length=None, logprobs=None, echo=False
+    ):
+        """
+        Compute log probabilities for all tokens in the given text.
+
+        Args:
+            text: The full text to analyze (e.g., prompt + completion)
+            prompt_length: Number of tokens in the prompt. If provided and echo=False,
+                only completion tokens after this position will be returned.
+            logprobs: If not None, return log probabilities. Value indicates how many top
+                alternatives to return. If True but not an integer, defaults to 5 alternatives.
+            echo: If True, include logprobs for prompt tokens. If False, only return logprobs
+                for completion tokens.
+
+        Returns:
+            - text_offset: Character offsets for each token in the text
+            - token_logprobs: Log probability for each token
+            - tokens: The actual tokens used
+            - top_logprobs: Top alternative log probabilities for each position
+        """
+        if tokenizer is None:
+            raise ValueError("Tokenizer is required for logprob calculation")
+
+        # Encode the full text
+        tokens = tokenizer(text).input_ids
+
+        # Track character offsets for each token
+        text_offset = []
+        start_idx = 0
+
+        token_strings = []
+        for token_id in tokens:
+            token_str = tokenizer.decode([token_id])
+            token_strings.append(token_str)
+
+            # Calculate character offsets for tokens - handles cases where tokens
+            # may not directly match in the original text due to encoding differences,
+            # special characters, or tokenization artifacts
+            try:
+                pos = text[start_idx:].find(token_str)
+                if pos != -1:
+                    text_offset.append(start_idx + pos)
+                    start_idx += pos + len(token_str)
+                else:
+                    text_offset.append(start_idx)
+            except (TypeError, ValueError, UnicodeError):
+                # Fallback to current position when matching fails due to encoding issues
+                text_offset.append(start_idx)
+
+        # Convert to tensor and get model output
+        input_tensor = torch.tensor([tokens], dtype=torch.long, device=self.device)
+        logits = self._model_call(input_tensor)[0]
+
+        # Calculate log probabilities for each token
+        all_log_probs = torch.log_softmax(logits, dim=-1)
+
+        # The first token doesn't have a conditional probability
+        # For tokens after the first, get the predicted probability
+        token_log_probs = []
+        top_logprobs_list = []
+
+        # For each position, get the actual token probability and top alternatives
+        for i in range(len(tokens)):
+            # Get previous token position logits
+            if i > 0:  # First token has no preceding context
+                prev_logits = all_log_probs[i - 1]
+                curr_token_id = tokens[i]
+                # Get probability of the actual token that appeared
+                token_logprob = prev_logits[curr_token_id].item()
+                token_log_probs.append(token_logprob)
+
+                # Get top-k alternatives if requested
+                if logprobs is not None:
+                    num_alternatives = logprobs if isinstance(logprobs, int) else 5
+                    topk_values, topk_indices = torch.topk(
+                        prev_logits, min(num_alternatives, prev_logits.size(-1))
+                    )
+
+                    # Create dictionary of token: logprob
+                    position_logprobs = {}
+                    for val, idx in zip(topk_values.tolist(), topk_indices.tolist()):
+                        token_str = tokenizer.decode([idx])
+                        position_logprobs[token_str] = val
+
+                    top_logprobs_list.append(position_logprobs)
+            else:
+                # For the first token, we don't have a conditional probability
+                token_log_probs.append(None)
+                top_logprobs_list.append({})
+
+        # If we don't want to echo prompt tokens, filter them out
+        if not echo and prompt_length is not None:
+            # Ensure prompt_length is within bounds
+            prompt_length = min(prompt_length, len(tokens))
+
+            # Filter results to only include completion tokens
+            if prompt_length < len(tokens):
+                filtered_text_offset = text_offset[prompt_length:]
+                filtered_token_logprobs = token_log_probs[prompt_length:]
+                filtered_tokens = token_strings[prompt_length:]
+                filtered_top_logprobs = top_logprobs_list[prompt_length:]
+
+                return (
+                    filtered_text_offset,
+                    filtered_token_logprobs,
+                    filtered_tokens,
+                    filtered_top_logprobs,
+                )
+            else:
+                # No completion tokens
+                return [], [], [], []
+
+        return text_offset, token_log_probs, token_strings, top_logprobs_list
+
+
+def benchmark_huggingface_llm(
+    model: torch.nn.Module,
+    tokenizer,
+    input_ids,
+    dtype,
+    num_beams: int,
+    target_output_tokens: int,
+    iterations: int,
+    warmup_iterations: int,
+    report_progress_fn,
+) -> List[Tuple[float, int]]:
+
+    amp_enabled = True if (dtype == torch.float16 or dtype == torch.bfloat16) else False
+    # The "if amp_enabled else nullcontext()" is to get around a bug in PyTorch 2.1
+    # where torch.cpu.amp.autocast(enabled=False) does nothing
+    with (
+        torch.cpu.amp.autocast(enabled=amp_enabled, dtype=dtype)
+        if amp_enabled
+        else nullcontext()
+    ):
+
+        per_iteration_result = []
+        tokens_out_len_list = []
+
+        # Early stopping is only a valid parameter with multiple beams
+        early_stopping = num_beams > 1
+
+        with torch.no_grad(), torch.inference_mode():
+            # Don't capture time for warmup
+            for count in range(warmup_iterations):
+                outputs = model.generate(
+                    input_ids,
+                    num_beams=num_beams,
+                    max_new_tokens=target_output_tokens,
+                    min_new_tokens=target_output_tokens,
+                    early_stopping=early_stopping,
+                    pad_token_id=tokenizer.eos_token_id,
+                )
+                tokens_out_len_list.append(outputs.shape[1] - input_ids.shape[1])
+                report_progress_fn((count + 1) / (warmup_iterations + iterations))
+
+            for count in range(iterations):
+                # CUDA synchronization is required prior to GPU benchmarking
+                # This has no negative effect on CPU-only benchmarks, and is more robust than
+                # checking `model.device == "cuda"` since it applies to multi-GPU environments
+                # Synchronization is done before collecting the start time because this will
+                # ensure that the GPU has finished initialization tasks such as loading weights
+                if torch.cuda.is_available():
+                    torch.cuda.synchronize()
+                start_time = time.perf_counter()
+
+                outputs = model.generate(
+                    input_ids,
+                    num_beams=num_beams,
+                    max_new_tokens=target_output_tokens,
+                    min_new_tokens=target_output_tokens,
+                    early_stopping=early_stopping,
+                    pad_token_id=tokenizer.eos_token_id,
+                )
+
+                if torch.cuda.is_available():
+                    torch.cuda.synchronize()
+                end_time = time.perf_counter()
+
+                latency = end_time - start_time
+
+                tokens_out_len_list.append(model.response_tokens)
+
+                # Only count an iteration if it produced enough tokens
+                if model.response_tokens >= target_output_tokens:
+                    per_iteration_result.append((latency, model.response_tokens))
+
+                report_progress_fn(
+                    (warmup_iterations + count + 1) / (warmup_iterations + iterations)
+                )
+
+        if not per_iteration_result:
+            raise Bench.not_enough_tokens(target_output_tokens)
+
+    return per_iteration_result, tokens_out_len_list

lemonade/tools/humaneval.py

@@ -0,0 +1,264 @@
+import argparse
+import os
+import csv
+from typing import Dict, Optional, Any
+
+
+from lemonade.state import State
+from lemonade.tools import Tool
+import lemonade.common.printing as printing
+import lemonade.common.build as build
+
+
+class AccuracyHumaneval(Tool):
+    """
+    HumanEval accuracy measurement tool.
+
+    This tool evaluates language models on the HumanEval dataset, which consists of
+    Python programming problems. It measures the model's ability to:
+    1. Generate functionally correct code completions
+    2. Pass unit tests for each programming problem
+
+    Metrics:
+    - pass@1: Percentage of problems solved with 1 generation attempt
+    - pass@10: Percentage of problems solved within 10 generation attempts
+    - pass@100: Percentage of problems solved within 100 generation attempts
+
+    See docs/dev_cli/humaneval_accuracy.md for more details
+    """
+
+    unique_name = "accuracy-humaneval"
+    DATASET = "https://github.com/openai/human-eval/blob/master/data/HumanEval.jsonl.gz?raw=true"
+    TOTAL_PROBLEMS = 164  # Total number of problems in the HumanEval dataset
+
+    def __init__(self):
+        super().__init__(monitor_message="Measuring accuracy with HumanEval")
+        self.status_stats = []
+        # Enable code evaluation for HumanEval
+        os.environ["HF_ALLOW_CODE_EVAL"] = "1"
+
+    @staticmethod
+    def parser(add_help: bool = True) -> argparse.ArgumentParser:
+        parser = __class__.helpful_parser(
+            short_description="Measure coding accuracy with HumanEval",
+            add_help=add_help,
+        )
+        parser.add_argument(
+            "--k-samples",
+            type=int,
+            default=1,
+            help="Number of completions to generate per prompt for pass@k calculation"
+            " (default: %(default)s)",
+        )
+        parser.add_argument(
+            "--first-n-samples",
+            type=int,
+            default=AccuracyHumaneval.TOTAL_PROBLEMS,
+            help=f"Evaluate only the first N problems from the dataset (default: "
+            f"%(default)s, evaluates all {AccuracyHumaneval.TOTAL_PROBLEMS} problems)",
+        )
+        parser.add_argument(
+            "--timeout",
+            type=float,
+            default=30.0,
+            help="Timeout in seconds for each test case (default: %(default)s)",
+        )
+        parser.add_argument(
+            "--data-dir",
+            type=str,
+            default=None,
+            help="Custom directory for dataset storage (default: %(default)s, "
+            "uses <lemonade_cache_dir>/data/humaneval)",
+        )
+        return parser
+
+    def run(
+        self,
+        state: State,
+        data_dir: Optional[str] = None,
+        k_samples: int = 1,
+        first_n_samples: Optional[int] = TOTAL_PROBLEMS,
+        timeout: float = 30.0,
+    ) -> State:
+        """
+        Run HumanEval evaluation on the model.
+
+        Args:
+            state: Current state containing model and tokenizer
+            data_dir: Optional custom directory for dataset storage
+            k_samples: Number of completions to generate per prompt for pass@k calculation
+            first_n_samples: Number of first N problems to evaluate
+            timeout: Timeout in seconds for each test case
+
+        Returns:
+            Updated state with evaluation results
+        """
+
+        # Validate required state components
+        if not hasattr(state, "model") or not hasattr(state, "tokenizer"):
+            raise ValueError("State must contain both 'model' and 'tokenizer'")
+
+        # Setup directories
+        data_dir_to_use = data_dir or os.path.join(state.cache_dir, "data", "humaneval")
+        data_path = os.path.join(data_dir_to_use, "HumanEval.jsonl.gz")
+        model_results_dir = os.path.join(
+            build.output_dir(state.cache_dir, state.build_name), "humaneval"
+        )
+        os.makedirs(model_results_dir, exist_ok=True)
+
+        # Download dataset if needed
+        self._download_dataset(data_path)
+
+        # Run evaluation
+        results = self._evaluate_model(
+            state.model,
+            state.tokenizer,
+            data_path,
+            k_samples,
+            timeout,
+            model_results_dir,
+            first_n_samples,
+        )
+
+        # Save metrics
+        self._save_metrics(state, results)
+
+        return state
+
+    def _download_dataset(self, output_path: str) -> None:
+        """Download HumanEval dataset if not already present."""
+
+        import requests
+
+        if os.path.exists(output_path):
+            printing.log_info(f"Dataset already exists at: {output_path}")
+            return
+
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        response = requests.get(self.DATASET, stream=True)
+
+        if response.status_code == 200:
+            with open(output_path, "wb") as file:
+                for chunk in response.iter_content(chunk_size=8192):
+                    file.write(chunk)
+            printing.log_info(f"Dataset downloaded successfully to: {output_path}")
+        else:
+            raise RuntimeError(
+                f"Failed to download dataset. Status code: {response.status_code}"
+            )
+
+    def _evaluate_model(
+        self,
+        model: Any,
+        tokenizer: Any,
+        data_path: str,
+        k_samples: int,
+        timeout: float,
+        results_dir: str,
+        first_n_samples: Optional[int] = TOTAL_PROBLEMS,
+    ) -> Dict[str, float]:
+        """
+        Evaluate model on HumanEval dataset.
+
+        Args:
+            model: The language model to evaluate
+            tokenizer: The tokenizer for the model
+            data_path: Path to the HumanEval dataset
+            k_samples: Number of completions per prompt for pass@k calculation
+            timeout: Test case timeout in seconds
+            results_dir: Directory to save results
+            first_n_samples: Number of first N problems to evaluate
+
+        Returns:
+            Dictionary containing evaluation metrics
+        """
+
+        from human_eval.data import write_jsonl, read_problems
+        from human_eval.evaluation import evaluate_functional_correctness
+
+        dataset = read_problems(data_path)
+
+        # Limit to first N problems
+        dataset_keys = list(dataset.keys())[:first_n_samples]
+        ignore_incomplete = True
+
+        samples = []
+
+        # Update Tool progress monitor
+        self.set_percent_progress(0.0)
+        questions_completed = 0
+        number_of_questions = first_n_samples * k_samples
+
+        # Save completions and expected answers
+        csv_path = os.path.join(results_dir, "evaluation_results.csv")
+        with open(
+            csv_path, mode="w", newline="", encoding="utf-8", errors="replace"
+        ) as file:
+            writer = csv.writer(file)
+            writer.writerow(["Prompt", "Completion", "Expected Answer"])
+
+            for task_id in dataset_keys:
+                try:
+                    for _ in range(k_samples):
+                        prompt = dataset[task_id]["prompt"]
+                        expected = dataset[task_id]["canonical_solution"]
+
+                        # Generate completion
+                        input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+                        completion = model.generate(
+                            input_ids,
+                            max_new_tokens=512,
+                            do_sample=False,
+                        )
+                        completion_text = tokenizer.decode(
+                            completion[0], skip_special_tokens=True
+                        )
+
+                        # Save results
+                        samples.append(
+                            {"task_id": task_id, "completion": completion_text}
+                        )
+                        writer.writerow([prompt, completion_text, expected])
+
+                        # Update progress monitor after completing all samples for a question
+                        questions_completed = questions_completed + 1
+                        percent_completed = (
+                            questions_completed / number_of_questions * 100
+                        )
+                        self.set_percent_progress(percent_completed)
+
+                # pylint: disable=W0718
+                except Exception as e:
+                    printing.log_info(f"Error processing task {task_id}: {str(e)}")
+                    continue
+
+        # Save predictions and evaluate
+        pred_path = os.path.join(results_dir, "humaneval_predictions.jsonl")
+        write_jsonl(pred_path, samples)
+        printing.log_info(f"Results saved in: {results_dir}")
+
+        # Run functional correctness evaluation
+        k_values = [k_samples]
+        results = evaluate_functional_correctness(
+            pred_path,
+            k_values,
+            n_workers=1,
+            timeout=timeout,
+            problem_file=data_path,
+            ignore_incomplete=ignore_incomplete,
+        )
+        return results
+
+    def _save_metrics(self, state: State, results: Dict[str, float]) -> None:
+        """Save evaluation metrics to state."""
+        for metric, value in results.items():
+            metric_name = f"humaneval_{metric}"
+            state.save_stat(
+                metric_name, float(value) * 100 if value is not None else None
+            )
+            state.save_stat(f"{metric_name}_units", "%")
+            self.status_stats.append(metric_name)
+
+
+# This file was originally licensed under Apache 2.0. It has been modified.
+# Modifications Copyright (c) 2025 AMD