New generate function

custom_generate/generate.py

@@ -14,21 +14,25 @@ In addition, I will modify this using my own method which is immaturely called s
 import copy
 import torch
 import logging
-from warnings import warn
+import inspect
+import warnings
 from transformers import GenerationConfig, LogitsProcessorList, StoppingCriteriaList
 from transformers.generation.utils import (
-    GenerationMixin,
+    GenerateOutput,
     GenerateNonBeamOutput,
     GenerateEncoderDecoderOutput,
     GenerateDecoderOnlyOutput,
-    _split_model_outputs
+    _split_model_outputs,
 )
+import logging
 from transformers.generation.candidate_generator import (
     _prepare_attention_mask, 
     _prepare_token_type_ids,
 )
 from typing import Optional, Union
 
+logger = logging.getLogger(__name__)
+
 def _speculative_sampling(
     candidate_input_ids,
     candidate_logits,
@@ -123,7 +127,7 @@ def _speculative_sampling(
 
     return valid_tokens, n_matches
 
-def generate(
+def _speculative_cascades(
     model,
     input_ids: torch.LongTensor,
     logits_processor: LogitsProcessorList,
@@ -192,14 +196,14 @@ def generate(
     # The cache must be dynamic for assisted generation, and the check must happen AFTER preparing cache
     if not model_kwargs["use_cache"]:
         model_kwargs["use_cache"] = True
-        warn("assisted generate requires `use_cache=True`. It has been set automatically.")
+        raise ValueError("assisted generate requires `use_cache=True`")
     if generation_config.cache_implementation in ["static", "hybrid", "sliding_window"] or (
         "past_key_values" in model_kwargs
         and hasattr(model_kwargs["past_key_values"], "layers")
         and any(getattr(l, "is_compileable", False) for l in model_kwargs["past_key_values"].layers)
     ):
         generation_config.cache_implementation = "dynamic"
-        warn("assisted generate is not supported with Static cache classes`. Chagned `cache_implementation` to `dynamic`")
+        raise ValueError("assisted generate is not supported with Static cache classes`")
     # Get the candidate generator, given the parameterization
     candidate_generator = model._get_candidate_generator(
         generation_config=generation_config,
@@ -428,4 +432,296 @@ def generate(
                 past_key_values=model_kwargs.get("past_key_values"),
             )
     else:
-        return input_ids
+        return input_ids
+    
+@torch.no_grad()
+def generate(
+    model,
+    inputs: torch.Tensor | None = None,
+    generation_config: GenerationConfig | None = None,
+    logits_processor: LogitsProcessorList | None = None,
+    stopping_criteria: StoppingCriteriaList | None = None,
+    prefix_allowed_tokens_fn: Callable[[int, torch.Tensor], list[int]] | None = None,
+    synced_gpus: bool | None = None,
+    assistant_model: Optional["PreTrainedModel"] = None,
+    streamer: Optional["BaseStreamer"] = None,
+    negative_prompt_ids: torch.Tensor | None = None,
+    negative_prompt_attention_mask: torch.Tensor | None = None,
+    **kwargs,
+) -> Union[GenerateOutput, torch.LongTensor]:
+    r"""
+
+    Generates sequences of token ids for models with a language modeling head.
+
+    <Tip warning={true}>
+
+    Most generation-controlling parameters are set in `generation_config` which, if not passed, will be set to the
+    model's default generation configuration. You can override any `generation_config` by passing the corresponding
+    parameters to generate(), e.g. `.generate(inputs, num_beams=4, do_sample=True)`.
+
+    For an overview of generation strategies and code examples, check out the [following
+    guide](../generation_strategies).
+
+    </Tip>
+
+    Parameters:
+        inputs (`torch.Tensor` of varying shape depending on the modality, *optional*):
+            The sequence used as a prompt for the generation or as model inputs to the encoder. If `None` the
+            method initializes it with `bos_token_id` and a batch size of 1. For decoder-only models `inputs`
+            should be in the format of `input_ids`. For encoder-decoder models *inputs* can represent any of
+            `input_ids`, `input_values`, `input_features`, or `pixel_values`.
+        generation_config ([`~generation.GenerationConfig`], *optional*):
+            The generation configuration to be used as base parametrization for the generation call. `**kwargs`
+            passed to generate matching the attributes of `generation_config` will override them. If
+            `generation_config` is not provided, the default will be used, which has the following loading
+            priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model
+            configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s
+            default values, whose documentation should be checked to parameterize generation.
+        logits_processor (`LogitsProcessorList`, *optional*):
+            Custom logits processors that complement the default logits processors built from arguments and
+            generation config. If a logit processor is passed that is already created with the arguments or a
+            generation config an error is thrown. This feature is intended for advanced users.
+        stopping_criteria (`StoppingCriteriaList`, *optional*):
+            Custom stopping criteria that complements the default stopping criteria built from arguments and a
+            generation config. If a stopping criteria is passed that is already created with the arguments or a
+            generation config an error is thrown. If your stopping criteria depends on the `scores` input, make
+            sure you pass `return_dict_in_generate=True, output_scores=True` to `generate`. This feature is
+            intended for advanced users.
+        prefix_allowed_tokens_fn (`Callable[[int, torch.Tensor], list[int]]`, *optional*):
+            If provided, this function constraints the beam search to allowed tokens only at each step. If not
+            provided no constraint is applied. This function takes 2 arguments: the batch ID `batch_id` and
+            `input_ids`. It has to return a list with the allowed tokens for the next generation step conditioned
+            on the batch ID `batch_id` and the previously generated tokens `inputs_ids`. This argument is useful
+            for constrained generation conditioned on the prefix, as described in [Autoregressive Entity
+            Retrieval](https://huggingface.co/papers/2010.00904).
+        synced_gpus (`bool`, *optional*):
+            Whether to continue running the while loop until max_length. Unless overridden, this flag will be set
+            to `True` if using `FullyShardedDataParallel` or DeepSpeed ZeRO Stage 3 with multiple GPUs to avoid
+            deadlocking if one GPU finishes generating before other GPUs. Otherwise, defaults to `False`.
+        assistant_model (`PreTrainedModel`, *optional*):
+            An assistant model that can be used to accelerate generation. The assistant model must have the exact
+            same tokenizer. The acceleration is achieved when forecasting candidate tokens with the assistant model
+            is much faster than running generation with the model you're calling generate from. As such, the
+            assistant model should be much smaller.
+        streamer (`BaseStreamer`, *optional*):
+            Streamer object that will be used to stream the generated sequences. Generated tokens are passed
+            through `streamer.put(token_ids)` and the streamer is responsible for any further processing.
+        negative_prompt_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            The negative prompt needed for some processors such as CFG. The batch size must match the input batch
+            size. This is an experimental feature, subject to breaking API changes in future versions.
+        negative_prompt_attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Attention_mask for `negative_prompt_ids`.
+        custom_generate (`str` or `Callable`, *optional*):
+            One of the following:
+            - `str` (Hugging Face Hub repository name): runs the custom `generate` function defined at
+                `custom_generate/generate.py` in that repository instead of the standard `generate` method. The
+                repository fully replaces the generation logic, and the return type may differ.
+            - `str` (local repository path): same as above but from a local path, `trust_remote_code` not required.
+            - `Callable`: `generate` will perform the usual input preparation steps, then call the provided callable to
+                run the decoding loop.
+            For more information, see [the docs](../../generation_strategies#custom-generation-methods).
+        kwargs (`dict[str, Any]`, *optional*):
+            Ad hoc parametrization of `generation_config` and/or additional model-specific kwargs that will be
+            forwarded to the `forward` function of the model. If the model is an encoder-decoder model, encoder
+            specific kwargs should not be prefixed and decoder specific kwargs should be prefixed with *decoder_*.
+
+    Return:
+        [`~utils.ModelOutput`] or `torch.LongTensor`: A [`~utils.ModelOutput`] (if `return_dict_in_generate=True`
+        or when `config.return_dict_in_generate=True`) or a `torch.LongTensor`.
+
+            If the model is *not* an encoder-decoder model (`model.config.is_encoder_decoder=False`), the possible
+            [`~utils.ModelOutput`] types are:
+
+                - [`~generation.GenerateDecoderOnlyOutput`],
+                - [`~generation.GenerateBeamDecoderOnlyOutput`]
+
+            If the model is an encoder-decoder model (`model.config.is_encoder_decoder=True`), the possible
+            [`~utils.ModelOutput`] types are:
+
+                - [`~generation.GenerateEncoderDecoderOutput`],
+                - [`~generation.GenerateBeamEncoderDecoderOutput`]
+    """
+    # 1. Handle kwargs, `generation_config`, validate them and obtain generation mode
+    generation_mode_kwargs = model._extract_generation_mode_kwargs(
+        None,
+        kwargs,
+        synced_gpus,
+        assistant_model,
+        streamer,
+    )
+
+    # Check length values before updating the config with defaults. We'll use it later to define the final min/max length (# 6)
+    has_default_max_length = kwargs.get("max_length") is None and (
+        generation_config is None or generation_config.max_length is None
+    )
+    has_default_min_length = kwargs.get("min_length") is None and (
+        generation_config is None or generation_config.min_length is None
+    )
+    generation_config, model_kwargs = model._prepare_generation_config(generation_config, **kwargs)
+
+    generation_mode = generation_config.get_generation_mode(assistant_model)
+    decoding_method = _speculative_cascades
+
+    model._validate_model_kwargs(model_kwargs.copy())
+    model._validate_generation_mode(generation_mode, generation_config, generation_mode_kwargs)
+
+    # 2. Set generation parameters if not already defined
+    logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
+    stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
+
+    accepts_attention_mask = "attention_mask" in set(inspect.signature(model.forward).parameters.keys())
+    requires_attention_mask = "encoder_outputs" not in model_kwargs
+    kwargs_has_attention_mask = model_kwargs.get("attention_mask", None) is not None
+
+    # 3. Define model inputs
+    inputs_tensor, model_input_name, model_kwargs = model._prepare_model_inputs(
+        inputs, generation_config.bos_token_id, model_kwargs
+    )
+    # Some generation modes (e.g. assisted) need `inputs_tensor` to rerun encoder.forward()
+    if "inputs_tensor" in inspect.signature(decoding_method).parameters.keys():
+        generation_mode_kwargs["inputs_tensor"] = inputs_tensor
+    batch_size = inputs_tensor.shape[0]
+
+    device = inputs_tensor.device
+    model._prepare_special_tokens(generation_config, kwargs_has_attention_mask, device=device)
+
+    # decoder-only models must use left-padding for batched generation.
+    if not model.config.is_encoder_decoder:
+        # If `input_ids` was given, check if the last id in any sequence is `pad_token_id`
+        # Note: If using, `inputs_embeds` this check does not work, because we want to be more hands-off.
+        if (
+            generation_config._pad_token_tensor is not None
+            and batch_size > 1
+            and len(inputs_tensor.shape) == 2
+            and torch.sum(inputs_tensor[:, -1] == generation_config._pad_token_tensor) > 0
+        ):
+            logger.warning(
+                "A decoder-only architecture is being used, but right-padding was detected! For correct "
+                "generation results, please set `padding_side='left'` when initializing the tokenizer."
+            )
+
+    # 4. Define other model kwargs
+    # decoder-only models with inputs_embeds forwarding must use caching (otherwise we can't detect whether we are
+    # generating the first new token or not, and we only want to use the embeddings for the first new token)
+    if not model.config.is_encoder_decoder and model_input_name == "inputs_embeds":
+        generation_config.use_cache = True
+
+    if not kwargs_has_attention_mask and requires_attention_mask and accepts_attention_mask:
+        model_kwargs["attention_mask"] = model._prepare_attention_mask_for_generation(
+            inputs_tensor, generation_config, model_kwargs
+        )
+    elif kwargs_has_attention_mask:
+        # TODO (joao): generalize this check with other types of inputs
+        if model_input_name == "input_ids" and len(model_kwargs["attention_mask"].shape) > 2:
+            raise ValueError("`attention_mask` passed to `generate` must be 2D.")
+
+    if model.config.is_encoder_decoder and "encoder_outputs" not in model_kwargs:
+        # if model is encoder decoder encoder_outputs are created and added to `model_kwargs`
+        model_kwargs = model._prepare_encoder_decoder_kwargs_for_generation(
+            inputs_tensor, model_kwargs, model_input_name, generation_config
+        )
+
+    # 5. Prepare `input_ids` which will be used for auto-regressive generation
+    if model.config.is_encoder_decoder:
+        input_ids, model_kwargs = model._prepare_decoder_input_ids_for_generation(
+            batch_size=batch_size,
+            model_input_name=model_input_name,
+            model_kwargs=model_kwargs,
+            decoder_start_token_id=generation_config._decoder_start_token_tensor,
+            device=inputs_tensor.device,
+        )
+    else:
+        input_ids = inputs_tensor if model_input_name == "input_ids" else model_kwargs.pop("input_ids")
+
+    # Expand inputs depending on the generation mode
+    input_ids, model_kwargs = model._expand_inputs_for_generation(
+        input_ids=input_ids,
+        expand_size=max(generation_config.num_beams, generation_config.num_return_sequences),
+        is_encoder_decoder=model.config.is_encoder_decoder,
+        **model_kwargs,
+    )
+
+    if generation_config.token_healing:
+        input_ids = model.heal_tokens(input_ids, generation_mode_kwargs.get("tokenizer"))
+
+    if streamer is not None:
+        streamer.put(input_ids.cpu())
+
+    # 6. Prepare `max_length` depending on other stopping criteria.
+    input_ids_length = input_ids.shape[1]
+    generation_config = model._prepare_generated_length(
+        generation_config=generation_config,
+        has_default_max_length=has_default_max_length,
+        has_default_min_length=has_default_min_length,
+        model_input_name=model_input_name,
+        inputs_tensor=inputs_tensor,
+        input_ids_length=input_ids_length,
+    )
+
+    # If the model supports `logits_to_keep` in forward(), set it to 1 to avoid computing the whole
+    # logit matrix. This can save a lot of memory during the first forward pass. Note that assisted decoding
+    # dynamically overrides this value as it can need more than the last token logits
+    if model._supports_logits_to_keep() and "logits_to_keep" not in model_kwargs:
+        model_kwargs["logits_to_keep"] = 1
+
+    model._validate_generated_length(generation_config, input_ids_length, has_default_max_length)
+
+    # 7. Prepare the cache.
+    # - `model_kwargs` may be updated in place with a cache as defined by the parameters in `generation_config`.
+    # - different models have a different cache name expected by the model (default = "past_key_values")
+    # - `max_length`, prepared above, is used to determine the maximum cache length
+    max_cache_length = generation_config.max_length - 1
+    if (
+        inputs_tensor.shape[1] != input_ids_length
+        and model_input_name == "inputs_embeds"
+        and not model.config.is_encoder_decoder
+    ):
+        max_cache_length += inputs_tensor.shape[1]
+    model._prepare_cache_for_generation(
+        generation_config, model_kwargs, generation_mode, batch_size, max_cache_length
+    )
+
+    if model.device.type != input_ids.device.type:
+        warnings.warn(
+            "You are calling .generate() with the `input_ids` being on a device type different"
+            f" than your model's device. `input_ids` is on {input_ids.device.type}, whereas the model"
+            f" is on {model.device.type}. You may experience unexpected behaviors or slower generation."
+            " Please make sure that you have put `input_ids` to the"
+            f" correct device by calling for example input_ids = input_ids.to('{model.device.type}') before"
+            " running `.generate()`.",
+            UserWarning,
+        )
+
+    # 8. Prepare logits processors and stopping criteria
+    prepared_logits_processor = model._get_logits_processor(
+        generation_config=generation_config,
+        input_ids_seq_length=input_ids_length,
+        encoder_input_ids=inputs_tensor,
+        prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
+        logits_processor=logits_processor,
+        device=inputs_tensor.device,
+        model_kwargs=model_kwargs,
+        negative_prompt_ids=negative_prompt_ids,
+        negative_prompt_attention_mask=negative_prompt_attention_mask,
+    )
+    prepared_stopping_criteria = model._get_stopping_criteria(
+        generation_config=generation_config,
+        stopping_criteria=stopping_criteria,
+        tokenizer=generation_mode_kwargs.get("tokenizer"),
+    )
+
+    # Set model_kwargs `use_cache` so we can use it later in forward runs
+    model_kwargs["use_cache"] = generation_config.use_cache
+
+    # 9. Call generation mode
+    result = decoding_method(
+        model,
+        input_ids,
+        logits_processor=prepared_logits_processor,
+        stopping_criteria=prepared_stopping_criteria,
+        generation_config=generation_config,
+        **generation_mode_kwargs,
+        **model_kwargs,
+    )
+
+    return result