# embedding_metadata

key: chroma:document
string_value:
CODE:

def insert_model_in_doc_toc(old_model_patterns, new_model_patterns):
    """
    Insert the new model in the doc TOC, in the same section as the old model.

    Args:
        old_model_patterns (`ModelPatterns`): The patterns for the old model.
        new_model_patterns (`ModelPatterns`): The patterns for the new model.
    """
    toc_file = REPO_PATH / "docs" / "source" / "en" / "_toctree.yml"
    with open(toc_file, "r", encoding="utf8") as f:
        content = yaml.safe_load(f)

    # Get to the model API doc
    api_idx = 0
    while content[api_idx]["title"] != "API":
        api_idx += 1
    api_doc = content[api_idx]["sections"]

    model_idx = 0
    while api_doc[model_idx]["title"] != "Models":
        model_idx += 1
    model_doc = api_doc[model_idx]["sections"]

    # Find the base model in the Toc
    old_model_type = old_model_patterns.model_type
    section_idx = 0
    while section_idx < len(model_doc):
        sections = [entry["local"] for entry in model_doc[section_idx]["sections"]]
        if f"model_doc/{old_model_type}" in sections:
            break

        section_idx += 1

    if section_idx == len(model_doc):
        old_model = old_model_patterns.model_name
        new_model = new_model_patterns.model_name
        print(f"Did not find {old_model} in the table of content, so you will need to add {new_model} manually.")
        return

    # Add the new model in the same toc
    toc_entry = {"local": f"model_doc/{new_model_patterns.model_type}", "title": new_model_patterns.model_name}
    model_doc[section_idx]["sections"].append(toc_entry)
    model_doc[section_idx]["sections"] = sorted(model_doc[section_idx]["sections"], key=lambda s: s["title"].lower())
    api_doc[model_idx]["sections"] = model_doc
    content[api_idx]["sections"] = api_doc

    with open(toc_file, "w", encoding="utf-8") as f:


SUMMARY:
**Function Summary**

```python
def insert_model_in_doc_toc(old_model_patterns, new_model_patterns)
```

**Purpose and Functionality**
Inserts the `new_model_patterns` into the doc TOC in the same section as the `old_model_patterns`, updating the table of contents (`toc_file`) accordingly.

**Key Implementation Details**

* Iterates through TOC sections to find the correct position for the new model.
* Updates the `model_doc` by adding a new entry with the `new_model_patterns.model_name`.
* Re-constructs the updated content and writes it back to `toc_file`.

**Notable Dependencies and Requirements**

* `REPO_PATH`: Path to repository root directory.
* `yaml`: YAML parsing library for loading/treating TOC file content.
* `ModelPatterns`: Class containing patterns for old model and new model (not shown in this snippet).

**External Dependencies**

* None mentioned; no external libraries besides YAML are used.


key: chroma:document
string_value:
CODE:

def create_new_model_like(
    model_type: str,
    new_model_patterns: ModelPatterns,
    add_copied_from: bool = True,
    frameworks: Optional[List[str]] = None,
    old_checkpoint: Optional[str] = None,
):
    """
    Creates a new model module like a given model of the Transformers library.

    Args:
        model_type (`str`): The model type to duplicate (like "bert" or "gpt2")
        new_model_patterns (`ModelPatterns`): The patterns for the new model.
        add_copied_from (`bool`, *optional*, defaults to `True`):
            Whether or not to add "Copied from" statements to all classes in the new model modeling files.
        frameworks (`List[str]`, *optional*):
            If passed, will limit the duplicate to the frameworks specified.
        old_checkpoint (`str`, *optional*):
            The name of the base checkpoint for the old model. Should be passed along when it can't be automatically
            recovered from the `model_type`.
    """
    # Retrieve all the old model info.
    model_info = retrieve_info_for_model(model_type, frameworks=frameworks)
    model_files = model_info["model_files"]
    old_model_patterns = model_info["model_patterns"]
    if old_checkpoint is not None:
        old_model_patterns.checkpoint = old_checkpoint
    if len(old_model_patterns.checkpoint) == 0:
        raise ValueError(
            "The old model checkpoint could not be recovered from the model type. Please pass it to the "
            "`old_checkpoint` argument."
        )

    keep_old_processing = True
    for processing_attr in ["image_processor_class", "feature_extractor_class", "processor_class", "tokenizer_class"]:
        if getattr(old_model_patterns, processing_attr) != getattr(new_model_patterns, processing_attr):
            keep_old_processing = False

    model_classes = model_info["model_classes"]

    # 1. We create the module for our new model.
    old_module_name = model_files["module_name"]
    module_folder = TRANSFORMERS_PATH / "models" / new_model_patterns.model_lower_cased
    os.makedirs(module_folder, exist_ok=True)

    files_to_adapt = model_files["model_files"]
    if keep_old_processing:
        files_to_adapt = [
            f
            for f in files_to_adapt
            if "tokenization" not in str(f)
            and "processing" not in str(f)
            and "feature_extraction" not in str(f)
            and "image_processing" not in str(f)
        ]

    os.makedirs(module_folder, exist_ok=True)
    for module_file in files_to_adapt:
        new_module_name = module_file.name.replace(
            old_model_patterns.model_lower_cased, new_model_patterns.model_lower_cased
        )
        dest_file = module_folder / new_module_name
        duplicate_module(
            module_file,
            old_model_patterns,
            new_model_patterns,
            dest_file=dest_file,
            add_copied_from=add_copied_from and "modeling" in new_module_name,
        )

    clean_frameworks_in_init(
        module_folder / "__init__.py", frameworks=frameworks, keep_processing=not keep_old_processing
    )

    # 2. We add our new model to the models init and the main init
    add_content_to_file(
        TRANSFORMERS_PATH / "models" / "__init__.py",
        f"    {new_model_patterns.model_lower_cased},",
        add_after=f"    {old_module_name},",
        exact_match=True,
    )
    add_model_to_main_init(
        old_model_patterns, new_model_patterns, frameworks=frameworks, with_processing=not keep_old_processing
    )

    # 3. Add test files
    files_to_adapt = model_files["test_files"]
    if keep_old_processing:
        files_to_adapt = [
            f
            for f in files_to_adapt
            if "tokenization" not in str(f)
            and "processor" not in str(f)
            and "feature_extraction" not in str(f)
            and "image_processing" not in str(f)
        ]

    def disable_fx_test(filename: Path) -> bool:
        with open(filename) as fp:
            content = fp.read()
        new_content = re.sub(r"fx_compatible\s*=\s*True", "fx_compatible = False", content)
        with open(filename, "w") as fp:
            fp.write(new_content)
        return content != new_content

    disabled_fx_test = False

    tests_folder = REPO_PATH / "tests" / "models" / new_model_patterns.model_lower_cased
    os.makedirs(tests_folder, exist_ok=True)
    with open(tests_folder / "__init__.py", "w"):
        pass

    for test_file in files_to_adapt:
        new_test_file_name = test_file.name.replace(
            old_model_patterns.model_lower_cased, new_model_patterns.model_lower_cased
        )
        dest_file = test_file.parent.parent / new_model_patterns.model_lower_cased / new_test_file_name
        duplicate_module(
            test_file,
            old_model_patterns,
            new_model_patterns,
            dest_file=dest_file,
            add_copied_from=False,
            attrs_to_remove=["pipeline_model_mapping", "is_pipeline_test_to_skip"],
        )
        disabled_fx_test = disabled_fx_test | disable_fx_test(dest_file)

    if disabled_fx_test:
        print(
            "The tests for symbolic tracing with torch.fx were disabled, you can add those once symbolic tracing works"
            " for your new model."
        )

    # 4. Add model to auto classes
    add_model_to_auto_classes(old_model_patterns, new_model_patterns, model_classes)

    # 5. Add doc file
    doc_file = REPO_PATH / "docs" / "source" / "en" / "model_doc" / f"{old_model_patterns.model_type}.md"
    duplicate_doc_file(doc_file, old_model_patterns, new_model_patterns, frameworks=frameworks)
    insert_model_in_doc_toc(old_model_patterns, new_model_patterns)

    # 6. Warn the user for duplicate patterns
    if old_model_patterns.model_type == old_model_patterns.checkpoint:
        print(
            "The model you picked has the same name for the model type and the checkpoint name "
            f"({old_model_patterns.model_type}). As a result, it's possible some places where the new checkpoint "
            f"should be, you have {new_model_patterns.model_type} instead. You should search for all instances of "
            f"{new_model_patterns.model_type} in the new files and check they're not badly used as checkpoints."
        )
    elif old_model_patterns.model_lower_cased == old_model_patterns.checkpoint:
        print(
            "The model you picked has the same name for the model type and the checkpoint name "
            f"({old_model_patterns.model_lower_cased}). As a result, it's possible some places where the new "
            f"checkpoint should be, you have {new_model_patterns.model_lower_cased} instead. You should search for "
            f"all instances of {new_model_patterns.model_lower_cased} in the new files and check they're not badly "
            "used as checkpoints."
        )
    if (
        old_model_patterns.model_type == old_model_patterns.model_lower_cased
        and new_model_patterns.model_type != new_model_patterns.model_lower_cased
    ):
        print(
            "The model you picked has the same name for the model type and the lowercased model name "
            f"({old_model_patterns.model_lower_cased}). As a result, it's possible some places where the new "
            f"model type should be, you have {new_model_patterns.model_lower_cased} instead. You should search for "
            f"all instances of {new_model_patterns.model_lower_cased} in the new files and check they're not badly "
            "used as the model type."
        )

    if not keep_old_processing and old_model_patterns.tokenizer_class is not None:
        print(
            "The constants at the start of the new tokenizer file created needs to be manually fixed. If your new "
            "model has a tokenizer fast, you will also need to manually add the converter in the "
            "`SLOW_TO_FAST_CONVERTERS` constant of `convert_slow_tokenizer.py`."


SUMMARY:
**Summary**

This function, `create_new_model_like`, generates a new model similar to an existing Transformer-based model. It:

1.  Duplicates files and classes from the old model based on specified patterns.
2.  Updates framework configurations and removes framework dependencies where possible.
3.  Adds the new model to the main init file and models init file.
4.  Copies test files and tests their compatibility with torch fx.
5.  Creates an auto class for the new model.
6.  Inserts a doc file in the documentation source directory.

**Functionality Breakdown**

The function can be categorized into six main tasks:

*   **File Duplications**: Duplicate specific files from the old model to create corresponding files for the new model, including tokenizer classes with updated constants.
*   **Framework Configuration Updates**: Update frameworks and remove unnecessary framework dependencies where possible.
*   **Auto Class Creation**: Create an auto class for the new model using data from both files.
*   **Doc File Insertion**: Copy a doc file from a source directory to the repository's documentation source directory.


key: chroma:document
string_value:
CODE:

def add_new_model_like_command_factory(args: Namespace):


SUMMARY:
**Summary**

Creates an instance of `AddNewModelLikeCommand` with configuration settings from the provided namespace.

**Key Details**

- Uses `Args.Namespace` for type hinting
- Conveys functionality through `config_file` and `path_to_repo` arguments, which are passed to the constructor of `AddNewModelLikeCommand`

**Notable Dependencies/Requirements**

- `Namespace` from an external library (e.g., `argparse`)
- `AddNewModelLikeCommand` class implementation not shown in this code snippet


key: chroma:document
string_value:
CODE:

class AddNewModelLikeCommand(BaseTransformersCLICommand):
    @staticmethod
    def register_subcommand(parser: ArgumentParser):
        add_new_model_like_parser = parser.add_parser("add-new-model-like")
        add_new_model_like_parser.add_argument(
            "--config_file", type=str, help="A file with all the information for this model creation."
        )
        add_new_model_like_parser.add_argument(
            "--path_to_repo", type=str, help="When not using an editable install, the path to the Transformers repo."
        )
        add_new_model_like_parser.set_defaults(func=add_new_model_like_command_factory)

    def __init__(self, config_file=None, path_to_repo=None, *args):
        if config_file is not None:
            with open(config_file, "r", encoding="utf-8") as f:
                config = json.load(f)
            self.old_model_type = config["old_model_type"]
            self.model_patterns = ModelPatterns(**config["new_model_patterns"])
            self.add_copied_from = config.get("add_copied_from", True)
            self.frameworks = config.get("frameworks", get_default_frameworks())
            self.old_checkpoint = config.get("old_checkpoint", None)
        else:
            (
                self.old_model_type,
                self.model_patterns,
                self.add_copied_from,
                self.frameworks,
                self.old_checkpoint,
            ) = get_user_input()

        self.path_to_repo = path_to_repo

    def run(self):
        if self.path_to_repo is not None:
            # Adapt constants
            global TRANSFORMERS_PATH
            global REPO_PATH

            REPO_PATH = Path(self.path_to_repo)
            TRANSFORMERS_PATH = REPO_PATH / "src" / "transformers"

        create_new_model_like(
            model_type=self.old_model_type,
            new_model_patterns=self.model_patterns,
            add_copied_from=self.add_copied_from,
            frameworks=self.frameworks,
            old_checkpoint=self.old_checkpoint,


SUMMARY:
**Class Summary: AddNewModelLikeCommand**

The `AddNewModelLikeCommand` class is a subclass of `BaseTransformersCLICommand`. It handles the "add-new-model-like" subcommand in a command-line interface (CLI).

**Primary Purpose and Functionality**

This class provides a CLI entry point for creating new models, similar to existing models in the Transformers framework. The primary function is to initialize and configure model creation parameters from either a configuration file or user input.

**Key Implementation Details**

* The `register_subcommand` method registers the "add-new-model-like" subcommand with the CLI parser.
* The `__init__` method loads configuration data from a file, if provided, and sets default values if no configuration is found. It also retrieves framework dependencies automatically.
* The `run` method adapts global constants and creates a new model using the configured parameters.

**Notable Dependencies or Requirements**

* The `ModelPatterns` class is required to parse patterns for new model creation.
* The `get_user_input` function returns default values if no configuration file is provided.
* External dependencies are:
	+ `json` for loading configuration data from files (.config).
	+ `Path` and ` glob` (implicit, as part of the Python standard library) for managing repository paths.

**Interfaces and Inherited Behavior**

This class inherits behavior from its parent class, `BaseTransformersCLICommand`, which likely handles common CLI tasks such as argument parsing and error handling.


key: chroma:document
string_value:
CODE:

def get_user_field(
    question: str,
    default_value: Optional[str] = None,
    is_valid_answer: Optional[Callable] = None,
    convert_to: Optional[Callable] = None,
    fallback_message: Optional[str] = None,
) -> Any:
    """
    A utility function that asks a question to the user to get an answer, potentially looping until it gets a valid
    answer.

    Args:
        question (`str`): The question to ask the user.
        default_value (`str`, *optional*): A potential default value that will be used when the answer is empty.
        is_valid_answer (`Callable`, *optional*):
            If set, the question will be asked until this function returns `True` on the provided answer.
        convert_to (`Callable`, *optional*):
            If set, the answer will be passed to this function. If this function raises an error on the procided
            answer, the question will be asked again.
        fallback_message (`str`, *optional*):
            A message that will be displayed each time the question is asked again to the user.

    Returns:
        `Any`: The answer provided by the user (or the default), passed through the potential conversion function.
    """
    if not question.endswith(" "):
        question = question + " "
    if default_value is not None:
        question = f"{question} [{default_value}] "

    valid_answer = False
    while not valid_answer:
        answer = input(question)
        if default_value is not None and len(answer) == 0:
            answer = default_value
        if is_valid_answer is not None:
            valid_answer = is_valid_answer(answer)
        elif convert_to is not None:
            try:
                answer = convert_to(answer)
                valid_answer = True
            except Exception:
                valid_answer = False
        else:
            valid_answer = True

        if not valid_answer:
            print(fallback_message)



SUMMARY:
**Summary: `get_user_field` Function**

### Purpose and Functionality

The `get_user_field` function is a utility tool that prompts a user to provide an input field, potentially looping until a valid response is obtained.

### Key Implementation Details

* Takes in a question string and optional parameters for:
	+ `default_value`: default value to use if answer is empty
	+ `is_valid_answer`: callback function to validate answer
	+ `convert_to`: callback function to convert input into desired format
	+ `fallback_message`: message displayed when looping back to user with the same question
* Uses a while loop to continuously ask the user for input until a valid response is provided
* Handles potential errors by skipping to the next iteration or displaying the fallback message

### Dependencies and Requirements

* `Callable` type hint for callback functions (`is_valid_answer`, `convert_to`)
* No external dependencies noted.


key: chroma:document
string_value:
CODE:

def convert_to_bool(x: str) -> bool:
    """
    Converts a string to a bool.
    """
    if x.lower() in ["1", "y", "yes", "true"]:
        return True
    if x.lower() in ["0", "n", "no", "false"]:
        return False


SUMMARY:
**Function Summary**

`convert_to_bool(x: str) -> bool`

Converts a string input `x` to a boolean value.

**Implementation Details**

- Uses case-insensitive comparison (`x.lower()`) for matching string values.
- Provides a list of approved values ("1", "y", "yes", "true") that are converted to `True`.
- Provides a list of rejected values ("0", "n", "no", "false") that result in an exception.

**Notable Dependencies or Requirements**

- None
- No external dependencies
- Follows standard Python convention for boolean conversion.


key: chroma:document
string_value:
CODE:

def get_user_input():
    """
    Ask the user for the necessary inputs to add the new model.
    """
    model_types = list(auto_module.configuration_auto.MODEL_NAMES_MAPPING.keys())

    # Get old model type
    valid_model_type = False
    while not valid_model_type:
        old_model_type = input(
            "What is the model you would like to duplicate? Please provide the lowercase `model_type` (e.g. roberta): "
        )
        if old_model_type in model_types:
            valid_model_type = True
        else:
            print(f"{old_model_type} is not a valid model type.")
            near_choices = difflib.get_close_matches(old_model_type, model_types)
            if len(near_choices) >= 1:
                if len(near_choices) > 1:
                    near_choices = " or ".join(near_choices)
                print(f"Did you mean {near_choices}?")

    old_model_info = retrieve_info_for_model(old_model_type)
    old_tokenizer_class = old_model_info["model_patterns"].tokenizer_class
    old_image_processor_class = old_model_info["model_patterns"].image_processor_class
    old_feature_extractor_class = old_model_info["model_patterns"].feature_extractor_class
    old_processor_class = old_model_info["model_patterns"].processor_class
    old_frameworks = old_model_info["frameworks"]

    old_checkpoint = None
    if len(old_model_info["model_patterns"].checkpoint) == 0:
        old_checkpoint = get_user_field(
            "We couldn't find the name of the base checkpoint for that model, please enter it here."
        )

    model_name = get_user_field(
        "What is the name (with no special casing) for your new model in the paper (e.g. RoBERTa)? "
    )
    default_patterns = ModelPatterns(model_name, model_name)

    model_type = get_user_field(
        "What identifier would you like to use for the `model_type` of this model? ",
        default_value=default_patterns.model_type,
    )
    model_lower_cased = get_user_field(
        "What lowercase name would you like to use for the module (folder) of this model? ",
        default_value=default_patterns.model_lower_cased,
    )
    model_camel_cased = get_user_field(
        "What prefix (camel-cased) would you like to use for the model classes of this model (e.g. Roberta)? ",
        default_value=default_patterns.model_camel_cased,
    )
    model_upper_cased = get_user_field(
        "What prefix (upper-cased) would you like to use for the constants relative to this model? ",
        default_value=default_patterns.model_upper_cased,
    )
    config_class = get_user_field(
        "What will be the name of the config class for this model? ", default_value=f"{model_camel_cased}Config"
    )
    checkpoint = get_user_field(
        "Please give a checkpoint identifier (on the model Hub) for this new model (e.g. facebook/FacebookAI/roberta-base): "
    )

    old_processing_classes = [
        c if not isinstance(c, tuple) else c[0]
        for c in [old_image_processor_class, old_feature_extractor_class, old_tokenizer_class, old_processor_class]
        if c is not None
    ]
    old_processing_classes = ", ".join(old_processing_classes)
    keep_processing = get_user_field(
        f"Will your new model use the same processing class as {old_model_type} ({old_processing_classes}) (yes/no)? ",
        convert_to=convert_to_bool,
        fallback_message="Please answer yes/no, y/n, true/false or 1/0. ",
    )
    if keep_processing:
        image_processor_class = old_image_processor_class
        feature_extractor_class = old_feature_extractor_class
        processor_class = old_processor_class
        tokenizer_class = old_tokenizer_class
    else:
        if old_tokenizer_class is not None:
            tokenizer_class = get_user_field(
                "What will be the name of the tokenizer class for this model? ",
                default_value=f"{model_camel_cased}Tokenizer",
            )
        else:
            tokenizer_class = None
        if old_image_processor_class is not None:
            image_processor_class = get_user_field(
                "What will be the name of the image processor class for this model? ",
                default_value=f"{model_camel_cased}ImageProcessor",
            )
        else:
            image_processor_class = None
        if old_feature_extractor_class is not None:
            feature_extractor_class = get_user_field(
                "What will be the name of the feature extractor class for this model? ",
                default_value=f"{model_camel_cased}FeatureExtractor",
            )
        else:
            feature_extractor_class = None
        if old_processor_class is not None:
            processor_class = get_user_field(
                "What will be the name of the processor class for this model? ",
                default_value=f"{model_camel_cased}Processor",
            )
        else:
            processor_class = None

    model_patterns = ModelPatterns(
        model_name,
        checkpoint,
        model_type=model_type,
        model_lower_cased=model_lower_cased,
        model_camel_cased=model_camel_cased,
        model_upper_cased=model_upper_cased,
        config_class=config_class,
        tokenizer_class=tokenizer_class,
        image_processor_class=image_processor_class,
        feature_extractor_class=feature_extractor_class,
        processor_class=processor_class,
    )

    add_copied_from = get_user_field(
        "Should we add # Copied from statements when creating the new modeling file (yes/no)? ",
        convert_to=convert_to_bool,
        default_value="yes",
        fallback_message="Please answer yes/no, y/n, true/false or 1/0.",
    )

    all_frameworks = get_user_field(
        "Should we add a version of your new model in all the frameworks implemented by"
        f" {old_model_type} ({old_frameworks}) (yes/no)? ",
        convert_to=convert_to_bool,
        default_value="yes",
        fallback_message="Please answer yes/no, y/n, true/false or 1/0.",
    )
    if all_frameworks:
        frameworks = None
    else:
        frameworks = get_user_field(
            "Please enter the list of framworks you want (pt, tf, flax) separated by spaces",
            is_valid_answer=lambda x: all(p in ["pt", "tf", "flax"] for p in x.split(" ")),
        )
        frameworks = list(set(frameworks.split(" ")))



SUMMARY:
**Function Summary: `get_user_input`**

### Purpose and Functionality

This function asks the user for input data to create a new model. It collects user input for various parameters, such as model name, identifier, checkpoint, frameworks, and additional fields.

### Key Implementation Details

*   Asks the user to select a valid old model type using `difflib.get_close_matches`.
*   Retrieves information for the selected old model type.
*   Asks the user to input fields such as:
    *   Model name
    *   Identifier (model type)
    *   Lowercase identifier (module folder)
    *   Camel-cased identifier (model class prefix)
    *   Upper-cased identifier (constant prefix)
    *   Config class name
    *   Checkpoint identifier
*   Determines which processing classes (image, feature extractor, tokenizer, processor) to use based on user input.
*   Asks the user if additional fields should be added:
    *   Copied from statements when creating a new modeling file
    *   Version of the model for each framework implemented by the old model type

### Notable Dependencies and Requirements

*   `difflib` for close matching functionality
*   `auto_module.configuration_auto.MODEL_NAMES_MAPPING` for old model types
*   `retrieve_info_for_model` function to retrieve information for the selected old model type
*   `gets_user_field` functions to collect user input for various fields


key: chroma:document
string_value:
CODE:

def download_command_factory(args):


SUMMARY:
**Summary:**

The `download_command_factory` function creates and returns an instance of `DownloadCommand`, which is used to initiate the download process. It takes four arguments:

- `model`
- `cache_dir`
- `force`
- `trust_remote_code`

 **Main Purpose:** Creates a factory for DownloadCommand instances.

**Key Implementation Details:**

* Initializes DownloadCommand with specific parameters (`model`, `cacheDir`, `force`, and `trustRemoteCode`).

**Notable Dependencies/Requirements:**

* The function relies on the existence of `DownloadCommand`, which is not shown in this snippet.
* No external dependencies are listed.


key: chroma:document
string_value:
CODE:

class DownloadCommand(BaseTransformersCLICommand):
    @staticmethod
    def register_subcommand(parser: ArgumentParser):
        download_parser = parser.add_parser("download")
        download_parser.add_argument(
            "--cache-dir", type=str, default=None, help="Path to location to store the models"
        )
        download_parser.add_argument(
            "--force", action="store_true", help="Force the model to be download even if already in cache-dir"
        )
        download_parser.add_argument(
            "--trust-remote-code",
            action="store_true",
            help="Whether or not to allow for custom models defined on the Hub in their own modeling files. Use only if you've reviewed the code as it will execute on your local machine",
        )
        download_parser.add_argument("model", type=str, help="Name of the model to download")
        download_parser.set_defaults(func=download_command_factory)

    def __init__(self, model: str, cache: str, force: bool, trust_remote_code: bool):
        self._model = model
        self._cache = cache
        self._force = force
        self._trust_remote_code = trust_remote_code

    def run(self):
        from ..models.auto import AutoModel, AutoTokenizer

        AutoModel.from_pretrained(
            self._model, cache_dir=self._cache, force_download=self._force, trust_remote_code=self._trust_remote_code
        )
        AutoTokenizer.from_pretrained(
            self._model, cache_dir=self._cache, force_download=self._force, trust_remote_code=self._trust_remote_code


SUMMARY:
**DownloadCommand Class Summary**

### Purpose and Functionality

The `DownloadCommand` class is a CLI command that allows users to download pre-trained models from the Hub. It provides options for customizing model storage and remote code execution.

### Key Implementation Details

*   The `download_command_factory` function is set as the default action for the "download" subcommand, which retrieves the AutoModel and AutoTokenizer instances.
*   The class uses `BaseTransformersCLICommand` as its base class to extend CLIParser functionality.

    ### Notable Dependencies or Requirements
    *   `ArgumentParser`
    *   `BaseTransformersCLICommand`

    The class inherits from `BaseTransformersCLICommand`, and utilizes the `ArgumentParser` to define the command-line interface for the "download" subcommand.


key: chroma:document
string_value:
CODE:

def train_command_factory(args: Namespace):
    """
    Factory function used to instantiate training command from provided command line arguments.

    Returns: TrainCommand
    """


SUMMARY:
**Summary**

`train_command_factory`: Factory function to create `TrainCommand` instance from provided argument namespace (`args`) with primary purpose of instantiating training commands from command line arguments. 

### Key Details:

* Returns an instance of `TrainCommand`
* Uses `args` namespace as input

### Notable Dependencies:

* None


key: chroma:document
string_value:
CODE:

class TrainCommand(BaseTransformersCLICommand):
    @staticmethod
    def register_subcommand(parser: ArgumentParser):
        """
        Register this command to argparse so it's available for the transformer-cli

        Args:
            parser: Root parser to register command-specific arguments
        """
        train_parser = parser.add_parser("train", help="CLI tool to train a model on a task.")

        train_parser.add_argument(
            "--train_data",
            type=str,
            required=True,
            help="path to train (and optionally evaluation) dataset as a csv with tab separated labels and sentences.",
        )
        train_parser.add_argument(
            "--column_label", type=int, default=0, help="Column of the dataset csv file with example labels."
        )
        train_parser.add_argument(
            "--column_text", type=int, default=1, help="Column of the dataset csv file with example texts."
        )
        train_parser.add_argument(
            "--column_id", type=int, default=2, help="Column of the dataset csv file with example ids."
        )
        train_parser.add_argument(
            "--skip_first_row", action="store_true", help="Skip the first row of the csv file (headers)."
        )

        train_parser.add_argument("--validation_data", type=str, default="", help="path to validation dataset.")
        train_parser.add_argument(
            "--validation_split",
            type=float,
            default=0.1,
            help="if validation dataset is not provided, fraction of train dataset to use as validation dataset.",
        )

        train_parser.add_argument("--output", type=str, default="./", help="path to saved the trained model.")

        train_parser.add_argument(
            "--task", type=str, default="text_classification", help="Task to train the model on."
        )
        train_parser.add_argument(
            "--model", type=str, default="google-bert/bert-base-uncased", help="Model's name or path to stored model."
        )
        train_parser.add_argument("--train_batch_size", type=int, default=32, help="Batch size for training.")
        train_parser.add_argument("--valid_batch_size", type=int, default=64, help="Batch size for validation.")
        train_parser.add_argument("--learning_rate", type=float, default=3e-5, help="Learning rate.")
        train_parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon for Adam optimizer.")
        train_parser.set_defaults(func=train_command_factory)

    def __init__(self, args: Namespace):
        self.logger = logging.get_logger("transformers-cli/training")

        self.framework = "tf" if is_tf_available() else "torch"

        os.makedirs(args.output, exist_ok=True)
        self.output = args.output

        self.column_label = args.column_label
        self.column_text = args.column_text
        self.column_id = args.column_id

        self.logger.info(f"Loading {args.task} pipeline for {args.model}")
        if args.task == "text_classification":
            self.pipeline = TextClassificationPipeline.from_pretrained(args.model)
        elif args.task == "token_classification":
            raise NotImplementedError
        elif args.task == "question_answering":
            raise NotImplementedError

        self.logger.info(f"Loading dataset from {args.train_data}")
        self.train_dataset = Processor.create_from_csv(
            args.train_data,
            column_label=args.column_label,
            column_text=args.column_text,
            column_id=args.column_id,
            skip_first_row=args.skip_first_row,
        )
        self.valid_dataset = None
        if args.validation_data:
            self.logger.info(f"Loading validation dataset from {args.validation_data}")
            self.valid_dataset = Processor.create_from_csv(
                args.validation_data,
                column_label=args.column_label,
                column_text=args.column_text,
                column_id=args.column_id,
                skip_first_row=args.skip_first_row,
            )

        self.validation_split = args.validation_split
        self.train_batch_size = args.train_batch_size
        self.valid_batch_size = args.valid_batch_size
        self.learning_rate = args.learning_rate
        self.adam_epsilon = args.adam_epsilon

    def run(self):
        if self.framework == "tf":
            return self.run_tf()
        return self.run_torch()

    def run_torch(self):
        raise NotImplementedError

    def run_tf(self):
        self.pipeline.fit(
            self.train_dataset,
            validation_data=self.valid_dataset,
            validation_split=self.validation_split,
            learning_rate=self.learning_rate,
            adam_epsilon=self.adam_epsilon,
            train_batch_size=self.train_batch_size,
            valid_batch_size=self.valid_batch_size,
        )

        # Save trained pipeline


SUMMARY:
**TrainCommand Class Summary**

### Purpose and Functionality

The `TrainCommand` class is a CLI tool for training a machine learning model using a dataset. It takes in various arguments, including the model architecture (`google-bert/bert-base-uncased` or custom), the task type (`text_classification`, `token_classification`, or `question_answering`), and hyperparameters such as batch size, learning rate, and Adam optimizer epsilon.

### Key Implementation Details

* Registers command-line arguments using `ArgumentParser`
* Creates `Pipeline` instance for each model architecture
* Loads dataset from CSV file(s) and processes it into a `Processor` object
* Trains the `Pipeline` on the processed dataset with validation data (if provided)
* Saves the trained pipeline to a directory

### Notable Dependencies and Requirements

* `BaseTransformersCLICommand` from `transformers-cli`
* `ArgumentParser` from `argparse`
* `Logging` for logging
* `Processor` for loading and processing datasets
* `TextClassificationPipeline` or other custom pipelines for specific tasks

### Class Members

* `logger`: instance of Logger class
* `framework`: either "tf" (TensorFlow) or "torch"
* `output`: path to saved trained model directory
* `column_label`, `column_text`, and `column_id`: CSV column indices
* `task`: task type (`text_classification`, `token_classification`, or `question_answering`)
* `model`: model architecture or path to stored model
* `validation_data` and `validation_split`: validation dataset path and split fraction for training


key: chroma:document
string_value:
CODE:

class BaseTransformersCLICommand(ABC):
    @staticmethod
    @abstractmethod
    def register_subcommand(parser: ArgumentParser):
        raise NotImplementedError()

    @abstractmethod
    def run(self):


SUMMARY:
**Class Summary**

`BaseTransformersCLICommand`: Abstract Base Class for Command-Line Interface Transformers.

**Primary Purpose**: Define the base structure for command-line interface transformers in a CLI system, requiring registration of subcommands and implementation of the run method.

**Key Implementation Details**:

*   **`register_subcommand` Method**: Register a subcommand with an `ArgumentParser`.
    *   Method Signature: `@staticmethod def register_subcommand(parser: ArgumentParser)`
*   **`run` Method**: Perform the actual transformation or execution.
    *   Method Signature: `@abstractmethod def run(self):`

**Notable Dependencies and Requirements**:

*   Uses the `ArgumentParser` from the `argparse` library for subcommand registration.
*   Requires inheritance from `ABC` (Abstract Base Class) to define abstract methods.


key: chroma:document
string_value:
CODE:

def convert_command_factory(args: Namespace):
    """
    Factory function used to convert a model TF 1.0 checkpoint in a PyTorch checkpoint.

    Returns: ServeCommand
    """
    return ConvertCommand(
        args.model_type, args.tf_checkpoint, args.pytorch_dump_output, args.config, args.finetuning_task_name


SUMMARY:
**Function Summary**

### Purpose and Functionality

The `convert_command_factory` function is a factory function that creates and returns an instance of the `ConvertCommand` class. Its primary purpose is to convert a TF 1.0 checkpoint to a PyTorch checkpoint.

### Key Implementation Details

* The function takes a `Namespace` object named `args` as input.
* It uses this input `args` to create a new instance of the `ConvertCommand` class, passing in several attributes (model type, TensorFlow checkpoint, PyTorch dump output, configuration, and fine-tuning task name).

### Notable Dependencies or Requirements

* The function relies on the `Namespace` object from the `argparse` library.
* It requires an instance of `ConvertCommand`.

**Summary Output**

`convert_command_factory`: Factory function to convert TF 1.0 checkpoint to PyTorch checkpoint, utilizing `ConvertCommand` class with input attributes (model type, TensorFlow checkpoint, etc.).


key: chroma:document
string_value:
CODE:

class ConvertCommand(BaseTransformersCLICommand):
    @staticmethod
    def register_subcommand(parser: ArgumentParser):
        """
        Register this command to argparse so it's available for the transformer-cli

        Args:
            parser: Root parser to register command-specific arguments
        """
        train_parser = parser.add_parser(
            "convert",
            help="CLI tool to run convert model from original author checkpoints to Transformers PyTorch checkpoints.",
        )
        train_parser.add_argument("--model_type", type=str, required=True, help="Model's type.")
        train_parser.add_argument(
            "--tf_checkpoint", type=str, required=True, help="TensorFlow checkpoint path or folder."
        )
        train_parser.add_argument(
            "--pytorch_dump_output", type=str, required=True, help="Path to the PyTorch saved model output."
        )
        train_parser.add_argument("--config", type=str, default="", help="Configuration file path or folder.")
        train_parser.add_argument(
            "--finetuning_task_name",
            type=str,
            default=None,
            help="Optional fine-tuning task name if the TF model was a finetuned model.",
        )
        train_parser.set_defaults(func=convert_command_factory)

    def __init__(
        self,
        model_type: str,
        tf_checkpoint: str,
        pytorch_dump_output: str,
        config: str,
        finetuning_task_name: str,
        *args,
    ):
        self._logger = logging.get_logger("transformers-cli/converting")

        self._logger.info(f"Loading model {model_type}")
        self._model_type = model_type
        self._tf_checkpoint = tf_checkpoint
        self._pytorch_dump_output = pytorch_dump_output
        self._config = config
        self._finetuning_task_name = finetuning_task_name

    def run(self):
        if self._model_type == "albert":
            try:
                from ..models.albert.convert_albert_original_tf_checkpoint_to_pytorch import (
                    convert_tf_checkpoint_to_pytorch,
                )
            except ImportError:
                raise ImportError(IMPORT_ERROR_MESSAGE)

            convert_tf_checkpoint_to_pytorch(self._tf_checkpoint, self._config, self._pytorch_dump_output)
        elif self._model_type == "bert":
            try:
                from ..models.bert.convert_bert_original_tf_checkpoint_to_pytorch import (
                    convert_tf_checkpoint_to_pytorch,
                )
            except ImportError:
                raise ImportError(IMPORT_ERROR_MESSAGE)

            convert_tf_checkpoint_to_pytorch(self._tf_checkpoint, self._config, self._pytorch_dump_output)
        elif self._model_type == "funnel":
            try:
                from ..models.funnel.convert_funnel_original_tf_checkpoint_to_pytorch import (
                    convert_tf_checkpoint_to_pytorch,
                )
            except ImportError:
                raise ImportError(IMPORT_ERROR_MESSAGE)

            convert_tf_checkpoint_to_pytorch(self._tf_checkpoint, self._config, self._pytorch_dump_output)
        elif self._model_type == "t5":
            try:
                from ..models.t5.convert_t5_original_tf_checkpoint_to_pytorch import convert_tf_checkpoint_to_pytorch
            except ImportError:
                raise ImportError(IMPORT_ERROR_MESSAGE)

            convert_tf_checkpoint_to_pytorch(self._tf_checkpoint, self._config, self._pytorch_dump_output)
        elif self._model_type == "gpt":
            from ..models.openai.convert_openai_original_tf_checkpoint_to_pytorch import (
                convert_openai_checkpoint_to_pytorch,
            )

            convert_openai_checkpoint_to_pytorch(self._tf_checkpoint, self._config, self._pytorch_dump_output)
        elif self._model_type == "gpt2":
            try:
                from ..models.gpt2.convert_gpt2_original_tf_checkpoint_to_pytorch import (
                    convert_gpt2_checkpoint_to_pytorch,
                )
            except ImportError:
                raise ImportError(IMPORT_ERROR_MESSAGE)

            convert_gpt2_checkpoint_to_pytorch(self._tf_checkpoint, self._config, self._pytorch_dump_output)
        elif self._model_type == "xlnet":
            try:
                from ..models.xlnet.convert_xlnet_original_tf_checkpoint_to_pytorch import (
                    convert_xlnet_checkpoint_to_pytorch,
                )
            except ImportError:
                raise ImportError(IMPORT_ERROR_MESSAGE)

            convert_xlnet_checkpoint_to_pytorch(
                self._tf_checkpoint, self._config, self._pytorch_dump_output, self._finetuning_task_name
            )
        elif self._model_type == "xlm":
            from ..models.xlm.convert_xlm_original_pytorch_checkpoint_to_pytorch import (
                convert_xlm_checkpoint_to_pytorch,
            )

            convert_xlm_checkpoint_to_pytorch(self._tf_checkpoint, self._pytorch_dump_output)
        elif self._model_type == "lxmert":
            from ..models.lxmert.convert_lxmert_original_tf_checkpoint_to_pytorch import (
                convert_lxmert_checkpoint_to_pytorch,
            )

            convert_lxmert_checkpoint_to_pytorch(self._tf_checkpoint, self._pytorch_dump_output)
        elif self._model_type == "rembert":
            from ..models.rembert.convert_rembert_tf_checkpoint_to_pytorch import (
                convert_rembert_tf_checkpoint_to_pytorch,
            )

            convert_rembert_tf_checkpoint_to_pytorch(self._tf_checkpoint, self._config, self._pytorch_dump_output)
        else:


SUMMARY:
**Class Summary: ConvertCommand**

The `ConvertCommand` class is a CLI command to convert original author checkpoints to Transformers PyTorch checkpoints.

**Purpose and Functionality**
---------------------------

* The primary purpose of this class is to enable the conversion of TensorFlow checkpoints to PyTorch models for various pre-trained language models, including ALBERT, BERT, Funnel, T5, GPT, GPT2, XLM, LXMert, and REMBERT.
* The class takes in several command-line arguments, including model type, TensorFlow checkpoint path, PyTorch dump output path, configuration file path or folder, and optional fine-tuning task name.

**Key Implementation Details**
--------------------------------

* The `convert_tf_checkpoint_to_pytorch` function is called dynamically based on the model type.
* The class uses try-except blocks to handle missing import errors for specific models.

**Notable Dependencies or Requirements**
-----------------------------------------

* `base/transformers/cli/BaseTransformersCLICommand`: Base class for CLI commands
* `argparse`: Library for parsing command-line arguments
* Various PyTorch checkpoint conversion functions (e.g., `convert_albert_original_tf_checkpoint_to_pytorch`, `convert_bert_original_tf_checkpoint_to_pytorch`)

**Interfaces and Functions**
------------------------------

* `__init__`: Initializes the class with model type, TensorFlow checkpoint path, PyTorch dump output path, configuration file path or folder, and optional fine-tuning task name.
* `run`: Runs the conversion process based on the specified model type.
* `register_subcommand`: Registers the command with argparse.


key: chroma:document
string_value:
CODE:

def convert_command_factory(args: Namespace):
    """
    Factory function used to convert a model PyTorch checkpoint in a TensorFlow 2 checkpoint.

    Returns: ServeCommand
    """
    return PTtoTFCommand(
        args.model_name,
        args.local_dir,
        args.max_error,
        args.new_weights,
        args.no_pr,
        args.push,
        args.extra_commit_description,
        args.override_model_class,


SUMMARY:
**Summary**

`convert_command_factory` Factory Function:

Converts a PyTorch checkpoint to a TensorFlow 2 checkpoint.

**Key Details**

- Converts a PyTorch model checkpoint to a TensorFlow 2-friendly format using `PTtoTFCommand`.

**Method Signature**
```python
def convert_command_factory(args: Namespace)
```
- Takes `Namespace` arguments (`model_name`, `local_dir`, `max_error`, `new_weights`, `no_pr`, `push`, `extra_commit_description`, and `override_model_class`) to create a command for the conversion process.

**Notable Dependencies**

- Uses `PTtoTFCommand` class, which is not part of this snippet but assumed to be defined elsewhere in the codebase.


key: chroma:document
string_value:
CODE:

class PTtoTFCommand(BaseTransformersCLICommand):
    @staticmethod
    def register_subcommand(parser: ArgumentParser):
        """
        Register this command to argparse so it's available for the transformer-cli

        Args:
            parser: Root parser to register command-specific arguments
        """
        train_parser = parser.add_parser(
            "pt-to-tf",
            help=(
                "CLI tool to run convert a transformers model from a PyTorch checkpoint to a TensorFlow checkpoint."
                " Can also be used to validate existing weights without opening PRs, with --no-pr."
            ),
        )
        train_parser.add_argument(
            "--model-name",
            type=str,
            required=True,
            help="The model name, including owner/organization, as seen on the hub.",
        )
        train_parser.add_argument(
            "--local-dir",
            type=str,
            default="",
            help="Optional local directory of the model repository. Defaults to /tmp/{model_name}",
        )
        train_parser.add_argument(
            "--max-error",
            type=float,
            default=MAX_ERROR,
            help=(
                f"Maximum error tolerance. Defaults to {MAX_ERROR}. This flag should be avoided, use at your own risk."
            ),
        )
        train_parser.add_argument(
            "--new-weights",
            action="store_true",
            help="Optional flag to create new TensorFlow weights, even if they already exist.",
        )
        train_parser.add_argument(
            "--no-pr", action="store_true", help="Optional flag to NOT open a PR with converted weights."
        )
        train_parser.add_argument(
            "--push",
            action="store_true",
            help="Optional flag to push the weights directly to `main` (requires permissions)",
        )
        train_parser.add_argument(
            "--extra-commit-description",
            type=str,
            default="",
            help="Optional additional commit description to use when opening a PR (e.g. to tag the owner).",
        )
        train_parser.add_argument(
            "--override-model-class",
            type=str,
            default=None,
            help="If you think you know better than the auto-detector, you can specify the model class here. "
            "Can be either an AutoModel class or a specific model class like BertForSequenceClassification.",
        )
        train_parser.set_defaults(func=convert_command_factory)

    def __init__(
        self,
        model_name: str,
        local_dir: str,
        max_error: float,
        new_weights: bool,
        no_pr: bool,
        push: bool,
        extra_commit_description: str,
        override_model_class: str,
        *args,
    ):
        self._logger = logging.get_logger("transformers-cli/pt_to_tf")
        self._model_name = model_name
        self._local_dir = local_dir if local_dir else os.path.join("/tmp", model_name)
        self._max_error = max_error
        self._new_weights = new_weights
        self._no_pr = no_pr
        self._push = push
        self._extra_commit_description = extra_commit_description
        self._override_model_class = override_model_class

    def run(self):
        # TODO (joao): delete file in v4.47
        raise NotImplementedError(
            "\n\nConverting PyTorch weights to TensorFlow weights was removed in v4.43. "
            "Instead, we recommend that you convert PyTorch weights to Safetensors, an improved "
            "format that can be loaded by any framework, including TensorFlow. For more information, "
            "please see the Safetensors conversion guide: "
            "https://huggingface.co/docs/safetensors/en/convert-weights\n\n"


SUMMARY:
**Class Summary**

`PTtoTFCommand`: A command-line interface (CLI) tool in `transformers-cli` for converting PyTorch models to TensorFlow checkpoints.

**Key Functionality and Implementation Details**

* Registers the "pt-to-tf" subcommand with arguments:
	+ `model-name`: required string model name
	+ `local-dir`: optional string local directory of the model repository
	+ `max-error`: maximum error tolerance for conversion (defaults to MAX_ERROR)
	+ `new-weights`: flag to create new TensorFlow weights even if they already exist
	+ `no-pr`: flag to NOT open a PR with converted weights
	+ `push`, `extra-commit-description`, and `override-model-class` optional flags
* Initializes the command-line tool with model name, local directory, max error tolerance, flag values, and override model class (if specified)
* Provides a `run()` method that raises a NotImplementedError, indicating that converting PyTorch weights to TensorFlow weights is no longer supported in this version

**Notable Dependencies and Requirements**

* `transformers-cli`
* `base-transformers-clicommand` (inherited interface)
* External dependencies not listed

**Context**
This command-line tool is part of the `transformers-cli` package, which provides a CLI for working with Transformers models. The "pt-to-tf" subcommand allows users to convert PyTorch checkpoints to TensorFlow checkpoints, with optional flags to customize the conversion process.


key: chroma:document
string_value:
CODE:

def serve_command_factory(args: Namespace):
    """
    Factory function used to instantiate serving server from provided command line arguments.

    Returns: ServeCommand
    """
    nlp = pipeline(
        task=args.task,
        model=args.model if args.model else None,
        config=args.config,
        tokenizer=args.tokenizer,
        device=args.device,
    )


SUMMARY:
**Summary:**

 Creates a serving server instance from command line arguments.

**Main purpose and functionality:**
`serve_command_factory` is a factory function that instantiates a `ServeCommand` object based on provided command line arguments. It configures an NLP pipeline using the passed arguments and returns the `ServeCommand` instance.

**Key implementation details:**

* Configures an NLP pipeline with the following parameters:
	+ `task`
	+ `model` (optional)
	+ `config`
	+ `tokenizer`
	+ `device`

* Returns a `ServeCommand` object, which is initialized with the loaded NLP pipeline and other server configuration parameters (`host`, `port`, and `workers`).

**Notable dependencies and requirements:**

* The `pipeline` function and its related classes are not explicitly shown in this code chunk.
* External dependency: This code assumes that a namespace class (`Namespace`) is available, which is likely defined in an external module (e.g., `argparse`).


key: chroma:document
string_value:
CODE:

class ServeModelInfoResult(BaseModel):
    """
    Expose model information
    """



SUMMARY:
**Summary**

The `ServeModelInfoResult` class is a data model that exposes metadata about a resource. It inherits from the `BaseModel` and contains a single attribute `infos`, which is a dictionary.

**Primary Purpose**: Exposes metadata as key-value pairs in a structured format.

**Key Implementation Details**:

* Uses a dictionary (`dict`) to store metadata.
* Inheriting from `BaseModel` suggests that other properties or behavior may be part of this result class.

**Notable Dependencies or Requirements**:

* Inherits from `BaseModel`, implying the use of another data model or framework.
* No external dependencies listed, but may require additional libraries or frameworks depending on the used `BaseModel`.


key: chroma:document
string_value:
CODE:

class ServeTokenizeResult(BaseModel):
    """
    Tokenize result model
    """

    tokens: List[str]


SUMMARY:
**Class Summary**

`ServeTokenizeResult`
-----------

### Purpose and Functionality

The `ServeTokenizeResult` class represents the outcome of tokenization, containing lists of **tokens** (`str`) and optional **token IDs** (`int`).

### Key Implementation Details

* Uses `BaseModel` as its base class for structuring data.
* No additional specific algorithms or design choices are present in this snippet.

### Notable Dependencies or Requirements

 None.

### Interface

Inherited from: `BaseModel`

 Interface: `ServeTokenizeResult`
    Model
  - tokens [List[str]]
  - tokens_ids [Optional[List[int]]]


key: chroma:document
string_value:
CODE:

class ServeDeTokenizeResult(BaseModel):
    """
    DeTokenize result model
    """



SUMMARY:
**Class Summary: ServeDeTokenizeResult**

* **Purpose**: A data model representing the de-tokenization result of a serve operation.

**Key Implementation Details:**

* The class has one property, `text`, which is a string representing the processed text.

**Notable Dependencies/Requirements:**

* Reuses `BaseModel` from an external library or module (assuming it's named exactly as defined).
* No explicit dependencies other than `BaseModel`.
* Does not inherit any interfaces.
* Implements no specific interfaces.


key: chroma:document
string_value:
CODE:

class ServeForwardResult(BaseModel):
    """
    Forward result model
    """



SUMMARY:
**Concise Summary**

`ServeForwardResult`: A data model representing the forward result of a serving operation, containing a single attribute `output` with type `Any`, indicating the result of the request processing.


key: chroma:document
string_value:
CODE:

class ServeCommand(BaseTransformersCLICommand):
    @staticmethod
    def register_subcommand(parser: ArgumentParser):
        """
        Register this command to argparse so it's available for the transformer-cli

        Args:
            parser: Root parser to register command-specific arguments
        """
        serve_parser = parser.add_parser(
            "serve", help="CLI tool to run inference requests through REST and GraphQL endpoints."
        )
        serve_parser.add_argument(
            "--task",
            type=str,
            choices=get_supported_tasks(),
            help="The task to run the pipeline on",
        )
        serve_parser.add_argument("--host", type=str, default="localhost", help="Interface the server will listen on.")
        serve_parser.add_argument("--port", type=int, default=8888, help="Port the serving will listen to.")
        serve_parser.add_argument("--workers", type=int, default=1, help="Number of http workers")
        serve_parser.add_argument("--model", type=str, help="Model's name or path to stored model.")
        serve_parser.add_argument("--config", type=str, help="Model's config name or path to stored model.")
        serve_parser.add_argument("--tokenizer", type=str, help="Tokenizer name to use.")
        serve_parser.add_argument(
            "--device",
            type=int,
            default=-1,
            help="Indicate the device to run onto, -1 indicates CPU, >= 0 indicates GPU (default: -1)",
        )
        serve_parser.set_defaults(func=serve_command_factory)

    def __init__(self, pipeline: Pipeline, host: str, port: int, workers: int):
        self._pipeline = pipeline

        self.host = host
        self.port = port
        self.workers = workers

        if not _serve_dependencies_installed:
            raise RuntimeError(
                "Using serve command requires FastAPI and uvicorn. "
                'Please install transformers with [serving]: pip install "transformers[serving]". '
                "Or install FastAPI and uvicorn separately."
            )
        else:
            logger.info(f"Serving model over {host}:{port}")
            self._app = FastAPI(
                routes=[
                    APIRoute(
                        "/",
                        self.model_info,
                        response_model=ServeModelInfoResult,
                        response_class=JSONResponse,
                        methods=["GET"],
                    ),
                    APIRoute(
                        "/tokenize",
                        self.tokenize,
                        response_model=ServeTokenizeResult,
                        response_class=JSONResponse,
                        methods=["POST"],
                    ),
                    APIRoute(
                        "/detokenize",
                        self.detokenize,
                        response_model=ServeDeTokenizeResult,
                        response_class=JSONResponse,
                        methods=["POST"],
                    ),
                    APIRoute(
                        "/forward",
                        self.forward,
                        response_model=ServeForwardResult,
                        response_class=JSONResponse,
                        methods=["POST"],
                    ),
                ],
                timeout=600,
            )

    def run(self):
        run(self._app, host=self.host, port=self.port, workers=self.workers)

    def model_info(self):
        return ServeModelInfoResult(infos=vars(self._pipeline.model.config))

    def tokenize(self, text_input: str = Body(None, embed=True), return_ids: bool = Body(False, embed=True)):
        """
        Tokenize the provided input and eventually returns corresponding tokens id: - **text_input**: String to
        tokenize - **return_ids**: Boolean flags indicating if the tokens have to be converted to their integer
        mapping.
        """
        try:
            tokens_txt = self._pipeline.tokenizer.tokenize(text_input)

            if return_ids:
                tokens_ids = self._pipeline.tokenizer.convert_tokens_to_ids(tokens_txt)
                return ServeTokenizeResult(tokens=tokens_txt, tokens_ids=tokens_ids)
            else:
                return ServeTokenizeResult(tokens=tokens_txt)

        except Exception as e:
            raise HTTPException(status_code=500, detail={"model": "", "error": str(e)})

    def detokenize(
        self,
        tokens_ids: List[int] = Body(None, embed=True),
        skip_special_tokens: bool = Body(False, embed=True),
        cleanup_tokenization_spaces: bool = Body(True, embed=True),
    ):
        """
        Detokenize the provided tokens ids to readable text: - **tokens_ids**: List of tokens ids -
        **skip_special_tokens**: Flag indicating to not try to decode special tokens - **cleanup_tokenization_spaces**:
        Flag indicating to remove all leading/trailing spaces and intermediate ones.
        """
        try:
            decoded_str = self._pipeline.tokenizer.decode(tokens_ids, skip_special_tokens, cleanup_tokenization_spaces)
            return ServeDeTokenizeResult(model="", text=decoded_str)
        except Exception as e:
            raise HTTPException(status_code=500, detail={"model": "", "error": str(e)})

    async def forward(self, inputs=Body(None, embed=True)):
        """
        **inputs**: **attention_mask**: **tokens_type_ids**:
        """

        # Check we don't have empty string
        if len(inputs) == 0:
            return ServeForwardResult(output=[], attention=[])

        try:
            # Forward through the model
            output = self._pipeline(inputs)
            return ServeForwardResult(output=output)
        except Exception as e:


SUMMARY:
**Summary**

`ServeCommand` is a `BaseTransformersCLICommand` class responsible for serving the transformer pipeline through REST and GraphQL endpoints.

**Main Purpose and Functionality**

The class allows users to run inference requests on a pre-defined task, model, and tokenizer. It listens on specific ports and handles worker processes to serve the pipeline.

**Key Implementation Details**

* Registers subcommands with `argparse` for task selection and configuration
* Installs necessary dependencies (FastAPI and uvicorn) at runtime if not already installed
* Creates a FastAPI app with routes for model info, tokenization, detokenization, and forward requests
* Defines methods for each route:
	+ `model_info`: returns model config information
	+ `tokenize`: tokenizes input text
	+ `detokenize`: detokenizes tokens back to readable text
	+ `forward`: forwards input through the model

**Notable Dependencies and Requirements**

* FastAPI
* uvicorn
* transformers (with serving package)
* Pipelin (assuming a custom pipeline class)


key: chroma:document
string_value:
CODE:

def main():
    parser = ArgumentParser("Transformers CLI tool", usage="transformers-cli <command> [<args>]")
    commands_parser = parser.add_subparsers(help="transformers-cli command helpers")

    # Register commands
    ConvertCommand.register_subcommand(commands_parser)
    DownloadCommand.register_subcommand(commands_parser)
    EnvironmentCommand.register_subcommand(commands_parser)
    RunCommand.register_subcommand(commands_parser)
    ServeCommand.register_subcommand(commands_parser)
    UserCommands.register_subcommand(commands_parser)
    AddNewModelLikeCommand.register_subcommand(commands_parser)
    LfsCommands.register_subcommand(commands_parser)
    PTtoTFCommand.register_subcommand(commands_parser)

    # Let's go
    args = parser.parse_args()

    if not hasattr(args, "func"):
        parser.print_help()
        exit(1)

    # Run
    service = args.func(args)


SUMMARY:
**Main Function Summary**

The `main` function is the entry point of a command-line interface (CLI) tool for transforming and managing various types of data. It defines a parser with subcommand registrations for various transformers, and parses command-line arguments to determine which transformation function to execute.

**Key Implementation Details**

* Defines a parser with subcommands using `ArgumentParser` and registers multiple transformer commands.
* Parses command-line arguments and checks if the `func` attribute is present before executing the corresponding function.

**Notable Dependencies or Requirements**

* `Argparse` library for parsing command-line arguments
* Multiple transformer packages (e.g., ConvertCommand, DownloadCommand) that implement subcommands and registered with the parser.


key: chroma:document
string_value:
CODE:

class LfsCommands(BaseTransformersCLICommand):
    """
    Implementation of a custom transfer agent for the transfer type "multipart" for git-lfs. This lets users upload
    large files >5GB 🔥. Spec for LFS custom transfer agent is:
    https://github.com/git-lfs/git-lfs/blob/master/docs/custom-transfers.md

    This introduces two commands to the CLI:

    1. $ transformers-cli lfs-enable-largefiles

    This should be executed once for each model repo that contains a model file >5GB. It's documented in the error
    message you get if you just try to git push a 5GB file without having enabled it before.

    2. $ transformers-cli lfs-multipart-upload

    This command is called by lfs directly and is not meant to be called by the user.
    """

    @staticmethod
    def register_subcommand(parser: ArgumentParser):
        enable_parser = parser.add_parser(
            "lfs-enable-largefiles",
            help=(
                "Deprecated: use `huggingface-cli` instead. Configure your repository to enable upload of files > 5GB."
            ),
        )
        enable_parser.add_argument("path", type=str, help="Local path to repository you want to configure.")
        enable_parser.set_defaults(func=lambda args: LfsEnableCommand(args))

        upload_parser = parser.add_parser(
            LFS_MULTIPART_UPLOAD_COMMAND,
            help=(
                "Deprecated: use `huggingface-cli` instead. "
                "Command will get called by git-lfs, do not call it directly."
            ),
        )


SUMMARY:
**LfsCommands Class Summary**

*   **Purpose**: Custom transfer agent for the `multiPart` transfer type in Git LFS, enabling large file uploads >5GB.
*   **Functionality**:
    *   Provides two CLI commands: `lfs-enable-largefiles` and `lfs-multipart-upload`.
    *   `lfs-enable-largefiles`: Configures a repository to enable large file uploads, deprecated in favor of `huggingface-cli`.
    *   `lfs-multipart-upload`: Called by Git LFS directly, not meant for direct user invocation.
*   **Key Implementation Details**:
    *   Uses the `ArgumentParser` from `Argparse` module to define CLI subcommands and arguments.
    *   Deploys two subcommands: `LfsEnableCommand` and `LfsUploadCommand`.
*   **Dependencies**: Inherited from `BaseTransformersCLICommand`, uses the `ArgumentParser` from `Argparse`. No external dependencies mentioned.

**Key Method/Function Signatures:**

*   `register_subcommand(parser)` - Registers two CLI subcommands.
*   `lfs-enable-largefiles` and `lfs-multipart-upload` functions (lambda expressions) - Invoked by the registered commands.


key: chroma:document
string_value:
CODE:

class LfsEnableCommand:
    def __init__(self, args):
        self.args = args

    def run(self):
        warnings.warn(
            "Managing repositories through transformers-cli is deprecated. Please use `huggingface-cli` instead."
        )
        local_path = os.path.abspath(self.args.path)
        if not os.path.isdir(local_path):
            print("This does not look like a valid git repo.")
            exit(1)
        subprocess.run(
            "git config lfs.customtransfer.multipart.path transformers-cli".split(), check=True, cwd=local_path
        )
        subprocess.run(
            f"git config lfs.customtransfer.multipart.args {LFS_MULTIPART_UPLOAD_COMMAND}".split(),
            check=True,
            cwd=local_path,
        )


SUMMARY:
**Class Summary: LfsEnableCommand**

*   **Purpose:** This class enables Local File System (LFS) support for a Git repository by configuring the `customtransfer` settings.
*

*   **Key Implementation Details:**
    *   Runs `git config` commands to set up LFS multipart upload transfer using `huggingface-cli`.
    *   Applies the uploaded `LFS_MULTIPART_UPLOAD_COMMAND` as the custom transfer argument.

       ```python
class LfsEnableCommand:
    def __init__(self, args):
        self.args = args

    def run(self):
        # Set up LFS multipart upload
        subprocess.run(
            "git config lfs.customtransfer.multipart.path transformers-cli".split(), check=True, cwd=self.args.path
        )
        subprocess.run(
            f"git config lfs.customtransfer.multipart.args {LFS_MULTIPART_UPLOAD_COMMAND}".split(),
            check=True,
            cwd=self.args.path,
        )
```
*   **Notable Dependencies:**
    *   `subprocess` module for executing shell commands.
    *   `os` module for working with file paths and directories.
    *   ` Git commands`

       ```python
import subprocess
import os

class LfsEnableCommand:
    def __init__(self, args):
        self.args = args

# Usage of the class would involve invoking this function with a path,
# like so...

args = SomeArguments
repo_path = PathToRepository
lfs_command = LfsEnableCommand(args)
```
*   **Interfaces and Inheritance:** 
    None


key: chroma:document
string_value:
CODE:

def write_msg(msg: Dict):
    """Write out the message in Line delimited JSON."""
    msg = json.dumps(msg) + "\n"
    sys.stdout.write(msg)


SUMMARY:
**Function Summary**

`write_msg` function writes out a given message `msg` to the standard output, formatting it as line-delimited JSON.

**Key Details:**

* Formats input `msg` into a JSON string using `json.dumps()`
* Appends a newline character and writes to the standard output
* Flushes stdout buffer after writing

**Note:** This function assumes that the input `msg` is of type `Dict`.


key: chroma:document
string_value:
CODE:

def read_msg() -> Optional[Dict]:
    """Read Line delimited JSON from stdin."""
    msg = json.loads(sys.stdin.readline().strip())

    if "terminate" in (msg.get("type"), msg.get("event")):
        # terminate message received
        return None

    if msg.get("event") not in ("download", "upload"):
        logger.critical("Received unexpected message")
        sys.exit(1)



SUMMARY:
**Summary**

Function `read_msg()` reads line-delimited JSON from stdin and processes it. Its primary purpose is to validate messages received for download or upload events.

*   **Functionality**: Reads JSON input, checks message type and event, and logs critical errors if invalid.

**Key Details**

*   Reads JSON from stdin using `json.loads()` and `sys.stdin.readline()`.
*   Checks if the message indicates termination (`terminate` in type or event).
*   Validates expected events for download or upload and logs an error if unexpected.
*   Returns the validated message (`None` if terminate).

**Notable Dependencies/Requirements**

*   JSON parsing: `json`
*   Logging functions: `logger.critical()`
*   Standard input streams: `sys.stdin`
*   System exit handling: `sys.exit()`


key: chroma:document
string_value:
CODE:

class FileSlice(AbstractContextManager):
    """
    File-like object that only reads a slice of a file

    Inspired by stackoverflow.com/a/29838711/593036
    """

    def __init__(self, filepath: str, seek_from: int, read_limit: int):
        self.filepath = filepath
        self.seek_from = seek_from
        self.read_limit = read_limit
        self.n_seen = 0

    def __enter__(self):
        self.f = open(self.filepath, "rb")
        self.f.seek(self.seek_from)
        return self

    def __len__(self):
        total_length = os.fstat(self.f.fileno()).st_size
        return min(self.read_limit, total_length - self.seek_from)

    def read(self, n=-1):
        if self.n_seen >= self.read_limit:
            return b""
        remaining_amount = self.read_limit - self.n_seen
        data = self.f.read(remaining_amount if n < 0 else min(n, remaining_amount))
        self.n_seen += len(data)
        return data

    def __iter__(self):
        yield self.read(n=4 * 1024 * 1024)

    def __exit__(self, *args):


SUMMARY:
**FileSlice Class Summary**

The `FileSlice` class is an iterable file object that reads a limited slice of a file. It provides a "file-like" interface for reading a subset of a file's contents without loading the entire file into memory.

**Key Features:**

* Reads a fixed-size slice from a file, specified by `seek_from` and `read_limit`.
* Returns an empty byte string if the read limit is reached.
* Uses `os.fstat` to estimate remaining file size.
* Iterates over the read data in chunks of 4MB.

**Notable Dependencies:**

* `os`
* `abstract`

**Implementation Details:**

* `__enter__`: Opens the file and seeks to the specified position.
* `read`: Reads up to `n` bytes from the current position, updating `self.n_seen`.
* `__iter__`: Generates chunks of 4MB using `read(n=4 * 1024 * 1024)`.

**Inherited Interfaces:**

* `AbstractContextManager` (not explicitly listed, but implied by class definition)


key: chroma:document
string_value:
CODE:

class LfsUploadCommand:
    def __init__(self, args):
        self.args = args

    def run(self):
        # Immediately after invoking a custom transfer process, git-lfs
        # sends initiation data to the process over stdin.
        # This tells the process useful information about the configuration.
        init_msg = json.loads(sys.stdin.readline().strip())
        if not (init_msg.get("event") == "init" and init_msg.get("operation") == "upload"):
            write_msg({"error": {"code": 32, "message": "Wrong lfs init operation"}})
            sys.exit(1)

        # The transfer process should use the information it needs from the
        # initiation structure, and also perform any one-off setup tasks it
        # needs to do. It should then respond on stdout with a simple empty
        # confirmation structure, as follows:
        write_msg({})

        # After the initiation exchange, git-lfs will send any number of
        # transfer requests to the stdin of the transfer process, in a serial sequence.
        while True:
            msg = read_msg()
            if msg is None:
                # When all transfers have been processed, git-lfs will send
                # a terminate event to the stdin of the transfer process.
                # On receiving this message the transfer process should
                # clean up and terminate. No response is expected.
                sys.exit(0)

            oid = msg["oid"]
            filepath = msg["path"]
            completion_url = msg["action"]["href"]
            header = msg["action"]["header"]
            chunk_size = int(header.pop("chunk_size"))
            presigned_urls: List[str] = list(header.values())

            parts = []
            for i, presigned_url in enumerate(presigned_urls):
                with FileSlice(filepath, seek_from=i * chunk_size, read_limit=chunk_size) as data:
                    r = requests.put(presigned_url, data=data)
                    r.raise_for_status()
                    parts.append(
                        {
                            "etag": r.headers.get("etag"),
                            "partNumber": i + 1,
                        }
                    )
                    # In order to support progress reporting while data is uploading / downloading,
                    # the transfer process should post messages to stdout
                    write_msg(
                        {
                            "event": "progress",
                            "oid": oid,
                            "bytesSoFar": (i + 1) * chunk_size,
                            "bytesSinceLast": chunk_size,
                        }
                    )
                    # Not precise but that's ok.

            r = requests.post(
                completion_url,
                json={
                    "oid": oid,
                    "parts": parts,
                },
            )
            r.raise_for_status()



SUMMARY:
**Class Summary: LfsUploadCommand**

### Purpose and Functionality
The `LfsUploadCommand` class represents a Git LFS (Large File Storage) upload process. It handles the transfer of large files from a source to a destination, exchanging information through stdin.

### Key Implementation Details

*   Initiates the file transfer process by parsing initiation data from stdin.
*   Verifies the operation type and exchanges useful information between processes.
*   Receives transfer requests and sends corresponding responses over stdin.
*   Updates progress messages to stdout during file uploading/downloading.

### Notable Dependencies or Requirements

*   `json` for initialization message parsing
*   `requests` for making HTTP POST requests
*   `write_msg` function (not defined in the provided code) for sending progress updates and completion notifications
*   Custom `FileSlice` class for handling file slice operations
*   `sys.stdin.readline()` and `sys.exit()` functions for reading input from stdin and exiting the process

### Method Signatures and Interfaces

#### run method
```python
run(self):
    # Transfer process logic
```

This summary provides a concise overview of the key aspects of the `LfsUploadCommand` class, focusing on its purpose, implementation details, and dependencies.


key: chroma:document
string_value:
CODE:

def try_infer_format_from_ext(path: str):
    if not path:
        return "pipe"

    for ext in PipelineDataFormat.SUPPORTED_FORMATS:
        if path.endswith(ext):
            return ext

    raise Exception(
        f"Unable to determine file format from file extension {path}. "
        f"Please provide the format through --format {PipelineDataFormat.SUPPORTED_FORMATS}"


SUMMARY:
**Summary**

tries_to_infer_format_from_ext Function:

```markdown
Tries to determine the file format from its extension. 
If no specific format is specified, it infers based on supported extensions. 
Raises exception if unknown format.

Signature:
    try_infer_format_from_ext(path: str) -> str

Returns:
  Format extension (str)
```

**Key Implementation Details**

- Iterates over `PipelineDataFormat.SUPPORTED_FORMATS` to check for matching file extensions.
- Raises an exception with a hint when no match is found.

**Dependencies/Notable Requirements**

- `PipelineDataFormat`
- `Exception` class for handling unknown formats
```

Let me know if you want any further summary or clarification


key: chroma:document
string_value:
CODE:

def run_command_factory(args):
    nlp = pipeline(
        task=args.task,
        model=args.model if args.model else None,
        config=args.config,
        tokenizer=args.tokenizer,
        device=args.device,
    )
    format = try_infer_format_from_ext(args.input) if args.format == "infer" else args.format
    reader = PipelineDataFormat.from_str(
        format=format,
        output_path=args.output,
        input_path=args.input,
        column=args.column if args.column else nlp.default_input_names,
        overwrite=args.overwrite,
    )


SUMMARY:
**Summary:**

The `run_command_factory` function creates an instance of `RunCommand` by initializing two pipeline components: `nlp` (natural language processing) and `reader` (pipeline data format). The function takes input arguments to configure both pipelines, including model selection, output paths, and input formatting. It returns a fully configured `RunCommand` object.

**Key Implementation Details:**

* NLP pipeline initialization with customizable tasks, models, and configuration.
* Format inference from file extension for the reader pipeline.
* Usage of PipelineDataFormat.from_str to create reader components.

**Notable Dependencies/Requirements:**

* No external dependencies are explicitly referenced in this snippet. However, it assumes presence of `pipeline` and `RunCommand` functions/classes, likely from the NLTK library.


key: chroma:document
string_value:
CODE:

class RunCommand(BaseTransformersCLICommand):
    def __init__(self, nlp: Pipeline, reader: PipelineDataFormat):
        self._nlp = nlp
        self._reader = reader

    @staticmethod
    def register_subcommand(parser: ArgumentParser):
        run_parser = parser.add_parser("run", help="Run a pipeline through the CLI")
        run_parser.add_argument("--task", choices=get_supported_tasks(), help="Task to run")
        run_parser.add_argument("--input", type=str, help="Path to the file to use for inference")
        run_parser.add_argument("--output", type=str, help="Path to the file that will be used post to write results.")
        run_parser.add_argument("--model", type=str, help="Name or path to the model to instantiate.")
        run_parser.add_argument("--config", type=str, help="Name or path to the model's config to instantiate.")
        run_parser.add_argument(
            "--tokenizer", type=str, help="Name of the tokenizer to use. (default: same as the model name)"
        )
        run_parser.add_argument(
            "--column",
            type=str,
            help="Name of the column to use as input. (For multi columns input as QA use column1,columns2)",
        )
        run_parser.add_argument(
            "--format",
            type=str,
            default="infer",
            choices=PipelineDataFormat.SUPPORTED_FORMATS,
            help="Input format to read from",
        )
        run_parser.add_argument(
            "--device",
            type=int,
            default=-1,
            help="Indicate the device to run onto, -1 indicates CPU, >= 0 indicates GPU (default: -1)",
        )
        run_parser.add_argument("--overwrite", action="store_true", help="Allow overwriting the output file.")
        run_parser.set_defaults(func=run_command_factory)

    def run(self):
        nlp, outputs = self._nlp, []

        for entry in self._reader:
            output = nlp(**entry) if self._reader.is_multi_columns else nlp(entry)
            if isinstance(output, dict):
                outputs.append(output)
            else:
                outputs += output

        # Saving data
        if self._nlp.binary_output:
            binary_path = self._reader.save_binary(outputs)
            logger.warning(f"Current pipeline requires output to be in binary format, saving at {binary_path}")
        else:


SUMMARY:
**Summary**

`RunCommand` is a CLI command for running a natural language processing (NLP) pipeline. Its primary purpose is to execute the pipeline through the provided `PipelineDataFormat` and save the results.

**Key Implementation Details**

* Registers subcommand `run` with several optional arguments:
	+ `task`, `input`, `output`, `model`, `config`, `tokenizer`, `column`, `format`, `device`
		- Corresponding default values for some arguments
* Uses a predefined factory to create the actual execution process (`run_command_factory`)
* Executes the pipeline by iterating over the input data and saving the results
* Supports binary output format when specified

**Notable Dependencies or Requirements**

* `Pipeline` class from `BaseTransformersCLICommand`
* `PipelineDataFormat` interface or class
* `TokenizationFactory` (assuming a tokenization setup for tokenizer usage)
* External dependency on logger and binary saving functionality

**Interfaces/Classes Used:**

* `Pipeline`
* `PipelineDataFormat`
* `ArgumentParser`
* `BaseTransformersCLICommand`


key: chroma:document
string_value:
CODE:

def info_command_factory(_):


SUMMARY:
**Summary**

Creates and returns an instance of `EnvironmentCommand`, a command that likely interacts with an environment-specific configuration.

**Key Details**

* Function name: `info_command_factory`
* Return type: `EnvironmentCommand` instance
* Primary purpose: Returns a command instance for information-related operations

Note: The actual implementation of the returned `EnvironmentCommand` is not shown in this code chunk.


key: chroma:document
string_value:
CODE:

def download_command_factory(args):


SUMMARY:
**Summary**

`download_command_factory` Function:

* Creates an instance of `EnvironmentCommand` with the provided `accelerate_config_file`.
* Main purpose: Factory function for generating EnvironmentCommand instances based on acceleration configuration.

**Implementation Details**

* Uses the `args` object to extract the `accelerate_config_file` parameter.
* Directly returns an instance of `EnvironmentCommand`, bypassing a separate constructor call.

**Notable Dependencies/Requirements**

* None explicitly mentioned; `EnvironmentCommand` is assumed to be defined elsewhere in the codebase.


key: chroma:document
string_value:
CODE:

class EnvironmentCommand(BaseTransformersCLICommand):
    @staticmethod
    def register_subcommand(parser: ArgumentParser):
        download_parser = parser.add_parser("env")
        download_parser.set_defaults(func=info_command_factory)
        download_parser.add_argument(
            "--accelerate-config_file",
            default=None,
            help="The accelerate config file to use for the default values in the launching script.",
        )
        download_parser.set_defaults(func=download_command_factory)

    def __init__(self, accelerate_config_file, *args) -> None:
        self._accelerate_config_file = accelerate_config_file

    def run(self):
        safetensors_version = "not installed"
        if is_safetensors_available():
            import safetensors

            safetensors_version = safetensors.__version__
        elif importlib.util.find_spec("safetensors") is not None:
            import safetensors

            safetensors_version = f"{safetensors.__version__} but is ignored because of PyTorch version too old."

        accelerate_version = "not installed"
        accelerate_config = accelerate_config_str = "not found"
        if is_accelerate_available():
            import accelerate
            from accelerate.commands.config import default_config_file, load_config_from_file

            accelerate_version = accelerate.__version__
            # Get the default from the config file.
            if self._accelerate_config_file is not None or os.path.isfile(default_config_file):
                accelerate_config = load_config_from_file(self._accelerate_config_file).to_dict()

            accelerate_config_str = (
                "\n".join([f"\t- {prop}: {val}" for prop, val in accelerate_config.items()])
                if isinstance(accelerate_config, dict)
                else f"\t{accelerate_config}"
            )

        pt_version = "not installed"
        pt_cuda_available = "NA"
        if is_torch_available():
            import torch

            pt_version = torch.__version__
            pt_cuda_available = torch.cuda.is_available()
            pt_npu_available = is_torch_npu_available()

        tf_version = "not installed"
        tf_cuda_available = "NA"
        if is_tf_available():
            import tensorflow as tf

            tf_version = tf.__version__
            try:
                # deprecated in v2.1
                tf_cuda_available = tf.test.is_gpu_available()
            except AttributeError:
                # returns list of devices, convert to bool
                tf_cuda_available = bool(tf.config.list_physical_devices("GPU"))

        flax_version = "not installed"
        jax_version = "not installed"
        jaxlib_version = "not installed"
        jax_backend = "NA"
        if is_flax_available():
            import flax
            import jax
            import jaxlib

            flax_version = flax.__version__
            jax_version = jax.__version__
            jaxlib_version = jaxlib.__version__
            jax_backend = jax.lib.xla_bridge.get_backend().platform

        info = {
            "`transformers` version": version,
            "Platform": platform.platform(),
            "Python version": platform.python_version(),
            "Huggingface_hub version": huggingface_hub.__version__,
            "Safetensors version": f"{safetensors_version}",
            "Accelerate version": f"{accelerate_version}",
            "Accelerate config": f"{accelerate_config_str}",
            "PyTorch version (GPU?)": f"{pt_version} ({pt_cuda_available})",
            "Tensorflow version (GPU?)": f"{tf_version} ({tf_cuda_available})",
            "Flax version (CPU?/GPU?/TPU?)": f"{flax_version} ({jax_backend})",
            "Jax version": f"{jax_version}",
            "JaxLib version": f"{jaxlib_version}",
            "Using distributed or parallel set-up in script?": "<fill in>",
        }
        if is_torch_available():
            if pt_cuda_available:
                info["Using GPU in script?"] = "<fill in>"
                info["GPU type"] = torch.cuda.get_device_name()
            elif pt_npu_available:
                info["Using NPU in script?"] = "<fill in>"
                info["NPU type"] = torch.npu.get_device_name()
                info["CANN version"] = torch.version.cann

        print("\nCopy-and-paste the text below in your GitHub issue and FILL OUT the two last points.\n")
        print(self.format_dict(info))

        return info

    @staticmethod
    def format_dict(d):


SUMMARY:
**EnvironmentCommand Class Summary**

### Purpose and Functionality

The `EnvironmentCommand` class is a CLI command that provides information about the environment and its installed dependencies. It uses various external libraries (e.g., `safetensors`, `accelerate`, `torch`, `tensorflow`, `flax`, `jax`) to gather version numbers, platform information, and GPU availability.

### Implementation Details

* The `register_subcommand` method adds a parser for the "env" command, which runs when the "--env" flag is used.
* The `run` method checks for available libraries, collects their versions and configurations, and formats the results as a string.

### Key Method Signatures and Interfaces

* `is_safetensors_available()`
* `importlib.util.find_spec("safetensors")`
* `download_command_factory`
* `info_command_factory`
* `load_config_from_file`

### Notable Dependencies or Requirements

* External libraries: `safetensors`, `accelerate`, `torch`, `tensorflow`, `flax`, `jax`
* Inherited interfaces:
	+ `ArgumentParser` (from ` Click`)
	+ `BaseTransformersCLICommand`


key: chroma:document
string_value:
CODE:

class UserCommands(BaseTransformersCLICommand):
    @staticmethod
    def register_subcommand(parser: ArgumentParser):
        login_parser = parser.add_parser("login", help="Log in using the same credentials as on huggingface.co")
        login_parser.set_defaults(func=lambda args: LoginCommand(args))
        whoami_parser = parser.add_parser("whoami", help="Find out which huggingface.co account you are logged in as.")
        whoami_parser.set_defaults(func=lambda args: WhoamiCommand(args))
        logout_parser = parser.add_parser("logout", help="Log out")
        logout_parser.set_defaults(func=lambda args: LogoutCommand(args))

        # new system: git-based repo system
        repo_parser = parser.add_parser(
            "repo",
            help="Deprecated: use `huggingface-cli` instead. Commands to interact with your huggingface.co repos.",
        )
        repo_subparsers = repo_parser.add_subparsers(
            help="Deprecated: use `huggingface-cli` instead. huggingface.co repos related commands"
        )
        repo_create_parser = repo_subparsers.add_parser(
            "create", help="Deprecated: use `huggingface-cli` instead. Create a new repo on huggingface.co"
        )
        repo_create_parser.add_argument(
            "name",
            type=str,
            help="Name for your model's repo. Will be namespaced under your username to build the model id.",
        )
        repo_create_parser.add_argument("--organization", type=str, help="Optional: organization namespace.")
        repo_create_parser.add_argument("-y", "--yes", action="store_true", help="Optional: answer Yes to the prompt")


SUMMARY:
**Summary:**

`UserCommands` class enables subcommands for user interaction, including logging in/out and managing Hugging Face repositories. It defines:

*   Two basic commands (`login` and `whoami`) that delegate to specific `LoginCommand` and `WhoamiCommand` respectively.
*   A deprecated legacy repository-related command group (`repo`) with a create parser.

**Key Details:**

*   Subcommands:
    *   `login`: Logs in using the same credentials as on Hugging Face Co (`func=lambda args: LoginCommand(args)`)
    *   `whoami`: Finds out which Hugging Face Co account is logged in as (`func=lambda args: WhoamiCommand(args)`)
    *   `logout`: Logs out (`func=lambda args: LogoutCommand(args)`)

*   Legacy Repository Commands:
    *   `repo create` (deprecated, use Hugging Face CLI instead): 
        +   Creates a new repository on Hugging Face Co (`repo_create_parser.set_defaults(func=lambda args: RepoCreateCommand(args))`)
        +   Parameters:
            -   `name`: The name of the model's repo
            -   `-y` or `--yes`: Optional, answer Yes to the prompt

*   External Dependencies:

    +   `BaseTransformersCLICommand`
    +   `ArgumentParser`
    +   `LoginCommand`
    +   `WhoamiCommand`
    +   `LogoutCommand`
    +   `RepoCreateCommand`


key: chroma:document
string_value:
CODE:

class ANSI:
    """
    Helper for en.wikipedia.org/wiki/ANSI_escape_code
    """

    _bold = "\u001b[1m"
    _red = "\u001b[31m"
    _gray = "\u001b[90m"
    _reset = "\u001b[0m"

    @classmethod
    def bold(cls, s):
        return f"{cls._bold}{s}{cls._reset}"

    @classmethod
    def red(cls, s):
        return f"{cls._bold}{cls._red}{s}{cls._reset}"

    @classmethod
    def gray(cls, s):


SUMMARY:
**Summary: ANSI Color Code Helper Class**

The `ANSI` class provides a simple way to apply colored text in terminal outputs. It uses ANSI escape codes to color text and includes methods for bold, red, and gray text.

**Key Implementation Details:**

* `bold(s)`, `red(s)`, and `gray(s)` methods use ANSI escape codes (`\u001b[1m`, `\u001b[31m`, and `\u001b[90m`) to format the input string with the corresponding color.
* `_reset` code is used to reset the text color to default.

**Notable Dependencies or Requirements:**

* None. This class only relies on the terminal's support for ANSI escape codes.

```
class ANSI:
    @classmethod
    def bold(cls, s):
        return f"{cls._bold}{s}{cls._reset}"

    @classmethod
    def red(cls, s):
        return f"{cls._bold}{cls._red}{s}{cls._reset}"

    @classmethod
    def gray(cls, s):
        return f"{cls._gray}{s}{cls._reset}"
```


key: chroma:document
string_value:
CODE:

def tabulate(rows: List[List[Union[str, int]]], headers: List[str]) -> str:
    """
    Inspired by:

    - stackoverflow.com/a/8356620/593036
    - stackoverflow.com/questions/9535954/printing-lists-as-tabular-data
    """
    col_widths = [max(len(str(x)) for x in col) for col in zip(*rows, headers)]
    row_format = ("{{:{}}} " * len(headers)).format(*col_widths)
    lines = []
    lines.append(row_format.format(*headers))
    lines.append(row_format.format(*["-" * w for w in col_widths]))
    for row in rows:
        lines.append(row_format.format(*row))


SUMMARY:
**Function Summary: `tabulate`**

* **Purpose:** Formats a table from tabular data into a string representation.

* **Key Implementation Details:**
	+ Calculates column widths using the maximum length of values in each column.
	+ Uses Python's string formatting to create a tabular layout with specified columns.
* **Notable Dependencies/Requirements:**
	+ `List` and `Union` types from the standard library.
	+ No external dependencies, but it references stack overflow links for inspiration.

**Key Method Signature:**

`tabulate(rows: List[List[Union[str, int]]], headers: List[str]) -> str`

This function takes in a list of rows (where each row is a list of strings or integers) and a list of column headers. It returns a string representation of the tabulated data.

**Design Choices:**

* Uses Python's built-in string formatting to create a tabular layout, which provides flexibility and ease of use.
* Calculates column widths using `zip` and list comprehension, which is concise and efficient.
* Appends multiple lines to the result string, which demonstrates a clear understanding of how to build up a complex data structure.


key: chroma:document
string_value:
CODE:

class BaseUserCommand:
    def __init__(self, args):


SUMMARY:
**Summary:**

The `BaseUserCommand` class serves as the foundation for user commands. It initializes with an `args` parameter, storing it as an instance variable, allowing for data transfer between command instances.

**Key Details:**

* **Implementation:** Initializer function (`__init__`) sets up the base structure.
* **Primary Purpose:** Facilitates inheritance and polymorphism for user commands.
* **Important Interface:** None.

**Notable Dependencies/Requirements:**

* `args` parameter (instance variable)
* No external dependencies explicitly stated.


key: chroma:document
string_value:
CODE:

class LoginCommand(BaseUserCommand):
    def run(self):
        print(
            ANSI.red(
                "ERROR! `huggingface-cli login` uses an outdated login mechanism "
                "that is not compatible with the Hugging Face Hub backend anymore. "
                "Please use `huggingface-cli login instead."
            )


SUMMARY:
**Summary:**

`LoginCommand` class informs the user about the incompatibility of using an outdated login mechanism with the current Hugging Face Hub backend, and recommends updating to `huggingface-cli login`.

**Key Details:**

* Class: `LoginCommand`
* Uses: outdated login mechanism (incompatible with Hugging Face Hub backend)
* Recommendation: use `huggingface-cli login` instead
* Output:
	+ Error message printed in red ANSI color

**Interfaces and Dependencies:**

* `BaseUserCommand`: inherited from
* None noted


key: chroma:document
string_value:
CODE:

class WhoamiCommand(BaseUserCommand):
    def run(self):
        print(
            ANSI.red(
                "WARNING! `transformers-cli whoami` is deprecated and will be removed in v5. Please use "
                "`huggingface-cli whoami` instead."
            )
        )
        token = HfFolder.get_token()
        if token is None:
            print("Not logged in")
            exit()
        try:
            user, orgs = whoami(token)
            print(user)
            if orgs:
                print(ANSI.bold("orgs: "), ",".join(orgs))
        except HTTPError as e:
            print(e)
            print(ANSI.red(e.response.text))


SUMMARY:
**Class Summary**

`WhoamiCommand` is a command class that retrieves information about the current user on the Hugging Face platform. Its primary purpose is to display the user's details, including their ID and organizations they are part of (if applicable), while logged in using the `HfFolder.get_token()` method.

**Key Implementation Details**

* Displays deprecation warning for using deprecated `transformers-cli whoami` command.
* Retrieves user information using the `whoami` function and logs them to the console.
* Handles exceptions, such as HTTP errors, by printing error messages and exiting with a non-zero status code.

**Notable Dependencies/Requirements**

* `HfFolder`: referenced for getting the logged-in token.
* `transformers-cli` (deprecated): referenced in deprecation warning.
* `HTTPError`: exception type handled by capturing it from `e`.


key: chroma:document
string_value:
CODE:

class LogoutCommand(BaseUserCommand):
    def run(self):
        print(
            ANSI.red(
                "ERROR! `transformers-cli logout` uses an outdated logout mechanism "
                "that is not compatible with the Hugging Face Hub backend anymore. "
                "Please use `huggingface-cli logout instead."
            )


SUMMARY:
**Summary:**
```python
LogoutCommand(BaseUserCommand):
    Runs a command to log out, printing an error message due to outdated logout mechanism compatibility issues with the Hugging Face Hub backend. 
```

**Implementation Details:**

* Calls `print()` with an ANSI-red colored error message.

**Notable Dependencies:**

* `BaseUserCommand` class
* `ANSI` library for colorized output

Note: This summary captures the essence of the code, focusing on its purpose and key implementation details.


key: chroma:document
string_value:
CODE:

class RepoCreateCommand(BaseUserCommand):
    def run(self):
        print(
            ANSI.red(
                "WARNING! Managing repositories through transformers-cli is deprecated. "
                "Please use `huggingface-cli` instead."
            )
        )
        token = HfFolder.get_token()
        if token is None:
            print("Not logged in")
            exit(1)
        try:
            stdout = subprocess.check_output(["git", "--version"]).decode("utf-8")
            print(ANSI.gray(stdout.strip()))
        except FileNotFoundError:
            print("Looks like you do not have git installed, please install.")

        try:
            stdout = subprocess.check_output(["git-lfs", "--version"]).decode("utf-8")
            print(ANSI.gray(stdout.strip()))
        except FileNotFoundError:
            print(
                ANSI.red(
                    "Looks like you do not have git-lfs installed, please install."
                    " You can install from https://git-lfs.github.com/."
                    " Then run `git lfs install` (you only have to do this once)."
                )
            )
        print("")

        user, _ = whoami(token)
        namespace = self.args.organization if self.args.organization is not None else user
        full_name = f"{namespace}/{self.args.name}"
        print(f"You are about to create {ANSI.bold(full_name)}")

        if not self.args.yes:
            choice = input("Proceed? [Y/n] ").lower()
            if not (choice == "" or choice == "y" or choice == "yes"):
                print("Abort")
                exit()
        try:
            url = create_repo(repo_id=full_name, token=token)
        except HTTPError as e:
            print(e)
            print(ANSI.red(e.response.text))
            exit(1)
        print("\nYour repo now lives at:")
        print(f"  {ANSI.bold(url)}")
        print("\nYou can clone it locally with the command below, and commit/push as usual.")
        print(f"\n  git clone {url}")


SUMMARY:
**RepoCreateCommand Class Summary**

**Purpose and Functionality**
The `RepoCreateCommand` class is used to create a new repository using the Git Large File Storage (Git LFS) version control system. It authenticates the user, checks for installation of required tools, and creates a new repository.

**Key Implementation Details**

* Uses `HfFolder.get_token()` to authenticate the user
* Displays version information for `git` and `git-lfs`
* Calculates the repository full name based on namespace and repository name
* Prompts the user to confirm creation before proceeding
* Calls the `create_repo` function to create the new repository

**Notable Dependencies and Requirements**

* `git` and `git-lfs` are required tools
* `HfFolder` class is used for authentication
* `whoami` function is used to authenticate the user
* `create_repo` function is called to create a new repository, which may throw an `HTTPError`

**Interface/Method Signatures**

* No public methods; all implementation details are within the `__init__` method (not shown in this code snippet).


key: chroma:document
string_value:
CODE:

def list_field(default=None, metadata=None):


SUMMARY:
**Function Summary**

`list_field` creates and returns a Django `ModelField`, providing a default value through a lambda function. It also accepts optional metadata parameters.

**Key Implementation Details**

* Uses Django's `default_factory` parameter to set the default value of the field
* Applies metadata to the created field using the `metadata` parameter

**Notable Dependencies**

* Django (specifically, its ORM and Field classes)


key: chroma:document
string_value:
CODE:

def run_with_tf_optimizations(do_eager_mode: bool, use_xla: bool):
    def run_func(func):
        @wraps(func)
        def run_in_eager_mode(*args, **kwargs):
            return func(*args, **kwargs)

        @wraps(func)
        @tf.function(experimental_compile=use_xla)
        def run_in_graph_mode(*args, **kwargs):
            return func(*args, **kwargs)

        if do_eager_mode is True:
            if use_xla is not False:
                raise ValueError(
                    "Cannot run model in XLA, if `args.eager_mode` is set to `True`. Please set `args.eager_mode=False`."
                )
            return run_in_eager_mode
        else:
            return run_in_graph_mode



SUMMARY:
**Summary:**

The `run_with_tf_optimizations` function enables or disables TensorFlow optimizations (XLA) in a given model execution. It takes two boolean inputs: `do_eager_mode` and `use_xla`.

**Key Functionality:**

- Runs the provided function (`func`) either in Eager mode or Graph mode based on the input value of `do_eager_mode`.
- When `do_eager_mode` is True, it raises an error if `use_xla` is set to True.
- Uses TensorFlow's `tf.function` decorator with `experimental_compile=use_xla` to enable XLA optimizations in Graph mode.

**Implementation Details:**

- The function utilizes a wrapper (`run_func`) that takes the input function (`func`) and returns another wrapper (`run_in_eager_mode` or `run_in_graph_mode`).
- The `@wraps(func)` decorator is used to preserve the original function's name and docstring.
- The `tf.function` decorator with the `experimental_compile=use_xla` argument enables XLA optimizations in Graph mode.

**Dependencies and Requirements:**

- TensorFlow (`tf`) with support for XLA (Experimental).


key: chroma:document
string_value:
CODE:

def random_input_ids(batch_size: int, sequence_length: int, vocab_size: int) -> ["tf.Tensor"]:
    rng = random.Random()
    values = [rng.randint(0, vocab_size - 1) for i in range(batch_size * sequence_length)]


SUMMARY:
**Function Summary:**

`random_input_ids`: generates a batch of randomly sampled integer indices from a vocabulary size, shaped as input tensor for transformer model.

**Key Implementation Details:**

* Utilizes `tf.constant` to convert list of integers into a TensorFlow tensor with shape `(batch_size, sequence_length)` and dtype `tf.int32`.
* Random integers are generated using `random.Random()` to ensure reproducibility.
* Input size defaults to batch size x sequence length.

**Notable Dependencies/Requirements:**

* `tensorflow (tf)`
* `random` module


key: chroma:document
string_value:
CODE:

class TensorFlowBenchmark(Benchmark):
    args: TensorFlowBenchmarkArguments
    configs: PretrainedConfig
    framework: str = "TensorFlow"

    @property
    def framework_version(self):
        return tf.__version__

    def _inference_speed(self, model_name: str, batch_size: int, sequence_length: int) -> float:
        # initialize GPU on separate process
        strategy = self.args.strategy
        if strategy is None:
            raise ValueError("A device strategy has to be initialized before using TensorFlow.")
        _inference = self._prepare_inference_func(model_name, batch_size, sequence_length)
        return self._measure_speed(_inference)

    def _train_speed(self, model_name: str, batch_size: int, sequence_length: int) -> float:
        strategy = self.args.strategy
        if strategy is None:
            raise ValueError("A device strategy has to be initialized before using TensorFlow.")
        _train = self._prepare_train_func(model_name, batch_size, sequence_length)
        return self._measure_speed(_train)

    def _inference_memory(
        self, model_name: str, batch_size: int, sequence_length: int
    ) -> [Memory, Optional[MemorySummary]]:
        # initialize GPU on separate process
        if self.args.is_gpu:
            tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx], True)
        strategy = self.args.strategy
        if strategy is None:
            raise ValueError("A device strategy has to be initialized before using TensorFlow.")
        _inference = self._prepare_inference_func(model_name, batch_size, sequence_length)
        return self._measure_memory(_inference)

    def _train_memory(
        self, model_name: str, batch_size: int, sequence_length: int
    ) -> [Memory, Optional[MemorySummary]]:
        if self.args.is_gpu:
            tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx], True)
        strategy = self.args.strategy
        if strategy is None:
            raise ValueError("A device strategy has to be initialized before using TensorFlow.")

        _train = self._prepare_train_func(model_name, batch_size, sequence_length)
        return self._measure_memory(_train)

    def _prepare_inference_func(self, model_name: str, batch_size: int, sequence_length: int) -> Callable[[], None]:
        config = self.config_dict[model_name]

        if self.args.fp16:
            raise NotImplementedError("Mixed precision is currently not supported.")

        has_model_class_in_config = (
            hasattr(config, "architectures")
            and isinstance(config.architectures, list)
            and len(config.architectures) > 0
        )
        if not self.args.only_pretrain_model and has_model_class_in_config:
            try:
                model_class = "TF" + config.architectures[0]  # prepend 'TF' for tensorflow model
                transformers_module = __import__("transformers", fromlist=[model_class])
                model_cls = getattr(transformers_module, model_class)
                model = model_cls(config)
            except ImportError:
                raise ImportError(
                    f"{model_class} does not exist. If you just want to test the pretrained model, you might want to"
                    " set `--only_pretrain_model` or `args.only_pretrain_model=True`."
                )
        else:
            model = TF_MODEL_MAPPING[config.__class__](config)

        # encoder-decoder has vocab size saved differently
        vocab_size = config.vocab_size if hasattr(config, "vocab_size") else config.encoder.vocab_size
        input_ids = random_input_ids(batch_size, sequence_length, vocab_size)

        @run_with_tf_optimizations(self.args.eager_mode, self.args.use_xla)
        def encoder_decoder_forward():
            return model(input_ids, decoder_input_ids=input_ids, training=False)

        @run_with_tf_optimizations(self.args.eager_mode, self.args.use_xla)
        def encoder_forward():
            return model(input_ids, training=False)

        _inference = encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

        return _inference

    def _prepare_train_func(self, model_name: str, batch_size: int, sequence_length: int) -> Callable[[], None]:
        config = self.config_dict[model_name]

        if self.args.eager_mode is not False:
            raise ValueError("Training cannot be done in eager mode. Please make sure that `args.eager_mode = False`.")

        if self.args.fp16:
            raise NotImplementedError("Mixed precision is currently not supported.")

        has_model_class_in_config = (
            hasattr(config, "architectures")
            and isinstance(config.architectures, list)
            and len(config.architectures) > 0
        )
        if not self.args.only_pretrain_model and has_model_class_in_config:
            try:
                model_class = "TF" + config.architectures[0]  # prepend 'TF' for tensorflow model
                transformers_module = __import__("transformers", fromlist=[model_class])
                model_cls = getattr(transformers_module, model_class)
                model = model_cls(config)
            except ImportError:
                raise ImportError(
                    f"{model_class} does not exist. If you just want to test the pretrained model, you might want to"
                    " set `--only_pretrain_model` or `args.only_pretrain_model=True`."
                )
        else:
            model = TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](config)

        # encoder-decoder has vocab size saved differently
        vocab_size = config.vocab_size if hasattr(config, "vocab_size") else config.encoder.vocab_size
        input_ids = random_input_ids(batch_size, sequence_length, vocab_size)

        @run_with_tf_optimizations(self.args.eager_mode, self.args.use_xla)
        def encoder_decoder_train():
            loss = model(input_ids, decoder_input_ids=input_ids, labels=input_ids, training=True)[0]
            gradients = tf.gradients(loss, model.trainable_variables)
            return gradients

        @run_with_tf_optimizations(self.args.eager_mode, self.args.use_xla)
        def encoder_train():
            loss = model(input_ids, labels=input_ids, training=True)[0]
            gradients = tf.gradients(loss, model.trainable_variables)
            return gradients

        _train = encoder_decoder_train if config.is_encoder_decoder else encoder_train

        return _train

    def _measure_speed(self, func) -> float:
        with self.args.strategy.scope():
            try:
                if self.args.is_tpu or self.args.use_xla:
                    # run additional 10 times to stabilize compilation for tpu
                    logger.info("Do inference on TPU. Running model 5 times to stabilize compilation")
                    timeit.repeat(func, repeat=1, number=5)

                # as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
                runtimes = timeit.repeat(
                    func,
                    repeat=self.args.repeat,
                    number=10,
                )

                return min(runtimes) / 10.0
            except ResourceExhaustedError as e:
                self.print_fn(f"Doesn't fit on GPU. {e}")

    def _measure_memory(self, func: Callable[[], None]) -> [Memory, MemorySummary]:
        logger.info(
            "Note that TensorFlow allocates more memory than "
            "it might need to speed up computation. "
            "The memory reported here corresponds to the memory "
            "reported by `nvidia-smi`, which can vary depending "
            "on total available memory on the GPU that is used."
        )
        with self.args.strategy.scope():
            try:
                if self.args.trace_memory_line_by_line:
                    if not self.args.eager_mode:
                        raise ValueError(
                            "`args.eager_mode` is set to `False`. Make sure to run model in eager mode to measure memory"
                            " consumption line by line."
                        )
                    trace = start_memory_tracing("transformers")

                if self.args.is_tpu:
                    # tpu
                    raise NotImplementedError(
                        "Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking"
                        " with `args.memory=False`"
                    )
                elif self.args.is_gpu:
                    # gpu
                    if not is_py3nvml_available():
                        logger.warning(
                            "py3nvml not installed, we won't log GPU memory usage. "
                            "Install py3nvml (pip install py3nvml) to log information about GPU."
                        )
                        memory = "N/A"
                    else:
                        logger.info(
                            "Measuring total GPU usage on GPU device. Make sure to not have additional processes"
                            " running on the same GPU."
                        )
                        # init nvml
                        nvml.nvmlInit()
                        func()
                        handle = nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx)
                        meminfo = nvml.nvmlDeviceGetMemoryInfo(handle)
                        max_bytes_in_use = meminfo.used
                        memory = Memory(max_bytes_in_use)
                        # shutdown nvml
                        nvml.nvmlShutdown()
                else:
                    # cpu
                    if self.args.trace_memory_line_by_line:
                        logger.info(
                            "When enabling line by line tracing, the max peak memory for CPU is inaccurate in"
                            " TensorFlow."
                        )
                        memory = None
                    else:
                        memory_bytes = measure_peak_memory_cpu(func)
                        memory = Memory(memory_bytes) if isinstance(memory_bytes, int) else memory_bytes
                if self.args.trace_memory_line_by_line:
                    summary = stop_memory_tracing(trace)
                    if memory is None:
                        memory = summary.total
                else:
                    summary = None

                return memory, summary
            except ResourceExhaustedError as e:
                self.print_fn(f"Doesn't fit on GPU. {e}")


SUMMARY:
This code appears to be part of a larger class or module for benchmarking and measuring the performance of transformer models, likely in a deep learning context. The specific methods of the class are:

1. **Model-related methods**: These include various functions related to training, inference, and benchmarking of transformer models.
2. **Benchmarking methods**:
	* `_measure_speed`: Measures the execution time of a given function.
	* `_measure_memory`: Measures the memory usage of a given function.

**Key takeaways:**

1. The code uses various libraries such as `timeit`, `nvidia-smi`, and `py3nvml` to measure performance and memory usage.
2. It supports both GPU and CPU benchmarking, with some specific configurations only applicable to certain devices or environments (e.g., TPU, multiple GPUs).
3. The code includes various logging statements to provide feedback on the progress of the benchmarks.

**Suggestions:**

1. **Code organization**: Some methods seem duplicated in terms of functionality or scope. Consider refactoring these parts for better modularity and reusability.
2. **Type hints and documentation**: Adding explicit type hints and comments (preferably with Dox) to describe each method's purpose, input parameters, output values, and any dependencies would make the code more readable and maintainable.
3. **Exception handling**: While some exceptions are caught and handled, it might be beneficial to include additional logging statements for non-catchable or unhandled exceptions to provide a clearer error message when applicable.

**Additional information:**

If possible, include any specific requirements, constraints, or optimizations that the code should adhere to (e.g., PyTorch, TensorFlow, performance metrics). This would enable an individual to better understand the context and implications of the provided code.


key: chroma:document
string_value:
CODE:

def separate_process_wrapper_fn(func: Callable[[], None], do_multi_processing: bool) -> Callable[[], None]:
    """
    This function wraps another function into its own separated process. In order to ensure accurate memory
    measurements it is important that the function is executed in a separate process

    Args:
        - `func`: (`callable`): function() -> ... generic function which will be executed in its own separate process
        - `do_multi_processing`: (`bool`) Whether to run function on separate process or not
    """

    def multi_process_func(*args, **kwargs):
        # run function in an individual
        # process to get correct memory
        def wrapper_func(queue: Queue, *args):
            try:
                result = func(*args)
            except Exception as e:
                logger.error(e)
                print(e)
                result = "N/A"
            queue.put(result)

        queue = Queue()
        p = Process(target=wrapper_func, args=[queue] + list(args))
        p.start()
        result = queue.get()
        p.join()
        return result

    if do_multi_processing:
        logger.info(f"Function {func} is executed in its own process...")
        return multi_process_func
    else:


SUMMARY:
**Summary: Separate Function Execution**

The `separate_process_wrapper_fn` function wraps another function (`func`) to execute it in a separate process. It ensures accurate memory measurements by running the function individually.

**Key Details:**

* Wraps `func` in a separate process if `do_multi_processing` is True
* Executes `func` with individual arguments in a separate process using multiprocessing

**Notable Function Signatures:**

* `wrapper_func(Queue, *args)`: wrapper function for executing `func` in a separate process
* `multi_process_func(*args, **kwargs)`: main function that runs the wrapped function in a separate process

**External Dependencies:**

* `Queue` and `Process` from multiprocessing library for parallel processing
* `logger` module for logging errors


key: chroma:document
string_value:
CODE:

def is_memory_tracing_enabled():
    global _is_memory_tracing_enabled


SUMMARY:
**Code Summary:**

`is_memory_tracing_enabled()` Function

*   **Purpose**: Returns whether memory tracing is currently enabled in the application.

*   **Key Implementation Details**:
    *   Uses a global variable `_is_memory_tracing_enabled` to store the enablement status.
    *   Simply returns the value of this global variable without any logic or calculation.

*   **Notable Dependencies or Requirements**:
    +   Global variable `_is_memory_tracing_enabled`

This function has no complex logic and solely relies on a global variable for its return value.


key: chroma:document
string_value:
CODE:

class Frame(NamedTuple):
    """
    `Frame` is a NamedTuple used to gather the current frame state. `Frame` has the following fields:

        - 'filename' (string): Name of the file currently executed
        - 'module' (string): Name of the module currently executed
        - 'line_number' (int): Number of the line currently executed
        - 'event' (string): Event that triggered the tracing (default will be "line")
        - 'line_text' (string): Text of the line in the python script
    """

    filename: str
    module: str
    line_number: int
    event: str


SUMMARY:
**Class Summary**

`Frame`: A NamedTuple representing the current frame state, comprising file and module information, execution metadata, and Python script text. The primary purpose is to gather and store frame data for tracing purposes.

**Key Implementation Details**

- `event`: An optional string field defaulting to "line" to track event triggers.
- Four specific fields: `filename`, `module`, `line_number`, `line_text` (strings), providing essential frame state information.

**Notable Dependencies or Requirements**

- Referenced as a NamedTuple from the `dataclasses` module using `namedtuple()`.
- No explicit external dependencies.


key: chroma:document
string_value:
CODE:

class UsedMemoryState(NamedTuple):
    """
    `UsedMemoryState` are named tuples with the following fields:

        - 'frame': a `Frame` namedtuple (see below) storing information on the current tracing frame (current file,
          location in current file)
        - 'cpu_memory': CPU RSS memory state *before* executing the line
        - 'gpu_memory': GPU used memory *before* executing the line (sum for all GPUs or for only `gpus_to_trace` if
          provided)
    """

    frame: Frame
    cpu_memory: int


SUMMARY:
**UsedMemoryState Class Summary**

*   **Purpose and Functionality**: The `UsedMemoryState` class represents the memory state before executing a line of code, consisting of:
    *   `frame`: current tracing frame (file and location)
    *   `cpu_memory`: CPU RSS memory usage
    *   `gpu_memory`: GPU used memory (combined or per-GPU, depending on `gpus_to_trace`)

*   **Key Implementation Details**:
    *   The class utilizes a named tuple to store memory state information.

*   **Notable Dependencies or Requirements**:
    *   `Frame` namedtuple is referenced.
    *   `gpus_to_trace` (an optional parameter) influences GPU memory calculation.
    
**Example Code**
```python
from collections import namedtuple

class UsedMemoryState(NamedTuple):
    pass  # implementation details omitted for brevity
```
Note: The actual implementation of the `Frame` namedtuple and `gpus_to_trace` variable are not provided in the given code chunk.


key: chroma:document
string_value:
CODE:

class Memory(NamedTuple):
    """
    `Memory` NamedTuple have a single field `bytes` and you can get a human readable str of the number of mega bytes by
    calling `__repr__`

        - `byte` (integer): number of bytes,
    """

    bytes: int

    def __repr__(self) -> str:


SUMMARY:
**Class Summary: Memory**

The `Memory` class represents a memory allocation with a single field `bytes`. It utilizes the `NamedTuple` data structure.

**Key Features:**

- Single field `bytes` (integer)
- Implements custom `__repr__` method to return human-readable representation of memory size in MB
- Utilizes a separate function `byte_to_mega_bytes` internally

**Implementation Details:**

- Field: `bytes`
- Implementation Method Signature: `str(bytes_to_mega_bytes(self.bytes))`

**Notable Dependencies and Requirements:**

- Reuses the `NamedTuple` data structure from the `typing` module
- Depends on the functionality of the `byte_to_mega_bytes` function (not shown in this snippet)


key: chroma:document
string_value:
CODE:

class MemoryState(NamedTuple):
    """
    `MemoryState` are namedtuples listing frame + CPU/GPU memory with the following fields:

        - `frame` (`Frame`): the current frame (see above)
        - `cpu`: CPU memory consumed at during the current frame as a `Memory` named tuple
        - `gpu`: GPU memory consumed at during the current frame as a `Memory` named tuple
        - `cpu_gpu`: CPU + GPU memory consumed at during the current frame as a `Memory` named tuple
    """

    frame: Frame
    cpu: Memory
    gpu: Memory


SUMMARY:
**Summary**

`MemoryState` class is an immutable namedtuple representing the current memory state, containing:
*   `frame`: the current frame
*   CPU and GPU memory usage (`cpu` and `gpu`) as separate `Memory` tuples
*   A combined `CPU + GPU` memory usage (`cpu_gpu`) tuple.

**Implementation Details**

- Uses an inner class `Frame`, referencing `Memory`, likely used for memory allocation details.
- Utilizes namedtuples from the `typing` module to define and access fields without ambiguity.


key: chroma:document
string_value:
CODE:

class MemorySummary(NamedTuple):
    """
    `MemorySummary` namedtuple otherwise with the fields:

        - `sequential`: a list of `MemoryState` namedtuple (see below) computed from the provided `memory_trace` by
          subtracting the memory after executing each line from the memory before executing said line.
        - `cumulative`: a list of `MemoryState` namedtuple (see below) with cumulative increase in memory for each line
          obtained by summing repeated memory increase for a line if it's executed several times. The list is sorted
          from the frame with the largest memory consumption to the frame with the smallest (can be negative if memory
          is released)
        - `total`: total memory increase during the full tracing as a `Memory` named tuple (see below). Line with
          memory release (negative consumption) are ignored if `ignore_released_memory` is `True` (default).
    """

    sequential: List[MemoryState]
    cumulative: List[MemoryState]
    current: List[MemoryState]


SUMMARY:
**Summary:**

The `MemorySummary` class is a named tuple representing the memory state of a program during tracing. It contains three lists and one total value:

*   **sequential**: A list of `MemoryState` objects, where each represents the memory change after executing a line.
*   **cumulative**: A list of `MemoryState` objects, showing cumulative increases in memory for each line, sorted from highest to lowest consumption.
*   **current**: The current memory state (empty by default).
*   **total**: The total memory increase during tracing, excluding lines with released memory if `ignore_released_memory` is True.

**Key Implementation Details:**

*   Uses `MemoryState`(named tuple) objects to represent the memory changes.
*   Sorts `cumulative` list in descending order by memory consumption.

**Notable Dependencies and Requirements:**

*   Refers to `MemoryState` namedtuple.
*   Assumes existence of a `Memory` named tuple.


key: chroma:document
string_value:
CODE:

def measure_peak_memory_cpu(function: Callable[[], None], interval=0.5, device_idx=None) -> int:
    """
    measures peak cpu memory consumption of a given `function` running the function for at least interval seconds and
    at most 20 * interval seconds. This function is heavily inspired by: `memory_usage` of the package
    `memory_profiler`:
    https://github.com/pythonprofilers/memory_profiler/blob/895c4ac7a08020d66ae001e24067da6dcea42451/memory_profiler.py#L239

    Args:
        - `function`: (`callable`): function() -> ... function without any arguments to measure for which to measure
          the peak memory

        - `interval`: (`float`, `optional`, defaults to `0.5`) interval in second for which to measure the memory usage

        - `device_idx`: (`int`, `optional`, defaults to `None`) device id for which to measure gpu usage

    Returns:

        - `max_memory`: (`int`) consumed memory peak in Bytes
    """

    def get_cpu_memory(process_id: int) -> int:
        """
        measures current cpu memory usage of a given `process_id`

        Args:
            - `process_id`: (`int`) process_id for which to measure memory

        Returns

            - `memory`: (`int`) consumed memory in Bytes
        """
        process = psutil.Process(process_id)
        try:
            meminfo_attr = "memory_info" if hasattr(process, "memory_info") else "get_memory_info"
            memory = getattr(process, meminfo_attr)()[0]
        except psutil.AccessDenied:
            raise ValueError("Error with Psutil.")
        return memory

    if not is_psutil_available():
        logger.warning(
            "Psutil not installed, we won't log CPU memory usage. "
            "Install Psutil (pip install psutil) to use CPU memory tracing."
        )
        max_memory = "N/A"
    else:

        class MemoryMeasureProcess(Process):
            """
            `MemoryMeasureProcess` inherits from `Process` and overwrites its `run()` method. Used to measure the
            memory usage of a process
            """

            def __init__(self, process_id: int, child_connection: Connection, interval: float):
                super().__init__()
                self.process_id = process_id
                self.interval = interval
                self.connection = child_connection
                self.num_measurements = 1
                self.mem_usage = get_cpu_memory(self.process_id)

            def run(self):
                self.connection.send(0)
                stop = False
                while True:
                    self.mem_usage = max(self.mem_usage, get_cpu_memory(self.process_id))
                    self.num_measurements += 1

                    if stop:
                        break

                    stop = self.connection.poll(self.interval)

                # send results to parent pipe
                self.connection.send(self.mem_usage)
                self.connection.send(self.num_measurements)

        while True:
            # create child, parent connection
            child_connection, parent_connection = Pipe()

            # instantiate process
            mem_process = MemoryMeasureProcess(os.getpid(), child_connection, interval)
            mem_process.start()

            # wait until we get memory
            parent_connection.recv()

            try:
                # execute function
                function()

                # start parent connection
                parent_connection.send(0)

                # receive memory and num measurements
                max_memory = parent_connection.recv()
                num_measurements = parent_connection.recv()
            except Exception:
                # kill process in a clean way
                parent = psutil.Process(os.getpid())
                for child in parent.children(recursive=True):
                    os.kill(child.pid, SIGKILL)
                mem_process.join(0)
                raise RuntimeError("Process killed. Error in Process")

            # run process at least 20 * interval or until it finishes
            mem_process.join(20 * interval)

            if (num_measurements > 4) or (interval < 1e-6):
                break

            # reduce interval
            interval /= 10



SUMMARY:
Here is a concise summary of the `measure_peak_memory_cpu` function:

**Purpose and Functionality**

Measures peak CPU memory consumption of a given function, running it for an extended period and capturing peak memory usage.

**Key Implementation Details**

* Uses `psutil` to measure current CPU memory usage.
* Creates a custom process (`MemoryMeasureProcess`) to execute the given function while measuring memory usage.
* Uses a pipe-based approach to communicate with the child process.

**Notable Dependencies and Requirements**

* `psutil`
* `memory_profiler` (inspiration)
* External dependencies:
	+ `os` for pid and signal handling
	+ `psutil` for OS-level processes functionality

**Primary Purpose**

Calculates the maximum amount of memory consumed by a given function during its execution.

**Notable Method Signatures and Interfaces**

* `get_cpu_memory(process_id: int) -> int`: measures current CPU memory usage.
* `is_psutil_available() -> bool`: checks if Psutil is installed.
* `MemoryMeasureProcess(os.getpid(), child_connection, interval)` : custom process class for measuring memory usage.


key: chroma:document
string_value:
CODE:

def start_memory_tracing(
    modules_to_trace: Optional[Union[str, Iterable[str]]] = None,
    modules_not_to_trace: Optional[Union[str, Iterable[str]]] = None,
    events_to_trace: str = "line",
    gpus_to_trace: Optional[List[int]] = None,
) -> MemoryTrace:
    """
    Setup line-by-line tracing to record rss mem (RAM) at each line of a module or sub-module. See `./benchmark.py` for
    usage examples. Current memory consumption is returned using psutil and in particular is the RSS memory "Resident
    Set Size” (the non-swapped physical memory the process is using). See
    https://psutil.readthedocs.io/en/latest/#psutil.Process.memory_info

    Args:
        - `modules_to_trace`: (None, string, list/tuple of string) if None, all events are recorded if string or list
          of strings: only events from the listed module/sub-module will be recorded (e.g. 'fairseq' or
          'transformers.models.gpt2.modeling_gpt2')
        - `modules_not_to_trace`: (None, string, list/tuple of string) if None, no module is avoided if string or list
          of strings: events from the listed module/sub-module will not be recorded (e.g. 'torch')
        - `events_to_trace`: string or list of string of events to be recorded (see official python doc for
          `sys.settrace` for the list of events) default to line
        - `gpus_to_trace`: (optional list, default None) list of GPUs to trace. Default to tracing all GPUs

    Return:

        - `memory_trace` is a list of `UsedMemoryState` for each event (default each line of the traced script).

            - `UsedMemoryState` are named tuples with the following fields:

                - 'frame': a `Frame` namedtuple (see below) storing information on the current tracing frame (current
                  file, location in current file)
                - 'cpu_memory': CPU RSS memory state *before* executing the line
                - 'gpu_memory': GPU used memory *before* executing the line (sum for all GPUs or for only
                  `gpus_to_trace` if provided)

    `Frame` is a namedtuple used by `UsedMemoryState` to list the current frame state. `Frame` has the following
    fields: - 'filename' (string): Name of the file currently executed - 'module' (string): Name of the module
    currently executed - 'line_number' (int): Number of the line currently executed - 'event' (string): Event that
    triggered the tracing (default will be "line") - 'line_text' (string): Text of the line in the python script

    """
    if is_psutil_available():
        process = psutil.Process(os.getpid())
    else:
        logger.warning(
            "Psutil not installed, we won't log CPU memory usage. "
            "Install psutil (pip install psutil) to use CPU memory tracing."
        )
        process = None

    if is_py3nvml_available():
        try:
            nvml.nvmlInit()
            devices = list(range(nvml.nvmlDeviceGetCount())) if gpus_to_trace is None else gpus_to_trace
            nvml.nvmlShutdown()
        except (OSError, nvml.NVMLError):
            logger.warning("Error while initializing communication with GPU. We won't perform GPU memory tracing.")
            log_gpu = False
        else:
            log_gpu = is_torch_available() or is_tf_available()
    else:
        logger.warning(
            "py3nvml not installed, we won't log GPU memory usage. "
            "Install py3nvml (pip install py3nvml) to use GPU memory tracing."
        )
        log_gpu = False

    memory_trace = []

    def traceit(frame, event, args):
        """
        Tracing method executed before running each line in a module or sub-module Record memory allocated in a list
        with debugging information
        """
        global _is_memory_tracing_enabled

        if not _is_memory_tracing_enabled:
            return traceit

        # Filter events
        if events_to_trace is not None:
            if isinstance(events_to_trace, str) and event != events_to_trace:
                return traceit
            elif isinstance(events_to_trace, (list, tuple)) and event not in events_to_trace:
                return traceit

        if "__name__" not in frame.f_globals:
            return traceit

        # Filter modules
        name = frame.f_globals["__name__"]
        if not isinstance(name, str):
            return traceit
        else:
            # Filter whitelist of modules to trace
            if modules_to_trace is not None:
                if isinstance(modules_to_trace, str) and modules_to_trace not in name:
                    return traceit
                elif isinstance(modules_to_trace, (list, tuple)) and all(m not in name for m in modules_to_trace):
                    return traceit

            # Filter blacklist of modules not to trace
            if modules_not_to_trace is not None:
                if isinstance(modules_not_to_trace, str) and modules_not_to_trace in name:
                    return traceit
                elif isinstance(modules_not_to_trace, (list, tuple)) and any(m in name for m in modules_not_to_trace):
                    return traceit

        # Record current tracing state (file, location in file...)
        lineno = frame.f_lineno
        filename = frame.f_globals["__file__"]
        if filename.endswith(".pyc") or filename.endswith(".pyo"):
            filename = filename[:-1]
        line = linecache.getline(filename, lineno).rstrip()
        traced_state = Frame(filename, name, lineno, event, line)

        # Record current memory state (rss memory) and compute difference with previous memory state
        cpu_mem = 0
        if process is not None:
            mem = process.memory_info()
            cpu_mem = mem.rss

        gpu_mem = 0
        if log_gpu:
            # Clear GPU caches
            if is_torch_available():
                torch_empty_cache()
            if is_tf_available():
                tf_context.context()._clear_caches()  # See https://github.com/tensorflow/tensorflow/issues/20218#issuecomment-416771802

            # Sum used memory for all GPUs
            nvml.nvmlInit()

            for i in devices:
                handle = nvml.nvmlDeviceGetHandleByIndex(i)
                meminfo = nvml.nvmlDeviceGetMemoryInfo(handle)
                gpu_mem += meminfo.used

            nvml.nvmlShutdown()

        mem_state = UsedMemoryState(traced_state, cpu_mem, gpu_mem)
        memory_trace.append(mem_state)

        return traceit

    sys.settrace(traceit)

    global _is_memory_tracing_enabled
    _is_memory_tracing_enabled = True



SUMMARY:
**Summary: `start_memory_tracing` Function**

**Purpose and Functionality**: This function sets up line-by-line tracing to record RSS mem (RAM) at each line of a module or sub-module. It uses Psutil for CPU memory and py3nvml/torch/tf for GPU memory allocation, if available.

**Key Implementation Details**:

*   Uses `sys.settrace` with a custom tracing function (`traceit`) that records memory state before executing each line.
*   Filters events to trace using `events_to_trace`, modules to trace using `modules_to_trace`, and modules not to trace using `modules_not_to_trace`.
*   Records current tracing state (file, location in file...) and computes difference with previous memory state.

**Notable Dependencies and Requirements**:

*   Psutil for CPU memory
*   py3nvml/torch/tf for GPU memory allocation (if available)
*   sys.settrace for setting up tracing
*   linecache for accessing Python source files

**Method Signatures and Interfaces**:

*   `sys.settrace(traceit)` to set up tracing
*   `traceit` function to record memory state before executing each line


key: chroma:document
string_value:
CODE:

def stop_memory_tracing(
    memory_trace: Optional[MemoryTrace] = None, ignore_released_memory: bool = True
) -> Optional[MemorySummary]:
    """
    Stop memory tracing cleanly and return a summary of the memory trace if a trace is given.

    Args:
        `memory_trace` (optional output of start_memory_tracing, default: None):
            memory trace to convert in summary
        `ignore_released_memory` (boolean, default: None):
            if True we only sum memory increase to compute total memory

    Return:

        - None if `memory_trace` is None
        - `MemorySummary` namedtuple otherwise with the fields:

            - `sequential`: a list of `MemoryState` namedtuple (see below) computed from the provided `memory_trace` by
              subtracting the memory after executing each line from the memory before executing said line.
            - `cumulative`: a list of `MemoryState` namedtuple (see below) with cumulative increase in memory for each
              line obtained by summing repeated memory increase for a line if it's executed several times. The list is
              sorted from the frame with the largest memory consumption to the frame with the smallest (can be negative
              if memory is released)
            - `total`: total memory increase during the full tracing as a `Memory` named tuple (see below). Line with
              memory release (negative consumption) are ignored if `ignore_released_memory` is `True` (default).

    `Memory` named tuple have fields

        - `byte` (integer): number of bytes,
        - `string` (string): same as human readable string (ex: "3.5MB")

    `Frame` are namedtuple used to list the current frame state and have the following fields:

        - 'filename' (string): Name of the file currently executed
        - 'module' (string): Name of the module currently executed
        - 'line_number' (int): Number of the line currently executed
        - 'event' (string): Event that triggered the tracing (default will be "line")
        - 'line_text' (string): Text of the line in the python script

    `MemoryState` are namedtuples listing frame + CPU/GPU memory with the following fields:

        - `frame` (`Frame`): the current frame (see above)
        - `cpu`: CPU memory consumed at during the current frame as a `Memory` named tuple
        - `gpu`: GPU memory consumed at during the current frame as a `Memory` named tuple
        - `cpu_gpu`: CPU + GPU memory consumed at during the current frame as a `Memory` named tuple
    """
    global _is_memory_tracing_enabled
    _is_memory_tracing_enabled = False

    if memory_trace is not None and len(memory_trace) > 1:
        memory_diff_trace = []
        memory_curr_trace = []

        cumulative_memory_dict = defaultdict(lambda: [0, 0, 0])

        for (
            (frame, cpu_mem, gpu_mem),
            (next_frame, next_cpu_mem, next_gpu_mem),
        ) in zip(memory_trace[:-1], memory_trace[1:]):
            cpu_mem_inc = next_cpu_mem - cpu_mem
            gpu_mem_inc = next_gpu_mem - gpu_mem
            cpu_gpu_mem_inc = cpu_mem_inc + gpu_mem_inc
            memory_diff_trace.append(
                MemoryState(
                    frame=frame,
                    cpu=Memory(cpu_mem_inc),
                    gpu=Memory(gpu_mem_inc),
                    cpu_gpu=Memory(cpu_gpu_mem_inc),
                )
            )

            memory_curr_trace.append(
                MemoryState(
                    frame=frame,
                    cpu=Memory(next_cpu_mem),
                    gpu=Memory(next_gpu_mem),
                    cpu_gpu=Memory(next_gpu_mem + next_cpu_mem),
                )
            )

            cumulative_memory_dict[frame][0] += cpu_mem_inc
            cumulative_memory_dict[frame][1] += gpu_mem_inc
            cumulative_memory_dict[frame][2] += cpu_gpu_mem_inc

        cumulative_memory = sorted(
            cumulative_memory_dict.items(), key=lambda x: x[1][2], reverse=True
        )  # order by the total CPU + GPU memory increase
        cumulative_memory = [
            MemoryState(
                frame=frame,
                cpu=Memory(cpu_mem_inc),
                gpu=Memory(gpu_mem_inc),
                cpu_gpu=Memory(cpu_gpu_mem_inc),
            )
            for frame, (cpu_mem_inc, gpu_mem_inc, cpu_gpu_mem_inc) in cumulative_memory
        ]

        memory_curr_trace = sorted(memory_curr_trace, key=lambda x: x.cpu_gpu.bytes, reverse=True)

        if ignore_released_memory:
            total_memory = sum(max(0, step_trace.cpu_gpu.bytes) for step_trace in memory_diff_trace)
        else:
            total_memory = sum(step_trace.cpu_gpu.bytes for step_trace in memory_diff_trace)

        total_memory = Memory(total_memory)

        return MemorySummary(
            sequential=memory_diff_trace,
            cumulative=cumulative_memory,
            current=memory_curr_trace,
            total=total_memory,
        )



SUMMARY:
**Function Summary: `stop_memory_tracing`**

### Purpose and Functionality

The `stop_memory_tracing` function stops memory tracing cleanly and returns a summary of the memory trace if a trace is provided. It takes two optional inputs:

*   `memory_trace`: an output from `start_memory_tracing`, used to compute a memory summary.
*   `ignore_released_memory`: a boolean flag (default: True) that determines whether to ignore lines with released memory.

### Implementation Details

The function performs the following key steps:

1.  Updates global `_is_memory_tracing_enabled` to False.
2.  If a memory trace is provided, computes two lists: `memory_diff_trace` and `memory_curr_trace`.
    *   `memory_diff_trace`: contains frames with CPU/GPU memory consumed at each step of the tracing process.
    *   `memory_curr_trace`: contains the current frame state at each step.
3.  Computes cumulative memory values for each frame using a dictionary.
4.  Sorts `cumulative_memory` and `memory_curr_trace` based on specific criteria.
5.  If `ignore_released_memory` is True, computes total memory increase excluding released memory.

### Output

The function returns an optional `MemorySummary` object if a valid memory trace is provided, or None otherwise.

**Key Method Signatures:**

*   `cpu_mem_inc = next_cpu_mem - cpu_mem`
*   `cumulative_memory_dict[frame][0] += cpu_mem_inc`

**Notable Dependencies and Requirements:**

*   `_is_memory_tracing_enabled` (global variable)
*   `MemoryTrace` class
*   `Frame`, `MemoryState`, `Memory` classes
*   `defaultdict`
*   `sorted` function


key: chroma:document
string_value:
CODE:

def bytes_to_mega_bytes(memory_amount: int) -> int:
    """Utility to convert a number of bytes (int) into a number of mega bytes (int)"""


SUMMARY:
**Summary**

This function converts an integer amount of bytes to its equivalent in mega bytes. It performs a bitwise right shift operation on the input `memory_amount` by 20 places, effectively dividing by 1,048,576.

**Key Details**

* Function signature: `bytes_to_mega_bytes(memory_amount: int) -> int`
* Operation: Right shift by 20 places (`memory_amount >> 20`)
* Purpose: Convert bytes to mega bytes

**Note:** This function assumes that the input is a non-negative integer, as byte amounts are inherently positive.


key: chroma:document
string_value:
CODE:

class Benchmark(ABC):
    """
    Benchmarks is a simple but feature-complete benchmarking script to compare memory and time performance of models in
    Transformers.
    """

    args: BenchmarkArguments
    configs: PretrainedConfig
    framework: str

    def __init__(self, args: BenchmarkArguments = None, configs: PretrainedConfig = None):
        self.args = args
        if configs is None:
            self.config_dict = {
                model_name: AutoConfig.from_pretrained(model_name) for model_name in self.args.model_names
            }
        else:
            self.config_dict = dict(zip(self.args.model_names, configs))

        warnings.warn(
            f"The class {self.__class__} is deprecated. Hugging Face Benchmarking utils"
            " are deprecated in general and it is advised to use external Benchmarking libraries "
            " to benchmark Transformer models.",
            FutureWarning,
        )

        if self.args.memory and os.getenv("TRANSFORMERS_USE_MULTIPROCESSING") == 0:
            logger.warning(
                "Memory consumption will not be measured accurately if `args.multi_process` is set to `False.` The"
                " flag 'TRANSFORMERS_USE_MULTIPROCESSING' should only be disabled for debugging / testing."
            )

        self._print_fn = None
        self._framework_version = None
        self._environment_info = None

    @property
    def print_fn(self):
        if self._print_fn is None:
            if self.args.log_print:

                def print_and_log(*args):
                    with open(self.args.log_filename, "a") as log_file:
                        log_file.write("".join(args) + "\n")
                    print(*args)

                self._print_fn = print_and_log
            else:
                self._print_fn = print
        return self._print_fn

    @property
    @abstractmethod
    def framework_version(self):
        pass

    @abstractmethod
    def _inference_speed(self, model_name: str, batch_size: int, sequence_length: int) -> float:
        pass

    @abstractmethod
    def _train_speed(self, model_name: str, batch_size: int, sequence_length: int) -> float:
        pass

    @abstractmethod
    def _inference_memory(
        self, model_name: str, batch_size: int, sequence_length: int
    ) -> [Memory, Optional[MemorySummary]]:
        pass

    @abstractmethod
    def _train_memory(
        self, model_name: str, batch_size: int, sequence_length: int
    ) -> [Memory, Optional[MemorySummary]]:
        pass

    def inference_speed(self, *args, **kwargs) -> float:
        return separate_process_wrapper_fn(self._inference_speed, self.args.do_multi_processing)(*args, **kwargs)

    def train_speed(self, *args, **kwargs) -> float:
        return separate_process_wrapper_fn(self._train_speed, self.args.do_multi_processing)(*args, **kwargs)

    def inference_memory(self, *args, **kwargs) -> [Memory, Optional[MemorySummary]]:
        return separate_process_wrapper_fn(self._inference_memory, self.args.do_multi_processing)(*args, **kwargs)

    def train_memory(self, *args, **kwargs) -> [Memory, Optional[MemorySummary]]:
        return separate_process_wrapper_fn(self._train_memory, self.args.do_multi_processing)(*args, **kwargs)

    def run(self):
        result_dict = {model_name: {} for model_name in self.args.model_names}
        inference_result_time = copy.deepcopy(result_dict)
        inference_result_memory = copy.deepcopy(result_dict)
        train_result_time = copy.deepcopy(result_dict)
        train_result_memory = copy.deepcopy(result_dict)

        for c, model_name in enumerate(self.args.model_names):
            self.print_fn(f"{c + 1} / {len(self.args.model_names)}")

            model_dict = {
                "bs": self.args.batch_sizes,
                "ss": self.args.sequence_lengths,
                "result": {i: {} for i in self.args.batch_sizes},
            }
            inference_result_time[model_name] = copy.deepcopy(model_dict)
            inference_result_memory[model_name] = copy.deepcopy(model_dict)
            train_result_time[model_name] = copy.deepcopy(model_dict)
            train_result_memory[model_name] = copy.deepcopy(model_dict)

            inference_summary = train_summary = None

            for batch_size in self.args.batch_sizes:
                for sequence_length in self.args.sequence_lengths:
                    if self.args.inference:
                        if self.args.memory:
                            memory, inference_summary = self.inference_memory(model_name, batch_size, sequence_length)
                            inference_result_memory[model_name]["result"][batch_size][sequence_length] = memory
                        if self.args.speed:
                            time = self.inference_speed(model_name, batch_size, sequence_length)
                            inference_result_time[model_name]["result"][batch_size][sequence_length] = time

                    if self.args.training:
                        if self.args.memory:
                            memory, train_summary = self.train_memory(model_name, batch_size, sequence_length)
                            train_result_memory[model_name]["result"][batch_size][sequence_length] = memory
                        if self.args.speed:
                            time = self.train_speed(model_name, batch_size, sequence_length)
                            train_result_time[model_name]["result"][batch_size][sequence_length] = time

        if self.args.inference:
            if self.args.speed:
                self.print_fn("\n" + 20 * "=" + ("INFERENCE - SPEED - RESULT").center(40) + 20 * "=")
                self.print_results(inference_result_time, type_label="Time in s")
                self.save_to_csv(inference_result_time, self.args.inference_time_csv_file)
                if self.args.is_tpu:
                    self.print_fn(
                        "TPU was used for inference. Note that the time after compilation stabilized (after ~10"
                        " inferences model.forward(..) calls) was measured."
                    )

            if self.args.memory:
                self.print_fn("\n" + 20 * "=" + ("INFERENCE - MEMORY - RESULT").center(40) + 20 * "=")
                self.print_results(inference_result_memory, type_label="Memory in MB")
                self.save_to_csv(inference_result_memory, self.args.inference_memory_csv_file)

            if self.args.trace_memory_line_by_line:
                self.print_fn("\n" + 20 * "=" + ("INFERENCE - MEMOMRY - LINE BY LINE - SUMMARY").center(40) + 20 * "=")
                self.print_memory_trace_statistics(inference_summary)

        if self.args.training:
            if self.args.speed:
                self.print_fn("\n" + 20 * "=" + ("TRAIN - SPEED - RESULTS").center(40) + 20 * "=")
                self.print_results(train_result_time, "Time in s")
                self.save_to_csv(train_result_time, self.args.train_time_csv_file)
                if self.args.is_tpu:
                    self.print_fn(
                        "TPU was used for training. Note that the time after compilation stabilized (after ~10 train"
                        " loss=model.forward(...) + loss.backward() calls) was measured."
                    )

            if self.args.memory:
                self.print_fn("\n" + 20 * "=" + ("TRAIN - MEMORY - RESULTS").center(40) + 20 * "=")
                self.print_results(train_result_memory, type_label="Memory in MB")
                self.save_to_csv(train_result_memory, self.args.train_memory_csv_file)

            if self.args.trace_memory_line_by_line:
                self.print_fn("\n" + 20 * "=" + ("TRAIN - MEMOMRY - LINE BY LINE - SUMMARY").center(40) + 20 * "=")
                self.print_memory_trace_statistics(train_summary)

        if self.args.env_print:
            self.print_fn("\n" + 20 * "=" + ("ENVIRONMENT INFORMATION").center(40) + 20 * "=")
            self.print_fn("\n".join([f"- {prop}: {val}" for prop, val in self.environment_info.items()]) + "\n")

        if self.args.save_to_csv:
            with open(self.args.env_info_csv_file, mode="w", newline="") as csv_file:
                writer = csv.writer(csv_file)
                for key, value in self.environment_info.items():
                    writer.writerow([key, value])

        return BenchmarkOutput(
            inference_result_time,
            inference_result_memory,
            train_result_time,
            train_result_memory,
            inference_summary,
            train_summary,
        )

    @property
    def environment_info(self):
        if self._environment_info is None:
            info = {}
            info["transformers_version"] = version
            info["framework"] = self.framework
            if self.framework == "PyTorch":
                info["use_torchscript"] = self.args.torchscript
            if self.framework == "TensorFlow":
                info["eager_mode"] = self.args.eager_mode
                info["use_xla"] = self.args.use_xla
            info["framework_version"] = self.framework_version
            info["python_version"] = platform.python_version()
            info["system"] = platform.system()
            info["cpu"] = platform.processor()
            info["architecture"] = platform.architecture()[0]
            info["date"] = datetime.date(datetime.now())
            info["time"] = datetime.time(datetime.now())
            info["fp16"] = self.args.fp16
            info["use_multiprocessing"] = self.args.do_multi_processing
            info["only_pretrain_model"] = self.args.only_pretrain_model

            if is_psutil_available():
                info["cpu_ram_mb"] = bytes_to_mega_bytes(psutil.virtual_memory().total)
            else:
                logger.warning(
                    "Psutil not installed, we won't log available CPU memory. "
                    "Install psutil (pip install psutil) to log available CPU memory."
                )
                info["cpu_ram_mb"] = "N/A"

            info["use_gpu"] = self.args.is_gpu
            if self.args.is_gpu:
                info["num_gpus"] = 1  # TODO(PVP) Currently only single GPU is supported
                if is_py3nvml_available():
                    nvml.nvmlInit()
                    handle = nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx)
                    info["gpu"] = nvml.nvmlDeviceGetName(handle)
                    info["gpu_ram_mb"] = bytes_to_mega_bytes(nvml.nvmlDeviceGetMemoryInfo(handle).total)
                    info["gpu_power_watts"] = nvml.nvmlDeviceGetPowerManagementLimit(handle) / 1000
                    info["gpu_performance_state"] = nvml.nvmlDeviceGetPerformanceState(handle)
                    nvml.nvmlShutdown()
                else:
                    logger.warning(
                        "py3nvml not installed, we won't log GPU memory usage. "
                        "Install py3nvml (pip install py3nvml) to log information about GPU."
                    )
                    info["gpu"] = "N/A"
                    info["gpu_ram_mb"] = "N/A"
                    info["gpu_power_watts"] = "N/A"
                    info["gpu_performance_state"] = "N/A"

            info["use_tpu"] = self.args.is_tpu
            # TODO(PVP): See if we can add more information about TPU
            # see: https://github.com/pytorch/xla/issues/2180

            self._environment_info = info
        return self._environment_info

    def print_results(self, result_dict, type_label):
        self.print_fn(80 * "-")
        self.print_fn(
            "Model Name".center(30) + "Batch Size".center(15) + "Seq Length".center(15) + type_label.center(15)
        )
        self.print_fn(80 * "-")
        for model_name in self.args.model_names:
            for batch_size in result_dict[model_name]["bs"]:
                for sequence_length in result_dict[model_name]["ss"]:
                    result = result_dict[model_name]["result"][batch_size][sequence_length]
                    if isinstance(result, float):
                        result = round(1000 * result) / 1000
                        result = "< 0.001" if result == 0.0 else str(result)
                    else:
                        result = str(result)
                    self.print_fn(
                        model_name[:30].center(30) + str(batch_size).center(15),
                        str(sequence_length).center(15),
                        result.center(15),
                    )
        self.print_fn(80 * "-")

    def print_memory_trace_statistics(self, summary: MemorySummary):
        self.print_fn(
            "\nLine by line memory consumption:\n"
            + "\n".join(
                f"{state.frame.filename}:{state.frame.line_number}: mem {state.cpu_gpu}: {state.frame.line_text}"
                for state in summary.sequential
            )
        )
        self.print_fn(
            "\nLines with top memory consumption:\n"
            + "\n".join(
                f"=> {state.frame.filename}:{state.frame.line_number}: mem {state.cpu_gpu}: {state.frame.line_text}"
                for state in summary.cumulative[:6]
            )
        )
        self.print_fn(
            "\nLines with lowest memory consumption:\n"
            + "\n".join(
                f"=> {state.frame.filename}:{state.frame.line_number}: mem {state.cpu_gpu}: {state.frame.line_text}"
                for state in summary.cumulative[-6:]
            )
        )
        self.print_fn(f"\nTotal memory increase: {summary.total}")

    def save_to_csv(self, result_dict, filename):
        if not self.args.save_to_csv:
            return
        self.print_fn("Saving results to csv.")
        with open(filename, mode="w") as csv_file:
            if len(self.args.model_names) <= 0:
                raise ValueError(f"At least 1 model should be defined, but got {self.model_names}")

            fieldnames = ["model", "batch_size", "sequence_length"]
            writer = csv.DictWriter(csv_file, fieldnames=fieldnames + ["result"])
            writer.writeheader()

            for model_name in self.args.model_names:
                result_dict_model = result_dict[model_name]["result"]
                for bs in result_dict_model:
                    for ss in result_dict_model[bs]:
                        result_model = result_dict_model[bs][ss]
                        writer.writerow(
                            {
                                "model": model_name,
                                "batch_size": bs,
                                "sequence_length": ss,
                                "result": ("{}" if not isinstance(result_model, float) else "{:.4f}").format(
                                    result_model
                                ),
                            }


SUMMARY:
This is a Python code snippet that appears to be part of a benchmarking or testing framework for machine learning models. Here's a concise summary:

**Purpose:** The code measures the performance of various machine learning models under different conditions (e.g., batch sizes, sequence lengths) and saves the results to a CSV file.

**Key functionality:**

1. Iterates over a list of model names (`self.args.model_names`) and collects performance metrics for each model.
2. For each model, it iterates over various batch sizes and sequence lengths, running experiments and collecting timing data.
3. Saves the results to a CSV file using `csv.DictWriter`.

**Notable aspects:**

* The code uses a custom `MemorySummary` class to analyze memory usage during testing.
* It has support for GPU acceleration and TPU inference, depending on the `--gpu` or `--gpu-tpu` command-line flags.
* There are utility functions (`print_fn`, `print_results`) that facilitate printing human-friendly output.

**Possible improvements:**

1. Error handling and logging mechanisms can be improved to handle exceptions and provide more informative error messages.
2. The test framework could benefit from additional features, such as support for multiple test environments, load balancing, or asynchronous testing.
3. Code readability and organization could be optimized through reusing modules or creating a hierarchical structure.

**Example usage:**

To use this code snippet in your own project, you would need to:

1. Create an instance of the `Benchmark` class, passing in the necessary command-line flags (e.g., `--model-model-name`, `--gpu-device-index`).
2. Run the benchmarking tests using a loop or script that iterates over various model names and configurations.
3. Finally, save the results to a CSV file using the `save_to_csv` method.

```python
# Example usage:
benchmark = Benchmark(
    --model-model-name="my_model",
    --gpu-device-index=0,
)
benchmark.run()
results = benchmark.get_results()
benchmark.save_to_csv("results.csv", results)
```


key: chroma:document
string_value:
CODE:

class PyTorchBenchmark(Benchmark):
    args: PyTorchBenchmarkArguments
    configs: PretrainedConfig
    framework: str = "PyTorch"

    @property
    def framework_version(self):
        return torch.__version__

    def _inference_speed(self, model_name: str, batch_size: int, sequence_length: int) -> float:
        _inference = self._prepare_inference_func(model_name, batch_size, sequence_length)
        return self._measure_speed(_inference)

    def _inference_memory(
        self, model_name: str, batch_size: int, sequence_length: int
    ) -> [Memory, Optional[MemorySummary]]:
        _inference = self._prepare_inference_func(model_name, batch_size, sequence_length)
        return self._measure_memory(_inference)

    def _train_speed(self, model_name: str, batch_size: int, sequence_length: int) -> float:
        _train = self._prepare_train_func(model_name, batch_size, sequence_length)
        return self._measure_speed(_train)

    def _train_memory(
        self, model_name: str, batch_size: int, sequence_length: int
    ) -> [Memory, Optional[MemorySummary]]:
        _train = self._prepare_train_func(model_name, batch_size, sequence_length)
        return self._measure_memory(_train)

    def _prepare_inference_func(self, model_name: str, batch_size: int, sequence_length: int) -> Callable[[], None]:
        config = self.config_dict[model_name]

        if self.args.torchscript:
            config.torchscript = True

        has_model_class_in_config = (
            hasattr(config, "architectures")
            and isinstance(config.architectures, list)
            and len(config.architectures) > 0
        )
        if not self.args.only_pretrain_model and has_model_class_in_config:
            try:
                model_class = config.architectures[0]
                transformers_module = __import__("transformers", fromlist=[model_class])
                model_cls = getattr(transformers_module, model_class)
                model = model_cls(config)
            except ImportError:
                raise ImportError(
                    f"{model_class} does not exist. If you just want to test the pretrained model, you might want to"
                    " set `--only_pretrain_model` or `args.only_pretrain_model=True`."
                )
        else:
            model = MODEL_MAPPING[config.__class__](config)

        model.eval()
        model.to(self.args.device)

        # encoder-decoder has vocab size saved differently
        vocab_size = config.vocab_size if hasattr(config, "vocab_size") else config.encoder.vocab_size
        input_ids = torch.randint(vocab_size, (batch_size, sequence_length), dtype=torch.long, device=self.args.device)

        if self.args.fp16:
            logger.info("Running training in Mixed Precision...")
            if not self.args.is_gpu:
                raise ValueError("Mixed precision is possible only for GPU.")
            # amp seems to have memory leaks so that memory usage
            # is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
            model.half()

        if self.args.torchscript:
            with torch.no_grad():
                inference_model = torch.jit.trace(model, input_ids)
        else:
            inference_model = model

        def encoder_decoder_forward():
            with torch.no_grad():
                outputs = inference_model(input_ids, decoder_input_ids=input_ids)
            return outputs

        def encoder_forward():
            with torch.no_grad():
                outputs = inference_model(input_ids)
            return outputs

        _forward = encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
        return _forward

    def _prepare_train_func(self, model_name: str, batch_size: int, sequence_length: int) -> Callable[[], None]:
        config = self.config_dict[model_name]

        has_model_class_in_config = (
            hasattr(config, "architectures")
            and isinstance(config.architectures, list)
            and len(config.architectures) > 0
        )
        if not self.args.only_pretrain_model and has_model_class_in_config:
            try:
                model_class = config.architectures[0]
                transformers_module = __import__("transformers", fromlist=[model_class])
                model_cls = getattr(transformers_module, model_class)
                model = model_cls(config)
            except ImportError:
                raise ImportError(
                    f"{model_class} does not exist. If you just want to test the pretrained model, you might want to"
                    " set `--only_pretrain_model` or `args.only_pretrain_model=True`."
                )
        else:
            model = MODEL_WITH_LM_HEAD_MAPPING[config.__class__](config)

        if self.args.torchscript:
            raise NotImplementedError("Training for torchscript is currently not implemented")
        else:
            train_model = model

        model.train()
        model.to(self.args.device)

        # encoder-decoder has vocab size saved differently
        vocab_size = config.vocab_size if hasattr(config, "vocab_size") else config.encoder.vocab_size
        input_ids = torch.randint(vocab_size, (batch_size, sequence_length), dtype=torch.long, device=self.args.device)

        if self.args.fp16:
            logger.info("Running training in Mixed Precision...")
            if not self.args.is_gpu:
                raise ValueError("Mixed precision is possible only for GPU.")

            # amp seems to have memory leaks so that memory usage
            # is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
            model.half()

        def compute_loss_and_backprob_encoder():
            loss = train_model(input_ids, labels=input_ids)[0]
            loss.backward()
            return loss

        def compute_loss_and_backprob_encoder_decoder():
            loss = train_model(input_ids, decoder_input_ids=input_ids, labels=input_ids)[0]
            loss.backward()
            return loss

        _train = (
            compute_loss_and_backprob_encoder_decoder
            if config.is_encoder_decoder
            else compute_loss_and_backprob_encoder
        )
        return _train

    def _measure_speed(self, func) -> float:
        try:
            if self.args.is_tpu or self.args.torchscript:
                # run additional 10 times to stabilize compilation for tpu and torchscript
                logger.info("Do inference on TPU or torchscript. Running model 5 times to stabilize compilation")
                timeit.repeat(
                    func,
                    repeat=1,
                    number=5,
                )

            # as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
            runtimes = timeit.repeat(
                func,
                repeat=self.args.repeat,
                number=10,
            )

            if self.args.is_tpu and self.args.torch_xla_tpu_print_metrics:
                import torch_xla.debug.metrics as met

                self.print_fn(met.metrics_report())

            return min(runtimes) / 10.0
        except RuntimeError as e:
            self.print_fn(f"Doesn't fit on GPU. {e}")
            return "N/A"

    def _measure_memory(self, func: Callable[[], None]) -> [Memory, MemorySummary]:
        try:
            if self.args.trace_memory_line_by_line:
                trace = start_memory_tracing("transformers")

            if self.args.is_tpu:
                # tpu
                raise NotImplementedError(
                    "Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking with"
                    " `--no-memory` or `args.memory=False`"
                )
            elif self.args.is_gpu:
                if not is_py3nvml_available():
                    logger.warning(
                        "py3nvml not installed, we won't log GPU memory usage. "
                        "Install py3nvml (pip install py3nvml) to log information about GPU."
                    )
                    memory = "N/A"
                else:
                    logger.info(
                        "Measuring total GPU usage on GPU device. Make sure to not have additional processes running"
                        " on the same GPU."
                    )
                    # init nvml
                    nvml.nvmlInit()
                    func()
                    handle = nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx)
                    meminfo = nvml.nvmlDeviceGetMemoryInfo(handle)
                    max_bytes_in_use = meminfo.used
                    memory = Memory(max_bytes_in_use)
                    # shutdown nvml
                    nvml.nvmlShutdown()
            else:
                # cpu
                memory_bytes = measure_peak_memory_cpu(func)
                memory = Memory(memory_bytes) if isinstance(memory_bytes, int) else memory_bytes

            if self.args.trace_memory_line_by_line:
                summary = stop_memory_tracing(trace)
            else:
                summary = None

            return memory, summary
        except RuntimeError as e:
            self.print_fn(f"Doesn't fit on GPU. {e}")


SUMMARY:
Here's a concise summary of the code:

**Class Overview**
PyTorchBenchmark is a class that benchmarks models using PyTorch. It provides methods to measure inference and training speed, memory usage, and other metrics.

**Key Methods**

1. `_prepare_inference_func` and `_prepare_train_func`: These methods prepare the model and its inputs for inference and training respectively.
2. `_measure_speed` and `_train_speed`: These methods measure the speed of inference and training, taking into account various factors such as TPU or mixed precision training.
3. `_measure_memory`: This method measures the memory usage during inference, using NVML (Nvidia Performance Library) on GPU devices.

**Additional Features**

* Support for TPU and torchscript
* Memory line-by-line tracing if enabled
* Print function-like behavior with `self.print_fn`
* Handling of exceptions and errors

This class is designed to provide a robust framework for benchmarking PyTorch models, covering various aspects such as speed, memory usage, and performance.


key: chroma:document
string_value:
CODE:

class WatermarkDetector:
    r"""
    Detector for detection of watermark generated text. The detector needs to be given the exact same settings that were
    given during text generation to replicate the watermark greenlist generation and so detect the watermark. This includes
    the correct device that was used during text generation, the correct watermarking arguments and the correct tokenizer vocab size.
    The code was based on the [original repo](https://github.com/jwkirchenbauer/lm-watermarking/tree/main).

    See [the paper](https://arxiv.org/abs/2306.04634) for more information.

    Args:
        model_config (`PretrainedConfig`):
            The model config that will be used to get model specific arguments used when generating.
        device (`str`):
            The device which was used during watermarked text generation.
        watermarking_config (Union[`WatermarkingConfig`, `Dict`]):
            The exact same watermarking config and arguments used when generating text.
        ignore_repeated_ngrams (`bool`, *optional*, defaults to `False`):
            Whether to count every unique ngram only once or not.
        max_cache_size (`int`, *optional*, defaults to 128):
            The max size to be used for LRU caching of seeding/sampling algorithms called for every token.

    Examples:

    ```python
    >>> from transformers import AutoTokenizer, AutoModelForCausalLM, WatermarkDetector, WatermarkingConfig

    >>> model_id = "openai-community/gpt2"
    >>> model = AutoModelForCausalLM.from_pretrained(model_id)
    >>> tok = AutoTokenizer.from_pretrained(model_id)
    >>> tok.pad_token_id = tok.eos_token_id
    >>> tok.padding_side = "left"

    >>> inputs = tok(["This is the beginning of a long story", "Alice and Bob are"], padding=True, return_tensors="pt")
    >>> input_len = inputs["input_ids"].shape[-1]

    >>> # first generate text with watermark and without
    >>> watermarking_config = WatermarkingConfig(bias=2.5, seeding_scheme="selfhash")
    >>> out_watermarked = model.generate(**inputs, watermarking_config=watermarking_config, do_sample=False, max_length=20)
    >>> out = model.generate(**inputs, do_sample=False, max_length=20)

    >>> # now we can instantiate the detector and check the generated text
    >>> detector = WatermarkDetector(model_config=model.config, device="cpu", watermarking_config=watermarking_config)
    >>> detection_out_watermarked = detector(out_watermarked, return_dict=True)
    >>> detection_out = detector(out, return_dict=True)
    >>> detection_out_watermarked.prediction
    array([ True,  True])

    >>> detection_out.prediction
    array([False,  False])
    ```
    """

    def __init__(
        self,
        model_config: PretrainedConfig,
        device: str,
        watermarking_config: Union[WatermarkingConfig, Dict],
        ignore_repeated_ngrams: bool = False,
        max_cache_size: int = 128,
    ):
        if isinstance(watermarking_config, WatermarkingConfig):
            watermarking_config = watermarking_config.to_dict()

        self.bos_token_id = (
            model_config.bos_token_id if not model_config.is_encoder_decoder else model_config.decoder_start_token_id
        )
        self.greenlist_ratio = watermarking_config["greenlist_ratio"]
        self.ignore_repeated_ngrams = ignore_repeated_ngrams
        self.processor = WatermarkLogitsProcessor(
            vocab_size=model_config.vocab_size, device=device, **watermarking_config
        )

        # Expensive re-seeding and sampling is cached.
        self._get_ngram_score_cached = lru_cache(maxsize=max_cache_size)(self._get_ngram_score)

    def _get_ngram_score(self, prefix: torch.LongTensor, target: int):
        greenlist_ids = self.processor._get_greenlist_ids(prefix)
        return target in greenlist_ids

    def _score_ngrams_in_passage(self, input_ids: torch.LongTensor):
        batch_size, seq_length = input_ids.shape
        selfhash = int(self.processor.seeding_scheme == "selfhash")
        n = self.processor.context_width + 1 - selfhash
        indices = torch.arange(n).unsqueeze(0) + torch.arange(seq_length - n + 1).unsqueeze(1)
        ngram_tensors = input_ids[:, indices]

        num_tokens_scored_batch = np.zeros(batch_size)
        green_token_count_batch = np.zeros(batch_size)
        for batch_idx in range(ngram_tensors.shape[0]):
            frequencies_table = collections.Counter(ngram_tensors[batch_idx])
            ngram_to_watermark_lookup = {}
            for ngram_example in frequencies_table.keys():
                prefix = ngram_example if selfhash else ngram_example[:-1]
                target = ngram_example[-1]
                ngram_to_watermark_lookup[ngram_example] = self._get_ngram_score_cached(prefix, target)

            if self.ignore_repeated_ngrams:
                # counts a green/red hit once per unique ngram.
                # num total tokens scored becomes the number unique ngrams.
                num_tokens_scored_batch[batch_idx] = len(frequencies_table.keys())
                green_token_count_batch[batch_idx] = sum(ngram_to_watermark_lookup.values())
            else:
                num_tokens_scored_batch[batch_idx] = sum(frequencies_table.values())
                green_token_count_batch[batch_idx] = sum(
                    freq * outcome
                    for freq, outcome in zip(frequencies_table.values(), ngram_to_watermark_lookup.values())
                )
        return num_tokens_scored_batch, green_token_count_batch

    def _compute_z_score(self, green_token_count: np.array, total_num_tokens: np.array) -> np.array:
        expected_count = self.greenlist_ratio
        numer = green_token_count - expected_count * total_num_tokens
        denom = np.sqrt(total_num_tokens * expected_count * (1 - expected_count))
        z = numer / denom
        return z

    def _compute_pval(self, x, loc=0, scale=1):
        z = (x - loc) / scale
        return 1 - (0.5 * (1 + np.sign(z) * (1 - np.exp(-2 * z**2 / np.pi))))

    def __call__(
        self,
        input_ids: torch.LongTensor,
        z_threshold: float = 3.0,
        return_dict: bool = False,
    ) -> Union[WatermarkDetectorOutput, np.array]:
        """
                Args:
                input_ids (`torch.LongTensor`):
                    The watermark generated text. It is advised to remove the prompt, which can affect the detection.
                z_threshold (`Dict`, *optional*, defaults to `3.0`):
                    Changing this threshold will change the sensitivity of the detector. Higher z threshold gives less
                    sensitivity and vice versa for lower z threshold.
                return_dict (`bool`,  *optional*, defaults to `False`):
                    Whether to return `~generation.WatermarkDetectorOutput` or not. If not it will return boolean predictions,
        ma
                Return:
                    [`~generation.WatermarkDetectorOutput`] or `np.array`: A [`~generation.WatermarkDetectorOutput`]
                    if `return_dict=True` otherwise a `np.array`.

        """

        # Let's assume that if one batch start with `bos`, all batched also do
        if input_ids[0, 0] == self.bos_token_id:
            input_ids = input_ids[:, 1:]

        if input_ids.shape[-1] - self.processor.context_width < 1:
            raise ValueError(
                f"Must have at least `1` token to score after the first "
                f"min_prefix_len={self.processor.context_width} tokens required by the seeding scheme."
            )

        num_tokens_scored, green_token_count = self._score_ngrams_in_passage(input_ids)
        z_score = self._compute_z_score(green_token_count, num_tokens_scored)
        prediction = z_score > z_threshold

        if return_dict:
            p_value = self._compute_pval(z_score)
            confidence = 1 - p_value

            return WatermarkDetectorOutput(
                num_tokens_scored=num_tokens_scored,
                num_green_tokens=green_token_count,
                green_fraction=green_token_count / num_tokens_scored,
                z_score=z_score,
                p_value=p_value,
                prediction=prediction,
                confidence=confidence,
            )


SUMMARY:
**WatermarkDetector Class Summary**

*   **Purpose/Functionality**: Detects watermark-generated text based on model-specific arguments, device used, and watermarking config.
*   **Main Features:**
    *   Re-seeding and sampling algorithms are cached for efficiency.
    *   Scores n-grams in the input sequence and compares with a reference list.
    *   Calculates z-scores and p-values to determine detection confidence.

**Method Summaries**

*   `__call__`: Takes in watermarked text, computes scores, and returns prediction or WatermarkDetectorOutput if return_dict is True.
*   `_score_ngrams_in_passage`: Scores n-grams in the input sequence, excluding leading tokens (if applicable).
*   `_compute_z_score`: Calculates z-score for green token counts vs. expected count.
*   `_compute_pval`: Computes p-value for detected watermarks.

**Key Attributes**

*   `bipartite_graph`: Reference graph containing ngram-watermark pairs (not explicitly defined in the given code snippet).
*   `seeding_scheme`: Determine if selfhashing is used or not (determines output length).


key: chroma:document
string_value:
CODE:

class BayesianDetectorConfig(PretrainedConfig):
    """
    This is the configuration class to store the configuration of a [`BayesianDetectorModel`]. It is used to
    instantiate a Bayesian Detector model according to the specified arguments.

    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
    documentation from [`PretrainedConfig`] for more information.

    Args:
        watermarking_depth (`int`, *optional*):
            The number of tournament layers.
        base_rate (`float1`, *optional*, defaults to 0.5):
            Prior probability P(w) that a text is watermarked.
    """

    def __init__(self, watermarking_depth: int = None, base_rate: float = 0.5, **kwargs):
        self.watermarking_depth = watermarking_depth
        self.base_rate = base_rate
        # These can be set later to store information about this detector.
        self.model_name = None
        self.watermarking_config = None

        super().__init__(**kwargs)

    def set_detector_information(self, model_name, watermarking_config):
        self.model_name = model_name


SUMMARY:
**BayesianDetectorConfig Class Summary**

This class configures a Bayesian Detector model. It inherits from `PretrainedConfig` and provides parameters for water marking depth (`watermarking_depth`) and base rate (`base_rate`). The class has two optional attributes: `model_name` and `watermarking_config`, which can be set later.

**Key Implementation Details**

* Sets the water marking depth using the `set_detector_information` method, taking `model_name` and `watermarking_config` as input.
* Initializes with `watermarking_depth` (default value), `base_rate` (default value: 0.5) and additional keyword arguments (`**kwargs`).

**Dependencies and Requirements**

* Inherits from `PretrainedConfig`
* Has optional attributes:
	+ `model_name`: None (settable)
	+ `watermarking_config`: None (settable)
* Does not mention explicit external dependencies.


key: chroma:document
string_value:
CODE:

class BayesianDetectorWatermarkedLikelihood(nn.Module):
    """Watermarked likelihood model for binary-valued g-values.

    This takes in g-values and returns p(g_values|watermarked).
    """

    def __init__(self, watermarking_depth: int):
        """Initializes the model parameters."""
        super().__init__()
        self.watermarking_depth = watermarking_depth
        self.beta = torch.nn.Parameter(-2.5 + 0.001 * torch.randn(1, 1, watermarking_depth))
        self.delta = torch.nn.Parameter(0.001 * torch.randn(1, 1, self.watermarking_depth, watermarking_depth))

    def _compute_latents(self, g_values: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Computes the unique token probability distribution given g-values.

        Args:
            g_values (`torch.Tensor` of shape `(batch_size, seq_len, watermarking_depth)`):
                PRF values.

        Returns:
            p_one_unique_token and p_two_unique_tokens, both of shape
            [batch_size, seq_len, watermarking_depth]. p_one_unique_token[i,t,l]
            gives the probability of there being one unique token in a tournament
            match on layer l, on timestep t, for batch item i.
            p_one_unique_token[i,t,l] + p_two_unique_token[i,t,l] = 1.
        """
        # Tile g-values to produce feature vectors for predicting the latents
        # for each layer in the tournament; our model for the latents psi is a
        # logistic regression model psi = sigmoid(delta * x + beta).

        # [batch_size, seq_len, watermarking_depth, watermarking_depth]
        x = torch.repeat_interleave(torch.unsqueeze(g_values, dim=-2), self.watermarking_depth, axis=-2)

        # mask all elements above -1 diagonal for autoregressive factorization
        x = torch.tril(x, diagonal=-1)

        # [batch_size, seq_len, watermarking_depth]
        # (i, j, k, l) x (i, j, k, l) -> (i, j, k) einsum equivalent
        logits = (self.delta[..., None, :] @ x.type(self.delta.dtype)[..., None]).squeeze() + self.beta

        p_two_unique_tokens = torch.sigmoid(logits)
        p_one_unique_token = 1 - p_two_unique_tokens
        return p_one_unique_token, p_two_unique_tokens

    def forward(self, g_values: torch.Tensor) -> torch.Tensor:
        """Computes the likelihoods P(g_values|watermarked).

        Args:
            g_values (`torch.Tensor` of shape `(batch_size, seq_len, watermarking_depth)`):
                g-values (values 0 or 1)

        Returns:
            p(g_values|watermarked) of shape [batch_size, seq_len, watermarking_depth].
        """
        p_one_unique_token, p_two_unique_tokens = self._compute_latents(g_values)

        # P(g_tl | watermarked) is equal to
        # 0.5 * [ (g_tl+0.5) * p_two_unique_tokens + p_one_unique_token].


SUMMARY:
**BayesianDetectorWatermarkedLikelihood Class Summary**

**Class Purpose:**
This class implements a watermarked likelihood model for binary-valued g-values, taking in g-values and returning the conditional probability `P(g_values|watermarked)`.

**Key Implementation Details:**

* `_compute_latents` method:
	+ Computes unique token probability distribution given PRF values.
	+ Uses logistic regression model to predict latents using delta and beta parameters.
	+ Tiles and masks input tensor for autoregressive factorization.
* `forward` method:
	+ Computes conditional probability using results from `_compute_latents`.
	+ Applies sigmoid function and scaling factors.

**Notable Dependencies/Requirements:**

* `torch.nn.Module`: The class inherits from PyTorch's neural network module.
* `torch.nn.Parameter`: `beta` and `delta` are initialized as learnable parameters.
* `torch.Tensor`: Input tensor is expected to have shape `(batch_size, seq_len, watermarking_depth)`.

**Primary Parameter:**
`watermarking_depth`: The depth of the watermarking scheme.


key: chroma:document
string_value:
CODE:

class BayesianDetectorModel(PreTrainedModel):
    r"""
    Bayesian classifier for watermark detection.

    This detector uses Bayes' rule to compute a watermarking score, which is the sigmoid of the log of ratio of the
    posterior probabilities P(watermarked|g_values) and P(unwatermarked|g_values). Please see the section on
    BayesianScore in the paper for further details.
    Paper URL: https://www.nature.com/articles/s41586-024-08025-4

    Note that this detector only works with non-distortionary Tournament-based watermarking using the Bernoulli(0.5)
    g-value distribution.

    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
    etc.)

    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
    and behavior.

    Parameters:
        config ([`BayesianDetectorConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
    """

    config_class = BayesianDetectorConfig
    base_model_prefix = "model"

    def __init__(self, config):
        super().__init__(config)

        self.watermarking_depth = config.watermarking_depth
        self.base_rate = config.base_rate
        self.likelihood_model_watermarked = BayesianDetectorWatermarkedLikelihood(
            watermarking_depth=self.watermarking_depth
        )
        self.prior = torch.nn.Parameter(torch.tensor([self.base_rate]))

    def _init_weights(self, module):
        """Initialize the weights."""
        if isinstance(module, nn.Parameter):
            module.weight.data.normal_(mean=0.0, std=0.02)

    def _compute_posterior(
        self,
        likelihoods_watermarked: torch.Tensor,
        likelihoods_unwatermarked: torch.Tensor,
        mask: torch.Tensor,
        prior: float,
    ) -> torch.Tensor:
        """
        Compute posterior P(w|g) given likelihoods, mask and prior.

        Args:
            likelihoods_watermarked (`torch.Tensor` of shape `(batch, length, depth)`):
                Likelihoods P(g_values|watermarked) of g-values under watermarked model.
            likelihoods_unwatermarked (`torch.Tensor` of shape `(batch, length, depth)`):
                Likelihoods P(g_values|unwatermarked) of g-values under unwatermarked model.
            mask (`torch.Tensor` of shape `(batch, length)`):
                A binary array indicating which g-values should be used. g-values with mask value 0 are discarded.
            prior (`float`):
                the prior probability P(w) that the text is watermarked.

        Returns:
            Posterior probability P(watermarked|g_values), shape [batch].
        """
        mask = torch.unsqueeze(mask, dim=-1)
        prior = torch.clamp(prior, min=1e-5, max=1 - 1e-5)
        log_likelihoods_watermarked = torch.log(torch.clamp(likelihoods_watermarked, min=1e-30, max=float("inf")))
        log_likelihoods_unwatermarked = torch.log(torch.clamp(likelihoods_unwatermarked, min=1e-30, max=float("inf")))
        log_odds = log_likelihoods_watermarked - log_likelihoods_unwatermarked

        # Sum relative surprisals (log odds) across all token positions and layers.
        relative_surprisal_likelihood = torch.einsum("i...->i", log_odds * mask)

        # Compute the relative surprisal prior
        relative_surprisal_prior = torch.log(prior) - torch.log(1 - prior)

        # Combine prior and likelihood.
        # [batch_size]
        relative_surprisal = relative_surprisal_prior + relative_surprisal_likelihood

        # Compute the posterior probability P(w|g) = sigmoid(relative_surprisal).
        return torch.sigmoid(relative_surprisal)

    def forward(
        self,
        g_values: torch.Tensor,
        mask: torch.Tensor,
        labels: Optional[torch.Tensor] = None,
        loss_batch_weight=1,
        return_dict=False,
    ) -> BayesianWatermarkDetectorModelOutput:
        """
        Computes the watermarked posterior P(watermarked|g_values).

        Args:
            g_values (`torch.Tensor` of shape `(batch_size, seq_len, watermarking_depth, ...)`):
                g-values (with values 0 or 1)
            mask:
                A binary array shape [batch_size, seq_len] indicating which g-values should be used. g-values with mask
                value 0 are discarded.

        Returns:
            p(watermarked | g_values), of shape [batch_size].
        """

        likelihoods_watermarked = self.likelihood_model_watermarked(g_values)
        likelihoods_unwatermarked = 0.5 * torch.ones_like(g_values)
        out = self._compute_posterior(
            likelihoods_watermarked=likelihoods_watermarked,
            likelihoods_unwatermarked=likelihoods_unwatermarked,
            mask=mask,
            prior=self.prior,
        )

        loss = None
        if labels is not None:
            loss_fct = BCELoss()
            loss_unwweight = torch.sum(self.likelihood_model_watermarked.delta**2)
            loss_weight = loss_unwweight * loss_batch_weight
            loss = loss_fct(torch.clamp(out, 1e-5, 1 - 1e-5), labels) + loss_weight

        if not return_dict:
            return (out,) if loss is None else (out, loss)



SUMMARY:
**BayesianDetectorModel**

**Main Purpose and Functionality:**
This class implements a Bayesian classifier for watermark detection using Bayes' rule to compute a watermarking score.

**Key Implementation Details:**

*   Initializes weights with normal distribution.
*   Computes posterior probability P(watermarked|g_values) given likelihoods, mask, and prior.
*   Applies sigmoid function to combine prior and likelihood to obtain the posterior probability.

**Notable Dependencies and Requirements:**

*   `BayesianDetectorWatermarkedLikelihood`: a likelihood model for watermarked g-values.
*   `BCELoss` loss function for training the model.
*   `PreTrainedModel`: the superclass of BayesianDetectorModel that provides generic methods for model handling.

**Class Summary (RAG System):**

|  | Field        | Purpose |
| --- | ------------ | -------- |
| `__init__`      | config       | Initializes model configuration.                     |
| `_compute_posterior` | likelihoods_watermarked, likelihoods_unwatermarked     | Calculates posterior probability for watermarked and unwatermarked g-values.         |
| `forward`        | g_values, mask    | Computes the watermarked posterior P(watermarked|g_values).              |

**Dependencies:**

*   `BayesianDetectorConfig`
*   `BayesianDetectorWatermarkedLikelihood`
*   `BCELoss`

**Parent Class:**
`PreTrainedModel`

This summary is designed to be concise and focused, providing essential details about the class and its key functionality without excessive information.


key: chroma:document
string_value:
CODE:

class SynthIDTextWatermarkDetector:
    r"""
    SynthID text watermark detector class.

    This class has to be initialized with the trained bayesian detector module check script
    in examples/synthid_text/detector_training.py for example in training/saving/loading this
    detector module. The folder also showcases example use case of this detector.

    Parameters:
        detector_module ([`BayesianDetectorModel`]):
            Bayesian detector module object initialized with parameters.
            Check examples/research_projects/synthid_text/detector_training.py for usage.
        logits_processor (`SynthIDTextWatermarkLogitsProcessor`):
            The logits processor used for watermarking.
        tokenizer (`Any`):
            The tokenizer used for the model.

    Examples:
    ```python
    >>> from transformers import (
    ...     AutoTokenizer, BayesianDetectorModel, SynthIDTextWatermarkLogitsProcessor, SynthIDTextWatermarkDetector
    ... )

    >>> # Load the detector. See examples/research_projects/synthid_text for training a detector.
    >>> detector_model = BayesianDetectorModel.from_pretrained("joaogante/dummy_synthid_detector")
    >>> logits_processor = SynthIDTextWatermarkLogitsProcessor(
    ...     **detector_model.config.watermarking_config, device="cpu"
    ... )
    >>> tokenizer = AutoTokenizer.from_pretrained(detector_model.config.model_name)
    >>> detector = SynthIDTextWatermarkDetector(detector_model, logits_processor, tokenizer)

    >>> # Test whether a certain string is watermarked
    >>> test_input = tokenizer(["This is a test input"], return_tensors="pt")
    >>> is_watermarked = detector(test_input.input_ids)
    ```
    """

    def __init__(
        self,
        detector_module: BayesianDetectorModel,
        logits_processor: SynthIDTextWatermarkLogitsProcessor,
        tokenizer: Any,
    ):
        self.detector_module = detector_module
        self.logits_processor = logits_processor
        self.tokenizer = tokenizer

    def __call__(self, tokenized_outputs: torch.Tensor):
        # eos mask is computed, skip first ngram_len - 1 tokens
        # eos_mask will be of shape [batch_size, output_len]
        eos_token_mask = self.logits_processor.compute_eos_token_mask(
            input_ids=tokenized_outputs,
            eos_token_id=self.tokenizer.eos_token_id,
        )[:, self.logits_processor.ngram_len - 1 :]

        # context repetition mask is computed
        context_repetition_mask = self.logits_processor.compute_context_repetition_mask(
            input_ids=tokenized_outputs,
        )
        # context repitition mask shape [batch_size, output_len - (ngram_len - 1)]

        combined_mask = context_repetition_mask * eos_token_mask

        g_values = self.logits_processor.compute_g_values(
            input_ids=tokenized_outputs,
        )
        # g values shape [batch_size, output_len - (ngram_len - 1), depth]


SUMMARY:
**SynthIDTextWatermarkDetector Class Summary**

The `SynthIDTextWatermarkDetector` class is designed to detect watermarks in Synchro-ID text inputs. It utilizes a trained Bayesian detector module and combines it with a logits processor for watermarking.

**Primary Functionality:**

* Initializes the detector module, logits processor, and tokenizer with input parameters
* Takes in tokenized outputs as input and computes EOS mask and context repetition mask
* Combines these masks to form a combined mask for detection
* Passes the combined mask and g-values to the detector module for watermarking

**Key Implementation Details:**

* `compute_eos_token_mask` and `compute_context_repetition_mask` methods of the logits processor
* Detection is performed by passing g-values and combined mask to the Bayesian detector module

**Notable Dependencies and Requirements:**

* `BayesianDetectorModel`
* `SynthIDTextWatermarkLogitsProcessor`
* `AutoTokenizer`

Note: The class assumes the detection module is already trained and loaded via `BayesianDetectorModel` from a specific script (`detector_training.py`).


key: chroma:document
string_value:
CODE:

class LogitsProcessor:
    """Abstract base class for all logit processors that can be applied during generation."""

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        raise NotImplementedError(
            f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."


SUMMARY:
**Concise Summary**

`LogitsProcessor` (Abstract Class): 
  - **Purpose**: Abstract base class for logit processors used during generation, ensuring only subclasses are callable.

**Key Implementation Details:**

* Takes `input_ids` and `scores` as input
* Returns a `torch.FloatTensor`
* Subclasses must implement this method to provide actual processing logic

**Notable Dependencies/Requirements:**

* Torch (`torch.LongTensor`, `torch.FloatTensor`)
* A decorator `add_start_docstrings` (not shown in provided code)


key: chroma:document
string_value:
CODE:

class LogitsWarper:
    """Abstract base class for all logit warpers that can be applied during generation with multinomial sampling."""

    def __init__(self):
        logger.warning_once(
            "`LogitsWarper` is deprecated and will be removed in v4.48. Your class should inherit `LogitsProcessor` "
            "instead, which has the same properties and interface."
        )

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        raise NotImplementedError(
            f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."


SUMMARY:
**LogitsWarper Class Summary**

The `LogitsWarper` class is an abstract base class for logit warpers used in multinomial sampling during model generation. It serves as a deprecated intermediate step, recommending inheritance of the `LogitsProcessor` class instead.

**Primary Purpose:**
Apply logits warping during generation to process input IDs and scores.

**Key Implementation Details:**

*   Takes `input_ids` (long tensor) and `scores` (float tensor) as inputs.
*   Raises a `NotImplementedError`, indicating that only subclasses can be instantiated.
*   Initializes logger warnings upon initiation due to deprecation.

**Notable Dependencies/Requirements:**

*   Inheriting classes should implement the `__call__` method, which takes `input_ids` and `scores`.
*   `LogitsProcessor` interface is recommended for subclass inheritance.


key: chroma:document
string_value:
CODE:

class LogitsProcessorList(list):
    """
    This class can be used to create a list of [`LogitsProcessor`] to subsequently process a `scores` input tensor.
    This class inherits from list and adds a specific *__call__* method to apply each [`LogitsProcessor`] to the
    inputs.
    """

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
        r"""
        Args:
            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
                Indices of input sequence tokens in the vocabulary. [What are input IDs?](../glossary#input-ids)
            scores (`torch.FloatTensor` of shape `(batch_size, config.vocab_size)`):
                Prediction scores of a language modeling head. These can be logits for each vocabulary when not using
                beam search or log softmax for each vocabulary token when using beam search
            kwargs (`Dict[str, Any]`, *optional*):
                Additional kwargs that are specific to a logits processor.

        Return:
            `torch.FloatTensor` of shape `(batch_size, config.vocab_size)`:
                The processed prediction scores.

        """
        for processor in self:
            function_args = inspect.signature(processor.__call__).parameters
            if len(function_args) > 2:
                if not all(arg in kwargs for arg in list(function_args.keys())[2:]):
                    raise ValueError(
                        f"Make sure that all the required parameters: {list(function_args.keys())} for "
                        f"{processor.__class__} are passed to the logits processor."
                    )
                scores = processor(input_ids, scores, **kwargs)
            else:
                scores = processor(input_ids, scores)



SUMMARY:
**LogitsProcessorList Class Summary**

The `LogitsProcessorList` class is a custom list implementation that enables processing multiple logits processors sequentially. Its primary purpose is to apply each processor to the input tensors and combine their results.

**Key Implementation Details:**

* The `__call__` method takes in `input_ids`, `scores`, and optional keyword arguments (`kwargs`) as inputs.
* Each processor in the list is applied to the input tensors using its own `__call__` method.
* The processors are inspected for required positional arguments using `inspect.signature`.
* If all required positional arguments are present, the scores are updated with the results of the processor. Otherwise, a `ValueError` is raised.

**Notable Dependencies and Requirements:**

* `torch`: PyTorch library for tensor operations.
* `inspect`: Python built-in module for inspecting function signatures.
* A list of [`LogitsProcessor`] instances to be processed sequentially.

**Method Signatures and Interfaces:**

* `__call__(input_ids, scores, **kwargs) -> torch.FloatTensor`: The main entry point of the class, which applies each processor to the input tensors.
* [`LogitsProcessor`]: An interface class that defines the behavior for processing logits.


key: chroma:document
string_value:
CODE:

class MinLengthLogitsProcessor(LogitsProcessor):
    r"""
    [`LogitsProcessor`] enforcing a min-length by setting EOS probability to 0. Note that, for decoder-only models
    like most LLMs, the length includes the prompt.

    Args:
        min_length (`int`):
            The minimum length below which the score of `eos_token_id` is set to `-float("Inf")`.
        eos_token_id (`Union[int, List[int], torch.Tensor]`):
            The id(s) of the *end-of-sequence* token.
        device (`str`, *optional*, defaults to `"cpu"`):
            The device to allocate the tensors.

    Examples:

    ```python
    >>> from transformers import AutoModelForCausalLM, AutoTokenizer

    >>> tokenizer = AutoTokenizer.from_pretrained("bigscience/bloomz-560m")
    >>> model = AutoModelForCausalLM.from_pretrained("bigscience/bloomz-560m")

    >>> inputs = tokenizer("A number:", return_tensors="pt")
    >>> gen_out = model.generate(**inputs)
    >>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0])
    A number: one

    >>> # setting `min_length` to a value smaller than the uncontrolled output length has no impact
    >>> gen_out = model.generate(**inputs, min_length=3)
    >>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0])
    A number: one

    >>> # setting a larger `min_length` will force the model to generate beyond its natural ending point, which is not
    >>> # necessarily incorrect
    >>> gen_out = model.generate(**inputs, min_length=10)
    >>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0])
    A number: one thousand, nine hundred and ninety-four
    ```
    """

    def __init__(self, min_length: int, eos_token_id: Union[int, List[int], torch.Tensor], device: str = "cpu"):
        if not isinstance(min_length, int) or min_length < 0:
            raise ValueError(f"`min_length` has to be a non-negative integer, but is {min_length}")

        if not isinstance(eos_token_id, torch.Tensor):
            if isinstance(eos_token_id, int):
                eos_token_id = [eos_token_id]
            eos_token_id = torch.tensor(eos_token_id, device=device)

        self.min_length = min_length
        self.eos_token_id = eos_token_id

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        vocab_tensor = torch.arange(scores.shape[-1], device=scores.device)
        eos_token_mask = isin_mps_friendly(vocab_tensor, self.eos_token_id)
        scores_processed = scores.clone()
        if input_ids.shape[-1] < self.min_length:
            scores_processed = torch.where(eos_token_mask, -math.inf, scores)


SUMMARY:
**Class Summary: MinLengthLogitsProcessor**

**Purpose and Functionality**
---------------------------

This class is a processor for logits in a sequence-to-sequence model. It enforces a minimum length by setting the end-of-sequence probability to 0.

* The `MinLengthLogitsProcessor` takes an input tensor (`input_ids`) and scores tensor (`scores`), and returns the processed scores.
* It checks if the minimum length is exceeded in the input sequence, and sets the model's output probability to -inf for EOS tokens beyond that length.

**Key Implementation Details**
-------------------------------

* The `__init__` method initializes theprocessor with input parameters:
	+ `min_length`: the minimum length below which the score of `eos_token_id` is set to -inf.
	+ `eos_token_id`: the id(s) of the *end-of-sequence* token.
	+ `device`: the device to allocate the tensors (defaults to `"cpu"`).
* The `__call__` method applies the processor logic:
	1. Create a vocabulary tensor of all possible scores.
	2. Compute an EOS token mask using `is_inmps_friendly`.
	3. Mask the model's output probability for scores beyond the minimum length.

**Notable Dependencies or Requirements**
--------------------------------------

* `torch`: a PyTorch library for efficient numerical computation.
* `math`: a Python library for mathematical functions.
* `[LOGITS_PROCESSOR_INPUTS_DOCSTRING]`: a custom docstring string.


key: chroma:document
string_value:
CODE:

class MinNewTokensLengthLogitsProcessor(LogitsProcessor):
    r"""
    [`LogitsProcessor`] enforcing a min-length of new tokens by setting EOS (End-Of-Sequence) token probability to 0.
    Contrarily to [`MinLengthLogitsProcessor`], this processor ignores the prompt.

    Args:
        prompt_length_to_skip (`int`):
            The input tokens length. Not a valid argument when used with `generate` as it will automatically assign the
            input length.
        min_new_tokens (`int`):
            The minimum *new* tokens length below which the score of `eos_token_id` is set to `-float("Inf")`.
        eos_token_id (`Union[int, List[int], torch.Tensor]`):
            The id(s) of the *end-of-sequence* token.
        device (`str`, *optional*, defaults to `"cpu"`):
            The device to allocate the tensors.

    Examples:

    ```python
    >>> from transformers import AutoModelForCausalLM, AutoTokenizer

    >>> tokenizer = AutoTokenizer.from_pretrained("bigscience/bloomz-560m")
    >>> model = AutoModelForCausalLM.from_pretrained("bigscience/bloomz-560m")

    >>> inputs = tokenizer(["A number:"], return_tensors="pt")
    >>> gen_out = model.generate(**inputs)
    >>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0])
    A number: one

    >>> # setting `min_new_tokens` will force the model to generate beyond its natural ending point, which is not
    >>> # necessarily incorrect
    >>> gen_out = model.generate(**inputs, min_new_tokens=2)
    >>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0])
    A number: one thousand
    ```
    """

    def __init__(
        self,
        prompt_length_to_skip: int,
        min_new_tokens: int,
        eos_token_id: Union[int, List[int], torch.Tensor],
        device: str = "cpu",
    ):
        for arg_name, arg_value in [
            ("prompt_length_to_skip", prompt_length_to_skip),
            ("min_new_tokens", min_new_tokens),
        ]:
            if not isinstance(arg_value, int) or arg_value < 0:
                raise ValueError(f"`{arg_name}` has to be a positive integer, but is {arg_value}")

        if not isinstance(eos_token_id, torch.Tensor):
            if isinstance(eos_token_id, int):
                eos_token_id = [eos_token_id]
            eos_token_id = torch.tensor(eos_token_id, device=device)

        self.prompt_length_to_skip = prompt_length_to_skip
        self.min_new_tokens = min_new_tokens
        self.eos_token_id = eos_token_id

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        new_tokens_length = input_ids.shape[-1] - self.prompt_length_to_skip
        scores_processed = scores.clone()
        vocab_tensor = torch.arange(scores.shape[-1], device=scores.device)
        eos_token_mask = isin_mps_friendly(vocab_tensor, self.eos_token_id)
        if new_tokens_length < self.min_new_tokens:
            scores_processed = torch.where(eos_token_mask, -math.inf, scores)



SUMMARY:
**MinNewTokensLengthLogitsProcessor Class Summary**

### Purpose and Functionality
This class extends `LogitsProcessor` to enforce a minimum length for new tokens by setting the End-Of-Sequence (EOS) token probability to 0. It ignores the prompt input when generating sequence.

### Implementation Details

* Sets EOS token probability to -inf if new tokens length is below `min_new_tokens`.
* Uses `is_in_mps_friendly` function to create an EOS token mask.
* Defaults device to "cpu" with type hinting for compatibility.

### Notable Dependencies and Requirements

* Torch ( library for numerical computations)
* Math (library for mathematical operations)

### Method Signatures

* `__init__(prompt_length_to_skip: int, min_new_tokens: int, eos_token_id: Union[int, List[int], torch.Tensor], device: str = "cpu"): Initializes the processor with prompt length skipping and EOS token ID.
* `__call__(input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor`: Processes logits to enforce minimum new tokens length.


key: chroma:document
string_value:
CODE:

class TemperatureLogitsWarper(LogitsProcessor):
    r"""
    [`LogitsProcessor`] for temperature (exponential scaling output probability distribution), which effectively means
    that it can control the randomness of the predicted tokens. Often used together with [`TopPLogitsWarper`] and
    [`TopKLogitsWarper`].

    <Tip>

    Make sure that `do_sample=True` is included in the `generate` arguments otherwise the temperature value won't have
    any effect.

    </Tip>

    Args:
        temperature (`float`):
            Strictly positive float value used to modulate the logits distribution. A value smaller than `1` decreases
            randomness (and vice versa), with `0` being equivalent to shifting all probability mass to the most likely
            token.

    Examples:

    ```python
    >>> import torch
    >>> from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed

    >>> set_seed(0)  # for reproducibility

    >>> tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
    >>> model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
    >>> model.config.pad_token_id = model.config.eos_token_id
    >>> inputs = tokenizer(["Hugging Face Company is"], return_tensors="pt")

    >>> # With temperature=1.0, the default, we consistently get random outputs due to random sampling.
    >>> generate_kwargs = {"max_new_tokens": 10, "do_sample": True, "temperature": 1.0, "num_return_sequences": 2}
    >>> outputs = model.generate(**inputs, **generate_kwargs)
    >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
    ['Hugging Face Company is one of these companies that is going to take a',
    "Hugging Face Company is a brand created by Brian A. O'Neil"]

    >>> # However, with temperature close to 0, it approximates greedy decoding strategies (invariant)
    >>> generate_kwargs["temperature"] = 0.0001
    >>> outputs = model.generate(**inputs, **generate_kwargs)
    >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
    ['Hugging Face Company is a company that has been around for over 20 years',
    'Hugging Face Company is a company that has been around for over 20 years']
    ```
    """

    def __init__(self, temperature: float):
        if not isinstance(temperature, float) or not (temperature > 0):
            except_msg = (
                f"`temperature` (={temperature}) has to be a strictly positive float, otherwise your next token "
                "scores will be invalid."
            )
            if isinstance(temperature, float) and temperature == 0.0:
                except_msg += " If you're looking for greedy decoding strategies, set `do_sample=False`."
            raise ValueError(except_msg)

        self.temperature = temperature

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        scores_processed = scores / self.temperature


SUMMARY:
**TemperatureLogitsWarper Class Summary**

The `TemperatureLogitsWarper` class is a `LogitsProcessor` that applies exponential scaling to the output probability distribution of a language model, effectively controlling randomness in predicted tokens.

**Main Purpose and Functionality:**
Modulates logits distribution using a temperature value (strictly positive float), adjusting randomness in output probabilities.

**Key Implementation Details:**

* Applies temperature scaling by dividing input scores by `temperature`.
* Raises a `ValueError` if `temperature` is not a strictly positive float.

**Notable Dependencies and Requirements:**

* `torch` for tensor operations (`LongTensor`, `FloatTensor`)
* `add_start_docstrings` function (not shown in code snippet, but assumed to be part of the class hierarchy)


key: chroma:document
string_value:
CODE:

class RepetitionPenaltyLogitsProcessor(LogitsProcessor):
    r"""
    [`LogitsProcessor`] that prevents the repetition of previous tokens through a penalty. This penalty is applied at
    most once per token. Note that, for decoder-only models like most LLMs, the considered tokens include the prompt.

    In the original [paper](https://arxiv.org/pdf/1909.05858.pdf), the authors suggest the use of a penalty of around
    1.2 to achieve a good balance between truthful generation and lack of repetition. To penalize and reduce
    repetition, use `penalty` values above 1.0, where a higher value penalizes more strongly. To reward and encourage
    repetition, use `penalty` values between 0.0 and 1.0, where a lower value rewards more strongly.

    Args:
        penalty (`float`):
            The parameter for repetition penalty. 1.0 means no penalty. Above 1.0 penalizes previously generated
            tokens. Between 0.0 and 1.0 rewards previously generated tokens.

    Examples:

    ```py
    >>> from transformers import AutoTokenizer, AutoModelForCausalLM

    >>> # Initializing the model and tokenizer for it
    >>> model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2")
    >>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2")
    >>> inputs = tokenizer(["I'm not going to"], return_tensors="pt")

    >>> # This shows a normal generate without any specific parameters
    >>> summary_ids = model.generate(**inputs)
    >>> print(tokenizer.batch_decode(summary_ids, skip_special_tokens=True)[0])
    I'm not going to be able to do that. I'm going to be able to do that

    >>> # This generates a penalty for repeated tokens
    >>> penalized_ids = model.generate(**inputs, repetition_penalty=1.1)
    >>> print(tokenizer.batch_decode(penalized_ids, skip_special_tokens=True)[0])
    I'm not going to be able to do that. I'll just have to go out and play
    ```
    """

    def __init__(self, penalty: float):
        if not isinstance(penalty, float) or not (penalty > 0):
            raise ValueError(f"`penalty` has to be a strictly positive float, but is {penalty}")

        self.penalty = penalty

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        score = torch.gather(scores, 1, input_ids)

        # if score < 0 then repetition penalty has to be multiplied to reduce the token probabilities
        score = torch.where(score < 0, score * self.penalty, score / self.penalty)

        scores_processed = scores.scatter(1, input_ids, score)


SUMMARY:
**Summary**

`RepetitionPenaltyLogitsProcessor` class implements a penalty-based `LogitsProcessor` to prevent repetition of previous tokens in sequence generation tasks. The penalty is applied up to once per token and can be adjusted using the `penalty` parameter.

**Key Implementation Details**

* Applies repeat penalty to tokens with negative scores
* Penalizes tokens up to once per token
* Reward repeats generated tokens using a multiplier (`self.penalty`)
* Accepts `penalty` value as input for adjustable repetition strength

**Notable Dependencies and Requirements**

* `torch`: used for tensor operations and inputs/outputs
* `LogitsProcessor`: base class that this processor inherits from, providing additional functionality

**Method Signatures and Interfaces**

* `__call__(input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor`: processes input tokens with repetition penalty using given scores


key: chroma:document
string_value:
CODE:

class EncoderRepetitionPenaltyLogitsProcessor(LogitsProcessor):
    r"""
    [`LogitsProcessor`] that works similarly to [`RepetitionPenaltyLogitsProcessor`], but with an *inverse* penalty
    that is applied to the tokens present in the prompt. In other words, a penalty above 1.0 increases the odds of
    selecting tokens that were present in the prompt.

    It was designed to avoid hallucination in input-grounded tasks, like summarization. Although originally intended
    for encoder-decoder models, it can also be used with decoder-only models like LLMs.

    Args:
        penalty (`float`):
            The parameter for repetition penalty. 1.0 means no penalty. Above 1.0 rewards prompt tokens. Between 0.0
            and 1.0 penalizes prompt tokens.
        encoder_input_ids (`torch.LongTensor`):
            The encoder_input_ids that should be repeated within the decoder ids.

    Examples:

    ```python
    >>> from transformers import AutoModelForCausalLM, AutoTokenizer

    >>> tokenizer = AutoTokenizer.from_pretrained("bigscience/bloomz-560m")
    >>> model = AutoModelForCausalLM.from_pretrained("bigscience/bloomz-560m")

    >>> inputs = tokenizer(["Alice and Bob. The third member's name was"], return_tensors="pt")
    >>> gen_out = model.generate(**inputs)
    >>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0])
    Alice and Bob. The third member's name was not mentioned.

    >>> # With the `encoder_repetition_penalty` argument we can trigger this logits processor in `generate`, which can
    >>> # promote the use of prompt tokens ("Bob" in this example)
    >>> gen_out = model.generate(**inputs, encoder_repetition_penalty=1.2)
    >>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0])
    Alice and Bob. The third member's name was Bob. The third member's name was Bob.
    ```
    """

    def __init__(self, penalty: float, encoder_input_ids: torch.LongTensor):
        if not isinstance(penalty, float) or not (penalty > 0):
            raise ValueError(f"`penalty` has to be a strictly positive float, but is {penalty}")

        self.penalty = 1 / penalty
        self.encoder_input_ids = encoder_input_ids

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        score = torch.gather(scores, 1, self.encoder_input_ids)

        # if score < 0 then hallucination penalty has to be multiplied to increase the token probabilities
        score = torch.where(score < 0, score * self.penalty, score / self.penalty)

        scores_processed = scores.scatter(1, self.encoder_input_ids, score)


SUMMARY:
**Class Summary: EncoderRepetitionPenaltyLogitsProcessor**

The `EncoderRepetitionPenaltyLogitsProcessor` is a custom logit processor designed to avoid hallucination in input-grounded tasks. It applies an inverse penalty to tokens present in the prompt.

**Key Functionality:**

* Inverts penalty above 1.0, increasing odds of selecting prompt tokens
* Penalizes prompt tokens between 0.0 and 1.0

**Implementation Details:**

* Multiplies hallucination penalty by `penalty` when score < 0 to increase token probabilities
* Applies inverse penalty to scores using `torch.where`

**Dependencies and Interfaces:**

* `LogitsProcessor`
* `RepetitionPenaltyLogitsProcessor` ( superclass )
* `torch.FloatTensor`
* `torch.LongTensor`
* Inherited from `add_start_docstrings` function

**Example:**
```python
inputs = tokenizer(["Alice and Bob. The third member's name was"], return_tensors="pt")
gen_out = model.generate(**inputs, encoder_repetition_penalty=1.2)
```
Note: This summary focuses on the core functionality, implementation details, and dependencies of the class.


key: chroma:document
string_value:
CODE:

class TopPLogitsWarper(LogitsProcessor):
    """
    [`LogitsProcessor`] that performs top-p, i.e. restricting to top tokens summing to prob_cut_off <= prob_cut_off.
    Often used together with [`TemperatureLogitsWarper`] and [`TopKLogitsWarper`].

    Args:
        top_p (`float`):
            If set to < 1, only the smallest set of most probable tokens with probabilities that add up to `top_p` or
            higher are kept for generation.
        filter_value (`float`, *optional*, defaults to -inf):
            All filtered values will be set to this float value.
        min_tokens_to_keep (`int`, *optional*, defaults to 1):
            Minimum number of tokens that cannot be filtered.

    Examples:

    ```python
    >>> from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed

    >>> set_seed(1)
    >>> model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2")
    >>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2")

    >>> inputs = tokenizer("A sequence: 1, 2", return_tensors="pt")

    >>> # With sampling, the output is unexpected -- sometimes too unexpected.
    >>> outputs = model.generate(**inputs, do_sample=True)
    >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
    A sequence: 1, 2, 3 | < 4 (left-hand pointer) ;
    <BLANKLINE>
    <BLANKLINE>

    >>> # With `top_p` sampling, the output gets restricted to high-probability tokens.
    >>> # Pro tip: In practice, LLMs use `top_p` in the 0.9-0.95 range.
    >>> outputs = model.generate(**inputs, do_sample=True, top_p=0.1)
    >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
    A sequence: 1, 2, 3, 4, 5, 6, 7, 8, 9
    ```
    """

    def __init__(self, top_p: float, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):
        top_p = float(top_p)
        if top_p < 0 or top_p > 1.0:
            raise ValueError(f"`top_p` has to be a float > 0 and < 1, but is {top_p}")
        if not isinstance(min_tokens_to_keep, int) or (min_tokens_to_keep < 1):
            raise ValueError(f"`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}")

        self.top_p = top_p
        self.filter_value = filter_value
        self.min_tokens_to_keep = min_tokens_to_keep

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        sorted_logits, sorted_indices = torch.sort(scores, descending=False)
        cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)

        # Remove tokens with cumulative top_p above the threshold (token with 0 are kept)
        sorted_indices_to_remove = cumulative_probs <= (1 - self.top_p)
        # Keep at least min_tokens_to_keep
        sorted_indices_to_remove[..., -self.min_tokens_to_keep :] = 0

        # scatter sorted tensors to original indexing
        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
        scores_processed = scores.masked_fill(indices_to_remove, self.filter_value)


SUMMARY:
**TopPLogitsWarper Class Summary**

**Class Purpose**: Restricts tokens to summing to `prob_cut_off` or higher, commonly used in sampling for Language Models.

**Key Implementation Details**:

* Sorts logits and indices in descending order
* Calculates cumulative probabilities
* Removes tokens with cumulative top-p above threshold (except those below min_tokens_to_keep)
* Scatter processed logits back to original indexing

**Method Signatures**:

* `__call__(input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor`: Main processing function

**Dependencies and Requirements**:

* `LOGITS_PROCESSOR_INPUTS_DOCSTRING` annotation
* `torch` tensor library (specifically `LongTensor` and `FloatTensor`)
* `add_start_docstrings` decorator


key: chroma:document
string_value:
CODE:

class TopKLogitsWarper(LogitsProcessor):
    r"""
    [`LogitsProcessor`] that performs top-k, i.e. restricting to the k highest probability elements. Often used
    together with [`TemperatureLogitsWarper`] and [`TopPLogitsWarper`].

    Args:
        top_k (`int`):
            The number of highest probability vocabulary tokens to keep for top-k-filtering.
        filter_value (`float`, *optional*, defaults to -inf):
            All filtered values will be set to this float value.
        min_tokens_to_keep (`int`, *optional*, defaults to 1):
            Minimum number of tokens that cannot be filtered.

    Examples:

    ```python
    >>> from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed

    >>> set_seed(1)
    >>> model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2")
    >>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2")

    >>> inputs = tokenizer("A sequence: A, B, C, D", return_tensors="pt")

    >>> # With sampling, the output is unexpected -- sometimes too unexpected.
    >>> outputs = model.generate(**inputs, do_sample=True)
    >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
    A sequence: A, B, C, D, E — S — O, P — R

    >>> # With `top_k` sampling, the output gets restricted the k most likely tokens.
    >>> # Pro tip: In practice, LLMs use `top_k` in the 5-50 range.
    >>> outputs = model.generate(**inputs, do_sample=True, top_k=2)
    >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
    A sequence: A, B, C, D, E, F, G, H, I
    ```
    """

    def __init__(self, top_k: int, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):
        if not isinstance(top_k, int) or top_k <= 0:
            raise ValueError(f"`top_k` has to be a strictly positive integer, but is {top_k}")

        self.top_k = max(top_k, min_tokens_to_keep)
        self.filter_value = filter_value

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        top_k = min(self.top_k, scores.size(-1))  # Safety check
        # Remove all tokens with a probability less than the last token of the top-k
        indices_to_remove = scores < torch.topk(scores, top_k)[0][..., -1, None]
        scores_processed = scores.masked_fill(indices_to_remove, self.filter_value)


SUMMARY:
**TopKLogitsWarper Class Summary**

The `TopKLogitsWarper` class is a subclass of `_LogitsProcessor_` that applies top-k filtering to logits, limiting the k highest probability vocabulary tokens. This processor is often used in conjunction with other logit warppers.

**Key Features:**

* Restricts to the k highest probability elements.
* Implements `TopKFiltering` algorithm.
* Accepts the following parameters:
	+ `top_k`: The number of highest probability vocabulary tokens to keep (required).
	+ `filter_value`: The value to set for filtered values (default is `-inf`).
	+ `min_tokens_to_keep`: Minimum number of tokens that cannot be filtered (default is 1).

**Method Signature:**

`__call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor`

* Processes logits by applying top-k filtering and setting filtered values to a specified filter value.

**Imports:**

None

**Interfaces:**

None


key: chroma:document
string_value:
CODE:

class MinPLogitsWarper(LogitsProcessor):
    """
    [`LogitsProcessor`] that performs min-p, i.e. keeps all tokens that are above a minimum probability, scaled by the
    probability of the most likely token. As a result, the filter becomes more agressive in the presence of
    high-probability tokens, which is a sign of a confident output that we shouldn't deviate from.

    Often used together with [`TemperatureLogitsWarper`]. Used as an alternative to [`TopPLogitsWarper`] and
    [`TopKLogitsWarper`].

    Created by @menhguin and @kalomaze (github handles). Code adapted from [this external PR](https://github.com/oobabooga/text-generation-webui/pull/4449/files)

    Args:
        min_p (`float`):
            Minimum token probability, which will be scaled by the probability of the most likely token. It must be a
            value between 0 and 1. Typical values are in the 0.01-0.2 range, comparably selective as setting `top_p` in
            the 0.99-0.8 range (use the opposite of normal `top_p` values).
        filter_value (`float`, *optional*, defaults to -inf):
            All filtered values will be set to this float value.
        min_tokens_to_keep (`int`, *optional*, defaults to 1):
            Minimum number of tokens that cannot be filtered.

    Examples:

    ```python
    >>> from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed

    >>> set_seed(1)
    >>> model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2")
    >>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2")

    >>> inputs = tokenizer("A sequence: 1, 2", return_tensors="pt")

    >>> # With sampling, the output is unexpected -- sometimes too unexpected.
    >>> outputs = model.generate(**inputs, do_sample=True)
    >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
    A sequence: 1, 2, 3 | < 4 (left-hand pointer) ;
    <BLANKLINE>
    <BLANKLINE>

    >>> # With `min_p` sampling, the output gets restricted to high-probability tokens.
    >>> # Pro tip: In practice, LLMs use `min_p` in the 0.01-0.2 range.
    >>> outputs = model.generate(**inputs, do_sample=True, min_p=0.1)
    >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
    A sequence: 1, 2, 3, 4, 5, 6, 7, 8, 9
    ```
    """

    def __init__(self, min_p: float, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):
        if not (0 <= min_p <= 1.0):
            raise ValueError(f"`min_p` has to be a float in the [0, 1] interval, but is {min_p}")
        if not isinstance(min_tokens_to_keep, int) or (min_tokens_to_keep < 1):
            raise ValueError(f"`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}")

        self.min_p = min_p
        self.filter_value = filter_value
        self.min_tokens_to_keep = min_tokens_to_keep

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        # Convert logits to probabilities
        probs = torch.softmax(scores, dim=-1)
        # Get the probability of the top token for each sequence in the batch
        top_probs, _ = probs.max(dim=-1, keepdim=True)
        # Calculate the actual min_p threshold by scaling min_p with the top token's probability
        scaled_min_p = self.min_p * top_probs
        # Create a mask for tokens that have a probability less than the scaled min_p
        tokens_to_remove = probs < scaled_min_p

        sorted_indices = torch.argsort(scores, descending=True, dim=-1)
        sorted_indices_to_remove = torch.gather(tokens_to_remove, dim=-1, index=sorted_indices)
        # Keep at least min_tokens_to_keep
        sorted_indices_to_remove[..., : self.min_tokens_to_keep] = False

        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
        scores_processed = scores.masked_fill(indices_to_remove, self.filter_value)


SUMMARY:
**Class Summary: MinPLogitsWarper**

A `LogitsProcessor` implementation that performs min-p logits warping, keeping all tokens with probabilities above a minimum threshold scaled by the probability of the most likely token.

### Main Purpose and Functionality

* Scale minimum logit values based on top-probability tokens
* Remove low-probability tokens while preserving high-probability ones

### Key Implementation Details

* Uses `torch.softmax` to convert logits to probabilities
* Calculates scaled minimum threshold (`min_p`) by scaling `min_p` with the probability of the top token
* Creates a mask for tokens below the scaled threshold and removes them

### Notable Dependencies or Requirements

* `torch`
* `numpy` (for numerical computations)
* External dependencies:
	+ `transformers` library for accessing pre-trained models

### Class Structure

* Initializes with `min_p`, `filter_value`, and `min_tokens_to_keep` parameters
* Takes an input tensor of `input_ids` and a scores tensor as input
* Returns the processed scores after applying min-p warping


key: chroma:document
string_value:
CODE:

class TypicalLogitsWarper(LogitsProcessor):
    r"""
    [`LogitsProcessor`] that performs typical decoding. Inspired on how humans use language, it prioritizes tokens
    whose log probability is close to the entropy of the token probability distribution. This means that the most
    likely tokens may be discarded in the process.

    See [Typical Decoding for Natural Language Generation](https://arxiv.org/abs/2202.00666) for more information.

    Args:
        mass (`float`, *optional*, defaults to 0.9):
            Value of typical_p between 0 and 1 inclusive, defaults to 0.9.
        filter_value (`float`, *optional*, defaults to -inf):
            All filtered values will be set to this float value.
        min_tokens_to_keep (`int`, *optional*, defaults to 1):
            Minimum number of tokens that cannot be filtered.

    Examples:

    ```python
    >>> from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed

    >>> model = AutoModelForCausalLM.from_pretrained("bigscience/bloomz-560m")
    >>> tokenizer = AutoTokenizer.from_pretrained("bigscience/bloomz-560m")

    >>> inputs = tokenizer("1, 2, 3", return_tensors="pt")

    >>> # We can see that greedy decoding produces a sequence of numbers
    >>> outputs = model.generate(**inputs)
    >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
    1, 2, 3, 4, 5, 6, 7, 8, 9, 10,

    >>> # For this particular seed, we can see that sampling produces nearly the same low-information (= low entropy)
    >>> # sequence
    >>> set_seed(18)
    >>> outputs = model.generate(**inputs, do_sample=True)
    >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
    1, 2, 3, 4, 5, 6, 7, 8, 9 and 10

    >>> # With `typical_p` set, the most obvious sequence is no longer produced, which may be good for your problem
    >>> set_seed(18)
    >>> outputs = model.generate(
    ...     **inputs, do_sample=True, typical_p=0.1, return_dict_in_generate=True, output_scores=True
    ... )
    >>> print(tokenizer.batch_decode(outputs.sequences, skip_special_tokens=True)[0])
    1, 2, 3 and 5

    >>> # We can see that the token corresponding to "4" (token 934) in the second position, the most likely token
    >>> # as seen with greedy decoding, was entirely blocked out
    >>> print(outputs.scores[1][0, 934])
    tensor(-inf)
    ```
    """

    def __init__(self, mass: float = 0.9, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):
        mass = float(mass)
        if not (mass > 0 and mass < 1):
            raise ValueError(f"`typical_p` has to be a float > 0 and < 1, but is {mass}")
        if not isinstance(min_tokens_to_keep, int) or (min_tokens_to_keep < 1):
            raise ValueError(f"`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}")

        self.filter_value = filter_value
        self.mass = mass
        self.min_tokens_to_keep = min_tokens_to_keep

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        # calculate entropy
        normalized = torch.nn.functional.log_softmax(scores, dim=-1)
        p = torch.exp(normalized)
        ent = -(normalized * p).nansum(-1, keepdim=True)

        # shift and sort
        shifted_scores = torch.abs((-normalized) - ent)
        sorted_scores, sorted_indices = torch.sort(shifted_scores, descending=False)
        sorted_logits = scores.gather(-1, sorted_indices)
        cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)

        # Remove tokens with cumulative mass above the threshold
        last_ind = (cumulative_probs < self.mass).sum(dim=1)
        last_ind.clamp_(max=sorted_scores.shape[-1] - 1)
        sorted_indices_to_remove = sorted_scores > sorted_scores.gather(1, last_ind.view(-1, 1))
        sorted_indices_to_remove[..., : self.min_tokens_to_keep] = 0
        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)

        scores_processed = scores.masked_fill(indices_to_remove, self.filter_value)


SUMMARY:
**TypicalLogitsWarper Class Summary**

The `TypicalLogitsWarper` class is a custom implementation of the LogitsProcessor for natural language generation tasks. It applies typical decoding strategy, prioritizing tokens with log probability close to entropy.

**Primary Purpose:**
Prioritizes tokens based on their information content, allowing for more efficient generation while discarding most likely tokens.

**Key Implementation Details:**

* Calculates entropy using `torch.nn.functional.log_softmax` and `torch.exp`
* Shifts scores by subtracting the negative entropy value
* Sorts scores in descending order and gathers logits from sorted indices
* Applies cumulative softmax to compute probabilities
* Removes tokens with cumulative mass above a threshold set by `mass` parameter

**Notable Dependencies:**

* `torch`
* `torch.nn.functional`

**External Dependencies:**

None

Note: The class inherits the `LOGITS_PROCESSOR_INPUTS_DOCSTRING` docstring from its parent class.


key: chroma:document
string_value:
CODE:

class EpsilonLogitsWarper(LogitsProcessor):
    r"""
    [`LogitsProcessor`] that performs epsilon-sampling, i.e. restricting to tokens with `prob >= epsilon`. Takes the
    largest min_tokens_to_keep tokens if no tokens satisfy this constraint. See [Truncation Sampling as Language Model
    Desmoothing](https://arxiv.org/abs/2210.15191) for more information.

    Args:
        epsilon (`float`):
            If set to > 0, only the most tokens with probabilities `epsilon` or higher are kept for generation.
        filter_value (`float`, *optional*, defaults to -inf):
            All filtered values will be set to this float value.
        min_tokens_to_keep (`int`, *optional*, defaults to 1):
            Minimum number of tokens that cannot be filtered.

    Examples:
    ```python
    >>> from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed

    >>> set_seed(1)
    >>> model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2")
    >>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2")

    >>> inputs = tokenizer("A sequence: 1, 2", return_tensors="pt")

    >>> # With sampling, the output is unexpected -- sometimes too unexpected.
    >>> outputs = model.generate(**inputs, do_sample=True)
    >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
    A sequence: 1, 2, 3 | < 4 (left-hand pointer) ;
    <BLANKLINE>
    <BLANKLINE>

    >>> # With epsilon sampling, the output gets restricted to high-probability tokens. Note that this is similar to
    >>> # Top P sampling, which restricts tokens based on their cumulative probability.
    >>> # Pro tip: The paper recomends using `epsilon_cutoff` values between 3e-4 and 9e-4
    >>> outputs = model.generate(**inputs, do_sample=True, epsilon_cutoff=0.1)
    >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
    A sequence: 1, 2, 3, 4, 5, 6, 7, 8, 9
    ```
    """

    def __init__(self, epsilon: float, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):
        epsilon = float(epsilon)
        if epsilon <= 0 or epsilon >= 1:
            raise ValueError(f"`epsilon_cutoff` has to be a float > 0 and < 1, but is {epsilon}")

        min_tokens_to_keep = int(min_tokens_to_keep)
        if min_tokens_to_keep < 1:
            raise ValueError(
                f"`min_tokens_to_keep` has to be a strictly positive integer, but is {min_tokens_to_keep}"
            )

        self.epsilon = epsilon
        self.filter_value = filter_value
        self.min_tokens_to_keep = min_tokens_to_keep

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        # Determine which indices to remove
        probabilities = scores.softmax(dim=-1)
        indices_to_remove = probabilities < self.epsilon

        # Keep the words with the 'min_tokens_to_keep'-highest probabilities
        top_k = min(self.min_tokens_to_keep, scores.size(-1))  # Safety check
        indices_to_remove = indices_to_remove & (scores < torch.topk(scores, top_k)[0][..., -1, None])

        scores_processed = scores.masked_fill(indices_to_remove, self.filter_value)


SUMMARY:
**EpsilonLogitsWarper Class Summary**

The `EpsilonLogitsWarper` class is a specialized logit processor that applies epsilon-sampling to restrict tokens with probabilities ≥ `epsilon`. It takes the largest `min_tokens_to_keep` tokens if no tokens satisfy this constraint.

**Primary Purpose:** Restrict tokens based on epsilon-sampling and keep minimum number of top-ranked tokens.

**Implementation Details:**

* Applies softmax to scores tensor along last dimension.
* Identifies indices to remove using epsilon threshold.
* Keeps top-k tokens with highest probabilities, where k is determined by `min_tokens_to_keep`.
* Masks filtered values with a specified filter value.

**Notable Dependencies/Requirements:**

* `torch` library for tensor operations
* `LOGITS_PROCESSOR_INPUTS_DOCSTRING` documentation string for initialization logic.
* External dependencies:
	+ `transformers` library for AutoTokenizer and AutoModelForCausalLM.


key: chroma:document
string_value:
CODE:

class EtaLogitsWarper(LogitsProcessor):
    r"""
    [`LogitsProcessor`] that performs eta-sampling, a technique to filter out tokens with probabilities below a dynamic
    cutoff value, `eta`, which is calculated based on a combination of the hyperparameter `epsilon` and the entropy of
    the token probabilities, i.e. `eta := min(epsilon, sqrt(epsilon * e^-entropy(probabilities)))`. Takes the largest
    min_tokens_to_keep tokens if no tokens satisfy this constraint. It addresses the issue of poor quality in long
    samples of text generated by neural language models leading to more coherent and fluent text. See [Truncation
    Sampling as Language Model Desmoothing](https://arxiv.org/abs/2210.15191) for more information. Note: `do_sample`
    must be set to `True` for this `LogitsProcessor` to work.


    Args:
        epsilon (`float`):
            A float value in the range (0, 1). Hyperparameter used to calculate the dynamic cutoff value, `eta`. The
            suggested values from the paper ranges from 3e-4 to 4e-3 depending on the size of the model.
        filter_value (`float`, *optional*, defaults to -inf):
            All values that are found to be below the dynamic cutoff value, `eta`, are set to this float value. This
            parameter is useful when logits need to be modified for very low probability tokens that should be excluded
            from generation entirely.
        min_tokens_to_keep (`int`, *optional*, defaults to 1):
            Specifies the minimum number of tokens that must be kept for generation, regardless of their probabilities.
            For example, if `min_tokens_to_keep` is set to 1, at least one token will always be kept for generation,
            even if all tokens have probabilities below the cutoff `eta`.
        device (`str`, *optional*, defaults to `"cpu"`):
            The device to allocate the tensors.

    Examples:
    ```python
    >>> from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed

    >>> set_seed(1)
    >>> model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2")
    >>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2")

    >>> inputs = tokenizer("A sequence: 1, 2", return_tensors="pt")

    >>> # With sampling, the output is unexpected -- sometimes too unexpected.
    >>> outputs = model.generate(**inputs, do_sample=True)
    >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
    A sequence: 1, 2, 3 | < 4 (left-hand pointer) ;
    <BLANKLINE>
    <BLANKLINE>

    >>> # With eta sampling, the output gets restricted to high-probability tokens. You can see it as a dynamic form of
    >>> # epsilon sampling that adapts its cutoff probability based on the entropy (high entropy = lower cutoff).
    >>> # Pro tip: The paper recomends using `eta_cutoff` values between 3e-4 to 4e-3
    >>> outputs = model.generate(**inputs, do_sample=True, eta_cutoff=0.1)
    >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
    A sequence: 1, 2, 3, 4, 5, 6, 7, 8, 9
    ```
    """

    def __init__(
        self, epsilon: float, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1, device: str = "cpu"
    ):
        epsilon = float(epsilon)
        if epsilon <= 0 or epsilon >= 1:
            raise ValueError(f"`eta_cutoff` has to be a float > 0 and < 1, but is {epsilon}")

        min_tokens_to_keep = int(min_tokens_to_keep)
        if min_tokens_to_keep < 1:
            raise ValueError(
                f"`min_tokens_to_keep` has to be a strictly positive integer, but is {min_tokens_to_keep}"
            )

        self.epsilon = torch.tensor(epsilon, device=device)
        self.filter_value = filter_value
        self.min_tokens_to_keep = min_tokens_to_keep

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        probabilities = scores.softmax(dim=-1)
        entropy = torch.distributions.Categorical(logits=scores).entropy()
        eta = torch.min(self.epsilon, torch.sqrt(self.epsilon) * torch.exp(-entropy))[..., None]
        indices_to_remove = probabilities < eta

        # Keep the words with the 'min_tokens_to_keep'-highest probabilities
        top_k = min(self.min_tokens_to_keep, scores.size(-1))  # Safety check
        indices_to_remove = indices_to_remove & (scores < torch.topk(scores, top_k)[0][..., -1, None])

        scores_processed = scores.masked_fill(indices_to_remove, self.filter_value)


SUMMARY:
**Class Summary: EtaLogitsWarper**

**Purpose and Functionality:** The `EtaLogitsWarper` class is a variant of the `LogitsProcessor` that applies eta-sampling to filter out tokens with probabilities below a dynamic cutoff value, calculated based on hyperparameter `epsilon` and entropy.

**Key Implementation Details:**

* Calculates `eta`: `min(epsilon, sqrt(epsilon * exp(-entropy)))`
* Removes tokens with probability below `eta` or lowest `k` top-probability tokens
* Sets filtered probabilities to a constant value (`filter_value`)

**Notable Dependencies and Requirements:**

* `torch` for tensor operations
* `torch.distributions.Categorical` for entropy calculation
* Hyperparameter `epsilon`: float in (0, 1) range

Example usage:
```python
_eta_logits_warper = EtaLogitsWarper(epsilon=0.5)
scores = ...
filtered_scores = eta_logits_warper(scores)
```
Note: The `do_sample` argument must be set to `True` for this class to work.


key: chroma:document
string_value:
CODE:

def _get_ngrams(ngram_size: int, prev_input_ids: torch.Tensor, num_hypos: int):
    """
    Assume ngram_size=2 and prev_input_ids=tensor([[40, 2883, 2712, 4346]]). The output of generated ngrams look like
    this {(40,): [2883], (2883,): [2712], (2712,): [4346]}.

    Args:
        ngram_size (`int`):
            The number sequential tokens taken as a group which may only occur once before being banned.
        prev_input_ids (`torch.Tensor`):
           Generated token ids for the current hypothesis.
        num_hypos (`int`):
            The number of hypotheses for which n-grams need to be generated.

    Returns:
        generated_ngrams (`dict`):
            Dictionary of generated ngrams.
    """
    # Initialize an empty list of dictionaries, one for each hypothesis (index) in the range of num_hypos
    generated_ngrams = [{} for _ in range(num_hypos)]
    for idx in range(num_hypos):
        gen_tokens = prev_input_ids[idx].tolist()
        generated_ngram = generated_ngrams[idx]
        # Loop through each n-gram of size ngram_size in the list of tokens (gen_tokens)
        for ngram in zip(*[gen_tokens[i:] for i in range(ngram_size)]):
            prev_ngram_tuple = tuple(ngram[:-1])
            generated_ngram[prev_ngram_tuple] = generated_ngram.get(prev_ngram_tuple, []) + [ngram[-1]]


SUMMARY:
**Summary**

 `_get_ngrams` function generates `n-grams` from previously generated token IDs.

**Purpose and Functionality**
 Creates a dictionary of generated `n-grams`, where each key is an `(n-1)` sequence tuple and the value is a list of tokens in that sequence, for `num_hypos` hypotheses.

**Key Implementation Details**

* Iterates through sequential tokens using `zip(*[gen_tokens[i:] for i in range(ngram_size)])`
* Updates dictionary values with new token additions to existing `_ngram_tuple`

**Dependencies and Requirements**

* `torch.Tensor` (input IDs)
* `int`: `ngram_size`, `num_hypos`
* Reference classes: `torch.Tensor`

Note: This summary focuses on the main purpose, key implementation details, and dependencies. It is concise and follows the required format.


key: chroma:document
string_value:
CODE:

def _get_generated_ngrams(banned_ngrams, prev_input_ids, ngram_size, cur_len):
    """
    Determines the banned tokens for the current hypothesis based on previously generated n-grams.

    Args:
        banned_ngrams (`dict`):
            A dictionary containing previously generated n-grams for each hypothesis.
        prev_input_ids (`torch.Tensor`):
            Generated token ids for the current hypothesis.
        ngram_size (`int`):
            The number sequential tokens taken as a group which may only occur once before being banned.
        cur_len (`int`):
            The current length of the token sequences for which the n-grams are being checked.

    Returns:
        List of tokens that are banned.
    """
    # Before decoding the next token, prevent decoding of ngrams that have already appeared
    start_idx = cur_len + 1 - ngram_size
    ngram_idx = tuple(prev_input_ids[start_idx:cur_len].tolist())


SUMMARY:
**Summary**

The `_get_generated_ngrams` function determines banned tokens for the current hypothesis based on previously generated n-grams. It takes as input:

*   `banned_ngrams`: a dictionary of previously generated n-grams
*   `prev_input_ids`: token ids for the current hypothesis
*   `ngram_size`: number of sequential tokens for group analysis
*   `cur_len`: length of token sequences for analysis

The function returns a list of banned tokens based on matching n-gram indices with existing banned n-grams.

**Key Implementation Details**

*   Uses PyTorch tensor for input IDs (`torch.Tensor prev_input_ids`)
*   Utilizes tuple indexing for n-gram creation and lookup (`ngram_idx = tuple(prev_input_ids[start_idx:cur_len].tolist())`)
*   Leverages dictionary lookups for retrieving banned tokens (with default return value of [])

**Notable Dependencies or Requirements**

*   `torch` library for tensor operations
*   Dictionary data structure (Python built-in)
*   Tuple creation and indexing


key: chroma:document
string_value:
CODE:

def _calc_banned_ngram_tokens(
    ngram_size: int, prev_input_ids: torch.Tensor, num_hypos: int, cur_len: int
) -> List[Iterable[int]]:
    """Copied from fairseq for no_repeat_ngram in beam_search"""
    if cur_len + 1 < ngram_size:
        # return no banned tokens if we haven't generated no_repeat_ngram_size tokens yet
        return [[] for _ in range(num_hypos)]
    generated_ngrams = _get_ngrams(ngram_size, prev_input_ids, num_hypos)
    banned_tokens = [
        _get_generated_ngrams(generated_ngrams[hypo_idx], prev_input_ids[hypo_idx], ngram_size, cur_len)
        for hypo_idx in range(num_hypos)
    ]


SUMMARY:
**Summary:**

 Calculates banned (non-repeating) N-gram tokens for beam search with `no_repeat_ngram`.

**Main Purpose and Functionality:**
The `_calc_banned_ngram_tokens` function generates list of iterative lists containing banned tokens based on the previous input IDs and their size, ensuring no repeat tokens in beam search algorithms (`ngram_size`).

**Key Implementation Details:**

* Utilizes `_get_ngrams` to generate n-grams from previous input IDs
* Creates a sequence of `_get_generated_ngrams` for each hypothesis
* Returns a list of lists containing banned tokens

**Notable Dependencies and Requirements:**
- `torch.Tensor` (input IDs)
- `_get_ngrams` function
- `_get_generated_ngrams` function
- `no_repeat_ngram` parameter

Note: The summary preserves class/function/variable names exactly, uses precise technical terms, and states the primary purpose.


key: chroma:document
string_value:
CODE:

class NoRepeatNGramLogitsProcessor(LogitsProcessor):
    r"""
    N-grams are groups of "n" consecutive words, characters, or tokens taken from a sequence of text. Given the
    sentence: "She runs fast", the bi-grams (n=2) would be ("she", "runs") and ("runs", "fast"). In text generation,
    avoiding repetitions of word sequences provides a more diverse output. This [`LogitsProcessor`] enforces no
    repetition of n-grams by setting the scores of banned tokens to negative infinity which eliminates those tokens
    from consideration when further processing the scores. Note that, for decoder-only models like most LLMs, the
    prompt is also considered to obtain the n-grams.
    [Fairseq](https://github.com/pytorch/fairseq/blob/a07cb6f40480928c9e0548b737aadd36ee66ac76/fairseq/sequence_generator.py#L345).

    <Tip>

    Use n-gram penalties with care. For instance, penalizing 2-grams (bigrams) in an article about the city of New York
    might lead to undesirable outcomes where the city's name appears only once in the entire text.
    [Reference](https://huggingface.co/blog/how-to-generate)

    </Tip>

    Args:
        ngram_size (`int`):
            All ngrams of size `ngram_size` can only occur once.

    Examples:

    ```py
    >>> from transformers import AutoTokenizer, AutoModelForCausalLM

    >>> model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2")
    >>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2")
    >>> inputs = tokenizer(["Today I"], return_tensors="pt")

    >>> output = model.generate(**inputs)
    >>> print(tokenizer.decode(output[0], skip_special_tokens=True))
    Today I’m not sure if I’m going to be able to do it.

    >>> # Now let's add ngram size using `no_repeat_ngram_size`. This stops the repetitions ("I’m") in the output.
    >>> output = model.generate(**inputs, no_repeat_ngram_size=2)
    >>> print(tokenizer.decode(output[0], skip_special_tokens=True))
    Today I’m not sure if I can get a better understanding of the nature of this issue
    ```
    """

    def __init__(self, ngram_size: int):
        if not isinstance(ngram_size, int) or ngram_size <= 0:
            raise ValueError(f"`ngram_size` has to be a strictly positive integer, but is {ngram_size}")
        self.ngram_size = ngram_size

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        num_batch_hypotheses = scores.shape[0]
        cur_len = input_ids.shape[-1]
        scores_processed = scores.clone()
        banned_batch_tokens = _calc_banned_ngram_tokens(self.ngram_size, input_ids, num_batch_hypotheses, cur_len)
        for i, banned_tokens in enumerate(banned_batch_tokens):
            scores_processed[i, banned_tokens] = -float("inf")



SUMMARY:
**Summary:**

`NoRepeatNGramLogitsProcessor` is a `LogitsProcessor` class that enforces no repetition of n-grams in text generation by setting the scores of banned tokens to negative infinity. This prevents undesirable outcomes of repeated word sequences.

**Key Implementation Details:**

* The class takes an `ngram_size` as input and raises a `ValueError` if it's not a strictly positive integer.
* It processes input `input_ids` and `scores` tensors to calculate banned n-grams using `_calc_banned_ngram_tokens`.
* For each banned token, it sets the corresponding score in `scores_processed` to negative infinity.
* The class returns the processed scores tensor.

**Notable Dependencies:**

* `LogitsProcessor`: A base class for logit processors that takes input and score tensors.
* `_calc_banned_ngram_tokens`: A function that calculates banned n-grams based on `ngram_size`, `input_ids`, `num_batch_hypotheses`, and `cur_len`.
* `torch` library is required for tensor operations.

**Interfaces:**

* `LogitsProcessor` interface: This class implements the `__call__` method to process input tensors.
* `_calc_banned_ngram_tokens` function: Not explicitly documented, but it's an important dependency that calculates banned n-grams.


key: chroma:document
string_value:
CODE:

class EncoderNoRepeatNGramLogitsProcessor(LogitsProcessor):
    r"""
    [`LogitsProcessor`] that works similarly to [`NoRepeatNGramLogitsProcessor`], but applied exclusively to prevent
    the repetition of n-grams present in the prompt.

    It was designed to promote chattiness in a language model, by preventing the generation of n-grams present in
    previous conversation rounds.

    Args:
        encoder_ngram_size (`int`):
            All ngrams of size `ngram_size` can only occur within the encoder input ids.
        encoder_input_ids (`int`):
            The encoder_input_ids that should not be repeated within the decoder ids.

    Examples:

    ```py
    >>> from transformers import AutoTokenizer, AutoModelForCausalLM

    >>> model = AutoModelForCausalLM.from_pretrained("bigscience/bloomz-560m")
    >>> tokenizer = AutoTokenizer.from_pretrained("bigscience/bloomz-560m")

    >>> inputs = tokenizer("Alice: I love cats. What do you love?\nBob:", return_tensors="pt")

    >>> # With greedy decoding, we see Bob repeating Alice's opinion. If Bob was a chatbot, it would be a poor one.
    >>> outputs = model.generate(**inputs)
    >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
    Alice: I love cats. What do you love?
    Bob: I love cats. What do you

    >>> # With this logits processor, we can prevent Bob from repeating Alice's opinion.
    >>> outputs = model.generate(**inputs, encoder_no_repeat_ngram_size=2)
    >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
    Alice: I love cats. What do you love?
    Bob: My cats are very cute.
    ```
    """

    def __init__(self, encoder_ngram_size: int, encoder_input_ids: torch.LongTensor):
        if not isinstance(encoder_ngram_size, int) or encoder_ngram_size <= 0:
            raise ValueError(
                f"`encoder_ngram_size` has to be a strictly positive integer, but is {encoder_ngram_size}"
            )
        self.ngram_size = encoder_ngram_size
        if len(encoder_input_ids.shape) == 1:
            encoder_input_ids = encoder_input_ids.unsqueeze(0)
        self.batch_size = encoder_input_ids.shape[0]
        self.generated_ngrams = _get_ngrams(encoder_ngram_size, encoder_input_ids, self.batch_size)

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        # B x num_beams
        num_hypos = scores.shape[0]
        num_beams = num_hypos // self.batch_size
        cur_len = input_ids.shape[-1]
        scores_processed = scores.clone()
        banned_batch_tokens = [
            _get_generated_ngrams(
                self.generated_ngrams[hypo_idx // num_beams], input_ids[hypo_idx], self.ngram_size, cur_len
            )
            for hypo_idx in range(num_hypos)
        ]

        for i, banned_tokens in enumerate(banned_batch_tokens):
            scores_processed[i, banned_tokens] = -float("inf")



SUMMARY:
**Class Summary: EncoderNoRepeatNGramLogitsProcessor**

### Purpose and Functionality
This class extends the `LogitsProcessor` class to prevent repetition of n-grams present in the prompt. It aims to promote chattiness in language models by preventing generation of repeated n-grams.

### Key Implementation Details

* Uses `_get_ngrams` to generate all n-grams of size `ngram_size` from `encoder_input_ids`.
* Applies a logits processor to scores, banning tokens that match previously generated n-grams.
* Returns processed scores with banned tokens replaced by `-inf`.

### Notable Dependencies and Requirements
* `LogitsProcessor`
* `_get_ngrams`
* `torch`

### Method Signatures

* `__call__(input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor`: Processed scores after applying n-gram filtering.

### Attributes

* `ngram_size` (int): Size of n-grams to filter.
* `batch_size` (int): Batch size for processing.
* `generated_ngrams` (list): Pre-generated n-grams for fast lookups.


key: chroma:document
string_value:
CODE:

class SequenceBiasLogitsProcessor(LogitsProcessor):
    """
    [`LogitsProcessor`] that applies an additive bias on sequences. The bias is applied to the last token of a sequence
    when the next generated token can complete it. Consequently, to take the most of biasing sequences with more than
    one token, consider using beam methods (to gracefully work around partially completed sequences that have a
    negative bias) and applying the bias to their prefixes (to ensure the bias is applied earlier).

    <Tip>

    In order to get the token ids of the sequences that you want to bias, make sure to set `add_prefix_space=True` when
    initializing the tokenizer, and use `tokenizer(bad_words, add_special_tokens=False).input_ids`. The
    `add_prefix_space` argument is only supported for some slow tokenizers, as fast tokenizers' prefixing behaviours
    come from `pre tokenizers`. Read more [here](https://huggingface.co/docs/tokenizers/api/pre-tokenizers).

    </Tip>

    Args:
        sequence_bias (`List[List[Union[List[int], float]]]`):
            List of lists that maps a sequence of tokens to its bias term (e.g. `[[[10, 45], -2.0],
            [[64], -7.5]]`). Positive biases increase the odds of the
            sequence being selected, while negative biases do the opposite. If a sequence has a length of 1, its bias
            will always be applied. Otherwise, the bias will only be applied if the sequence in question is about to be
            completed (in the token selection step after this processor is applied).

    Examples:

    ```python
    >>> from transformers import AutoTokenizer, AutoModelForCausalLM

    >>> model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
    >>> tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
    >>> inputs = tokenizer(["The full name of Donald is Donald"], return_tensors="pt")

    >>> summary_ids = model.generate(inputs["input_ids"], max_new_tokens=4)
    >>> print(tokenizer.batch_decode(summary_ids, skip_special_tokens=True)[0])
    The full name of Donald is Donald J. Trump Jr

    >>> # Now let's control generation through a bias. Please note that the tokenizer is initialized differently!
    >>> tokenizer_with_prefix_space = AutoTokenizer.from_pretrained("openai-community/gpt2", add_prefix_space=True)


    >>> def get_tokens(word):
    ...     return tokenizer_with_prefix_space([word], add_special_tokens=False).input_ids[0]


    >>> # If we add a negative bias without beam search, it may become "stuck" in a prefix without good continuations
    >>> sequence_bias = [get_tokens("Trump"), -10.0]
    >>> biased_ids = model.generate(inputs["input_ids"], max_new_tokens=4, sequence_bias=sequence_bias)
    >>> print(tokenizer.batch_decode(biased_ids, skip_special_tokens=True)[0])
    The full name of Donald is Donald J. Donald,

    >>> biased_ids = model.generate(inputs["input_ids"], max_new_tokens=4, num_beams=4, sequence_bias=sequence_bias)
    >>> print(tokenizer.batch_decode(biased_ids, skip_special_tokens=True)[0])
    The full name of Donald is Donald Rumsfeld,

    >>> # We can also add a positive bias to nudge the model towards specific tokens or continuations
    >>> sequence_bias = [get_tokens("Donald Duck"), 10.0]
    >>> biased_ids = model.generate(inputs["input_ids"], max_new_tokens=4, num_beams=4, sequence_bias=sequence_bias)
    >>> print(tokenizer.batch_decode(biased_ids, skip_special_tokens=True)[0])
    The full name of Donald is Donald Duck.
    ```
    """

    def __init__(self, sequence_bias: List[List[Union[List[int], float]]]):
        self.sequence_bias = sequence_bias
        self._validate_arguments()
        self._convert_list_arguments_into_dict()

        # Bias variables that will be populated on the first call (for retrocompatibility purposes, the vocabulary size
        # is infered in the first usage, which inhibits initializing here)
        self.length_1_bias = None
        self.prepared_bias_variables = False

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        # 1 - Prepares the bias tensors. This is only needed the first time the logit processor is called.
        if not self.prepared_bias_variables:
            self._prepare_bias_variables(scores)

        # 2 - prepares an empty bias to add
        bias = torch.zeros_like(scores)

        # 3 - include the bias from length = 1
        bias += self.length_1_bias

        # 4 - include the bias from length > 1, after determining which biased sequences may be completed.
        for sequence_ids, sequence_bias in self.sequence_bias.items():
            if len(sequence_ids) == 1:  # the sequence is of length 1, already applied
                continue
            if len(sequence_ids) > input_ids.shape[1]:  # the sequence is longer than the context, ignore
                continue
            prefix_length = len(sequence_ids) - 1
            last_token = sequence_ids[-1]
            matching_rows = torch.eq(
                input_ids[:, -prefix_length:],
                torch.tensor(sequence_ids[:-1], dtype=input_ids.dtype, device=input_ids.device),
            ).prod(dim=1)
            bias[:, last_token] += torch.where(
                matching_rows.bool(),
                torch.tensor(sequence_bias, device=input_ids.device),
                torch.tensor(0.0, device=input_ids.device),
            )

        # 5 - apply the bias to the scores
        scores_processed = scores + bias
        return scores_processed

    def _prepare_bias_variables(self, scores: torch.FloatTensor):
        vocabulary_size = scores.shape[-1]

        # Check biased tokens out of bounds
        invalid_biases = []
        for sequence_ids in self.sequence_bias:
            for token_id in sequence_ids:
                if token_id >= vocabulary_size:
                    invalid_biases.append(token_id)
        if len(invalid_biases) > 0:
            raise ValueError(
                f"The model vocabulary size is {vocabulary_size}, but the following tokens were being biased: "
                f"{invalid_biases}"
            )

        # Precompute the bias tensors to be applied. Sequences of length 1 are kept separately, as they can be applied
        # with simpler logic.
        self.length_1_bias = torch.zeros((vocabulary_size,), dtype=torch.float).to(scores.device)
        for sequence_ids, bias in self.sequence_bias.items():
            if len(sequence_ids) == 1:
                self.length_1_bias[sequence_ids[-1]] = bias

        self.prepared_bias_variables = True

    def _validate_arguments(self):
        sequence_bias = self.sequence_bias
        if not isinstance(sequence_bias, dict) and not isinstance(sequence_bias, list) or len(sequence_bias) == 0:
            raise ValueError(
                f"`sequence_bias` has to be a non-empty dictionary, or non-empty list of lists but is {sequence_bias}."
            )
        if isinstance(sequence_bias, dict) and any(
            not isinstance(sequence_ids, tuple) for sequence_ids in sequence_bias.keys()
        ):
            raise ValueError(f"`sequence_bias` has to be a dict with tuples as keys, but is {sequence_bias}.")
        if isinstance(sequence_bias, dict) and any(
            any((not isinstance(token_id, (int, np.integer)) or token_id < 0) for token_id in sequence_ids)
            or len(sequence_ids) == 0
            for sequence_ids in sequence_bias.keys()
        ):
            raise ValueError(
                f"Each key in `sequence_bias` has to be a non-empty tuple of positive integers, but is "
                f"{sequence_bias}."
            )

        def all_token_bias_pairs_are_valid(sequence):
            return (
                isinstance(sequence[0], list)
                and all(isinstance(token_id, (int, np.integer)) and token_id > 0 for token_id in sequence[0])
                and isinstance(sequence[1], float)
            )

        if isinstance(sequence_bias, list) and any(
            (not all_token_bias_pairs_are_valid(sequence)) or len(sequence) == 0 for sequence in sequence_bias
        ):
            raise ValueError(
                f"Each element in `sequence_bias` has to be a non-empty list of lists of positive integers and float, but is "
                f"{sequence_bias}."
            )
        if isinstance(sequence_bias, dict) and any(not isinstance(bias, float) for bias in sequence_bias.values()):
            raise ValueError(f"`sequence_bias` has to be a dict with floats as values, but is {sequence_bias}.")

    def _convert_list_arguments_into_dict(self):
        """BC: we used to accept `dict{tuple of tokens: float}` directly, now we expect a list"""
        if isinstance(self.sequence_bias, list):
            temp_sequence = self.sequence_bias


SUMMARY:
The provided class code chunk defines `SequenceBiasLogitsProcessor`, a custom logits processor that applies an additive bias to sequences. The bias is applied to the last token of a sequence when it can complete the sequence.

Key functions and features:

1.  **Initialization**: The class initializes with a list or dictionary containing sequences of tokens (e.g., tuples of integers) paired with biased floats. Sequences are ignored if they would be longer than the input context.

2.  **Bias Application**: A bias is applied to each element in scores where `sequence_ids` match an input element at position one before end

3.  **Data Validation**: The class includes checks to ensure that sequences of tokens are valid, meaning positive integers and a float, and that both sequence-based data and list-of-lists representation can be used correctly.

4.  **BC (Backward Compatibility)**: If using a dictionary, it is now converted from list format for backward compatibility purposes.


key: chroma:document
string_value:
CODE:

class NoBadWordsLogitsProcessor(SequenceBiasLogitsProcessor):
    """
    [`LogitsProcessor`] that enforces that specified sequences will never be selected.

    <Tip>

    In order to get the token ids of the words that should not appear in the generated text, make sure to set
    `add_prefix_space=True` when initializing the tokenizer, and use `tokenizer(bad_words,
    add_special_tokens=False).input_ids`. The `add_prefix_space` argument is only supported for some slow tokenizers,
    as fast tokenizers' prefixing behaviours come from `pre tokenizers`. Read more
    [here](https://huggingface.co/docs/tokenizers/api/pre-tokenizers).

    </Tip>

    Args:
        bad_words_ids (`List[List[int]]`):
            List of list of token ids that are not allowed to be generated.
        eos_token_id (`Union[int, List[int], torch.Tensor]`, *optional*):
            The id(s) of the *end-of-sequence* token.

    Examples:

    ```python
    >>> from transformers import AutoTokenizer, AutoModelForCausalLM

    >>> model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
    >>> tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
    >>> inputs = tokenizer(["In a word, the cake is a"], return_tensors="pt")

    >>> output_ids = model.generate(inputs["input_ids"], max_new_tokens=5, pad_token_id=tokenizer.eos_token_id)
    >>> print(tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0])
    In a word, the cake is a bit of a mess.

    >>> # Now let's take the bad words out. Please note that the tokenizer is initialized differently
    >>> tokenizer_with_prefix_space = AutoTokenizer.from_pretrained("openai-community/gpt2", add_prefix_space=True)


    >>> def get_tokens_as_list(word_list):
    ...     "Converts a sequence of words into a list of tokens"
    ...     tokens_list = []
    ...     for word in word_list:
    ...         tokenized_word = tokenizer_with_prefix_space([word], add_special_tokens=False).input_ids[0]
    ...         tokens_list.append(tokenized_word)
    ...     return tokens_list


    >>> bad_words_ids = get_tokens_as_list(word_list=["mess"])
    >>> output_ids = model.generate(
    ...     inputs["input_ids"], max_new_tokens=5, bad_words_ids=bad_words_ids, pad_token_id=tokenizer.eos_token_id
    ... )
    >>> print(tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0])
    In a word, the cake is a bit of a surprise.
    ```
    """

    def __init__(
        self, bad_words_ids: List[List[int]], eos_token_id: Optional[Union[int, List[int], torch.Tensor]] = None
    ):
        self.bad_word_ids = bad_words_ids
        self._validate_arguments()

        # Filter EOS token from bad_words_ids
        if eos_token_id is not None:
            if not isinstance(eos_token_id, torch.Tensor):
                if isinstance(eos_token_id, int):
                    eos_token_id = [eos_token_id]
                eos_token_id = torch.tensor(eos_token_id)

            bad_words_ids = list(
                filter(lambda bad_token_seq: all(bad_token_seq != [i] for i in eos_token_id), bad_words_ids)
            )

        # Forbidding a sequence is equivalent to setting its bias to -inf
        sequence_bias = {tuple(sequence): float("-inf") for sequence in bad_words_ids}
        super().__init__(sequence_bias=sequence_bias)

    def _validate_arguments(self):
        bad_words_ids = self.bad_word_ids
        if not isinstance(bad_words_ids, list) or len(bad_words_ids) == 0:
            raise ValueError(f"`bad_words_ids` has to be a non-empty list, but is {bad_words_ids}.")
        if any(not isinstance(bad_word_ids, list) for bad_word_ids in bad_words_ids):
            raise ValueError(f"`bad_words_ids` has to be a list of lists, but is {bad_words_ids}.")
        if any(
            any((not isinstance(token_id, (int, np.integer)) or token_id < 0) for token_id in bad_word_ids)
            for bad_word_ids in bad_words_ids
        ):
            raise ValueError(
                f"Each list in `bad_words_ids` has to be a list of positive integers, but is {bad_words_ids}."


SUMMARY:
**Class Summary: NoBadWordsLogitsProcessor**

### Main Purpose and Functionality
The `NoBadWordsLogitsProcessor` class is a custom logit processor that prevents specified sequences from being selected during text generation. It filters out token IDs of specified "bad words" and sets their bias to -inf, effectively forbidding their existence in the generated output.

### Key Implementation Details

*   Initializes with `bad_words_ids`, a list of lists containing token IDs of forbidden words.
*   Validates input arguments, raising errors if `bad_words_ids` is empty or contains non-integer values.
*   Applies prefixing to tokenizer (if required) and filters out end-of-sequence tokens from `bad_words_ids`.
*   Sets bias for each bad word sequence to -inf.

### Notable Dependencies or Requirements

*   `tokenizer`: an instance of a tokenization model with support for `add_prefix_space` feature.
*   `torch.Tensor`: used to represent EOS token IDs for filtering purposes.

### Example Usage

```python
model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2", add_prefix_space=True)
inputs = tokenizer(["In a word, the cake is a"], return_tensors="pt")

output_ids = model.generate(inputs["input_ids"], max_new_tokens=5, pad_token_id=tokenizer.eos_token_id,
                            bad_words_ids=[get_tokens_as_list({"mess"})], eos_token_id=tokenizer.eos_token_id)
print(tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0])
```

In this example, `NoBadWordsLogitsProcessor` is passed the list of forbidden words `["mess"]`, which prevents its occurrence in the generated output.


key: chroma:document
string_value:
CODE:

class PrefixConstrainedLogitsProcessor(LogitsProcessor):
    r"""
    [`LogitsProcessor`] that enforces constrained generation and is useful for prefix-conditioned constrained
    generation. See [Autoregressive Entity Retrieval](https://arxiv.org/abs/2010.00904) for more information.

    Args:
        prefix_allowed_tokens_fn (`Callable[[int, torch.Tensor], List[int]]`):
            This function constraints the beam search to allowed tokens only at each step. This function takes 2
            arguments `inputs_ids` and the batch ID `batch_id`. It has to return a list with the allowed tokens for the
            next generation step conditioned on the previously generated tokens `inputs_ids` and the batch ID
            `batch_id`.

    Examples:

    ```py
    >>> from transformers import AutoTokenizer, AutoModelForCausalLM

    >>> model = AutoModelForCausalLM.from_pretrained("bigscience/bloomz-560m")
    >>> tokenizer = AutoTokenizer.from_pretrained("bigscience/bloomz-560m")

    >>> inputs = tokenizer("Alice and Bob", return_tensors="pt")

    >>> # By default, it continues generating according to the model's logits
    >>> outputs = model.generate(**inputs, max_new_tokens=5)
    >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
    Alice and Bob are friends

    >>> # We can contrain it with `prefix_allowed_tokens_fn` to force a certain behavior based on a prefix.
    >>> # For instance, we can force an entire entity to be generated when its beginning is detected.
    >>> entity = tokenizer(" Bob Marley", return_tensors="pt").input_ids[0]  # 3 tokens
    >>> def prefix_allowed_tokens_fn(batch_id, input_ids):
    ...     '''
    ...     Attempts to generate 'Bob Marley' when 'Bob' is detected.
    ...     In this case, `batch_id` is not used, but you can set rules for each batch member.
    ...     '''
    ...     if input_ids[-1] == entity[0]:
    ...         return [entity[1].item()]
    ...     elif input_ids[-2] == entity[0] and input_ids[-1] == entity[1]:
    ...         return [entity[2].item()]
    ...     return list(range(tokenizer.vocab_size))  # If no match, allow all tokens

    >>> outputs = model.generate(**inputs, max_new_tokens=5, prefix_allowed_tokens_fn=prefix_allowed_tokens_fn)
    >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
    Alice and Bob Marley
    ```
    """

    def __init__(self, prefix_allowed_tokens_fn: Callable[[int, torch.Tensor], List[int]], num_beams: int):
        self._prefix_allowed_tokens_fn = prefix_allowed_tokens_fn
        self._num_beams = num_beams

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        mask = torch.full_like(scores, -math.inf)
        for batch_id, beam_sent in enumerate(input_ids.view(-1, self._num_beams, input_ids.shape[-1])):
            for beam_id, sent in enumerate(beam_sent):
                prefix_allowed_tokens = self._prefix_allowed_tokens_fn(batch_id, sent)
                if len(prefix_allowed_tokens) == 0:
                    raise ValueError(
                        f"`prefix_allowed_tokens_fn` returned an empty list for batch ID {batch_id}."
                        f"This means that the constraint is unsatisfiable. Please check your implementation"
                        f"of `prefix_allowed_tokens_fn` "
                    )
                mask[batch_id * self._num_beams + beam_id, prefix_allowed_tokens] = 0

        scores_processed = scores + mask


SUMMARY:
**PrefixConstrainedLogitsProcessor Class Summary**

This class (`PrefixConstrainedLogitsProcessor`) is a subtype of `LogitsProcessor` designed for prefix-conditioned constrained generation. Its primary purpose is to enforce a specific token sequence constraint, as defined by the user-provided `prefix_allowed_tokens_fn` function.

**Key Implementation Details:**

* The `_prefix_allowed_tokens_fn` function takes `inputs_ids` and `batch_id` as inputs and returns a list of allowed tokens for each step.
* A mask is created to filter out allowed tokens from the input scores using the `_prefix_allowed_tokens_fn`.
* Scores are updated with the mask, discarding any scores below -inf.

**Notable Dependencies/Requirements:**

* `Callable[[int, torch.Tensor], List[int]]`: The user-provided `prefix_allowed_tokens_fn` function.
* `torch.FloatTensor` and `torch.LongTensor`: Used to represent input scores and IDs, respectively.

**Internal Methods:**

* `__init__`: Initializes the processor with a `_prefix_allowed_tokens_fn` function and number of beams (`_num_beams`).
* `__call__`: Processes input scores, applying the constraint defined by `_prefix_allowed_tokens_fn`.


key: chroma:document
string_value:
CODE:

class HammingDiversityLogitsProcessor(LogitsProcessor):
    r"""
    [`LogitsProcessor`] that enforces diverse beam search.

    Note that this logits processor is only effective for [`PreTrainedModel.group_beam_search`]. See [Diverse Beam
    Search: Decoding Diverse Solutions from Neural Sequence Models](https://arxiv.org/pdf/1610.02424.pdf) for more
    details.

    Traditional beam search often generates very similar sequences across different beams.
    `HammingDiversityLogitsProcessor` addresses this by penalizing beams that generate tokens already chosen by other
    beams in the same time step.

    Args:
        diversity_penalty (`float`):
            This value is subtracted from a beam's score if it generates a token same as any beam from other group at a
            particular time. A higher `diversity_penalty` will enforce greater diversity among the beams. Adjusting
            this value can help strike a balance between diversity and natural likelihood.
        num_beams (`int`):
            Number of beams for beam search. 1 means no beam search.
        num_beam_groups (`int`):
            Number of groups to divide `num_beams` into in order to ensure diversity among different groups of beams.
            [this paper](https://arxiv.org/pdf/1610.02424.pdf) for more details.

    Examples:

    ```python
    >>> from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
    >>> import torch

    >>> # Initialize the model and tokenizer
    >>> tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-base")
    >>> model = AutoModelForSeq2SeqLM.from_pretrained("google-t5/t5-base")

    >>> # A long text about the solar system
    >>> text = (
    ...     "The Solar System is a gravitationally bound system comprising the Sun and the objects that orbit it, "
    ...     "either directly or indirectly. Of the objects that orbit the Sun directly, the largest are the eight "
    ...     "planets, with the remainder being smaller objects, such as the five dwarf planets and small Solar System "
    ...     "bodies. The Solar System formed 4.6 billion years ago from the gravitational collapse of a giant "
    ...     "interstellar molecular cloud."
    ... )
    >>> inputs = tokenizer("summarize: " + text, return_tensors="pt")

    >>> # Generate diverse summary
    >>> outputs_diverse = model.generate(
    ...     **inputs,
    ...     num_beam_groups=2,
    ...     diversity_penalty=10.0,
    ...     max_length=100,
    ...     num_beams=4,
    ...     num_return_sequences=2,
    ... )
    >>> summaries_diverse = tokenizer.batch_decode(outputs_diverse, skip_special_tokens=True)

    >>> # Generate non-diverse summary
    >>> outputs_non_diverse = model.generate(
    ...     **inputs,
    ...     max_length=100,
    ...     num_beams=4,
    ...     num_return_sequences=2,
    ... )
    >>> summary_non_diverse = tokenizer.batch_decode(outputs_non_diverse, skip_special_tokens=True)

    >>> # With `diversity_penalty`, the resulting beams are much more diverse
    >>> print(summary_non_diverse)
    ['the solar system formed 4.6 billion years ago from the collapse of a giant interstellar molecular cloud. of the objects that orbit the Sun directly, the largest are the eight planets.',
    'the Solar System formed 4.6 billion years ago from the collapse of a giant interstellar molecular cloud. of the objects that orbit the Sun directly, the largest are the eight planets.']

    >>> print(summaries_diverse)
    ['the solar system formed 4.6 billion years ago from the collapse of a giant interstellar molecular cloud. of the objects that orbit the Sun directly, the largest are the eight planets.',
    'the solar system formed 4.6 billion years ago from the collapse of a giant interstellar molecular cloud. of the objects that orbit the Sun directly, the largest are the eight planets. the rest of the objects are smaller objects, such as the five dwarf planets and small solar system bodies.']
    ```
    """

    def __init__(self, diversity_penalty: float, num_beams: int, num_beam_groups: int):
        if not isinstance(diversity_penalty, float) or (not diversity_penalty > 0.0):
            raise ValueError("`diversity_penalty` should be a float strictly larger than 0.")
        self._diversity_penalty = diversity_penalty
        if not isinstance(num_beams, int) or num_beams < 2:
            raise ValueError("`num_beams` should be an integer strictly larger than 1.")
        self._num_beams = num_beams
        if not isinstance(num_beam_groups, int) or num_beam_groups < 2:
            raise ValueError("`num_beam_groups` should be an integer strictly larger than 1.")
        if num_beam_groups > num_beams:
            raise ValueError("`beam_groups` has to be smaller or equal to `num_beams`.")
        self._num_sub_beams = num_beams // num_beam_groups

    def __call__(
        self,
        input_ids: torch.LongTensor,
        scores: torch.FloatTensor,
        current_tokens: torch.LongTensor,
        beam_group_idx: int,
    ) -> torch.FloatTensor:
        r"""
        Args:
            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
                Indices of input sequence tokens in the vocabulary. [What are input IDs?](../glossary#input-ids)
            scores (`torch.FloatTensor` of shape `(batch_size, config.vocab_size)`):
                Prediction scores of a language modeling head. These can be logits for each vocabulary when not using
                beam search or log softmax for each vocabulary token when using beam search
            current_tokens (`torch.LongTensor` of shape `(batch_size)`):
                Indices of input sequence tokens in the vocabulary, corresponding to the tokens selected by the other
                beam groups in the current generation step.
            beam_group_idx (`int`):
                The index of the beam group currently being processed.

        Return:
            `torch.FloatTensor` of shape `(batch_size, config.vocab_size)`:
                The processed prediction scores.
        """
        # hamming diversity: penalise using same token in current group which was used in previous groups at
        # the same time step
        batch_size = current_tokens.shape[0] // self._num_beams
        group_start_idx = beam_group_idx * self._num_sub_beams
        group_end_idx = min(group_start_idx + self._num_sub_beams, self._num_beams)
        group_size = group_end_idx - group_start_idx
        vocab_size = scores.shape[-1]

        if group_start_idx == 0:
            return scores

        scores_processed = scores.clone()
        for batch_idx in range(batch_size):
            # predicted tokens of last time step of previous groups
            previous_group_tokens = current_tokens[
                batch_idx * self._num_beams : batch_idx * self._num_beams + group_start_idx
            ]
            token_frequency = torch.bincount(previous_group_tokens, minlength=vocab_size).to(scores.device)
            scores_processed[batch_idx * group_size : (batch_idx + 1) * group_size] -= (
                self._diversity_penalty * token_frequency
            )



SUMMARY:
**Hamming Diversity Logits Processor**

This class implements a [`LogitsProcessor`] that enforces diverse beam search in sequence generation models. Its primary purpose is to address the issue of traditional beam search generating similar sequences across different beams.

### Key Features:

*   Enforces diversity among beam groups using `diversity_penalty`
*   Divides beams into groups (`num_beam_groups`) and calculates sub-beams' scores
*   Penalizes tokens used in previous groups at the same time step (`hamming_diversity`)

### Implementation Details:

*   Initialises with `diversity_penalty`, `num_beams`, and `num_beam_groups`
*   Processes prediction scores by subtracting divergence penalty for shared tokens across beam groups

### Notable Dependencies/Requirements:

*   (`LogitsProcessor`)
*   (`PreTrainedModel.group_beam_search`)

### Example Usage:

```python
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Initialize the model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-base")
model = AutoModelForSeq2SeqLM.from_pretrained("google-t5/t5-base")

# Generate diverse summary with Hamming Diversity Logits Processor
diversity_penalty = 10.0
num_beams = 4
num_beam_groups = 2

inputs = tokenizer("summarize: ", return_tensors="pt")
outputs_diverse = model(inputs['input_ids'], attention_mask=inputs['attention_mask'])
processed_scores = HammingDiversityLogitsProcessor(diversity_penalty, num_beams, num_beam_groups)(inputs['scores'], inputs['current_tokens'])

# Generate summary without Hamming Diversity Logits Processor
processes_without_hamming_diversity = model(inputs['input_ids'], attention_mask=inputs['attention_mask'])
```


key: chroma:document
string_value:
CODE:

class ForcedBOSTokenLogitsProcessor(LogitsProcessor):
    r"""
    [`LogitsProcessor`] that enforces the specified token as the first generated token. Used with encoder-decoder
    models.

    Args:
        bos_token_id (`int`):
            The id of the token to force as the first generated token.

    Examples:

    ```python
    >>> from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

    >>> model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
    >>> tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")

    >>> inputs = tokenizer("Translate from English to German: I love cats.", return_tensors="pt")

    >>> # By default, it continues generating according to the model's logits
    >>> outputs = model.generate(**inputs, max_new_tokens=10)
    >>> print(tokenizer.batch_decode(outputs)[0])
    <pad> Ich liebe Kitty.</s>

    >>> # We can use `forced_bos_token_id` to force the start of generation with an encoder-decoder model
    >>> # (including forcing it to end straight away with an EOS token)
    >>> outputs = model.generate(**inputs, max_new_tokens=10, forced_bos_token_id=tokenizer.eos_token_id)
    >>> print(tokenizer.batch_decode(outputs)[0])
    <pad></s>
    ```
    """

    def __init__(self, bos_token_id: int):
        self.bos_token_id = bos_token_id

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        cur_len = input_ids.shape[-1]
        scores_processed = scores
        if cur_len == 1:
            scores_processed = torch.full_like(scores, -math.inf)
            scores_processed[:, self.bos_token_id] = 0


SUMMARY:
**Class Summary: ForcedBOSTokenLogitsProcessor**

A custom LogitsProcessor class that ensures the first generated token is fixed to a specified BOS (beginning-of-sequence) token in encoder-decoder models.

**Key Functionality:**
- Takes an additional `bos_token_id` argument in its constructor, which specifies the ID of the token to enforce as the start of generation.
- Modifies the input scores by setting them to negative infinity (`-inf`) if the sequence length is 1, effectively freezing the first token. It then sets the score for the specified BOS token to zero.

**Notable Implementation Details:**

* Uses `torch.LongTensor` and `torch.FloatTensor` as input types.
* Implements a class method (`__call__`) that takes in `input_ids` and `scores`, returning modified scores.

**Dependencies/Requirements:**

* No external dependencies mentioned in the code snippet, but it is assumed to be part of a Transformers library setup.


key: chroma:document
string_value:
CODE:

class ForcedEOSTokenLogitsProcessor(LogitsProcessor):
    r"""
    [`LogitsProcessor`] that enforces the specified token as the last generated token when `max_length` is reached.

    Args:
        max_length (`int`):
            The maximum length of the sequence to be generated.
        eos_token_id (`Union[int, List[int], torch.Tensor]`):
            The id(s) of the *end-of-sequence* token.
        device (`str`, *optional*, defaults to `"cpu"`):
            The device to allocate the tensors.

    Examples:

    ```python
    >>> from transformers import AutoTokenizer, AutoModelForCausalLM

    >>> model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2")
    >>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2")

    >>> inputs = tokenizer("A sequence: 1, 2, 3", return_tensors="pt")

    >>> # By default, it continues generating according to the model's logits
    >>> outputs = model.generate(**inputs, max_new_tokens=10)
    >>> print(tokenizer.batch_decode(outputs)[0])
    A sequence: 1, 2, 3, 4, 5, 6, 7, 8

    >>> # `forced_eos_token_id` ensures the generation ends with a EOS token
    >>> outputs = model.generate(**inputs, max_new_tokens=10, forced_eos_token_id=tokenizer.eos_token_id)
    >>> print(tokenizer.batch_decode(outputs)[0])
    A sequence: 1, 2, 3, 4, 5, 6, 7,<|endoftext|>
    ```
    """

    def __init__(self, max_length: int, eos_token_id: Union[int, List[int], torch.Tensor], device: str = "cpu"):
        self.max_length = max_length

        if not isinstance(eos_token_id, torch.Tensor):
            if isinstance(eos_token_id, int):
                eos_token_id = [eos_token_id]
            eos_token_id = torch.tensor(eos_token_id, device=device)
        self.eos_token_id = eos_token_id

        if torch.is_floating_point(eos_token_id) or (eos_token_id < 0).any():
            raise ValueError(f"`eos_token_id` has to be a list of positive integers, but is {eos_token_id}")

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        cur_len = input_ids.shape[-1]
        scores_processed = scores
        if cur_len == self.max_length - 1:
            scores_processed = torch.full_like(scores, -math.inf)
            scores_processed[:, self.eos_token_id] = 0


SUMMARY:
**ForcedEOSTokenLogitsProcessor Class Summary**

### Purpose and Functionality

The `ForcedEOSTokenLogitsProcessor` is a custom class that extends the `[LogitsProcessor]` (Base) class to enforce a specified token as the last generated token when `max_length` is reached.

### Key Implementation Details

*   Sets the device for tensor allocation (`device= "cpu"` by default).
*   Applies a maximum length constraint on sequence generation.
*   Masks out scores corresponding to the end-of-sequence token using a special value of `-math.inf`.
*   Updates scores to reflect forced EOS token at the last position.

### Notable Dependencies or Requirements

*   Requires input IDs and logit tensor as inputs (`input_ids: torch.LongTensor`, `scores: torch.FloatTensor`).
*   Depends on custom functions like `add_start_docstrings()` for documentation update.
*   Uses external module imports such as `torch` and `add_start_docstrings()`.

### Interfaces

Extends the `[LogitsProcessor]` interface, which implies compatibility with a specific logging or processing system. Key methods are:

*   `__call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor`
    *   Accepts `input_ids` and `scores`, returning processed logit tensor (`torch.FloatTensor`).


key: chroma:document
string_value:
CODE:

class InfNanRemoveLogitsProcessor(LogitsProcessor):
    r"""
    [`LogitsProcessor`] that removes all `nan` and `inf` values to avoid the generation method to fail. Note that using
    the logits processor should only be used if necessary since it can slow down the generation method.

    This logits processor has no `generate` example, as there shouldn't be a correct combination of flags that warrants
    its use.
    """

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        # set all nan values to 0.0
        scores_processed = torch.where(scores != scores, 0.0, scores)

        # set all +/-inf values to max/min possible value
        scores_processed = torch.where(scores == float("inf"), torch.finfo(scores.dtype).max, scores_processed)
        scores_processed = torch.where(scores == -float("inf"), torch.finfo(scores.dtype).min, scores_processed)



SUMMARY:
**InfNanRemoveLogitsProcessor Class Summary**

* **Purpose**: Removes `nan` and `inf` values from input scores to prevent generation method failure.

**Implementation Details**

* Sets `nan` values to 0.0 using `torch.where`
* Sets `inf` values to max/min possible value based on data type using `torch.where`

**Primary Method Signature**
```python
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
```
**Notable Dependencies**

* `torch`: Used for tensor operations (e.g., `torch.where`, `torch.finfo`)
* `LOGITS_PROCESSOR_INPUTS_DOCSTRING`: Used for docstring purpose
* `InfNanRemoveLogitsProcessor` inherits from `LogitsProcessor`


key: chroma:document
string_value:
CODE:

class ExponentialDecayLengthPenalty(LogitsProcessor):
    r"""
    [`LogitsProcessor`] that exponentially increases the score of the `eos_token_id` after `start_index` has been
    reached. This allows generating shorter sequences without having a hard cutoff, allowing the `eos_token` to be
    predicted in a meaningful position.

    Args:
        exponential_decay_length_penalty (`tuple(int, float)`):
            This tuple shall consist of: `(start_index, decay_factor)` where `start_index` indicates where penalty
            starts and `decay_factor` represents the factor of exponential decay
        eos_token_id (`Union[int, List[int], torch.Tensor]`):
            The id(s) of the *end-of-sequence* token.
        input_ids_seq_length (`int`):
            The length of the input sequence.

    Examples:

    ```python
    >>> from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed

    >>> model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
    >>> tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")

    >>> text = "Just wanted to let you know, I"
    >>> inputs = tokenizer(text, return_tensors="pt")

    >>> # Let's consider that we want short sentences, so we limit `max_length=30`. However, we observe that the answer
    >>> # tends to end abruptly.
    >>> set_seed(1)
    >>> outputs = model.generate(**inputs, do_sample=True, temperature=0.9, max_length=30, pad_token_id=50256)
    >>> print(tokenizer.batch_decode(outputs)[0])
    Just wanted to let you know, I received a link to an ebook, the book How To Start A Social Network which was
    published in 2010. Although

    >>> # To promote the appearance of the EOS token at the right time, we add the `exponential_decay_length_penalty =
    >>> # (start_index, decay_factor)`. Instead of cutting at max_tokens, the output comes to an end before and usually
    >>> # with more meaning. What happens is that starting from `start_index` the EOS token score will be increased
    >>> # by `decay_factor` exponentially. However, if you set a high decay factor, you may also end up with abruptly
    >>> # ending sequences.
    >>> set_seed(1)
    >>> outputs = model.generate(
    ...     **inputs,
    ...     do_sample=True,
    ...     temperature=0.9,
    ...     max_length=30,
    ...     pad_token_id=50256,
    ...     exponential_decay_length_penalty=(15, 1.6),
    ... )
    >>> print(tokenizer.batch_decode(outputs)[0])
    Just wanted to let you know, I received a link to an ebook, the book How To Start A Social Network
    which<|endoftext|>

    >>> # With a small decay factor, you will have a higher chance of getting a meaningful sequence.
    >>> set_seed(1)
    >>> outputs = model.generate(
    ...     **inputs,
    ...     do_sample=True,
    ...     temperature=0.9,
    ...     max_length=30,
    ...     pad_token_id=50256,
    ...     exponential_decay_length_penalty=(15, 1.01),
    ... )
    >>> print(tokenizer.batch_decode(outputs)[0])
    Just wanted to let you know, I received a link to an ebook, the book How To Start A Social Network which was
    published in 2010.<|endoftext|>
    ```
    """

    def __init__(
        self,
        exponential_decay_length_penalty: Tuple[int, float],
        eos_token_id: Union[int, List[int], torch.Tensor],
        input_ids_seq_length: int,
    ):
        self.regulation_start = exponential_decay_length_penalty[0] + input_ids_seq_length
        self.regulation_factor = exponential_decay_length_penalty[1]

        if not isinstance(eos_token_id, torch.Tensor):
            if isinstance(eos_token_id, int):
                eos_token_id = [eos_token_id]
            eos_token_id = torch.tensor(eos_token_id)
        self.eos_token_id = eos_token_id

        if torch.is_floating_point(eos_token_id) or (eos_token_id < 0).any():
            raise ValueError(f"`eos_token_id` has to be a list of positive integers, but is {eos_token_id}")

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        cur_len = input_ids.shape[-1]
        self.eos_token_id = self.eos_token_id.to(scores.device)
        penalties = torch.zeros_like(scores)
        scores_processed = scores
        if cur_len > self.regulation_start:
            penalty_idx = cur_len - self.regulation_start
            # To support negative logits we compute the penalty of the absolute value and add to the original logit
            penalty = torch.abs(scores[:, self.eos_token_id]) * (pow(self.regulation_factor, penalty_idx) - 1)
            penalties[:, self.eos_token_id] = penalty
            scores_processed = scores + penalties


SUMMARY:
**Exponential Decay Length Penalty Class**

The `ExponentialDecayLengthPenalty` class extends the `LogitsProcessor` and applies an exponential decay penalty to the score of the EOS token (`eos_token_id`) after a specified start index (`exponential_decay_length_penalty`). This allows generating shorter sequences with a smooth cutoff, enabling meaningful prediction of the EOS token.

**Key Implementation Details:**

* The class is initialized with `exponential_decay_length_penalty`, `eos_token_id`, and `input_ids_seq_length`.
* In the `__call__` method:
	+ It first checks if the input `input_ids` has a length greater than the start index, then applies the exponential decay penalty to the EOS token.
	+ The calculated penalty is added to the original scores, resulting in `scores_processed`.

**Notable Dependencies:**

* `torch`: for tensor operations (e.g., `torch.zeros_like`, `torch.abs`, `pow`)
* `LOGITS_PROCESSOR_INPUTS_DOCSTRING`: for docstring and input validation

**Interface and Class Inheritance:**

Inherited from `LogitsProcessor`.


key: chroma:document
string_value:
CODE:

class LogitNormalization(LogitsProcessor):
    r"""
    [`LogitsProcessor`] for normalizing the scores using log-softmax. It's important to normalize
    the scores during beam search, after applying the logits processors or warpers, since the search algorithm used in
    this library doesn't do it (it only does it before, but they may need re-normalization) but it still supposes that
    the scores are normalized when comparing the hypotheses.

    Examples:

    ```python
    >>> from transformers import AutoTokenizer, AutoModelForCausalLM
    >>> import torch

    >>> model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2")
    >>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2")

    >>> inputs = tokenizer("A sequence: 1, 2, 3", return_tensors="pt")

    >>> # By default, the scores are not normalized -- the sum of their exponentials is NOT a normalized probability
    >>> # distribution, summing to 1
    >>> outputs = model.generate(**inputs, return_dict_in_generate=True, output_scores=True)
    >>> print(torch.allclose(torch.sum(torch.exp(outputs.scores[-1])), torch.Tensor((1.000,)), rtol=1e-4))
    False

    >>> # Normalizing them may have a positive impact on beam methods, or when using the scores on your application
    >>> outputs = model.generate(**inputs, renormalize_logits=True, return_dict_in_generate=True, output_scores=True)
    >>> print(torch.allclose(torch.sum(torch.exp(outputs.scores[-1])), torch.Tensor((1.000,)), rtol=1e-4))
    True
    ```
    """

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        scores_processed = scores.log_softmax(dim=-1)


SUMMARY:
**LogitNormalization Class Summary**

**Main Purpose and Functionality:** 
The `LogitNormalization` class normalizes scores using log-softmax, preserving their exponential sum to 1, essential for beam search algorithms.

### Implementation Details:
* The method applies `log_softmax` on the input `scores` tensor along the last dimension (`dim=-1`).
* The normalized scores are returned as a new tensor.

### Notable Dependencies and Requirements:
* Torch library for tensor operations.
* `LOGITS_PROCESSOR_INPUTS_DOCSTRING` docstring for adding start documentation.


key: chroma:document
string_value:
CODE:

class SuppressTokensAtBeginLogitsProcessor(LogitsProcessor):
    r"""
    [`SuppressTokensAtBeginLogitsProcessor`] supresses a list of tokens as soon as the `generate` function starts
    generating using `begin_index` tokens. This should ensure that the tokens defined by `begin_suppress_tokens` are
    not generated at the beginning. Originally created for
    [Whisper](https://huggingface.co/docs/transformers/model_doc/whisper).

    Examples:

    ```python
    >>> from transformers import AutoProcessor, WhisperForConditionalGeneration
    >>> from datasets import load_dataset

    >>> processor = AutoProcessor.from_pretrained("openai/whisper-tiny.en")
    >>> model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en")
    >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
    >>> inputs = processor(ds[0]["audio"]["array"], return_tensors="pt")

    >>> # Whisper has `begin_suppress_tokens` set by default (= `[220, 50256]`). 50256 is the EOS token, so this means
    >>> # it can't generate and EOS token in the first iteration, but it can in the others.
    >>> outputs = model.generate(**inputs, return_dict_in_generate=True, output_scores=True)
    >>> print(outputs.scores[0][0, 50256])
    tensor(-inf)
    >>> print(outputs.scores[-1][0, 50256])  # in other places we can see some probability mass for EOS
    tensor(29.9010)

    >>> # If we disable `begin_suppress_tokens`, we can generate EOS in the first iteration.
    >>> outputs = model.generate(
    ...     **inputs, return_dict_in_generate=True, output_scores=True, begin_suppress_tokens=None
    ... )
    >>> print(outputs.scores[0][0, 50256])
    tensor(11.2027)
    ```
    """

    def __init__(self, begin_suppress_tokens, begin_index, device: str = "cpu"):
        self.begin_suppress_tokens = torch.tensor(list(begin_suppress_tokens), device=device)
        self.begin_index = begin_index

    def set_begin_index(self, begin_index):
        self.begin_index = begin_index

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        vocab_tensor = torch.arange(scores.shape[-1], device=scores.device)
        suppress_token_mask = isin_mps_friendly(vocab_tensor, self.begin_suppress_tokens)
        scores_processed = scores
        if input_ids.shape[-1] == self.begin_index:
            scores_processed = torch.where(suppress_token_mask, -float("inf"), scores)



SUMMARY:
**SuppressTokensAtBeginLogitsProcessor Class Summary**

* **Purpose:** Suppress tokens at the beginning of logits generation to prevent default tokens from being generated prematurely.

**Key Features:**

* Suppresses specific tokens (`begin_suppress_tokens`) for the first `begin_index` tokens in the input sequence.
* Adjusts scores accordingly to reflect suppressed tokens as `-inf`.
* Accepts `device` parameter for device placement, defaults to "cpu".

**Notable Methods:**

* `__init__(self, begin_suppress_tokens: [int], begin_index: int)`: Initializes processor with suppression token list and start index.
* `set_begin_index(self, begin_index):`: Sets new beginning index for suppression.

**Imports/Dependencies:**

* `torch` (Tensor Library)
* `torch.nn.functional.isin_mps_friendly(vocab_tensor, tokens)` (Helper function to create masks)


key: chroma:document
string_value:
CODE:

class SuppressTokensLogitsProcessor(LogitsProcessor):
    r"""
    This processor can be used to suppress a list of tokens. The processor will set their log probs to `-inf` so
    that they are not generated. Originally created for
    [Whisper](https://huggingface.co/docs/transformers/model_doc/whisper).

    Examples:

    ```python
    >>> from transformers import AutoProcessor, WhisperForConditionalGeneration
    >>> from datasets import load_dataset

    >>> processor = AutoProcessor.from_pretrained("openai/whisper-tiny.en")
    >>> model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en")
    >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
    >>> inputs = processor(ds[0]["audio"]["array"], return_tensors="pt")

    >>> # Whisper has a long list of suppressed tokens. For instance, in this case, the token 1 is suppressed by default.
    >>> outputs = model.generate(**inputs, return_dict_in_generate=True, output_scores=True)
    >>> print(outputs.scores[1][0, 1])  # 1 (and not 0) is the first freely generated token
    tensor(-inf)

    >>> # If we disable `suppress_tokens`, we can generate it.
    >>> outputs = model.generate(**inputs, return_dict_in_generate=True, output_scores=True, suppress_tokens=None)
    >>> print(outputs.scores[1][0, 1])
    tensor(6.0678)
    ```
    """

    def __init__(self, suppress_tokens, device: str = "cpu"):
        self.suppress_tokens = torch.tensor(list(suppress_tokens), device=device)

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        vocab_tensor = torch.arange(scores.shape[-1], device=scores.device)
        suppress_token_mask = isin_mps_friendly(vocab_tensor, self.suppress_tokens)
        scores = torch.where(suppress_token_mask, -float("inf"), scores)


SUMMARY:
**Class Summary: SuppressTokensLogitsProcessor**

**Purpose:** Suppresses specified tokens in logits output by setting their log probabilities to `-inf`.

**Main Functionality:**

* Initializes the processor with a list of suppressed tokens and a device (defaulting to "cpu").
* Takes input IDs and scores tensors as input.
* Returns a new tensor where the values corresponding to the suppressed tokens are set to -inf.

**Key Implementation Details:**

* Uses `torch.arange` to create a vocabulary tensor.
* Applies `isin_mps_friendly` function to check for suppression token masks.
* Utilizes `where` to replace scores with -inf when suppress_token_mask is True.

**Notable Dependencies/Requirements:**

* `torch.tensor`, `torch.LongTensor`, `torch.FloatTensor`, `add_start_docstrings`
* `LogitsProcessor` (inherited from)
* No external dependencies other than the PyTorch library.


key: chroma:document
string_value:
CODE:

class WhisperTimeStampLogitsProcessor(LogitsProcessor):
    r"""

    [`LogitsProcessor`] that modifies the logits for the generation of timestamps in the transcription. When the input
    tokens are at a specific threshold, the processor sets the scores to negative infinity. The processor makes sure
    that timestamp tokens appear in pairs, by masking out the logits that would break this pairing pattern. This is
    done to maintain the consistency and structure of generated timestamps. It also ensures that when the predicted
    probability of sampling any of the timestamp token is greater than any individual non-timestamp token, those
    non-timestamp logits are set to negative infinity. This is done to ensure the generation of timestamps over other
    potential tokens.


    See [the paper](https://arxiv.org/abs/2212.04356) for more information.

    Args:
        generate_config (`GenerateConfig`):
            The generate config used to generate the output. The following parameters are required:
                eos_token_id (`int`, *optional*, defaults to 50257):
                    The id of the *end-of-sequence* token.
                no_timestamps_token_id (`int`, *optional*, defaults to 50363):
                    The id of the `"<|notimestamps|>"` token.
                max_initial_timestamp_index (`int`, *optional*, defaults to 1):
                    Used to set the maximum value of the initial timestamp. This is used to prevent the model from
                    predicting timestamps that are too far in the future.
        begin_index (`Optional`, *optional*): Token index of the first token that is generated by the model.
        _detect_timestamp_from_logprob (`bool`, *optional*): Whether timestamps can be predicted from logprobs over all timestamps.

    Examples:
    ``` python
    >>> import torch
    >>> from transformers import AutoProcessor, WhisperForConditionalGeneration, GenerationConfig
    >>> from datasets import load_dataset

    >>> processor = AutoProcessor.from_pretrained("openai/whisper-tiny.en")
    >>> model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en")
    >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
    >>> inputs = processor(ds[3]["audio"]["array"], return_tensors="pt")
    >>> input_features = inputs.input_features

    >>> #Displaying timestamps
    >>> generated_ids = model.generate(inputs=input_features, return_timestamps=True)
    >>> transcription = processor.batch_decode(generated_ids, decode_with_timestamps=True)[0]
    >>> print("Transcription:", transcription)
    Transcription: <|startoftranscript|><|0.00|> He has grave doubts whether Sir Frederick Layton's work is really Greek after all, and can<|6.44|><|6.44|> discover in it but little of rocky Ithaca.<|9.44|><|endoftext|>


    >>> #No timestamps & change EOS:
    >>> #This allows the user to select a specific token to terminate the sequence on, in this case it's the word "can"(460)
    >>> model.generation_config.eos_token_id = 460
    >>> generated_ids = model.generate(inputs=input_features,return_timestamps=False)
    >>> transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    >>> print("Transcription:", transcription)
    Transcription:  He has grave doubts whether Sir Frederick Layton's work is really Greek after all and can
    ```
    """

    def __init__(
        self,
        generate_config,
        begin_index: Optional[int] = None,
        _detect_timestamp_from_logprob: Optional[bool] = None,
    ):  # support for the kwargs
        self.no_timestamps_token_id = generate_config.no_timestamps_token_id
        self.timestamp_begin = generate_config.no_timestamps_token_id + 1
        self.eos_token_id = generate_config.eos_token_id or generate_config.bos_token_id

        # this variable is mostly just used for testing
        self._detect_timestamp_from_logprob = (
            _detect_timestamp_from_logprob
            if _detect_timestamp_from_logprob is not None
            else getattr(generate_config, "_detect_timestamp_from_logprob", True)
        )

        num_forced_ids = (
            len(generate_config.forced_decoder_ids) if generate_config.forced_decoder_ids is not None else 0
        )
        self.begin_index = begin_index or (num_forced_ids + 1)

        self.max_initial_timestamp_index = getattr(generate_config, "max_initial_timestamp_index", None)
        # TODO(Patrick): Make sure that official models have max_initial_timestamp_index set to 50
        # self.max_initial_timestamp_index = 50

    def set_begin_index(self, begin_index):
        self.begin_index = begin_index

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        # suppress <|notimestamps|> which is handled by without_timestamps
        scores_processed = scores.clone()
        scores_processed[:, self.no_timestamps_token_id] = -float("inf")

        # timestamps have to appear in pairs, except directly before eos_token; mask logits accordingly
        for k in range(input_ids.shape[0]):
            sampled_tokens = input_ids[k, self.begin_index :]
            seq = list(sampled_tokens.tolist())

            last_was_timestamp = len(seq) >= 1 and seq[-1] >= self.timestamp_begin
            penultimate_was_timestamp = len(seq) < 2 or seq[-2] >= self.timestamp_begin

            if last_was_timestamp:
                if penultimate_was_timestamp:  # has to be non-timestamp
                    scores_processed[k, self.timestamp_begin :] = -float("inf")
                else:  # cannot be normal text tokens
                    scores_processed[k, : self.eos_token_id] = -float("inf")

            timestamps = sampled_tokens[sampled_tokens.ge(self.timestamp_begin)]
            if timestamps.numel() > 0:
                # `timestamps` shouldn't decrease; forbid timestamp tokens smaller than the last
                # The following lines of code are copied from: https://github.com/openai/whisper/pull/914/files#r1137085090
                if last_was_timestamp and not penultimate_was_timestamp:
                    timestamp_last = timestamps[-1]
                else:
                    # Avoid to emit <|0.00|> again
                    timestamp_last = timestamps[-1] + 1

                scores_processed[k, self.timestamp_begin : timestamp_last] = -float("inf")

        # apply the `max_initial_timestamp` option
        if input_ids.shape[1] == self.begin_index:
            scores_processed[:, : self.timestamp_begin] = -float("inf")

            if self.max_initial_timestamp_index is not None:
                last_allowed = self.timestamp_begin + self.max_initial_timestamp_index
                scores_processed[:, last_allowed + 1 :] = -float("inf")

        # if sum of probability over timestamps is above any other token, sample timestamp
        logprobs = torch.nn.functional.log_softmax(scores_processed.float(), dim=-1)
        for k in range(input_ids.shape[0]):
            timestamp_logprob = logprobs[k, self.timestamp_begin :].logsumexp(dim=-1)
            max_text_token_logprob = logprobs[k, : self.timestamp_begin].max()
            if timestamp_logprob > max_text_token_logprob and self._detect_timestamp_from_logprob:
                scores_processed[k, : self.timestamp_begin] = -float("inf")



SUMMARY:
**WhisperTimeStampLogitsProcessor Class Summary**

**Purpose:**
A custom `LogitsProcessor` that modifies input logits to generate timestamps in transcription outputs. It ensures timestamp tokens appear in pairs, sets non-timestamp logits to negative infinity if predicted probability of sampling any timestamp token is higher.

**Main Functionality:**

1. Suppresses <|notimestamps|> and sets scores to negative infinity when input tokens are at a specific threshold.
2. Masks out logits that would break the pairing pattern between timestamps.
3. Sets non-timestamp logits to negative infinity if predicted probability of sampling any timestamp token is higher than individual non-timestamp tokens.

**Key Implementation Details:**

* Utilizes `torch.nn.functional.log_softmax` to compute probabilities.
* Features a conditional mask to ensure timestamp tokens appear in pairs, except directly before the eos_token ID.
* Checks for consistency between predicted probabilities and actual token values.

**Dependencies:**

* `LogitsProcessor`
* `GenerateConfig`
* `torch`

**Interfaces:**

* `self.no_timestamps_token_id` and `self.timestamp_begin` reference values from `GenerateConfig`.
* `torch.nn.functional.log_softmax` is used for computing probabilities.


key: chroma:document
string_value:
CODE:

class WhisperNoSpeechDetection(LogitsProcessor):
    r"""This processor can be used to detect silence when using Whisper. It should take as input unprocessed logits to follow the original implementation"""

    def __init__(self, no_speech_token: int, begin_index: int, scores_is_logprobs: bool = False):
        self.no_speech_token = no_speech_token
        # offset between <start-of-transcription> token, <SOT>, in paper and first generated token
        # is equal to the position of the first generated token index
        self.start_of_trans_offset = begin_index

        # `self.begin_index` is a running value that is changed on the fly
        self.begin_index = begin_index
        self._no_speech_prob = [0.0]
        self.is_scores_logprobs = scores_is_logprobs

        # overwritten dynamically
        self.model = None
        self.inputs = None

    def set_model(self, model):
        self.model = model

    def set_inputs(self, inputs):
        self.inputs = {**self.model.prepare_inputs_for_generation(**inputs), **inputs}
        self.inputs["input_features"] = self.inputs.pop("inputs")

    @property
    def no_speech_prob(self):
        return self._no_speech_prob

    def set_begin_index(self, begin_index):
        self.begin_index = begin_index

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        is_scores_logprobs = self.is_scores_logprobs

        if input_ids.shape[1] == self.begin_index:
            if self.start_of_trans_offset > 1:
                with torch.no_grad():
                    logits = self.model(**self.inputs).logits

                no_speech_index = self.begin_index - self.start_of_trans_offset
                no_speech_scores = logits[:, no_speech_index]
                is_scores_logprobs = False
            else:
                no_speech_scores = scores

            if is_scores_logprobs:
                probs = no_speech_scores.exp()
            else:
                probs = no_speech_scores.float().softmax(dim=-1)

            self._no_speech_prob = probs[:, self.no_speech_token]



SUMMARY:
**Class Summary: WhisperNoSpeechDetection**

The `WhisperNoSpeechDetection` class is a processor that detects silence in audio using the Whisper model. It takes unprocessed logits as input and returns the output.

**Key Functionality:**

1.  **Initialization**: The class takes three inputs:
    *   `no_speech_token`: an integer representing the index of the no-speech token
    *   `begin_index`: the start index for processing
    *   `scores_is_logprobs`: a boolean indicating whether scores are in log-probability form (default: False)
2.  **Model and Input Setup**: The class dynamically sets the model and inputs using the `set_model` and `set_inputs` methods, respectively.
3.  **Detection Logic**:
    *   When encountering an input ID at index `begin_index`, it checks if it's a start-of-transcription token (SOT).
    *   If it is, it extracts logits from the model using `self.model(**self.inputs).logits`.
    *   It then computes no-speech scores by selecting the corresponding element from the logits.
    *   The class returns either the original scores or the computed probabilities depending on the value of `is_scores_logprobs`.

**Implementation Details:**

1.  **Key Algorithm**: No-speech detection is performed using the extracted logits and the specified token index.

**Notable Dependencies and Requirements:**

*   `torch`: for numerical computations
*   `LOGITS_PROCESSOR_INPUTS_DOCSTRING`: a docstring indicating the expected input structure for the processor

**Notable Attributes and Properties:**

1.  `no_speech_token`
2.  `begin_index`
3.  `model` (dynamically set)
4.  `inputs` (dynamically set)
5.  `_no_speech_prob` (computationally expensive property)


key: chroma:document
string_value:
CODE:

class ClassifierFreeGuidanceLogitsProcessor(LogitsProcessor):
    r"""
    [`LogitsProcessor`] for classifier free guidance (CFG). The scores are split over the batch dimension,
    where the first half correspond to the conditional logits (predicted from the input prompt) and the second half
    correspond to the unconditional logits (predicted from an empty or 'null' prompt). The processor computes a
    weighted average across the conditional and unconditional logits, parameterised by the `guidance_scale`.

    See [the paper](https://arxiv.org/abs/2306.05284) for more information.

    <Tip warning={true}>

    This logits processor is exclusively compatible with
    [MusicGen](https://huggingface.co/docs/transformers/main/en/model_doc/musicgen)

    </Tip>

    Args:
        guidance_scale (float):
            The guidance scale for classifier free guidance (CFG). CFG is enabled by setting `guidance_scale > 1`.
            Higher guidance scale encourages the model to generate samples that are more closely linked to the input
            prompt, usually at the expense of poorer quality.

    Examples:

    ```python
    >>> from transformers import AutoProcessor, MusicgenForConditionalGeneration

    >>> processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
    >>> model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")

    >>> inputs = processor(
    ...     text=["80s pop track with bassy drums and synth", "90s rock song with loud guitars and heavy drums"],
    ...     padding=True,
    ...     return_tensors="pt",
    ... )
    >>> audio_values = model.generate(**inputs, do_sample=True, guidance_scale=3, max_new_tokens=256)
    ```
    """

    def __init__(self, guidance_scale):
        if guidance_scale > 1:
            self.guidance_scale = guidance_scale
        else:
            raise ValueError(
                "Require guidance scale >1 to use the classifier free guidance processor, got guidance scale "
                f"{guidance_scale}."
            )

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        # simple check to make sure we have compatible batch sizes between our
        # logits scores (cond + uncond) and input ids (cond only)
        if scores.shape[0] != 2 * input_ids.shape[0]:
            raise ValueError(
                f"Logits should have twice the batch size of the input ids, the first half of batches corresponding to "
                f"the conditional inputs, and the second half of batches corresponding to the unconditional inputs. Got "
                f"batch size {scores.shape[0]} for the logits and {input_ids.shape[0]} for the input ids."
            )
        unguided_bsz = scores.shape[0] // 2
        cond_logits, uncond_logits = scores.split(unguided_bsz, dim=0)
        scores_processed = uncond_logits + (cond_logits - uncond_logits) * self.guidance_scale


SUMMARY:
**Classifier Free Guidance Logits Processor**

This processor computes weighted averages of conditional and unconditional logits for classifier free guidance. It is designed to work with `MusicGen` models.

**Main Purpose and Functionality:**

The processor calculates a weighted average of the scores, where the weight is determined by the `guidance_scale`. This allows the model to generate samples that are more closely linked to the input prompt while encouraging poorer quality.

**Key Implementation Details:**

* The processor splits the scores into two halves based on the batch size.
* It computes the weighted average of the conditional and unconditional logits using the `guidance_scale`.
* The resulting processed scores are returned.

**Notable Dependencies or Requirements:**

* `MusicGen` model
* `transformers` library

**Method Signatures and Interfaces:**

* `__call__(input_ids: torch<LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor`
	+ Input: input IDs and logits scores
	+ Output: processed scores


key: chroma:document
string_value:
CODE:

class AlternatingCodebooksLogitsProcessor(LogitsProcessor):
    r"""
    [`LogitsProcessor`] enforcing alternated generation between the two codebooks of Bark.

    <Tip warning={true}>

    This logits processor is exclusively compatible with
    [Bark](https://huggingface.co/docs/transformers/en/model_doc/bark)'s fine submodel. See the model documentation
    for examples.

    </Tip>

    Args:
        input_start_len (`int`):
            The length of the initial input sequence.
        semantic_vocab_size (`int`):
            Vocabulary size of the semantic part, i.e number of tokens associated to the semantic vocabulary.
        codebook_size (`int`):
            Number of tokens associated to the codebook.
    """

    def __init__(self, input_start_len: int, semantic_vocab_size: int, codebook_size: int):
        if not isinstance(input_start_len, int) or input_start_len < 0:
            raise ValueError(f"`input_starting_length` has to be a non-negative integer, but is {input_start_len}")

        self.input_start_len = input_start_len
        self.semantic_vocab_size = semantic_vocab_size
        self.codebook_size = codebook_size

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        curr_len = input_ids.shape[-1]

        # even -> first codebook, odd -> second codebook
        is_first_codebook = ((curr_len - self.input_start_len) % 2) == 0

        scores_processed = scores.clone()
        if is_first_codebook:
            scores_processed[:, : self.semantic_vocab_size] = -float("inf")
            scores_processed[:, self.semantic_vocab_size + self.codebook_size :] = -float("inf")
        else:
            scores_processed[:, : self.semantic_vocab_size + self.codebook_size] = -float("inf")



SUMMARY:
**Alternating Codebooks Logits Processor**

**Purpose:** Enforces alternating generation between two codebooks of Bark's fine submodel.

**Functionality:**
This class inherits from `LogitsProcessor` and processes input logits. It alternates between the semantic vocabulary and a second codebook based on the input length.
```python
class AlternatingCodebooksLogitsProcessor(LogitsProcessor):
    def __init__(self, input_start_len: int, semantic_vocab_size: int, codebook_size: int):
        ...
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        # Process scores
        curr_len = input_ids.shape[-1]
        is_first_codebook = ((curr_len - self.input_start_len) % 2) == 0

        if is_first_codebook:
            scores_processed[:, :self.semantic_vocab_size] = -float("inf")
            scores_processed[:, self.semantic_vocab_size + self.codebook_size:] = -float("inf")
        else:
            scores_processed[:, : self.semantic_vocab_size + self.codebook_size] = -float("inf")

        return scores_processed
```
**Notable Details:**

* Applies masking to the input scores based on the input length and codebook sizes.
* Exclusively compatible with Bark's fine submodel.

**Dependencies:**
- `LogitsProcessor` (inherited)
- `torch.LongTensor`
- `torch.FloatTensor`


key: chroma:document
string_value:
CODE:

class UnbatchedClassifierFreeGuidanceLogitsProcessor(LogitsProcessor):
    r"""
    Logits processor for Classifier-Free Guidance (CFG). The processors computes a weighted average across scores
    from prompt conditional and prompt unconditional (or negative) logits, parameterized by the `guidance_scale`.
    The unconditional scores are computed internally by prompting `model` with the `unconditional_ids` branch.

    See [the paper](https://arxiv.org/abs/2306.17806) for more information.

    Args:
        guidance_scale (`float`):
            The guidance scale for classifier free guidance (CFG). CFG is enabled by setting `guidance_scale != 1`.
            Higher guidance scale encourages the model to generate samples that are more closely linked to the input
            prompt, usually at the expense of poorer quality. A value smaller than 1 has the opposite effect, while
            making the negative prompt provided with negative_prompt_ids (if any) act as a positive prompt.
        model (`PreTrainedModel`):
            The model computing the unconditional scores. Supposedly the same as the one computing the conditional
            scores. Both models must use the same tokenizer.
        unconditional_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Indices of input sequence tokens in the vocabulary for the unconditional branch. If unset, will default to
            the last token of the prompt.
        unconditional_attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Attention mask for unconditional_ids.
        use_cache (`bool`, *optional*, defaults to `True`):
            Whether to cache key/values during the negative prompt forward pass.


    Examples:

    ```python
    >>> from transformers import AutoTokenizer, AutoModelForCausalLM

    >>> model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
    >>> tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
    >>> inputs = tokenizer(["Today, a dragon flew over Paris, France,"], return_tensors="pt")
    >>> out = model.generate(inputs["input_ids"], guidance_scale=1.5)
    >>> tokenizer.batch_decode(out, skip_special_tokens=True)[0]
    'Today, a dragon flew over Paris, France, killing at least 50 people and injuring more than 100'

    >>> # with a negative prompt
    >>> neg_inputs = tokenizer(["A very happy event happened,"], return_tensors="pt")
    >>> out = model.generate(inputs["input_ids"], guidance_scale=2, negative_prompt_ids=neg_inputs["input_ids"])
    >>> tokenizer.batch_decode(out, skip_special_tokens=True)[0]
    'Today, a dragon flew over Paris, France, killing at least 130 people. French media reported that'

    >>> # with a positive prompt
    >>> neg_inputs = tokenizer(["A very happy event happened,"], return_tensors="pt")
    >>> out = model.generate(inputs["input_ids"], guidance_scale=0, negative_prompt_ids=neg_inputs["input_ids"])
    >>> tokenizer.batch_decode(out, skip_special_tokens=True)[0]
    "Today, a dragon flew over Paris, France, and I'm very happy to be here. I"
    ```
    """

    def __init__(
        self,
        guidance_scale: float,
        model,
        unconditional_ids: Optional[torch.LongTensor] = None,
        unconditional_attention_mask: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = True,
    ):
        self.guidance_scale = guidance_scale
        self.model = model
        self.unconditional_context = {
            "input_ids": unconditional_ids,
            "attention_mask": unconditional_attention_mask,
            "use_cache": use_cache,
            "past_key_values": None,
            "first_pass": True,
        }

    def get_unconditional_logits(self, input_ids):
        if self.unconditional_context["first_pass"]:
            if self.unconditional_context["input_ids"] is None:
                self.unconditional_context["input_ids"] = input_ids[:, -1:]
            if self.unconditional_context["attention_mask"] is None:
                self.unconditional_context["attention_mask"] = torch.ones_like(
                    self.unconditional_context["input_ids"], dtype=torch.long
                )
            input_ids = self.unconditional_context["input_ids"]
            attention_mask = self.unconditional_context["attention_mask"]
            self.unconditional_context["first_pass"] = False
        else:
            attention_mask = torch.cat(
                [
                    self.unconditional_context["attention_mask"],
                    torch.ones_like(input_ids[:, -1:], dtype=torch.long),
                ],
                dim=1,
            )
            if not self.unconditional_context["use_cache"]:
                input_ids = torch.cat([self.unconditional_context["input_ids"], input_ids[:, -1:]], dim=1)
            else:
                input_ids = input_ids[:, -1:]
            self.unconditional_context["input_ids"] = input_ids
            self.unconditional_context["attention_mask"] = attention_mask

        out = self.model(
            input_ids,
            attention_mask=attention_mask,
            use_cache=self.unconditional_context["use_cache"],
            past_key_values=self.unconditional_context["past_key_values"],
        )
        self.unconditional_context["past_key_values"] = out.get("past_key_values", None)

        return out.logits

    def __call__(self, input_ids, scores):
        scores = torch.nn.functional.log_softmax(scores, dim=-1)
        if self.guidance_scale == 1:
            return scores

        logits = self.get_unconditional_logits(input_ids)

        unconditional_logits = torch.nn.functional.log_softmax(logits[:, -1], dim=-1)
        scores_processed = self.guidance_scale * (scores - unconditional_logits) + unconditional_logits


SUMMARY:
**Class Summary: UnbatchedClassifierFreeGuidanceLogitsProcessor**

The `UnbatchedClassifierFreeGuidanceLogitsProcessor` class processes logits for Classifier-Free Guidance (CFG). It computes a weighted average of conditional and unconditional scores from prompt conditional and unconditional logits, parameterized by the `guidance_scale`.

**Key Details:**

* Computes weighted average of conditional and unconditional scores
* Uses `model` to compute unconditional scores
* Applies guidance scale (`guidance_scale`) to scores
* Supports negative prompt with `negative_prompt_ids`
* Caches key/values during negative prompt forward pass if `use_cache` is set to `True`

**Method Signatures:**

* `get_unconditional_logits(input_ids)`: computes unconditional scores from input IDs
* `__call__(input_ids, scores)`: processes logits for CFG

**Dependencies:**

* `PreTrainedModel`: model computing unconditional scores
* `torch.nn.functional.log_softmax`
* Optional dependencies:
	+ `torch.LongTensor` for input sequence tokens
	+ `torch.LongTensor` for attention mask (unconditional branch)
	+ `bool` for use_cache flag


key: chroma:document
string_value:
CODE:

class BarkEosPrioritizerLogitsProcessor(LogitsProcessor):
    r"""This processor ensures that the EOS token is selected if its probability is greater than the `min_eos_p`.

    <Tip warning={true}>

    This logits processor is exclusively compatible with
    [Bark](https://huggingface.co/docs/transformers/en/model_doc/bark). See the model documentation for examples.

    </Tip>

    Args:
        eos_token_id (`Union[int, List[int], torch.Tensor]`):
            The id(s) of the *end-of-sequence* token.
        min_eos_p (`float`, *optional*):
            Minimum end of speech threshold.
    """

    def __init__(self, eos_token_id: Union[int, List[int], torch.Tensor], min_eos_p: float, device: str = "cpu"):
        if not isinstance(eos_token_id, torch.Tensor):
            if isinstance(eos_token_id, int):
                eos_token_id = [eos_token_id]
            eos_token_id = torch.tensor(eos_token_id, device=device)
        self.eos_token_id = eos_token_id

        if torch.is_floating_point(eos_token_id) or (eos_token_id < 0).any():
            raise ValueError(f"`eos_token_id` has to be a list of positive integers, but is {eos_token_id}")

        if min_eos_p is not None and min_eos_p <= 0:
            raise ValueError(f"`min_eos_p` has to be a positive float, but is {min_eos_p}")
        self.min_eos_p = min_eos_p

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        scores_processed = scores
        if self.min_eos_p:
            probs = torch.nn.functional.softmax(scores.float(), dim=-1)
            # create scores full of -inf except for the eos_token_id
            early_stop_scores = torch.ones_like(scores) * -float("inf")
            early_stop_scores[:, self.eos_token_id] = scores[:, self.eos_token_id]

            do_early_stop = probs[:, self.eos_token_id] > self.min_eos_p
            do_early_stop = torch.any(do_early_stop, dim=1, keepdim=True)
            scores_processed = torch.where(do_early_stop, early_stop_scores, scores)



SUMMARY:
**Summary:**

The `BarkEosPrioritizerLogitsProcessor` class is a processor for logit outputs in the [Bark](https://huggingface.co/docs/transformers/en/model_doc/bark) model. It prioritizes the EOS token if its probability exceeds a minimum threshold (`min_eos_p`). The key functionality includes:

*   Filtering out tokens based on `min_eos_p`
*   Applying early stopping for the EOS token

**Key Implementation Details:**
-   Uses `torch.nn.functional.softmax` to normalize scores
-   Creates a mask to identify tokens where EOS probability exceeds `min_eos_p`

**Notable Dependencies and Requirements:**
-   [Bark](https://huggingface.co/docs/transformers/en/model_doc/bark)
-   `torch`
-   `torch.nn.functional.softmax`


key: chroma:document
string_value:
CODE:

class WatermarkLogitsProcessor(LogitsProcessor):
    r"""
    Logits processor for watermarking generated text. The processor modifies model output scores by adding a small bias to
    randomized set of "green" tokens before generating the next token. "Green" tokens selection process depends on the
    `seeding_scheme` used. The code was based on the [original repo](https://github.com/jwkirchenbauer/lm-watermarking/tree/main).

    The text generated by this `LogitsProcessor` can be detected using `WatermarkDetector`. See [`~WatermarkDetector.__call__`] for details,

    See [the paper](https://arxiv.org/abs/2306.04634) for more information.

    Args:
        vocab_size (`int`):
            The model tokenizer's vocab_size. Used to calculate "green" tokens ratio.
        device (`str`):
            The device where model is allocated.
        greenlist_ratio (`float`, optional, *optional*, defaults to 0.25):
            The ratio of "green" tokens used to the vocabulary size. Defaults to 0.25.
        bias (`float`, optional, *optional*, defaults to 2.0):
            The bias added to the selected "green" tokens' logits. Consider lowering the
            `bias` if the text generation quality degrades. Recommended values are in the
            range of [0.5, 2.0]. Defaults to 2.0.
        hashing_key (`int`, optional, *optional*, defaults to 15485863):
            Key used for hashing. If you deploy this watermark, we advise using another private key.
            Defaults to 15485863 (the millionth prime).
        seeding_scheme (`str`, optional, *optional*, defaults to `"lefthash"`):
            The seeding scheme used for selecting "green" tokens. Accepts values:
                - "lefthash" (default): "green" tokens selection depend on the last token (Algorithm 2 from paper)
                - "selfhash": "green" tokens selection depends on the current token itself (Algorithm 3 from paper)
                    The downside of this scheme is that it considers all possible next tokens and can be slower than "lefthash".
            The context length of previous tokens to use in seeding. Higher context length makes watermarking more robust.
        context_width (`int`, *optional*, defaults to 1):
            The number of previous tokens to use when setting the seed.

    Examples:

    ```python
    >>> from transformers import AutoTokenizer, AutoModelForCausalLM, WatermarkingConfig

    >>> model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
    >>> tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
    >>> inputs = tokenizer(["Alice and Bob are"], return_tensors="pt")

    >>> # normal generation
    >>> out = model.generate(inputs["input_ids"], max_length=20, do_sample=False)
    >>> tokenizer.batch_decode(out, skip_special_tokens=True)[0]
    'Alice and Bob are both in the same room.\n\n"I\'m not sure if you\'re'

    >>> # watermarked generation
    >>> watermarking_config = WatermarkingConfig(bias=2.5, context_width=2, seeding_scheme="selfhash")
    >>> out = model.generate(inputs["input_ids"], watermarking_config=watermarking_config, max_length=20, do_sample=False)
    >>> tokenizer.batch_decode(out, skip_special_tokens=True)[0]
    'Alice and Bob are both still alive and well and the story is pretty much a one-hour adventure'

    >>> # to detect watermarked text use the WatermarkDetector class
    >>> from transformers import WatermarkDetector
    >>> detector = WatermarkDetector(model_config=model.config, device="cpu", watermarking_config= watermarking_config)
    >>> detection_preds = detector(out)
    >>> detection_preds
    array([ True])
    ```
    """

    def __init__(
        self,
        vocab_size,
        device,
        greenlist_ratio: float = 0.25,
        bias: float = 2.0,
        hashing_key: int = 15485863,
        seeding_scheme: str = "lefthash",
        context_width: int = 1,
    ):
        if seeding_scheme not in ["selfhash", "lefthash"]:
            raise ValueError(f"seeding_scheme has to be one of [`selfhash`, `lefthash`], but found {seeding_scheme}")
        if greenlist_ratio >= 1.0 or greenlist_ratio <= 0.0:
            raise ValueError(
                f"greenlist_ratio has be in range between 0.0 and 1.0, exclusively. but found {greenlist_ratio}"
            )

        self.vocab_size = vocab_size
        self.greenlist_size = int(self.vocab_size * greenlist_ratio)
        self.bias = bias
        self.seeding_scheme = seeding_scheme
        self.rng = torch.Generator(device=device)
        self.hash_key = hashing_key
        self.context_width = context_width

        self.rng.manual_seed(hashing_key)
        self.table_size = 1_000_003
        self.fixed_table = torch.randperm(self.table_size, generator=self.rng, device=device)

    def set_seed(self, input_seq: torch.LongTensor):
        input_seq = input_seq[-self.context_width :]
        if self.seeding_scheme == "selfhash":
            a = self.fixed_table[input_seq % self.table_size] + 1
            b = self.fixed_table[input_seq[-1] % self.table_size] + 1
            seed = (self.hash_key * a * b).min().item()
        else:
            seed = self.hash_key * input_seq[-1].item()
        self.rng.manual_seed(seed % (2**64 - 1))

    def _get_greenlist_ids(self, input_seq: torch.LongTensor) -> torch.LongTensor:
        self.set_seed(input_seq)
        vocab_permutation = torch.randperm(self.vocab_size, device=input_seq.device, generator=self.rng)
        greenlist_ids = vocab_permutation[: self.greenlist_size]
        return greenlist_ids

    def _score_rejection_sampling(self, input_seq: torch.LongTensor, scores: torch.FloatTensor) -> torch.LongTensor:
        """
        Generate greenlist based on current candidate next token. Reject and move on if necessary.
        Runs for a fixed number of steps only for efficiency, since the methods is not batched.
        """
        final_greenlist = []
        _, greedy_predictions = scores.sort(dim=-1, descending=True)

        # 40 is an arbitrary number chosen to save compute and not run for long (taken from orig repo)
        for i in range(40):
            greenlist_ids = self._get_greenlist_ids(torch.cat([input_seq, greedy_predictions[i, None]], dim=-1))
            if greedy_predictions[i] in greenlist_ids:
                final_greenlist.append(greedy_predictions[i])
        return torch.tensor(final_greenlist, device=input_seq.device)

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        if input_ids.shape[-1] < self.context_width:
            logger.warning(
                f"`input_ids` should have at least `{self.context_width}` tokens but has {input_ids.shape[-1]}. "
                "The seeding will be skipped for this generation step!"
            )
            return scores

        scores_processed = scores.clone()
        for b_idx, input_seq in enumerate(input_ids):
            if self.seeding_scheme == "selfhash":
                greenlist_ids = self._score_rejection_sampling(input_seq, scores[b_idx])
            else:
                greenlist_ids = self._get_greenlist_ids(input_seq)
            scores_processed[b_idx, greenlist_ids] = scores_processed[b_idx, greenlist_ids] + self.bias



SUMMARY:
**WatermarkLogitsProcessor Class Summary**

**Purpose:** Modify model output scores by adding bias to randomized set of "green" tokens for watermarking generated text.

**Key Features:**

* Implemented as a subclass of `LogitsProcessor`
* `seeding_scheme` determines how "green" tokens are selected
* Supports two seeding schemes: `"lefthash"` and `"selfhash"`
* Greenlist ratio (`greenlist_ratio`) controls the number of "green" tokens in the vocabulary

**Implementation Details:**

* `_get_greenlist_ids` method selects "green" token indices based on input sequence and seeding scheme
* `_score_rejection_sampling` method generates "green" list by rejecting non-"green" tokens for each candidate next token
* `__call__` method processes input scores, adds bias to "green" tokens, and returns processed scores

**Dependencies:**

* `torch`
* `Autograd`
* Logging utilities (`logger`, `LogItsProcessorInput Docstring`)


key: chroma:document
string_value:
CODE:

class SynthIDTextWatermarkState:
    """SynthID watermarking state."""

    def __init__(
        self,
        batch_size: int,
        ngram_len: int,
        context_history_size: int,
        device: torch.device,
    ):
        """Initializes the state.

        Args:
            batch_size (`int`): Batch size.
            ngram_len (`int`): Ngram length.
            context_history_size (`int`): Size of the tensor to keep track of seen contexts.
            device (`int`): Device to use.
        """
        self.context = torch.zeros(
            (batch_size, ngram_len - 1),
            dtype=torch.int64,
            device=device,
        )
        self.context_history = torch.zeros(
            (batch_size, context_history_size),
            dtype=torch.int64,
            device=device,
        )


SUMMARY:
**SynthIDTextWatermarkState Class Summary**

The `SynthIDTextWatermarkState` class initializes the state for a SynthID watermarking algorithm. Its primary purpose is to keep track of the current and past contexts.

**Key Implementation Details:**

* Initializes two tensors (`context` and `context_history`) to store the sequence-based context information.
* Takes in four parameters:
	+ `batch_size`: The number of sequences processed.
	+ `ngram_len`: The length of each n-gram sequence segment.
	+ `context_history_size`: The size of the history tensor for tracking past contexts.
	+ `device`: The device (GPU or CPU) to use.

**Notable Dependencies and Requirements:**

* Utilizes `torch` as the primary library for tensor operations.
* Rely on the following tensors:
	+ `context`
	+ `context_history`
* No external dependencies are mentioned, but it is assumed to be used with a larger SynthID system.


key: chroma:document
string_value:
CODE:

class SynthIDTextWatermarkLogitsProcessor(LogitsProcessor):
    r"""
    Logits processor that implements watermarking techniques for text generation models.
    This class facilitates the application of SynthID text watermarking, a method for embedding imperceptible signals
    into generated text to aid in detecting synthetic content. It operates by subtly manipulating the probabilities of
    token selection during text generation in a manner that can be reliably recovered later for verification.

    Key Features:
    * **State Management:** Maintains internal state to track token sequences and generate watermarking keys
    dynamically.

    * **Key Generation:** Computes hashes based on token sequences and watermarking parameters to create unique keys
    for each position.

    * **G-Value Sampling:** Employs a pre-computed sampling table to sample watermarking values (g-values) based on
    the generated keys.

    * **Score Adjustment:** Applies calculated g-values to modify token probabilities during generation, embedding the
    watermark.

    * **Context Repetition Handling:** Incorporates logic to avoid watermarking tokens in repeated contexts,
    preserving naturalness.

    * **EOS Token Masking:** Supports masking end-of-sentence tokens to prevent their inclusion in watermarking
    calculations.

    * **Utility Functions:** Provides functions to compute g-values directly, check for context repetition, create
    EOS token masks, and estimate expected mean g-values.

    Refer to paper url: https://www.nature.com/articles/s41586-024-08025-4 for more details around this.

    Args:
        ngram_len (`int`):
            Ngram length.
        keys (`List[int]`):
            A sequence of watermarking keys, one for each depth.
        sampling_table_size (`int`):
            Size of the sampling table.
        sampling_table_seed (`int`):
            Random seed to generate the sampling table.
        context_history_size (`int`):
            Size of the tensor to keep track of seen contexts.
        device (`torch.device`):
            Device to use.
        skip_first_ngram_calls (`bool`, *optional*, defaults to `False`):
            Whether to skip first ngram calls.
        debug_mode (`bool`, optional, *optional*, defaults to `False`):
            Logits are modified to uniform one got before watermarking modification is applied. This is to test the
            implementation.

    Examples:
    ```python
    >>> from transformers import AutoModelForCausalLM, AutoTokenizer, SynthIDTextWatermarkingConfig

    >>> tokenizer = AutoTokenizer.from_pretrained('google/gemma-2-2b', padding_side="left")
    >>> model = AutoModelForCausalLM.from_pretrained('google/gemma-2-2b')

    >>> # SynthID Text configuration
    >>> watermarking_config = SynthIDTextWatermarkingConfig(
    ...     keys=[654, 400, 836, 123, 340, 443, 597, 160, 57],
    ...     ngram_len=5,
    ... )

    >>> # Generation with watermarking
    >>> tokenized_prompts = tokenizer(["Once upon a time, "], return_tensors="pt", padding=True)
    >>> output_sequences = model.generate(
    ...     **tokenized_prompts, watermarking_config=watermarking_config, do_sample=True, max_new_tokens=10
    ... )
    >>> watermarked_text = tokenizer.batch_decode(output_sequences, skip_special_tokens=True)
    ```
    """

    def __init__(
        self,
        ngram_len: int,
        keys: List[int],
        sampling_table_size: int,
        sampling_table_seed: int,
        context_history_size: int,
        device: torch.device,
        skip_first_ngram_calls: bool = False,
        debug_mode: bool = False,
    ):
        self.ngram_len = ngram_len
        self.keys = torch.tensor(keys, device=device)

        generator = torch.Generator(device=device).manual_seed(sampling_table_seed)
        # A random sampling table is pre-computed and modulo table size is applied to map from a hash of ngram keys to
        # g values, this is similar to the hashtable implementation used in
        # https://github.com/facebookresearch/three_bricks. We note that the hashing employed in this repository is
        # different from that used to watermark the Gemini App, and hence the detectors trained based on the
        # hashing in this repository will not transfer to text generated by the Gemini App.
        self.sampling_table = torch.randint(
            low=0,
            high=2,
            size=(sampling_table_size,),
            generator=generator,
            device=device,
        )
        self.context_history_size = context_history_size
        self.device = device
        self.state = None
        self.skip_first_ngram_calls = skip_first_ngram_calls
        self.debug_mode = debug_mode

    def _init_state(self, batch_size: int):
        """Initializes the state."""
        self.state = SynthIDTextWatermarkState(
            batch_size=batch_size,
            ngram_len=self.ngram_len,
            context_history_size=self.context_history_size,
            device=self.device,
        )

    def update_scores(self, scores: torch.FloatTensor, g_values: torch.FloatTensor) -> torch.FloatTensor:
        """Updates scores using the g values.

        We assume that the scores are in the log space.
        Args:
            scores (`torch.FloatTensor`): Scores (batch_size, vocab_size).
            g_values (`torch.FloatTensor`): G valus (batch_size, vocab_size, depth).

        Returns:
            Updated scores (batch_size, vocab_size).
        """
        _, _, depth = g_values.shape

        probs = torch.softmax(scores, dim=1)

        for i in range(depth):
            g_values_at_depth = g_values[:, :, i]
            g_mass_at_depth = (g_values_at_depth * probs).sum(axis=1, keepdims=True)
            probs = probs * (1 + g_values_at_depth - g_mass_at_depth)

        log_probs = torch.log(probs)
        log_probs = torch.where(torch.isfinite(log_probs), log_probs, torch.finfo(log_probs.dtype).min)
        return log_probs

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        self._check_input_ids_shape(input_ids)
        batch_size, vocab_size = scores.shape

        if self.debug_mode:
            scores = torch.ones_like(scores)

        # Currently indices is just a arange to compute watermarking on the desnse logits.
        all_indices = torch.stack([torch.arange(vocab_size, device=self.device) for _ in range(batch_size)])

        if self.state is None:
            # Initialize watermarking state if it does not exist.
            self._init_state(batch_size)
        else:
            # Append last input id (which is the input id added in last call) to the
            # previous context so we have the context to be used for current
            # watermarking.
            self.state.context = torch.concat(
                (self.state.context, input_ids[:, -1:]),
                dim=1,
            )
            self.state.context = self.state.context[:, 1:]

        if self.state is None:
            raise ValueError("self.state can't be None! Call `self._init_state` to initialize the state.")

        self.state.num_calls += 1

        # Don't watermark the first ngram_len - 1 tokens if set.
        if self.skip_first_ngram_calls and self.state.num_calls < self.ngram_len:
            return scores

        # 2. Generate random keys for each ngram key combination.
        ngram_keys, hash_result_with_just_context = self._compute_keys(self.state.context, all_indices)
        # ngram_keys shape [batch_size, top_k, depth]

        # 3. Sample g values.
        g_values = self.sample_g_values(ngram_keys)
        # g_values shape [batch_size, top_k, depth]

        # 4. Modify scores.
        updated_scores = self.update_scores(scores, g_values)
        # updated scores shape [batch_size, top_k]

        # 5. Check if the current watermarking context was previously used, if yes skip watermarking.
        hash_result_with_just_context = hash_result_with_just_context[:, None]
        is_repeated_context = (self.state.context_history == hash_result_with_just_context).any(
            dim=1,
            keepdim=True,
        )
        self.state.context_history = torch.concat(
            (hash_result_with_just_context, self.state.context_history),
            dim=1,
        )[:, :-1]

        updated_watermarked_scores = torch.where(
            is_repeated_context,
            input=scores,
            other=updated_scores,
        )
        return updated_watermarked_scores

    def accumulate_hash(
        self,
        current_hash: torch.LongTensor,
        data: torch.LongTensor,
        multiplier: int = 6364136223846793005,
        increment: int = 1,
    ) -> torch.LongTensor:
        """
        Accumulate hash of data on current hash.

        Method uses adapted linear congruential generator with newlib/musl parameters.

        This function has following property -
        f(x, data[T]) = f(f(x, data[:T - 1]), data[T])

        This function expects current_hash.shape and data.shape[:-1] to
        match/broadcastable.

        Args:
            current_hash (`torch.LongTensor`):
                (shape,)
            data (`torch.LongTensor`):
                (shape, tensor_len)
            multiplier (`int`, optional, *optional*, defaults to 6364136223846793005):
                multiplier of linear congruential generator
            increment (`int`, optional, *optional*, defaults to 1):
                increment of linear congruential generator

        Returns:
            updated hash (shape,)
        """
        for i in range(data.shape[-1]):
            current_hash = torch.add(current_hash, data[..., i])
            current_hash = torch.mul(current_hash, multiplier)
            current_hash = torch.add(current_hash, increment)
        return current_hash

    def compute_ngram_keys(self, ngrams: torch.LongTensor) -> torch.LongTensor:
        """Computes random keys for each ngram and depth.

        Args:
            ngrams (`torch.LongTensor`):
                Ngrams (batch_size, num_ngrams, ngram_len).

        Returns:
            ngram keys (batch_size, num_ngrams, depth).
        """
        if len(ngrams.shape) != 3:
            raise ValueError(
                "Ngrams should be of shape (batch_size, num_ngrams, ngram_len), but" f" is {ngrams.shape}"
            )
        if ngrams.shape[2] != self.ngram_len:
            raise ValueError(
                "Ngrams should be of shape (batch_size, num_ngrams, ngram_len),"
                f" where ngram_len is {self.ngram_len}, but is {ngrams.shape}"
            )
        batch_size, _, _ = ngrams.shape

        hash_result = torch.ones(batch_size, device=self.device, dtype=torch.long)
        # hash_result shape [batch_size,]
        # ngrams shape [batch_size, num_ngrams, ngram_len]
        hash_result = torch.vmap(self.accumulate_hash, in_dims=(None, 1), out_dims=1)(hash_result, ngrams)
        # hash_result shape [batch_size, num_ngrams]

        keys = self.keys[None, None, :, None]
        # hash_result shape [batch_size, num_ngrams]
        # keys shape [1, 1, depth, 1]
        hash_result = torch.vmap(self.accumulate_hash, in_dims=(None, 2), out_dims=2)(hash_result, keys)
        # hash_result shape [batch_size, num_ngrams, depth]

        return hash_result

    def _compute_keys(
        self, n_minus_1_grams: torch.LongTensor, indices: torch.LongTensor
    ) -> Tuple[torch.LongTensor, torch.LongTensor]:
        """Computes random keys for each ngram and depth.

        Args:
            n_minus_1_grams (`torch.LongTensor`):
                Ngrams (batch_size, ngram_len - 1).
            indices (`torch.LongTensor`):
                indices of the continuations (batch_size, num_indices)

        Returns:
            Ngram keys (batch_size, num_indices, depth).
        """
        batch_size, _ = n_minus_1_grams.shape

        hash_result = torch.ones(batch_size, device=self.device, dtype=torch.long)
        # First hash n_minus_1 gram, for each batch entry we have a single
        # n_minus_1 gram context.
        # hash_result shape [batch_size]
        # n_minus_1_gram shape [batch_size, ngram_len - 1]
        hash_result_with_just_context = self.accumulate_hash(hash_result, n_minus_1_grams)
        # hash_result shape [batch_size,]
        # Indices is of shape [batch_size, num_indices], so we make it
        # [batch_size, num_indices, 1] so we can vmap over num_indices dim.
        hash_result = torch.vmap(self.accumulate_hash, in_dims=(None, 1), out_dims=1)(
            hash_result_with_just_context, indices[:, :, None]
        )
        # hash_result shape [batch_size, num_indices]
        # Basically we have a hash for each batch entry and each indices
        # Now we add watermarking keys to this hash.
        # keys are of shape [depth,]
        # We add batch, num_indices and data dimension to this making it
        # [1, 1, depth, 1].
        # So we can vmap over the depth dimension for compute_hash
        keys = self.keys[None, None, :, None]
        hash_result = torch.vmap(self.accumulate_hash, in_dims=(None, 2), out_dims=2)(hash_result, keys)
        # hash_result shape should be [batch_size, num_indices, depth]
        return hash_result, hash_result_with_just_context

    def sample_g_values(self, ngram_keys: torch.LongTensor) -> torch.LongTensor:
        """
        Samples g values from Bernoulli distribution.

        It is not possible to pass random keys in a vectorized way in torch. Instead
        we pre-compute a random sampling table, and use apply modulo table size to
        map from ngram keys (int64) to g values.

        Args:
            ngram_keys (`torch.LongTensor`):
                Random keys (batch_size, num_ngrams, depth).

        Returns:
            G values (batch_size, num_ngrams, depth).
        """
        (sampling_table_size,) = self.sampling_table.shape
        sampling_table = self.sampling_table.reshape((1, 1, sampling_table_size))
        ngram_keys = ngram_keys % sampling_table_size
        return torch.take_along_dim(sampling_table, indices=ngram_keys, dim=2)

    def _check_input_ids_shape(self, input_ids: torch.LongTensor):
        """Checks the shape of input ids."""
        if len(input_ids.shape) != 2:
            raise ValueError("Input ids should be of shape (batch_size, input_len), but is" f" {input_ids.shape}")

    def compute_g_values(self, input_ids: torch.LongTensor) -> torch.LongTensor:
        """
        Computes g values for each ngram from the given sequence of tokens.

        Args:
            input_ids (`torch.LongTensor`):
                Input token ids (batch_size, input_len).

        Returns:
            G values (batch_size, input_len - (ngram_len - 1), depth).
        """
        self._check_input_ids_shape(input_ids)
        ngrams = input_ids.unfold(dimension=1, size=self.ngram_len, step=1)
        ngram_keys = self.compute_ngram_keys(ngrams)
        return self.sample_g_values(ngram_keys)

    def compute_context_repetition_mask(self, input_ids: torch.LongTensor) -> torch.LongTensor:
        """
        Computes repetition mask.

        0 and 1 stand for repeated and not repeated context n-1 grams respectively.

        Args:
            input_ids (`torch.LongTensor`):
                Input token ids (batch_size, input_len).

        Returns:
            Repetitions mask (batch_size, input_len - (ngram_len - 1)).
        """
        self._check_input_ids_shape(input_ids)
        batch_size, _ = input_ids.shape
        state = SynthIDTextWatermarkState(
            batch_size=batch_size,
            ngram_len=self.ngram_len,
            context_history_size=self.context_history_size,
            device=self.device,
        )
        contexts = input_ids[:, :-1].unfold(
            dimension=1,
            size=self.ngram_len - 1,
            step=1,
        )
        _, num_contexts, _ = contexts.shape

        are_repeated_contexts = []
        for i in range(num_contexts):
            context = contexts[:, i, :]
            hash_result = torch.ones(batch_size, device=self.device, dtype=torch.long)
            context_hash = self.accumulate_hash(hash_result, context)[:, None]
            is_repeated_context = (state.context_history == context_hash).any(
                dim=1,
                keepdim=True,
            )
            are_repeated_contexts.append(is_repeated_context)
            state.context_history = torch.concat(
                (context_hash, state.context_history),
                dim=1,
            )[:, :-1]
        are_repeated_contexts = torch.concat(are_repeated_contexts, dim=1)

        return torch.logical_not(are_repeated_contexts)

    def compute_eos_token_mask(self, input_ids: torch.LongTensor, eos_token_id: int) -> torch.LongTensor:
        """
        Computes repetitions mask.

        1 stands for ngrams that don't contain EOS tokens and vice versa.

        Args:
            input_ids (`torch.LongTensor`):
                Input token ids (batch_size, input_len).
            eos_token_id (`int`):
                EOS token ID.

        Returns:
            EOS token mask (batch_size, input_len).
        """
        self._check_input_ids_shape(input_ids)
        noneos_masks = []
        all_eos_equated = input_ids == eos_token_id
        for eos_equated in all_eos_equated:
            nonzero_idx = torch.nonzero(eos_equated)
            noneos_mask = torch.ones_like(eos_equated)
            if nonzero_idx.shape[0] != 0:
                noneos_mask[nonzero_idx[0][0] :] = 0
            noneos_masks.append(noneos_mask)
        return torch.stack(noneos_masks, dim=0)

    def expected_mean_g_value(self, vocab_size: int, coinflip_prob: float = 0.5) -> float:
        """
        Compute expected mean g-value after watermarking, assuming uniform LM dist.

        This is the theoretical expected value for single-layer watermarking.

        Args:
            vocab_size (`int`):
                The size of the vocabulary.
            coinflip_prob arg_name (`float`, *optional*, defaults to 0.5):
                Probability of 1 in boolean prf.

        Returns:
            The expected mean g-value for watermarked text.
        """


SUMMARY:
This is an implementation of a SynthIDTextWatermarking class, which appears to be related to watermarking in natural language processing (NLP). The primary function of this class likely serves as a model for watermarked sequences generated by a language model. Here's a high-level summary:

**Purpose:**

The SynthIDTextWatermarking class provides methods to compute various aspects of a watermarked text sequence, including the expected mean g-value.

**Key Methods:**

1. `compute_g_values`: Computes g-values for each n-gram in the input sequence.
2. `sample_g_values`: Samples g-values based on the computed values.
3. `accumulate_hash`: Computes the hash value of a given sequence.
4. `compute_ngram_keys`: Computes keys for an input sequence.

**Relevant Properties:**

1. `device`: The device (GPU or CPU) where computations will take place.
2. `ngram_len`: The length of each n-gram in the watermark.
3. `context_history_size`: The number of history steps used to compute context repetition masks.
4. `sampling_table`: A lookup table used for sampling g-values.

**Use Cases:**

This class can be used in various NLP applications, such as:

1. Generating watermarked sequences with controlled properties (e.g., expected mean g-value).
2. Evaluating the robustness of a language model to watermarking attacks.
3. Developing strategies for hiding or revealing sensitive information in text sequences.

Note that the implementation provides various methods and tools for analyzing and manipulating watermarked sequences, but its exact use cases and applications are not specified in this summary.


key: chroma:document
string_value:
CODE:

class BeamScorer(ABC):
    """
    Abstract base class for all beam scorers that are used for [`~PreTrainedModel.beam_search`] and
    [`~PreTrainedModel.beam_sample`].
    """

    @abstractmethod
    @add_start_docstrings(PROCESS_INPUTS_DOCSTRING)
    def process(
        self,
        input_ids: torch.LongTensor,
        next_scores: torch.FloatTensor,
        next_tokens: torch.LongTensor,
        next_indices: torch.LongTensor,
        **kwargs,
    ) -> Tuple[torch.Tensor]:
        raise NotImplementedError("This is an abstract method.")

    @abstractmethod
    @add_start_docstrings(FINALIZE_INPUTS_DOCSTRING)
    def finalize(
        self,
        input_ids: torch.LongTensor,
        next_scores: torch.FloatTensor,
        next_tokens: torch.LongTensor,
        next_indices: torch.LongTensor,
        max_length: int,
        **kwargs,
    ) -> torch.LongTensor:


SUMMARY:
**BeamScorer Class Summary**

The `BeamScorer` class is an abstract base class for implementing beam search algorithms in Natural Language Processing (NLP) models. It provides two main methods:

1. `process(input_ids, next_scores, next_tokens, next_indices, **kwargs) -> Tuple[torch.Tensor]`: This method receives input data and returns a tensor representing the updated score values.
2. `finalize(input_ids, next_scores, next_tokens, next_indices, max_length, **kwargs) -> torch.LongTensor`: This method finalizes the beam search process by returning the top scoring indices.

Key implementation details:

* The class requires implementing the `process` and `finalize` methods, which are marked as abstract methods to ensure concrete subclasses provide their own logic.
* The `add_start_docstrings` decorator is used to add docstring attributes for the input parameters and return values of these methods.

Notable dependencies or requirements:

* The class uses PyTorch tensors (`torch.LongTensor`, `torch.FloatTensor`) and requires access to the `PreTrainedModel.beam_search` and `PreTrainedModel.beam_sample` functions in its implementation.
* The `finalize` method inherits from `PreTrainedModel`, implying that concrete subclasses need to implement this based on a specific pre-trained model architecture.


key: chroma:document
string_value:
CODE:

class BeamSearchScorer(BeamScorer):
    r"""
    [`BeamScorer`] implementing standard beam search decoding.

    Adapted in part from [Facebook's XLM beam search
    code](https://github.com/facebookresearch/XLM/blob/9e6f6814d17be4fe5b15f2e6c43eb2b2d76daeb4/src/model/transformer.py#L529).

    Reference for the diverse beam search algorithm and implementation [Ashwin Kalyan's DBS
    implementation](https://github.com/ashwinkalyan/dbs/blob/master/dbs/beam_utils.lua)

    Args:
        batch_size (`int`):
            Batch Size of `input_ids` for which standard beam search decoding is run in parallel.
        num_beams (`int`):
            Number of beams for beam search.
        device (`torch.device`):
            Defines the device type (*e.g.*, `"cpu"` or `"cuda"`) on which this instance of `BeamSearchScorer` will be
            allocated.
        length_penalty (`float`, *optional*, defaults to 1.0):
            Exponential penalty to the length that is used with beam-based generation. It is applied as an exponent to
            the sequence length, which in turn is used to divide the score of the sequence. Since the score is the log
            likelihood of the sequence (i.e. negative), `length_penalty` > 0.0 promotes longer sequences, while
            `length_penalty` < 0.0 encourages shorter sequences.
        do_early_stopping (`bool` or `str`, *optional*, defaults to `False`):
            Controls the stopping condition for beam-based methods, like beam-search. It accepts the following values:
            `True`, where the generation stops as soon as there are `num_beams` complete candidates; `False`, where an
            heuristic is applied and the generation stops when is it very unlikely to find better candidates;
            `"never"`, where the beam search procedure only stops when there cannot be better candidates (canonical
            beam search algorithm).
        num_beam_hyps_to_keep (`int`, *optional*, defaults to 1):
            The number of beam hypotheses that shall be returned upon calling
            [`~transformers.BeamSearchScorer.finalize`].
        num_beam_groups (`int`, *optional*, defaults to 1):
            Number of groups to divide `num_beams` into in order to ensure diversity among different groups of beams.
            See [this paper](https://arxiv.org/pdf/1610.02424.pdf) for more details.
        max_length (`int`, *optional*):
            The maximum length of the sequence to be generated.
    """

    def __init__(
        self,
        batch_size: int,
        num_beams: int,
        device: torch.device,
        length_penalty: Optional[float] = 1.0,
        do_early_stopping: Optional[Union[bool, str]] = False,
        num_beam_hyps_to_keep: Optional[int] = 1,
        num_beam_groups: Optional[int] = 1,
        max_length: Optional[int] = None,
    ):
        self.num_beams = num_beams
        self.device = device
        self.length_penalty = length_penalty
        self.do_early_stopping = do_early_stopping
        self.num_beam_hyps_to_keep = num_beam_hyps_to_keep
        self.num_beam_groups = num_beam_groups
        self.group_size = self.num_beams // self.num_beam_groups

        self._is_init = False
        # self._beam_hyps[i*self.num_beam_groups+j] is the beam_hyps of the j-th group in the i-th mini-batch.
        # If group_beam_search is not used, the list consists of `batch_size` beam_hyps.
        self._beam_hyps = [
            BeamHypotheses(
                num_beams=self.group_size,
                length_penalty=self.length_penalty,
                early_stopping=self.do_early_stopping,
                max_length=max_length,
            )
            for _ in range(batch_size * self.num_beam_groups)
        ]
        # self._done[i*self.num_beam_groups+j] indicates whether the generation of the beam_hyps of the j-th group
        # in the i-th mini-batch is complete.
        self._done = torch.tensor(
            [False for _ in range(batch_size * self.num_beam_groups)], dtype=torch.bool, device=self.device
        )

        if not isinstance(num_beams, int) or num_beams <= 1:
            raise ValueError(
                f"`num_beams` has to be an integer strictly greater than 1, but is {num_beams}. For `num_beams` == 1,"
                " one should make use of `greedy_search` instead."
            )

        if not isinstance(num_beam_groups, int) or (num_beam_groups > num_beams) or (num_beams % num_beam_groups != 0):
            raise ValueError(
                "`num_beam_groups` has to be an integer smaller or equal than `num_beams` and `num_beams` has to be"
                f" divisible by `num_beam_groups`, but is {num_beam_groups} with `num_beams` being {num_beams}."
            )

    @property
    def is_done(self) -> bool:
        return self._done.all()

    def process(
        self,
        input_ids: torch.LongTensor,
        next_scores: torch.FloatTensor,
        next_tokens: torch.LongTensor,
        next_indices: torch.LongTensor,
        pad_token_id: Optional[Union[int, torch.Tensor]] = None,
        eos_token_id: Optional[Union[int, List[int], torch.Tensor]] = None,
        beam_indices: Optional[torch.LongTensor] = None,
        group_index: Optional[int] = 0,
        decoder_prompt_len: Optional[int] = 0,
    ) -> Dict[str, torch.Tensor]:
        # add up to the length which the next_scores is calculated on (including decoder prompt)
        cur_len = input_ids.shape[-1] + 1
        batch_size = len(self._beam_hyps) // self.num_beam_groups

        if not (batch_size == (input_ids.shape[0] // self.group_size)):
            if self.num_beam_groups > 1:
                raise ValueError(
                    f"A group beam size of {input_ids.shape[0]} is used as the input, but a group beam "
                    f"size of {self.group_size} is expected by the beam scorer."
                )
            else:
                raise ValueError(
                    f"A beam size of {input_ids.shape[0]} is used as the input, but a beam size of "
                    f"{self.group_size} is expected by the beam scorer."
                )

        device = input_ids.device
        next_beam_scores = torch.zeros((batch_size, self.group_size), dtype=next_scores.dtype, device=device)
        next_beam_tokens = torch.zeros((batch_size, self.group_size), dtype=next_tokens.dtype, device=device)
        next_beam_indices = torch.zeros((batch_size, self.group_size), dtype=next_indices.dtype, device=device)

        if eos_token_id is not None and not isinstance(eos_token_id, torch.Tensor):
            if isinstance(eos_token_id, int):
                eos_token_id = [eos_token_id]
            eos_token_id = torch.tensor(eos_token_id)

        for batch_idx in range(batch_size):
            batch_group_idx = batch_idx * self.num_beam_groups + group_index
            if self._done[batch_group_idx]:
                if self.num_beams < len(self._beam_hyps[batch_group_idx]):
                    raise ValueError(f"Batch can only be done if at least {self.num_beams} beams have been generated")
                if eos_token_id is None or pad_token_id is None:
                    raise ValueError("Generated beams >= num_beams -> eos_token_id and pad_token have to be defined")
                # pad the batch
                next_beam_scores[batch_idx, :] = 0
                next_beam_tokens[batch_idx, :] = pad_token_id
                next_beam_indices[batch_idx, :] = 0
                continue

            # next tokens for this sentence
            beam_idx = 0
            for beam_token_rank, (next_token, next_score, next_index) in enumerate(
                zip(next_tokens[batch_idx], next_scores[batch_idx], next_indices[batch_idx])
            ):
                batch_beam_idx = batch_idx * self.group_size + next_index
                # add to generated hypotheses if end of sentence
                if (eos_token_id is not None) and (next_token.item() in eos_token_id):
                    # if beam_token does not belong to top num_beams tokens, it should not be added
                    is_beam_token_worse_than_top_num_beams = beam_token_rank >= self.group_size
                    if is_beam_token_worse_than_top_num_beams:
                        continue
                    if beam_indices is not None:
                        beam_index = beam_indices[batch_beam_idx]
                        beam_index = beam_index + (batch_beam_idx,)
                    else:
                        beam_index = None

                    self._beam_hyps[batch_group_idx].add(
                        input_ids[batch_beam_idx].clone(),
                        next_score.item(),
                        beam_indices=beam_index,
                        generated_len=cur_len - decoder_prompt_len,
                    )
                else:
                    # add next predicted token since it is not eos_token
                    next_beam_scores[batch_idx, beam_idx] = next_score
                    next_beam_tokens[batch_idx, beam_idx] = next_token
                    next_beam_indices[batch_idx, beam_idx] = batch_beam_idx
                    beam_idx += 1

                # once the beam for next step is full, don't add more tokens to it.
                if beam_idx == self.group_size:
                    break

            if beam_idx < self.group_size:
                raise ValueError(
                    f"At most {self.group_size} tokens in {next_tokens[batch_idx]} can be equal to `eos_token_id:"
                    f" {eos_token_id}`. Make sure {next_tokens[batch_idx]} are corrected."
                )

            # Check if we are done so that we can save a pad step if all(done)
            self._done[batch_group_idx] = self._done[batch_group_idx] or self._beam_hyps[batch_group_idx].is_done(
                next_scores[batch_idx].max().item(), cur_len, decoder_prompt_len
            )

        return UserDict(
            {
                "next_beam_scores": next_beam_scores.view(-1),
                "next_beam_tokens": next_beam_tokens.view(-1),
                "next_beam_indices": next_beam_indices.view(-1),
            }
        )

    def finalize(
        self,
        input_ids: torch.LongTensor,
        final_beam_scores: torch.FloatTensor,
        final_beam_tokens: torch.LongTensor,
        final_beam_indices: torch.LongTensor,
        max_length: int,
        pad_token_id: Optional[Union[int, torch.Tensor]] = None,
        eos_token_id: Optional[Union[int, List[int], torch.Tensor]] = None,
        beam_indices: Optional[torch.LongTensor] = None,
        decoder_prompt_len: Optional[int] = 0,
    ) -> Tuple[torch.LongTensor]:
        batch_size = len(self._beam_hyps) // self.num_beam_groups

        if eos_token_id is not None and not isinstance(eos_token_id, torch.Tensor):
            if isinstance(eos_token_id, int):
                eos_token_id = [eos_token_id]
            eos_token_id = torch.tensor(eos_token_id)

        # finalize all open beam hypotheses and add to generated hypotheses
        for batch_group_idx, beam_hyp in enumerate(self._beam_hyps):
            if self._done[batch_group_idx]:
                continue

            # all open beam hypotheses are added to the beam hypothesis
            # beam hypothesis class automatically keeps the best beams
            for index_per_group in range(self.group_size):
                batch_beam_idx = batch_group_idx * self.group_size + index_per_group
                final_score = final_beam_scores[batch_beam_idx].item()
                final_tokens = input_ids[batch_beam_idx]
                beam_index = beam_indices[batch_beam_idx] if beam_indices is not None else None
                generated_len = final_tokens.shape[-1] - decoder_prompt_len
                beam_hyp.add(final_tokens, final_score, beam_indices=beam_index, generated_len=generated_len)

        # select the best hypotheses
        sent_lengths = input_ids.new(batch_size * self.num_beam_hyps_to_keep)
        best = []
        best_indices = []
        best_scores = torch.zeros(batch_size * self.num_beam_hyps_to_keep, device=self.device, dtype=torch.float32)

        # retrieve best hypotheses
        for i in range(batch_size):
            beam_hyps_in_batch = self._beam_hyps[i * self.num_beam_groups : (i + 1) * self.num_beam_groups]
            candidate_beams = [beam for beam_hyp in beam_hyps_in_batch for beam in beam_hyp.beams]
            sorted_hyps = sorted(candidate_beams, key=lambda x: x[0])
            for j in range(self.num_beam_hyps_to_keep):
                best_hyp_tuple = sorted_hyps.pop()
                best_score = best_hyp_tuple[0]
                best_hyp = best_hyp_tuple[1]
                best_index = best_hyp_tuple[2]
                sent_lengths[self.num_beam_hyps_to_keep * i + j] = len(best_hyp)

                # append hyp to lists
                best.append(best_hyp)

                # append indices to list
                best_indices.append(best_index)

                best_scores[i * self.num_beam_hyps_to_keep + j] = best_score

        # prepare for adding eos
        sent_lengths_max = sent_lengths.max().item() + 1
        sent_max_len = min(sent_lengths_max, max_length) if max_length is not None else sent_lengths_max
        decoded: torch.LongTensor = input_ids.new(batch_size * self.num_beam_hyps_to_keep, sent_max_len)

        if len(best_indices) > 0 and best_indices[0] is not None:
            indices: torch.LongTensor = input_ids.new(batch_size * self.num_beam_hyps_to_keep, sent_max_len)
        else:
            indices = None

        # shorter batches are padded if needed
        if sent_lengths.min().item() != sent_lengths.max().item():
            if pad_token_id is None:
                raise ValueError("`pad_token_id` has to be defined")
            decoded.fill_(pad_token_id)

        if indices is not None:
            indices.fill_(-1)

        # fill with hypotheses and eos_token_id if the latter fits in
        for i, (hypo, best_idx) in enumerate(zip(best, best_indices)):
            decoded[i, : sent_lengths[i]] = hypo

            if indices is not None:
                indices[i, : len(best_idx)] = torch.tensor(best_idx)

            if sent_lengths[i] < sent_max_len:
                # inserting only the first eos_token_id
                decoded[i, sent_lengths[i]] = eos_token_id[0]

        return UserDict(
            {
                "sequences": decoded,
                "sequence_scores": best_scores,
                "beam_indices": indices,
            }


SUMMARY:
This code appears to be part of the Fairseq model for natural language processing and machine translation. It is responsible for computing two key operations:
1. Generating final hypotheses: This involves iterating through open-ended beam search, aggregating scores, token sequences, and indices from generated hypotheses, filtering out those that do not meet conditions based on a maximum value of max_length, applying an end-of-sequence post-processing step using a set of tokens.
2. Finalizing the output: This includes selecting the top N best hypotheses for each mini-batch based on their scores, discarding those with lower beam indices and computing their token identities to form final sequences.

The essential functions involved in these operations include:

- `forward`: Invokes a forward pass using input IDs, `final_beam_scores`, `[final_beam_tokens, final_beam_indices]` output.
- `finalize`: Runs the entire model sequence on the provided inputs for the best scoring hypothesis up until max_length or completion of an N-th best hypotheses.

They use various operations like tensors concatenation, indexing & filtering, aggregation over elements. However, without seeing more of Fairseq and deeper into code implementation is difficult to describe further specific functionality


