# embedding_metadata

key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/seq2seq-distillation/callbacks.py


key: chroma:document
string_value:
CODE:

def count_trainable_parameters(model):
    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
    params = sum([np.prod(p.size()) for p in model_parameters])


SUMMARY:
**Summary**

The `count_trainable_parameters` function takes a PyTorch neural network model as input and returns the total number of trainable parameters.

**Key Details**

* Filters model parameters to only include those that require computational gradients using `model.parameters().filter(lambda p: p.requires_grad)`.
* Calculates the total number of trainable parameters by summing the product of each parameter's dimensions (`np.prod(p.size())`) using a list comprehension.
* Returns the total count.

**Implementation**

- Key algorithm: Filtering and summing model parameters to calculate total trainability.
- Method signature: `def count_trainable_parameters(model) -> int`

**Dependencies/Requirements**

* NumPy (`np`)
* PyTorch (specifically, model and parameter objects)
 
Note: No external dependencies or additional interfaces mentioned.


key: chroma:document
string_value:
CODE:

class Seq2SeqLoggingCallback(pl.Callback):
    def on_batch_end(self, trainer, pl_module):
        lrs = {f"lr_group_{i}": param["lr"] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups)}
        pl_module.logger.log_metrics(lrs)

    @rank_zero_only
    def _write_logs(
        self, trainer: pl.Trainer, pl_module: pl.LightningModule, type_path: str, save_generations=True
    ) -> None:
        logger.info(f"***** {type_path} results at step {trainer.global_step:05d} *****")
        metrics = trainer.callback_metrics
        trainer.logger.log_metrics({k: v for k, v in metrics.items() if k not in ["log", "progress_bar", "preds"]})
        # Log results
        od = Path(pl_module.hparams.output_dir)
        if type_path == "test":
            results_file = od / "test_results.txt"
            generations_file = od / "test_generations.txt"
        else:
            # this never gets hit. I prefer not to save intermediate generations, and results are in metrics.json
            # If people want this it will be easy enough to add back.
            results_file = od / f"{type_path}_results/{trainer.global_step:05d}.txt"
            generations_file = od / f"{type_path}_generations/{trainer.global_step:05d}.txt"
            results_file.parent.mkdir(exist_ok=True)
            generations_file.parent.mkdir(exist_ok=True)
        with open(results_file, "a+") as writer:
            for key in sorted(metrics):
                if key in ["log", "progress_bar", "preds"]:
                    continue
                val = metrics[key]
                if isinstance(val, torch.Tensor):
                    val = val.item()
                msg = f"{key}: {val:.6f}\n"
                writer.write(msg)

        if not save_generations:
            return

        if "preds" in metrics:
            content = "\n".join(metrics["preds"])
            generations_file.open("w+").write(content)

    @rank_zero_only
    def on_train_start(self, trainer, pl_module):
        try:
            npars = pl_module.model.model.num_parameters()
        except AttributeError:
            npars = pl_module.model.num_parameters()

        n_trainable_pars = count_trainable_parameters(pl_module)
        # mp stands for million parameters
        trainer.logger.log_metrics({"n_params": npars, "mp": npars / 1e6, "grad_mp": n_trainable_pars / 1e6})

    @rank_zero_only
    def on_test_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule):
        save_json(pl_module.metrics, pl_module.metrics_save_path)
        return self._write_logs(trainer, pl_module, "test")

    @rank_zero_only
    def on_validation_end(self, trainer: pl.Trainer, pl_module):
        save_json(pl_module.metrics, pl_module.metrics_save_path)
        # Uncommenting this will save val generations


SUMMARY:
**Sequence-to-Sequence Logging Callback Class Summary**

### Overview

This custom PyTorch Lightning `pl.Callback` class extends the built-in `pl.Callback` class to log training and evaluation metrics. It logs additional information such as the number of trainable parameters, learning rates (LRs), and model architecture details.

### Key Functions and Implementations

*   **_on_batch_end_**: Logs LR groups at each batch end.
    ```
    def on_batch_end(self, trainer, pl_module):
        lrs = {f"lr_group_{i}": param["lr"] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups)}
        pl_module.logger.log_metrics(lrs)
    """
*   `_write_logs_`: Writes test or validation results to files depending on the metric path.
    ```
    def _write_logs(
        self, trainer: pl.Trainer, pl_module: pl.LightningModule, type_path: str, save_generations=True
    ) -> None:
        ...
```
*   `_on_train_start_`: Logs model architecture details (number of trainable parameters and million parameters) before the training starts.
    ```
def _on_train_start(self, trainer, pl_module):
    try:
        npars = pl_module.model.model.num_parameters()
    except AttributeError:
        npars = pl_module.model.num_parameters()

    ...
```
*   `_on_test_end_` and `_on_validation_end_`: Calls the `_write_logs` method to save test or validation results to files.

### Dependencies and Requirements

The following dependencies are mentioned:

*   `pl.Trainer`, `pl.LightningModule`
*   External libraries: None explicitly mentioned
*   Inherited/Implemented interfaces: None explicitly mentioned


key: chroma:document
string_value:
CODE:

def get_checkpoint_callback(output_dir, metric, save_top_k=1, lower_is_better=False):
    """Saves the best model by validation ROUGE2 score."""
    if metric == "rouge2":
        exp = "{val_avg_rouge2:.4f}-{step_count}"
    elif metric == "bleu":
        exp = "{val_avg_bleu:.4f}-{step_count}"
    elif metric == "loss":
        exp = "{val_avg_loss:.4f}-{step_count}"
    else:
        raise NotImplementedError(
            f"seq2seq callbacks only support rouge2, bleu and loss, got {metric}, You can make your own by adding to"
            " this function."
        )

    checkpoint_callback = ModelCheckpoint(
        dirpath=output_dir,
        filename=exp,
        monitor=f"val_{metric}",
        mode="min" if "loss" in metric else "max",
        save_top_k=save_top_k,
    )


SUMMARY:
**Function Summary**

`get_checkpoint_callback`

Creates a model checkpoint callback function for saving the best model based on validation metrics.

**Key Implementation Details:**

* `dirpath` and `filename` are constructed dynamically using a ternary expression or conditionals depending on the metric type (`rouge2`, `bleu`, or `loss`)
* The `ModelCheckpoint` class is used to create a checkpoint callback, with parameters:
	+ `monitor`: uses a specific validation metric value
	+ `mode`: updates the mode based on the metric type ('min' for loss and 'max' otherwise)
	+ `save_top_k`: saves top k checkpoints

**Notable Dependencies or Requirements:**

* `ModelCheckpoint` class from a Keras/TensorFlow library (not explicitly stated, but assumed due to `dirpath`, `filename`, `monitor` usage)

**External Dependencies:**

* TensorFlow/Keras for the `ModelCheckpoint` class
* Keras/TensorFlow metrics ('val_avg_rouge2', 'val_avg_bleu', or 'val_avg_loss')


key: chroma:document
string_value:
CODE:

def get_early_stopping_callback(metric, patience):
    return EarlyStopping(
        monitor=f"val_{metric}",  # does this need avg?
        mode="min" if "loss" in metric else "max",
        patience=patience,
        verbose=True,


SUMMARY:
**Summary:**

The `get_early_stopping_callback` function returns an `EarlyStopping` callback configuration.

*   **Primary Purpose:** To monitor the early stopping of a metric (`metric`) with specified patience, adjusting for loss or maximization mode if present in the metric name.

**Key Implementation Details:**

*   `monitor=f"val_{metric}"`: Specifies the monitoring metric to track during validation.
*   `mode="min" if "loss" in metric else "max"` : Sets mode based on whether the metric is for loss (minimum) or maximization (maximum).
*   `patience=patience` : Accepts and configures the patience parameter.

**Notable Dependencies/Requirements:**

*   The function takes two parameters:
    *   `metric`: The evaluation metric to monitor during validation.
    *   `patience`: The number of epochs without improvement to restart training.

No external dependencies are mentioned in this code snippet.


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/seq2seq-distillation/_test_seq2seq_examples.py


key: chroma:document
string_value:
CODE:

def _dump_articles(path: Path, articles: list):
    content = "\n".join(articles)


SUMMARY:
**Summary**

`_dump_articles` function:

* Writes a list of article contents to a file at specified `path`.
* Takes two parameters: `path` (file location) and `articles` (list of content).

**Key implementation details**

* Uses `Path(path).open("w")` to open the file in write mode, ensuring the path is resolved correctly.
* Utilizes `writelines(content)` to efficiently write the combined article contents.

**Notable dependencies or requirements**

* `pathlib` library for working with file paths (specifically `Path` class).
* No external dependencies listed.
* Does not inherit from any interfaces.


key: chroma:document
string_value:
CODE:

def make_test_data_dir(tmp_dir):
    for split in ["train", "val", "test"]:
        _dump_articles(os.path.join(tmp_dir, f"{split}.source"), ARTICLES)
        _dump_articles(os.path.join(tmp_dir, f"{split}.target"), SUMMARIES)


SUMMARY:
**Summary**

This function, `make_test_data_dir`, creates test data directories for training and evaluation splits by dumping source and target articles to corresponding directories.

**Key Details**

*   Takes an optional temporary directory (`tmp_dir`) as input
*   Iterates over training (`"train"`), validation (`"val"`), and testing (`"test"` ) splits

**Implementation Highlights**

*   `_dump_articles` function is used to write source and target articles to directories in `ARTICLES` and `SUMMARIES`, respectively
*   Directory structure is expected with filenames having specific formats (`split.source` and `split.target`)

**External Dependencies and Requirements**

*   `os.path.join()` for constructing file paths
*   `_dump_articles()` function, which is not defined in this snippet ( likely resides elsewhere in the codebase)
*   `ARTICLES` and `SUMMARIES` datasets (not shown in this snippet)


key: chroma:document
string_value:
CODE:

class TestSummarizationDistiller(TestCasePlus):
    @classmethod
    def setUpClass(cls):
        logging.disable(logging.CRITICAL)  # remove noisy download output from tracebacks
        return cls

    @slow
    @require_torch_gpu
    def test_hub_configs(self):
        """I put require_torch_gpu cause I only want this to run with self-scheduled."""

        model_list = list_models()
        org = "sshleifer"
        model_ids = [x.modelId for x in model_list if x.modelId.startswith(org)]
        allowed_to_be_broken = ["sshleifer/blenderbot-3B", "sshleifer/blenderbot-90M"]
        failures = []
        for m in model_ids:
            if m in allowed_to_be_broken:
                continue
            try:
                AutoConfig.from_pretrained(m)
            except Exception:
                failures.append(m)
        assert not failures, f"The following models could not be loaded through AutoConfig: {failures}"

    def test_distill_no_teacher(self):
        updates = {"student_encoder_layers": 2, "student_decoder_layers": 1, "no_teacher": True}
        self._test_distiller_cli(updates)

    def test_distill_checkpointing_with_teacher(self):
        updates = {
            "student_encoder_layers": 2,
            "student_decoder_layers": 1,
            "max_epochs": 4,
            "val_check_interval": 0.25,
            "alpha_hid": 2.0,
            "model_name_or_path": "IGNORE_THIS_IT_DOESNT_GET_USED",
        }
        model = self._test_distiller_cli(updates, check_contents=False)

        ckpts = list(Path(model.output_dir).glob("*.ckpt"))
        self.assertEqual(1, len(ckpts))
        transformer_ckpts = list(Path(model.output_dir).glob("**/*.bin"))
        self.assertEqual(len(transformer_ckpts), 2)
        examples = lmap(str.strip, Path(model.hparams.data_dir).joinpath("test.source").open().readlines())
        out_path = tempfile.mktemp()  # XXX: not being cleaned up
        generate_summaries_or_translations(examples, out_path, str(model.output_dir / "best_tfmr"))
        self.assertTrue(Path(out_path).exists())

        out_path_new = self.get_auto_remove_tmp_dir()
        convert_pl_to_hf(ckpts[0], transformer_ckpts[0].parent, out_path_new)
        assert os.path.exists(os.path.join(out_path_new, "pytorch_model.bin"))

    def test_loss_fn(self):
        model = AutoModelForSeq2SeqLM.from_pretrained(BART_TINY)
        input_ids, mask = model.dummy_inputs["input_ids"], model.dummy_inputs["attention_mask"]
        target_ids = torch.tensor([[0, 4, 8, 2], [0, 8, 2, 1]], dtype=torch.long, device=model.device)
        decoder_input_ids = target_ids[:, :-1].contiguous()  # Why this line?
        lm_labels = target_ids[:, 1:].clone()  # why clone?
        model_computed_loss = model(
            input_ids, attention_mask=mask, decoder_input_ids=decoder_input_ids, labels=lm_labels, use_cache=False
        ).loss

        logits = model(input_ids, attention_mask=mask, decoder_input_ids=decoder_input_ids, use_cache=False).logits

        lprobs = nn.functional.log_softmax(logits, dim=-1)
        smoothed_loss, nll_loss = label_smoothed_nll_loss(
            lprobs, lm_labels, 0.1, ignore_index=model.config.pad_token_id
        )
        with self.assertRaises(AssertionError):
            # TODO: understand why this breaks
            self.assertEqual(nll_loss, model_computed_loss)

    def test_distill_mbart(self):
        updates = {
            "student_encoder_layers": 2,
            "student_decoder_layers": 1,
            "num_train_epochs": 4,
            "val_check_interval": 0.25,
            "alpha_hid": 2.0,
            "task": "translation",
            "model_name_or_path": "IGNORE_THIS_IT_DOESNT_GET_USED",
            "tokenizer_name": MBART_TINY,
            "teacher": MBART_TINY,
            "src_lang": "en_XX",
            "tgt_lang": "ro_RO",
        }
        model = self._test_distiller_cli(updates, check_contents=False)
        assert model.model.config.model_type == "mbart"

        ckpts = list(Path(model.output_dir).glob("*.ckpt"))
        self.assertEqual(1, len(ckpts))
        transformer_ckpts = list(Path(model.output_dir).glob("**/*.bin"))
        all_files = list(Path(model.output_dir).glob("best_tfmr/*"))
        assert len(all_files) > 2
        self.assertEqual(len(transformer_ckpts), 2)

    def test_distill_t5(self):
        updates = {
            "student_encoder_layers": 1,
            "student_decoder_layers": 1,
            "alpha_hid": 2.0,
            "teacher": T5_TINY,
            "model_name_or_path": T5_TINY,
            "tokenizer_name": T5_TINY,
        }
        self._test_distiller_cli(updates)

    def test_distill_different_base_models(self):
        updates = {
            "teacher": T5_TINY,
            "student": T5_TINIER,
            "model_name_or_path": T5_TINIER,
            "tokenizer_name": T5_TINIER,
        }
        self._test_distiller_cli(updates)

    def _test_distiller_cli(self, updates, check_contents=True):
        default_updates = {
            "label_smoothing": 0.0,
            "early_stopping_patience": -1,
            "train_batch_size": 1,
            "eval_batch_size": 2,
            "max_epochs": 2,
            "alpha_mlm": 0.2,
            "alpha_ce": 0.8,
            "do_predict": True,
            "model_name_or_path": "sshleifer/tinier_bart",
            "teacher": CHEAP_ARGS["model_name_or_path"],
            "val_check_interval": 0.5,
        }
        default_updates.update(updates)
        args_d: dict = CHEAP_ARGS.copy()
        tmp_dir = make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir())
        output_dir = self.get_auto_remove_tmp_dir()

        args_d.update(data_dir=tmp_dir, output_dir=output_dir, **default_updates)
        model = distill_main(argparse.Namespace(**args_d))
        if not check_contents:
            return model
        contents = os.listdir(output_dir)
        contents = {os.path.basename(p) for p in contents}
        ckpt_files = [p for p in contents if p.endswith("ckpt")]
        assert len(ckpt_files) > 0

        self.assertIn("test_generations.txt", contents)
        self.assertIn("test_results.txt", contents)

        metrics = load_json(model.metrics_save_path)
        last_step_stats = metrics["val"][-1]
        self.assertGreaterEqual(last_step_stats["val_avg_gen_time"], 0.01)
        self.assertGreaterEqual(1.0, last_step_stats["val_avg_gen_time"])
        self.assertIsInstance(last_step_stats[f"val_avg_{model.val_metric}"], float)
        desired_n_evals = int(args_d["max_epochs"] * (1 / args_d["val_check_interval"]) + 1)
        self.assertEqual(len(metrics["val"]), desired_n_evals)
        self.assertEqual(len(metrics["test"]), 1)


SUMMARY:
**Class: TestSummarizationDistiller**

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

The `TestSummarizationDistiller` class is a test suite for summarization distillation models, specifically designed to evaluate the performance of these models on various tasks. The main purpose is to test the model's ability to generate summaries, perform loss computation, and handle different baselines.

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

*   `test_hub_configs`: Tests the loading of hub configurations and checks if certain models can be loaded successfully.
*   `test_distill_no_teacher` and `test_distill_checkpointing_with_teacher`: Evaluate the model's performance when training without a teacher or with checkpointing, respectively.
*   `test_loss_fn`: Verifies the correct computation of loss during training.
*   `test_distill_mbart` and `test_distill_t5`: Test the models' performance on specific tasks, including translation and generation.

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

*   PyTorch: The class relies heavily on PyTorch for model architecture and operations.
*   TorchScript: Used for compiling PyTorch models to a fast and efficient byte-code format.
*   Hugging Face Transformers: Provides pre-trained models, tokenizers, and other utilities for natural language processing tasks.

**Code snippet**
```python
class TestSummarizationDistiller:
    def __init__(self):
        # Initialize test parameters and models
        self.model = None

    def _test_distiller_cli(self, updates):
        # Distill the model using the provided updates
        pass

    def test_hub_configs(self):
        # Load hub configurations and check successful loading
        pass

    def test_distill_no_teacher(self):
        # Train the model without a teacher
        pass

    def test_distill_checkpointing_with_teacher(self):
        # Train the model with checkpointing
        pass

    def test_loss_fn(self):
        # Compute loss during training
        pass

# Initialize the test class instance
test_summarization-distiller = TestSummarizationDistiller()
```

This summary captures the main purpose and functionality of the `TestSummarizationDistiller` class, as well as its notable dependencies and requirements. The class provides various tests for evaluating the model's performance on different tasks and scenarios.


key: chroma:document
string_value:
CODE:

class TestTheRest(TestCasePlus):
    @parameterized.expand(
        [T5_TINY, BART_TINY, MBART_TINY, MARIAN_TINY, FSMT_TINY],
    )
    def test_finetune(self, model):
        args_d: dict = CHEAP_ARGS.copy()
        task = "translation" if model in [MBART_TINY, MARIAN_TINY, FSMT_TINY] else "summarization"
        args_d["label_smoothing"] = 0.1 if task == "translation" else 0

        tmp_dir = make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir())
        output_dir = self.get_auto_remove_tmp_dir()
        args_d.update(
            data_dir=tmp_dir,
            model_name_or_path=model,
            tokenizer_name=None,
            train_batch_size=2,
            eval_batch_size=2,
            output_dir=output_dir,
            do_predict=True,
            task=task,
            src_lang="en_XX",
            tgt_lang="ro_RO",
            freeze_encoder=True,
            freeze_embeds=True,
        )
        assert "n_train" in args_d
        args = argparse.Namespace(**args_d)
        module = main(args)

        input_embeds = module.model.get_input_embeddings()
        assert not input_embeds.weight.requires_grad
        if model == T5_TINY:
            lm_head = module.model.lm_head
            assert not lm_head.weight.requires_grad
            assert (lm_head.weight == input_embeds.weight).all().item()
        elif model == FSMT_TINY:
            fsmt = module.model.model
            embed_pos = fsmt.decoder.embed_positions
            assert not embed_pos.weight.requires_grad
            assert not fsmt.decoder.embed_tokens.weight.requires_grad
            # check that embeds are not the same
            assert fsmt.decoder.embed_tokens != fsmt.encoder.embed_tokens
        else:
            bart = module.model.model
            embed_pos = bart.decoder.embed_positions
            assert not embed_pos.weight.requires_grad
            assert not bart.shared.weight.requires_grad
            # check that embeds are the same
            assert bart.decoder.embed_tokens == bart.encoder.embed_tokens
            assert bart.decoder.embed_tokens == bart.shared

        example_batch = load_json(module.output_dir / "text_batch.json")
        assert isinstance(example_batch, dict)
        assert len(example_batch) >= 4

    def test_finetune_extra_model_args(self):
        args_d: dict = CHEAP_ARGS.copy()

        task = "summarization"
        tmp_dir = make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir())

        args_d.update(
            data_dir=tmp_dir,
            tokenizer_name=None,
            train_batch_size=2,
            eval_batch_size=2,
            do_predict=False,
            task=task,
            src_lang="en_XX",
            tgt_lang="ro_RO",
            freeze_encoder=True,
            freeze_embeds=True,
        )

        # test models whose config includes the extra_model_args
        model = BART_TINY
        output_dir = self.get_auto_remove_tmp_dir()
        args_d1 = args_d.copy()
        args_d1.update(
            model_name_or_path=model,
            output_dir=output_dir,
        )
        extra_model_params = ("encoder_layerdrop", "decoder_layerdrop", "dropout", "attention_dropout")
        for p in extra_model_params:
            args_d1[p] = 0.5
        args = argparse.Namespace(**args_d1)
        model = main(args)
        for p in extra_model_params:
            assert getattr(model.config, p) == 0.5, f"failed to override the model config for param {p}"

        # test models whose config doesn't include the extra_model_args
        model = T5_TINY
        output_dir = self.get_auto_remove_tmp_dir()
        args_d2 = args_d.copy()
        args_d2.update(
            model_name_or_path=model,
            output_dir=output_dir,
        )
        unsupported_param = "encoder_layerdrop"
        args_d2[unsupported_param] = 0.5
        args = argparse.Namespace(**args_d2)
        with pytest.raises(Exception) as excinfo:
            model = main(args)
        assert str(excinfo.value) == f"model config doesn't have a `{unsupported_param}` attribute"

    def test_finetune_lr_schedulers(self):
        args_d: dict = CHEAP_ARGS.copy()

        task = "summarization"
        tmp_dir = make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir())

        model = BART_TINY
        output_dir = self.get_auto_remove_tmp_dir()

        args_d.update(
            data_dir=tmp_dir,
            model_name_or_path=model,
            output_dir=output_dir,
            tokenizer_name=None,
            train_batch_size=2,
            eval_batch_size=2,
            do_predict=False,
            task=task,
            src_lang="en_XX",
            tgt_lang="ro_RO",
            freeze_encoder=True,
            freeze_embeds=True,
        )

        # emulate finetune.py
        parser = argparse.ArgumentParser()
        parser = pl.Trainer.add_argparse_args(parser)
        parser = SummarizationModule.add_model_specific_args(parser, os.getcwd())
        args = {"--help": True}

        # --help test
        with pytest.raises(SystemExit) as excinfo:
            with CaptureStdout() as cs:
                args = parser.parse_args(args)
            assert False, "--help is expected to sys.exit"
        assert excinfo.type is SystemExit
        expected = lightning_base.arg_to_scheduler_metavar
        assert expected in cs.out, "--help is expected to list the supported schedulers"

        # --lr_scheduler=non_existing_scheduler test
        unsupported_param = "non_existing_scheduler"
        args = {f"--lr_scheduler={unsupported_param}"}
        with pytest.raises(SystemExit) as excinfo:
            with CaptureStderr() as cs:
                args = parser.parse_args(args)
            assert False, "invalid argument is expected to sys.exit"
        assert excinfo.type is SystemExit
        expected = f"invalid choice: '{unsupported_param}'"
        assert expected in cs.err, f"should have bailed on invalid choice of scheduler {unsupported_param}"

        # --lr_scheduler=existing_scheduler test
        supported_param = "cosine"
        args_d1 = args_d.copy()
        args_d1["lr_scheduler"] = supported_param
        args = argparse.Namespace(**args_d1)
        model = main(args)
        assert (
            getattr(model.hparams, "lr_scheduler") == supported_param


SUMMARY:
## TestTheRest Class Summary

### Overview
TestTheRest is a test class for finetuning Hugging Face models. It tests different models (T5, BART, FSMT) under various configurations.

### Main Purpose and Functionality
The class inherits from `TestCasePlus` and uses parameterized testing to run multiple tests with different model configurations. The primary purpose is to ensure that the models are correctly finetuned and behave as expected.

### Key Implementation Details

*   Used Hugging Face's transformers library for model loading and configuration.
*   Test cases include:
    *   `test_finetune`: Tests different models (T5, BART, FSMT) under various configurations, including tasks (translation and summarization), label smoothing, and output directories.
    *   `test_finetune_extra_model_args`: Tests models whose config includes the extra\_model\_args and those that don't.
    *   `test_finetune_lr_schedulers`: Tests different schedulers for finetuning.

### Notable Dependencies or Requirements

*   Hugging Face's transformers library (`transformers`)
*   PyTest
*   `argparse`
*   Plume (PyTorch Lightning)

### Class Structure and Interfaces

The class uses `argparse` to parse command-line arguments and `Plume` for Trainer configuration.

```python
class TestTheRest(TestCasePlus):
    @parameterized.expand(...)  # [T5_TINY, BART_TINY, MBART_TINY, MARIAN_TINY, FSMT_TINY],
    
    def test_finetune(self, model):
        ...
    ```


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/seq2seq-distillation/_test_bash_script.py


key: chroma:document
string_value:
CODE:

class TestMbartCc25Enro(TestCasePlus):
    def setUp(self):
        super().setUp()

        data_cached = cached_path(
            "https://cdn-datasets.huggingface.co/translation/wmt_en_ro-tr40k-va0.5k-te0.5k.tar.gz",
            extract_compressed_file=True,
        )
        self.data_dir = f"{data_cached}/wmt_en_ro-tr40k-va0.5k-te0.5k"

    @slow
    @require_torch_gpu
    def test_model_download(self):
        """This warms up the cache so that we can time the next test without including download time, which varies between machines."""
        MarianMTModel.from_pretrained(MARIAN_MODEL)

    # @timeout_decorator.timeout(1200)
    @slow
    @require_torch_gpu
    def test_train_mbart_cc25_enro_script(self):
        env_vars_to_replace = {
            "$MAX_LEN": 64,
            "$BS": 64,
            "$GAS": 1,
            "$ENRO_DIR": self.data_dir,
            "facebook/mbart-large-cc25": MARIAN_MODEL,
            # "val_check_interval=0.25": "val_check_interval=1.0",
            "--learning_rate=3e-5": "--learning_rate 3e-4",
            "--num_train_epochs 6": "--num_train_epochs 1",
        }

        # Clean up bash script
        bash_script = (self.test_file_dir / "train_mbart_cc25_enro.sh").open().read().split("finetune.py")[1].strip()
        bash_script = bash_script.replace("\\\n", "").strip().replace('"$@"', "")
        for k, v in env_vars_to_replace.items():
            bash_script = bash_script.replace(k, str(v))
        output_dir = self.get_auto_remove_tmp_dir()

        # bash_script = bash_script.replace("--fp16 ", "")
        args = f"""
            --output_dir {output_dir}
            --tokenizer_name Helsinki-NLP/opus-mt-en-ro
            --sortish_sampler
            --do_predict
            --gpus 1
            --freeze_encoder
            --n_train 40000
            --n_val 500
            --n_test 500
            --fp16_opt_level O1
            --num_sanity_val_steps 0
            --eval_beams 2
        """.split()
        # XXX: args.gpus > 1 : handle multi_gpu in the future

        testargs = ["finetune.py"] + bash_script.split() + args
        with patch.object(sys, "argv", testargs):
            parser = argparse.ArgumentParser()
            parser = pl.Trainer.add_argparse_args(parser)
            parser = SummarizationModule.add_model_specific_args(parser, os.getcwd())
            args = parser.parse_args()
            model = main(args)

        # Check metrics
        metrics = load_json(model.metrics_save_path)
        first_step_stats = metrics["val"][0]
        last_step_stats = metrics["val"][-1]
        self.assertEqual(len(metrics["val"]), (args.max_epochs / args.val_check_interval))
        assert isinstance(last_step_stats[f"val_avg_{model.val_metric}"], float)

        self.assertGreater(last_step_stats["val_avg_gen_time"], 0.01)
        # model hanging on generate. Maybe bad config was saved. (XXX: old comment/assert?)
        self.assertLessEqual(last_step_stats["val_avg_gen_time"], 1.0)

        # test learning requirements:

        # 1. BLEU improves over the course of training by more than 2 pts
        self.assertGreater(last_step_stats["val_avg_bleu"] - first_step_stats["val_avg_bleu"], 2)

        # 2. BLEU finishes above 17
        self.assertGreater(last_step_stats["val_avg_bleu"], 17)

        # 3. test BLEU and val BLEU within ~1.1 pt.
        self.assertLess(abs(metrics["val"][-1]["val_avg_bleu"] - metrics["test"][-1]["test_avg_bleu"]), 1.1)

        # check lightning ckpt can be loaded and has a reasonable statedict
        contents = os.listdir(output_dir)
        ckpt_path = [x for x in contents if x.endswith(".ckpt")][0]
        full_path = os.path.join(args.output_dir, ckpt_path)
        ckpt = torch.load(full_path, map_location="cpu")
        expected_key = "model.model.decoder.layers.0.encoder_attn_layer_norm.weight"
        assert expected_key in ckpt["state_dict"]
        assert ckpt["state_dict"]["model.model.decoder.layers.0.encoder_attn_layer_norm.weight"].dtype == torch.float32

        # TODO: turn on args.do_predict when PL bug fixed.
        if args.do_predict:
            contents = {os.path.basename(p) for p in contents}
            assert "test_generations.txt" in contents
            assert "test_results.txt" in contents
            # assert len(metrics["val"]) ==  desired_n_evals


SUMMARY:
**Summary**

The `TestMbartCc25Enro` class tests the MBART Large-CC25 model for machine translation from English to Romanian. It downloads and fine-tunes the model, checks its performance, and verifies its loading with a saved checkpoint.

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

*   Downloads and extracts the WMT En-Ro dataset using `cached_path`.
*   Creates a bash script to run the fine-tuning process.
*   Runs fine-tuning using `finetune.py`, which takes in parsed command-line arguments, executes the bash script, and performs hyperparameter tuning and testing.

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

*   The test uses the `MarianMTModel` class with a pre-trained checkpoint (`MARIAN_MODEL`) and applies transformations to the bash script using environment variables (`env_vars_to_replace`).
*   It verifies that specific metrics (e.g., BLEU score, generation time) are reported after fine-tuning.

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

*   The `torch` library for PyTorch models.
*   The `argparse` library for command-line argument parsing.
*   The `pl.Trainer` class from the Hugging Face Lightning Library.
*   The `SummarizationModule` from the Hugging Face Transformers Library.

```markdown
# Class Summary

## Class Definition
---------------

*   `TestMbartCc25Enro`: Test class for the MBART Large-CC25 model.

## Methods

*   `setUp`: Sets up the test environment by downloading and extracting the dataset.
*   `test_model_download`: Tests whether the MBART model can be loaded and pre-trained from a checkpoint.
*   `test_train_mbart_cc25_enro_script`: Tests the fine-tuning process, including execution and verification of metrics.

## Dependencies

*   `torch` for PyTorch models
*   `argparse` for command-line argument parsing
*   `pl.Trainer` from Lightning Library
*   `SummarizationModule` from Transformers Library
```


key: chroma:document
string_value:
CODE:

class TestDistilMarianNoTeacher(TestCasePlus):
    @timeout_decorator.timeout(600)
    @slow
    @require_torch_gpu
    def test_opus_mt_distill_script(self):
        data_dir = f"{self.test_file_dir_str}/test_data/wmt_en_ro"
        env_vars_to_replace = {
            "--fp16_opt_level=O1": "",
            "$MAX_LEN": 128,
            "$BS": 16,
            "$GAS": 1,
            "$ENRO_DIR": data_dir,
            "$m": "sshleifer/student_marian_en_ro_6_1",
            "val_check_interval=0.25": "val_check_interval=1.0",
        }

        # Clean up bash script
        bash_script = (
            (self.test_file_dir / "distil_marian_no_teacher.sh").open().read().split("distillation.py")[1].strip()
        )
        bash_script = bash_script.replace("\\\n", "").strip().replace('"$@"', "")
        bash_script = bash_script.replace("--fp16 ", " ")

        for k, v in env_vars_to_replace.items():
            bash_script = bash_script.replace(k, str(v))
        output_dir = self.get_auto_remove_tmp_dir()
        bash_script = bash_script.replace("--fp16", "")
        epochs = 6
        testargs = (
            ["distillation.py"]
            + bash_script.split()
            + [
                f"--output_dir={output_dir}",
                "--gpus=1",
                "--learning_rate=1e-3",
                f"--num_train_epochs={epochs}",
                "--warmup_steps=10",
                "--val_check_interval=1.0",
                "--do_predict",
            ]
        )
        with patch.object(sys, "argv", testargs):
            parser = argparse.ArgumentParser()
            parser = pl.Trainer.add_argparse_args(parser)
            parser = SummarizationDistiller.add_model_specific_args(parser, os.getcwd())
            args = parser.parse_args()
            # assert args.gpus == gpus THIS BREAKS for multi_gpu

            model = distill_main(args)

        # Check metrics
        metrics = load_json(model.metrics_save_path)
        first_step_stats = metrics["val"][0]
        last_step_stats = metrics["val"][-1]
        assert len(metrics["val"]) >= (args.max_epochs / args.val_check_interval)  # +1 accounts for val_sanity_check

        assert last_step_stats["val_avg_gen_time"] >= 0.01

        assert first_step_stats["val_avg_bleu"] < last_step_stats["val_avg_bleu"]  # model learned nothing
        assert 1.0 >= last_step_stats["val_avg_gen_time"]  # model hanging on generate. Maybe bad config was saved.
        assert isinstance(last_step_stats[f"val_avg_{model.val_metric}"], float)

        # check lightning ckpt can be loaded and has a reasonable statedict
        contents = os.listdir(output_dir)
        ckpt_path = [x for x in contents if x.endswith(".ckpt")][0]
        full_path = os.path.join(args.output_dir, ckpt_path)
        ckpt = torch.load(full_path, map_location="cpu")
        expected_key = "model.model.decoder.layers.0.encoder_attn_layer_norm.weight"
        assert expected_key in ckpt["state_dict"]
        assert ckpt["state_dict"]["model.model.decoder.layers.0.encoder_attn_layer_norm.weight"].dtype == torch.float32

        # TODO: turn on args.do_predict when PL bug fixed.
        if args.do_predict:
            contents = {os.path.basename(p) for p in contents}
            assert "test_generations.txt" in contents
            assert "test_results.txt" in contents
            # assert len(metrics["val"]) ==  desired_n_evals


SUMMARY:
**Class Summary:**

`TestDistilMarianNoTeacher`
==========================

A test class that verifies the `distil_marian_no_teacher` model's performance using OPUS-MT distillation.

*   Tests the model's ability to learn from teacher-free inference using a bash script.
    *   Verifies that the model improves over iterations by checking BLEU scores and generation times.
*   Validates the model's loadability from a checkpoint file (`.ckpt`).
*   Uses PyTorch `Trainer` and SummarizationDistiller utilities to add command-line arguments.

**Key Implementation Details:**

*   Modifies a bash script to remove unnecessary parameters and replaces environment variables with hardcoded values.
*   Creates a list of command-line arguments for the `distil_marian_no_teacher` model using `argparse`.
*   Verifies that the model improves over iterations by checking metrics (BLEU scores, generation times, model hanging).
*   Loads a checkpoint file (`.ckpt`) from an output directory and verifies its contents.

**Notable Dependencies:**

*   PyTorch
*   SummarizationDistiller
*   `argparse`
*   required environment variables (e.g., MAX_LEN, BS)
*   External dependencies:

    *   `distil_marian_no_teacher` model
    *   OPUS-MT distillation dataset

**Method Signatures and Interfaces:**

*   `test_opus_mt_distill_script`: The test method that runs the OPUS-MT distillation experiment.
*   `pl.Trainer.add_argparse_args`
*   `SummarizationDistiller.add_model_specific_args`

**Interface Dependencies:**

*   `os` for file system operations
*   `sys` for patching `sys.argv`
*   `argparse` for command-line argument parsing


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/seq2seq-distillation/utils.py


key: chroma:document
string_value:
CODE:

def label_smoothed_nll_loss(lprobs, target, epsilon, ignore_index=-100):
    """From fairseq"""
    if target.dim() == lprobs.dim() - 1:
        target = target.unsqueeze(-1)
    nll_loss = -lprobs.gather(dim=-1, index=target)
    smooth_loss = -lprobs.sum(dim=-1, keepdim=True)
    if ignore_index is not None:
        pad_mask = target.eq(ignore_index)
        nll_loss.masked_fill_(pad_mask, 0.0)
        smooth_loss.masked_fill_(pad_mask, 0.0)
    else:
        nll_loss = nll_loss.squeeze(-1)
        smooth_loss = smooth_loss.squeeze(-1)

    nll_loss = nll_loss.sum()  # mean()? Scared to break other math.
    smooth_loss = smooth_loss.sum()
    eps_i = epsilon / lprobs.size(-1)
    loss = (1.0 - epsilon) * nll_loss + eps_i * smooth_loss


SUMMARY:
**Function Summary: `label_smoothed_nll_loss`**

*Purpose:* Computes the smoothed NLL loss for sequence prediction.

**Main Purpose:**
The function calculates the smoothed Negative Log Likelihood (NLL) loss using label smoothing. It takes in predicted log probabilities (`lprobs`), target labels (`target`), a smoothing factor (`epsilon`), and an optional `ignore_index`.

**Key Implementation Details:**

* Uses `torch.gather` to calculate NLL loss.
* Mixture of true NLL loss and smoothed NLL loss using label smoothing.
* Scales epsilon by predicted logit dimension.

**Notable Dependencies:**

* Torch (`torch.gather`, `torch.unsqueeze`)

**Return Values:** `(loss, nll_loss)`

* `loss`: Smoothed NLL loss with added epsilon smoothing.
* `nll_loss`: True NLL loss without smoothing.


key: chroma:document
string_value:
CODE:

def lmap(f: Callable, x: Iterable) -> List:
    """list(map(f, x))"""


SUMMARY:
**Summary:**

This function, `lmap`, is an alias for the built-in `map` function that applies a given function `f` to each element of an iterable `x`. It returns a new list containing the results.

**Key Implementation Details:**

* Applies function `f` to each element in `x`
* Returns a list containing the mapped values
* Used as an alias for `list(map(f, x))`

**Notable Dependencies/Requirements:**

* `f`: a callable (function) that takes one argument
* `x`: an iterable (e.g., list, tuple, set)
* Built-in `map`, ` Callable`, `List`, and `Iterable` types


key: chroma:document
string_value:
CODE:

def calculate_bleu(output_lns, refs_lns, **kwargs) -> dict:
    """Uses sacrebleu's corpus_bleu implementation."""


SUMMARY:
**Summary:**

`calculate Bleu score`

Calculates the BLEU score using SacreBLEU's `corpus_bleu` implementation. The function takes two lists of lines (`output_lns` and `refs_lns`) as input, along with additional keyword arguments (`kwargs`). Returns a dictionary containing the calculated BLEU score.

**Key Details:**

* Uses `corpus_bleu` to calculate BLEU score
* Input: `output_lns`, `refs_lns`, and custom keyword arguments (`kwargs`)
* Output: Dictionary with BLEU score rounded to 4 decimal places

No external dependencies or notable interfaces are listed.


key: chroma:document
string_value:
CODE:

def build_compute_metrics_fn(task_name: str, tokenizer: PreTrainedTokenizer) -> Callable[[EvalPrediction], Dict]:
    def non_pad_len(tokens: np.ndarray) -> int:
        return np.count_nonzero(tokens != tokenizer.pad_token_id)

    def decode_pred(pred: EvalPrediction) -> Tuple[List[str], List[str]]:
        pred_str = tokenizer.batch_decode(pred.predictions, skip_special_tokens=True)
        label_str = tokenizer.batch_decode(pred.label_ids, skip_special_tokens=True)
        pred_str = lmap(str.strip, pred_str)
        label_str = lmap(str.strip, label_str)
        return pred_str, label_str

    def summarization_metrics(pred: EvalPrediction) -> Dict:
        pred_str, label_str = decode_pred(pred)
        rouge: Dict = calculate_rouge(pred_str, label_str)
        summ_len = np.round(np.mean(lmap(non_pad_len, pred.predictions)), 1)
        rouge.update({"gen_len": summ_len})
        return rouge

    def translation_metrics(pred: EvalPrediction) -> Dict:
        pred_str, label_str = decode_pred(pred)
        bleu: Dict = calculate_bleu(pred_str, label_str)
        gen_len = np.round(np.mean(lmap(non_pad_len, pred.predictions)), 1)
        bleu.update({"gen_len": gen_len})
        return bleu

    compute_metrics_fn = summarization_metrics if "summarization" in task_name else translation_metrics


SUMMARY:
**Compute Metrics Function Builder**

The `build_compute_metrics_fn` function constructs and returns a custom computing metrics function (`compute_metrics_fn`) for evaluating model performance on text classification tasks. The built function depends on the task type indicated by the input `task_name`.

**Key Implementation Details:**

1.  **Evaluation Metrics:** Two functions are implemented to compute metrics on evaluation predictions:
    *   `summarization_metrics` (default): Computed using ROUGE score with label lengths.
    *   `translation_metrics`: Computed using BLEU score with label lengths.
2.  **Input Decoding:** A shared decoding function (`decode_pred`) processes prediction values to transform them into strings, removing padding tokens and stripping extra whitespace.
3.  **Input Functions:**
    *   `non_pad_len` calculates the length of non-pad elements in a tensor.

**Notable Dependencies and Requirements:**

*   The `tokenizer` object from Hugging Face's Transformers library (`PreTrainedTokenizer`) is used to process input data.
*   Two external dependencies are not explicitly referenced:
    *   `calculate_rouge`: A function calculating ROUGE scores, though its implementation is not provided here.
    *   `calculate_bleu`: A function calculating BLEU scores, but again, its implementation details are not included in this summary.

**Return Type:**
The built computing metrics function returns a dictionary (`Dict`) containing evaluation results for the specified task type.


key: chroma:document
string_value:
CODE:

def trim_batch(
    input_ids,
    pad_token_id,
    attention_mask=None,
):
    """Remove columns that are populated exclusively by pad_token_id"""
    keep_column_mask = input_ids.ne(pad_token_id).any(dim=0)
    if attention_mask is None:
        return input_ids[:, keep_column_mask]
    else:


SUMMARY:
**Summary:**

The `trim_batch` function filters columns in input tensors based on the presence of a specific pad token (`pad_token_id`). It identifies which columns are exclusively populated with this token using mask manipulation.

**Key Details:**

* Filters input tensor (`input_ids`) to remove columns with only `pad_token_id` values.
* Accepts an optional argument `attention_mask`, if provided, filters both `input_ids` and its corresponding mask values.
* Returns two output tensors: filtered `input_ids` without the exclusive pad token columns, or both `input_ids` and `attention_mask`.

**Notable Dependencies/Requirements:**

* PyTorch library for tensor operations (`ne`, `any`, `dim=0`).
* Optional dependency: `attention_mask` tensor (similar type to `input_ids`).


key: chroma:document
string_value:
CODE:

class AbstractSeq2SeqDataset(Dataset):
    def __init__(
        self,
        tokenizer,
        data_dir,
        max_source_length,
        max_target_length,
        type_path="train",
        n_obs=None,
        prefix="",
        **dataset_kwargs,
    ):
        super().__init__()
        self.src_file = Path(data_dir).joinpath(type_path + ".source")
        self.tgt_file = Path(data_dir).joinpath(type_path + ".target")
        self.len_file = Path(data_dir).joinpath(type_path + ".len")
        if os.path.exists(self.len_file):
            self.src_lens = pickle_load(self.len_file)
            self.used_char_len = False
        else:
            self.src_lens = self.get_char_lens(self.src_file)
            self.used_char_len = True
        self.max_source_length = max_source_length
        self.max_target_length = max_target_length
        assert min(self.src_lens) > 0, f"found empty line in {self.src_file}"
        self.tokenizer = tokenizer
        self.prefix = prefix if prefix is not None else ""

        if n_obs is not None:
            self.src_lens = self.src_lens[:n_obs]
        self.pad_token_id = self.tokenizer.pad_token_id
        self.dataset_kwargs = dataset_kwargs
        dataset_kwargs.update({"add_prefix_space": True} if isinstance(self.tokenizer, BartTokenizer) else {})

    def __len__(self):
        return len(self.src_lens)

    @staticmethod
    def get_char_lens(data_file):
        return [len(x) for x in Path(data_file).open().readlines()]

    @cached_property
    def tgt_lens(self):
        """Length in characters of target documents"""
        return self.get_char_lens(self.tgt_file)

    def make_sortish_sampler(self, batch_size, distributed=False, shuffle=True, **kwargs):
        if distributed:
            return DistributedSortishSampler(self, batch_size, shuffle=shuffle, **kwargs)
        else:
            return SortishSampler(self.src_lens, batch_size, shuffle=shuffle)

    def make_dynamic_sampler(self, max_tokens_per_batch=1024, **kwargs):
        assert FAIRSEQ_AVAILABLE, "Dynamic batch size requires `pip install fairseq`"
        assert not self.used_char_len, "You must call  python make_len_file.py before calling make_dynamic_sampler"
        sorted_indices = list(self.make_sortish_sampler(1024, shuffle=False))

        def num_tokens_in_example(i):
            return min(self.src_lens[i], self.max_target_length)

        # call fairseq cython function
        batch_sampler: List[List[int]] = batch_by_size(
            sorted_indices,
            num_tokens_fn=num_tokens_in_example,
            max_tokens=max_tokens_per_batch,
            required_batch_size_multiple=64,
        )
        shuffled_batches = [batch_sampler[i] for i in np.random.permutation(range(len(batch_sampler)))]
        # move the largest batch to the front to OOM quickly (uses an approximation for padding)
        approximate_toks_per_batch = [max(self.src_lens[i] for i in batch) * len(batch) for batch in shuffled_batches]
        largest_batch_idx = np.argmax(approximate_toks_per_batch)
        shuffled_batches[0], shuffled_batches[largest_batch_idx] = (
            shuffled_batches[largest_batch_idx],
            shuffled_batches[0],
        )
        return shuffled_batches

    def __getitem__(self, item):
        raise NotImplementedError("You must implement this")

    def collate_fn(self, batch):


SUMMARY:
**Summary: AbstractSeq2SeqDataset Class**

The `AbstractSeq2SeqDataset` class is a custom dataset class for sequences-to-sequence tasks. It extends the PyTorch `Dataset` class.

**Primary Purpose:** Defines a sequence-to-sequence dataset with optional data preprocessing and shuffling.

**Key Implementation Details:**

1. Initializes dataset paths, lengths, and preprocesses data files.
2. Supports dynamic batch sampling using Fairseq (requires `pip install fairseq`).
3. Allows for filtering datasets to a specified observation count (`n_obs`).

**Notable Dependencies/Requirements:**

* Fairseq (for dynamic batch size)
* `BartTokenizer` (for adding prefix space)

**Interfaces Implemented:**

* PyTorch `Dataset`
* `DistributedSortishSampler`
* `SortishSampler`

**Notes:**

The class raises a `NotImplementedError` for accessing individual data points (`__getitem__`) and collating batches (`collate_fn`). Subclass implementations are required to fulfill these requirements.


key: chroma:document
string_value:
CODE:

class LegacySeq2SeqDataset(AbstractSeq2SeqDataset):
    def __getitem__(self, index) -> Dict[str, torch.Tensor]:
        """Call tokenizer on src and tgt_lines"""
        index = index + 1  # linecache starts at 1
        source_line = self.prefix + linecache.getline(str(self.src_file), index).rstrip("\n")
        tgt_line = linecache.getline(str(self.tgt_file), index).rstrip("\n")
        assert source_line, f"empty source line for index {index}"
        assert tgt_line, f"empty tgt line for index {index}"
        source_inputs = self.encode_line(self.tokenizer, source_line, self.max_source_length)
        target_inputs = self.encode_line(self.tokenizer, tgt_line, self.max_target_length)

        source_ids = source_inputs["input_ids"].squeeze()
        target_ids = target_inputs["input_ids"].squeeze()
        src_mask = source_inputs["attention_mask"].squeeze()
        return {
            "input_ids": source_ids,
            "attention_mask": src_mask,
            "labels": target_ids,
        }

    def encode_line(self, tokenizer, line, max_length, pad_to_max_length=True, return_tensors="pt"):
        """Only used by LegacyDataset"""
        return tokenizer(
            [line],
            max_length=max_length,
            padding="max_length" if pad_to_max_length else None,
            truncation=True,
            return_tensors=return_tensors,
            **self.dataset_kwargs,
        )

    def collate_fn(self, batch) -> Dict[str, torch.Tensor]:
        input_ids = torch.stack([x["input_ids"] for x in batch])
        masks = torch.stack([x["attention_mask"] for x in batch])
        target_ids = torch.stack([x["labels"] for x in batch])
        pad_token_id = self.pad_token_id
        y = trim_batch(target_ids, pad_token_id)
        source_ids, source_mask = trim_batch(input_ids, pad_token_id, attention_mask=masks)
        batch = {
            "input_ids": source_ids,
            "attention_mask": source_mask,
            "labels": y,
        }


SUMMARY:
**Summary: LegacySeq2SeqDataset Class**

**Main Purpose and Functionality:**
The `LegacySeq2SeqDataset` class is an implementation of the abstract `AbstractSeq2SeqDataset` class, providing a custom dataset for sequence-to-sequence tasks. It loads source and target text lines from files using `linecache`, tokenizes them with a tokenizer, and returns the encoded input IDs, attention masks, and labels.

**Key Implementation Details:**

* The `__getitem__` method retrieves a sample by index, processing the source and target lines using the `tokenizer`.
	+ Tokenization is performed on source and target lines, up to the maximum sequence length.
	+ Assert statements are used to ensure that both source and target lines exist at each index.
* The `encode_line` method tokenizes a single line of text using the tokenizer.
	+ Parameters: tokenizer, line, max_length, pad_to_max_length, return_tensors
	+ Returns: encoded input IDs, attention mask, and labels.
* The `collate_fn` method collates a batch of samples into a single dictionary.
	+ It creates tensor values for input IDs, attention masks, and labels using torch.stack.

**Notable Dependencies or Requirements:**

* Related classes:
	+ `AbstractSeq2SeqDataset`
	+ `tokenizer`
* External dependencies:
	+ `linecache`
* Interfaces implemented:
	+ Inherited from `AbstractSeq2SeqDataset`: `__getitem__`, `collate_fn`


key: chroma:document
string_value:
CODE:

class Seq2SeqDataset(AbstractSeq2SeqDataset):
    """A dataset that calls prepare_seq2seq_batch."""

    def __getitem__(self, index) -> Dict[str, str]:
        index = index + 1  # linecache starts at 1
        source_line = self.prefix + linecache.getline(str(self.src_file), index).rstrip("\n")
        tgt_line = linecache.getline(str(self.tgt_file), index).rstrip("\n")
        assert source_line, f"empty source line for index {index}"
        assert tgt_line, f"empty tgt line for index {index}"
        return {"tgt_texts": tgt_line, "src_texts": source_line, "id": index - 1}

    def collate_fn(self, batch) -> Dict[str, torch.Tensor]:
        """Call prepare_seq2seq_batch."""
        batch_encoding: Dict[str, torch.Tensor] = self.tokenizer.prepare_seq2seq_batch(
            [x["src_texts"] for x in batch],
            tgt_texts=[x["tgt_texts"] for x in batch],
            max_length=self.max_source_length,
            max_target_length=self.max_target_length,
            return_tensors="pt",
            **self.dataset_kwargs,
        ).data
        batch_encoding["ids"] = torch.tensor([x["id"] for x in batch])


SUMMARY:
**Class Summary: Seq2SeqDataset**

**Purpose**: A dataset class that loads and preprocesses sequential data for seq-to-seq tasks.

*   The primary purpose is to call the `prepare_seq2seq_batch` method from the tokenizer module to batch and encode input/output data.
    **Key Implementation Details**:
        *   Use of `linecache` to retrieve source and target lines from files (`self.src_file` and `self.tgt_file`).
        *   Tokenization using the `tokenizer.prepare_seq2seq_batch` method with specified maximum length parameters (`max_source_length` and `max_target_length`).
    **Notable Dependencies**:
        *   `linecache`: used for file-line retrieval.
        *   `tokenizer`: provides the `prepare_seq2seq_batch` method for tokenization.
*   The class also defines a custom collate function to batch and encoding input/output data.
    **Interface/Inherited Interfaces**: Exposes `__getitem__` and `collate_fn` methods that conform to PyTorch's data pipeline interface.


key: chroma:document
string_value:
CODE:

class Seq2SeqDataCollator:
    def __init__(self, tokenizer, data_args, tpu_num_cores=None):
        self.tokenizer = tokenizer
        self.pad_token_id = tokenizer.pad_token_id
        assert (
            self.pad_token_id is not None
        ), f"pad_token_id is not defined for ({self.tokenizer.__class__.__name__}), it must be defined."
        self.data_args = data_args
        self.tpu_num_cores = tpu_num_cores
        self.dataset_kwargs = {"add_prefix_space": True} if isinstance(tokenizer, BartTokenizer) else {}
        if data_args.src_lang is not None:
            self.dataset_kwargs["src_lang"] = data_args.src_lang
        if data_args.tgt_lang is not None:
            self.dataset_kwargs["tgt_lang"] = data_args.tgt_lang

    def __call__(self, batch) -> Dict[str, torch.Tensor]:
        if hasattr(self.tokenizer, "prepare_seq2seq_batch"):
            batch = self._encode(batch)
            input_ids, attention_mask, labels = (
                batch["input_ids"],
                batch["attention_mask"],
                batch["labels"],
            )
        else:
            input_ids = torch.stack([x["input_ids"] for x in batch])
            attention_mask = torch.stack([x["attention_mask"] for x in batch])
            labels = torch.stack([x["labels"] for x in batch])

            labels = trim_batch(labels, self.pad_token_id)
            input_ids, attention_mask = trim_batch(input_ids, self.pad_token_id, attention_mask=attention_mask)

        if isinstance(self.tokenizer, T5Tokenizer):
            decoder_input_ids = self._shift_right_t5(labels)
        else:
            decoder_input_ids = shift_tokens_right(labels, self.pad_token_id)

        batch = {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "decoder_input_ids": decoder_input_ids,
            "labels": labels,
        }
        return batch

    def _shift_right_t5(self, input_ids):
        # shift inputs to the right
        shifted_input_ids = input_ids.new_zeros(input_ids.shape)
        shifted_input_ids[..., 1:] = input_ids[..., :-1].clone()
        shifted_input_ids[..., 0] = self.pad_token_id
        return shifted_input_ids

    def _encode(self, batch) -> Dict[str, torch.Tensor]:
        batch_encoding = self.tokenizer.prepare_seq2seq_batch(
            [x["src_texts"] for x in batch],
            tgt_texts=[x["tgt_texts"] for x in batch],
            max_length=self.data_args.max_source_length,
            max_target_length=self.data_args.max_target_length,
            padding="max_length" if self.tpu_num_cores is not None else "longest",  # TPU hack
            return_tensors="pt",
            **self.dataset_kwargs,
        )


SUMMARY:
**Class Summary: Seq2SeqDataCollator**

**Purpose and Functionality**
The `Seq2SeqDataCollator` class collates batched data for sequence-to-sequence tasks, processing input and target text sequences.

**Key Implementation Details**

* Prepares batch encoding using `tokenizer.prepare_seq2seq_batch`, considering source and target language models (e.g., T5 or Bart).
* Applies padding to input sequences using the `pad_token_id` value.
* Shifts input IDs to the right for T5-based models, adding a special token at the beginning.
* Returns a dictionary containing preprocessed input IDs, attention masks, decoder input IDs, and labels.

**Notable Dependencies and Requirements**

* `tokenizer`: an instance of a sequence-to-sequence-aware tokenizer (e.g., `T5Tokenizer` or `BartTokenizer`).
* `data_args`: a dataset-related configuration object.
* `tpu_num_cores`: optional integer indicating the number of TPU cores.

**Method Signatures and Interfaces**

* `__init__(self, tokenizer, data_args, tpu_num_cores=None)`: constructor.
* `__call__(self, batch) -> Dict[str, torch.Tensor]`: collates batched data.
* `_encode(self, batch) -> Dict[str, torch tensor]` : internal encoding method.
* `_shift_right_t5(self, input_ids)` : shifts input IDs for T5-based models.


key: chroma:document
string_value:
CODE:

class SortishSampler(Sampler):
    "Go through the text data by order of src length with a bit of randomness. From fastai repo."

    def __init__(self, data, batch_size, shuffle=True):
        self.data, self.bs, self.shuffle = data, batch_size, shuffle

    def __len__(self) -> int:
        return len(self.data)

    def __iter__(self):


SUMMARY:
**Class Summary**

`SortishSampler` is a customized `Sampler` class for sampling data with a focus on source length order and added randomness. It takes in `data`, `batch_size`, and optional `shuffle` parameters.

**Key Implementation Details**

* Serves text data sorted by source length, introducing randomness via `sortish_sampler_indices`.
* Implements custom `length()` and `__iter__()` methods for sequence management.

**Notable Dependencies/Requirements**

* None explicitly listed; inherits from `Sampler` class (exact implementation not provided).
* External dependency: `fastai` repository, specifically the `SortishSampler` algorithm.


key: chroma:document
string_value:
CODE:

def sortish_sampler_indices(data: List, bs: int, shuffle=True) -> np.array:
    "Go through the text data by order of src length with a bit of randomness. From fastai repo."
    if not shuffle:
        return np.argsort(np.array(data) * -1)

    def key_fn(i):
        return data[i]

    idxs = np.random.permutation(len(data))
    sz = bs * 50
    ck_idx = [idxs[i : i + sz] for i in range(0, len(idxs), sz)]
    sort_idx = np.concatenate([sorted(s, key=key_fn, reverse=True) for s in ck_idx])
    sz = bs
    ck_idx = [sort_idx[i : i + sz] for i in range(0, len(sort_idx), sz)]
    max_ck = np.argmax([key_fn(ck[0]) for ck in ck_idx])  # find the chunk with the largest key,
    ck_idx[0], ck_idx[max_ck] = ck_idx[max_ck], ck_idx[0]  # then make sure it goes first.
    sort_idx = np.concatenate(np.random.permutation(ck_idx[1:])) if len(ck_idx) > 1 else np.array([], dtype=int)
    sort_idx = np.concatenate((ck_idx[0], sort_idx))


SUMMARY:
Here is a concise and focused summary of the provided function code chunk:

**Function Summary:**

`sortish_sampler_indices(data: List, bs: int, shuffle=True) -> np.array`

This function generates an array of indices (`sort_idx`) based on input `data`. It sorts the indices by source length in descending order with a mix of randomness. If `shuffle=False`, it returns an ascending ordered array.

**Key Implementation Details:**

* The algorithm combines both sorted and random permutation for index sorting.
* It uses `sorted` function with a custom key (`key_fn`) to sort data by source length.

**Notable Dependencies/Requirements:**

* Uses NumPy (`np.array`, `np.argsort`, etc.)
* Imports `List` from Python's built-in `typing` module or external library (e.g., Fastai repo).

Overall, the function is designed for sampling indices in a way that balances order and randomness.


key: chroma:document
string_value:
CODE:

class DistributedSortishSampler(Sampler):
    """Copied from torch DistributedSampler"""

    def __init__(self, dataset, batch_size, num_replicas=None, rank=None, add_extra_examples=True, shuffle=True):
        if num_replicas is None:
            if not dist.is_available():
                raise RuntimeError("Requires distributed package to be available")
            num_replicas = dist.get_world_size()
        if rank is None:
            if not dist.is_available():
                raise RuntimeError("Requires distributed package to be available")
            rank = dist.get_rank()
        self.dataset = dataset
        self.num_replicas = num_replicas
        self.rank = rank
        self.epoch = 0
        if add_extra_examples:
            self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.num_replicas))
            self.total_size = self.num_samples * self.num_replicas
        else:
            self.total_size = len(dataset)
            self.num_samples = len(self.available_indices)
        self.batch_size = batch_size
        self.add_extra_examples = add_extra_examples
        self.shuffle = shuffle

    def __iter__(self) -> Iterable:
        g = torch.Generator()
        g.manual_seed(self.epoch)

        sortish_data = [self.dataset.src_lens[i] for i in self.available_indices]
        sortish_indices = sortish_sampler_indices(sortish_data, self.batch_size, shuffle=self.shuffle)
        indices = [self.available_indices[i] for i in sortish_indices]
        assert len(indices) == self.num_samples
        return iter(indices)

    @cached_property
    def available_indices(self) -> np.array:
        indices = list(range(len(self.dataset)))
        # add extra samples to make it evenly divisible
        indices += indices[: (self.total_size - len(indices))]
        assert len(indices) == self.total_size
        # subsample
        available_indices = indices[self.rank : self.total_size : self.num_replicas]
        return available_indices

    def __len__(self):
        return self.num_samples

    def set_epoch(self, epoch):


SUMMARY:
**DistributedSortishSampler Class Summary**

### Purpose and Functionality

The `DistributedSortishSampler` class is a distributed sampler for PyTorch datasets, designed to efficiently sample data from multiple processes in a distributed environment.

### Key Implementation Details

*   The sampler uses the `torch.Generator` class to randomize seed values based on epoch.
*   The `available_indices` property subsets indices of the dataset dataset using rank and num_replicas parameters.
*   A custom sorting algorithm is applied in the `__iter__` method.

### Notable Dependencies and Requirements

*   PyTorch's distributed package (dist)
*   Torch Generator API (`torch.Generator`)
*   NumPy library (for `math.ceil` and array operations)


key: chroma:document
string_value:
CODE:

def use_task_specific_params(model, task):
    """Update config with summarization specific params."""
    task_specific_params = model.config.task_specific_params

    if task_specific_params is not None:
        pars = task_specific_params.get(task, {})
        logger.info(f"using task specific params for {task}: {pars}")


SUMMARY:
**Summary:**

`use_task_specific_params(model, task)` updates the model's configuration with task-specific summarization parameters.

**Key Implementation Details:**

* Retrieves `task_specific_params` from `model.config` using the `task` parameter to filter relevant params.
* Updates `model.config` with the updated parameters.
* Logs a message indicating which task-specific parameters are being used.

**Dependencies/Requirements:**

* `logger` object (not explicitly typed)
* `model.config` object
* `task_specific_params` dictionary
* External dependencies: None


key: chroma:document
string_value:
CODE:

def pickle_load(path):
    """pickle.load(path)"""
    with open(path, "rb") as f:


SUMMARY:
**Summary:**

`pickle_load()` Function Summary:

This function loads data from a file path using the `pickle.load()` method.

**Key Implementation Details:**

* Loads file in binary read mode (`"rb"`).
* Returns the loaded pickle object.

**Context:** This function appears to be part of a larger system that integrates with Python's built-in serialization library, `pickle`. It provides a simple way to load data from a file.

**Notable Dependencies/Requirements:**

* `pickle`: Built-in Python module for serializing and deserializing objects.
* No external dependencies noted.


key: chroma:document
string_value:
CODE:

def pickle_save(obj, path):
    """pickle.dump(obj, path)"""
    with open(path, "wb") as f:


SUMMARY:
**Summary:**

This functions writes an object to a binary file using `pickle.dump()`, but does not return the result like expected. Instead, it returns the function call itself.

```markdown
# Function: pickle_save
# Purpose: Saves an object in binary format

def pickle_save(obj, path):
  """pickle.dump(obj, path)"""
    with open(path, "wb") as f:
      # Implementation:
        return pickle.dump(obj, f)
```

**Key implementation details:**

* Uses `pickle.dump()` to serialize the object
* Writes to a binary file at the specified path

**Notable dependencies and requirements:**

* `pickle` module (part of Python Standard Library)


key: chroma:document
string_value:
CODE:

def flatten_list(summary_ids: List[List]):


SUMMARY:
**Function Summary**

### `flatten_list`
#### Purpose
Flattens a nested list of lists into a single list.

#### Functionality
Expands a hierarchical structure by iteratively unwraveling the innermost elements, returning a new flat list containing all elements from the input `summary_ids`.

#### Key Implementation Details

* Utilizes `itertools.chain.from_iterable` to flatten the list in-place.
* Takes advantage of mutable behavior using iteration.

#### Notable Dependencies
- **itertools.chain`
- **List**

Note: This summary focuses on the function's purpose, functionality, and key implementation details, while minimizing unnecessary information.


key: chroma:document
string_value:
CODE:

def save_git_info(folder_path: str) -> None:
    """Save git information to output_dir/git_log.json"""
    repo_infos = get_git_info()


SUMMARY:
**Summary**

 Saves Git log information from the repository at `folder_path` to a JSON file named `git_log.json`.

**Key Implementation Details**

* Uses `get_git_info()` function to retrieve repository metadata
* Utilizes `save_json()` function to persist data to a file located at `${folder_path}/git_log.json`
* Writes output to a single JSON file, storing repository log information

**Notable Dependencies/Requirements**

* `os` for joining path components
* `get_git_info()` and `save_json()` functions (external dependencies)

Note: No interfaces are implemented or inherited in this function.


key: chroma:document
string_value:
CODE:

def save_json(content, path, indent=4, **json_dump_kwargs):
    with open(path, "w") as f:


SUMMARY:
**Function Summary**

`save_json`: Saves JSON content to a file.

*   **Purpose and Functionality**: Writes the provided JSON payload (`content`) to a specified file path (`path`) with customizable indentation (`indent=4`) and additional keyword arguments (`json_dump_kwargs`).
*   **Key Implementation Details:**
    *   Uses `json.dump()` to serialize content into a JSON string.
    *   Opens the file in write mode (`"w"`).
*   **Notable Dependencies/Requirements:**
    *   Relies on the built-in `json` module for serialization.
    *   External dependency: The operating system needs permission to create and write files.


key: chroma:document
string_value:
CODE:

def load_json(path):
    with open(path) as f:


SUMMARY:
**Concise Summary**

* **Function Name:** `load_json`
* **Purpose:** Loads JSON data from a provided file path into a Python dictionary using the `json` module.
	+ Primary Purpose: Convert JSON string to a Python data structure.
* **Implementation Details:**
	+ Uses a `with` statement to ensure efficient file handling and resource closure.
	+ Utilizes the `load()` method of `json.load()` object from the `json` module for JSON parsing.

**Example Signature:** `json.load(file_object) -> dict`

**External Dependencies:**

* The `json` module is used to parse the JSON data.
* The `python` interpreter and the `os` system provides access to file handling mechanisms.


key: chroma:document
string_value:
CODE:

def get_git_info():
    try:
        repo = git.Repo(search_parent_directories=True)
        repo_infos = {
            "repo_id": str(repo),
            "repo_sha": str(repo.head.object.hexsha),
            "repo_branch": str(repo.active_branch),
            "hostname": str(socket.gethostname()),
        }
        return repo_infos
    except TypeError:
        return {
            "repo_id": None,
            "repo_sha": None,
            "repo_branch": None,
            "hostname": None,


SUMMARY:
**Function Summary:**

The `get_git_info` function retrieves and returns information about the current Git repository and the system's hostname.

*   Purpose:
    *   Retrieves Git repository information (id, SHA, branch) using `git.Repo`.
    *   Includes system host name via `socket.gethostname()`.
*   Return Values:
    *   A dictionary containing "repo_id", "repo_sha", "repo_branch", and "hostname" if successful.
    *   An empty dictionary with all values set to None in case of a TypeError exception.

**Key Implementation Details:**

*   Uses `git.Repo` to access Git repository information.
*   Utilizes `socket.gethostname()` for system hostname retrieval.
*   Handle exception using `try-except` block for TypeErrors.

**Notable Dependencies and Requirements:**

*   Requires the `gitpython` library (specifically, `git` module) for Git repository operations.
*   No known external dependencies.

To use this code in a RAG system, you can store it as follows:

```plaintext
get_git_info:
    Functionality: Retrieves and returns Git repository information and system host name.
    Purpose: Accesses current Git repository and system hostname data.
```

This concise summary highlights the function's primary purpose, key implementation details, and dependencies, making it easily understandable for developers familiar with the codebase.


key: chroma:document
string_value:
CODE:

def extract_rouge_mid_statistics(dct):
    new_dict = {}
    for k1, v1 in dct.items():
        mid = v1.mid
        new_dict[k1] = {stat: round(getattr(mid, stat), 4) for stat in ["precision", "recall", "fmeasure"]}


SUMMARY:
**Function Summary: `extract_rouge_mid_statistics`**

Extracts and computes ROUGE metrics (`precision`, `recall`, `fmeasure`) from `mid` objects in the input dictionary `dct`.

**Key Implementation Details:**

* Iterates through each key-value pair in the input dictionary.
* For each value, extracts the corresponding `mid` object and computes ROUGE metrics using setattr.
* Returns a new dictionary with the computed ROUGE metrics for each `mid` object.

**Notable Dependencies and Requirements:**
None identified. All necessary modules and functions are imported but not explicitly listed in this code snippet.


key: chroma:document
string_value:
CODE:

def calculate_rouge(
    pred_lns: List[str],
    tgt_lns: List[str],
    use_stemmer=True,
    rouge_keys=ROUGE_KEYS,
    return_precision_and_recall=False,
    bootstrap_aggregation=True,
    newline_sep=True,
) -> Dict:
    """Calculate rouge using rouge_scorer package.

    Args:
        pred_lns: list of summaries generated by model
        tgt_lns: list of groundtruth summaries (e.g. contents of val.target)
        use_stemmer:  Bool indicating whether Porter stemmer should be used to
        strip word suffixes to improve matching.
        rouge_keys:  which metrics to compute, defaults to rouge1, rouge2, rougeL, rougeLsum
        return_precision_and_recall: (False) whether to also return precision and recall.
        bootstrap_aggregation: whether to do the typical bootstrap resampling of scores. Defaults to True, if False
            this function returns a collections.defaultdict[metric: list of values for each observation for each subscore]``
        newline_sep:(default=True) whether to add newline between sentences. This is essential for calculation rougeL
        on multi sentence summaries (CNN/DM dataset).

    Returns:
         Dict[score: value] if aggregate else defaultdict(list) keyed by rouge_keys

    """
    scorer = rouge_scorer.RougeScorer(rouge_keys, use_stemmer=use_stemmer)
    aggregator = scoring.BootstrapAggregator()
    for pred, tgt in zip(tgt_lns, pred_lns):
        # rougeLsum expects "\n" separated sentences within a summary
        if newline_sep:
            pred = add_newline_to_end_of_each_sentence(pred)
            tgt = add_newline_to_end_of_each_sentence(tgt)
        scores = scorer.score(pred, tgt)
        aggregator.add_scores(scores)

    if bootstrap_aggregation:
        result = aggregator.aggregate()
        if return_precision_and_recall:
            return extract_rouge_mid_statistics(result)  # here we return dict
        else:
            return {k: round(v.mid.fmeasure * 100, 4) for k, v in result.items()}

    else:


SUMMARY:
**Summary**

Calculates ROUGE (Recall-Oriented Understudy for Gisting Evaluation Technique) scores using the `rouge_scorer` package. This function takes in predicted and target summaries, along with optional parameters to customize the calculation.

* Main purpose: Calculate ROUGE scores between predicted and target summaries.
* Key implementation details:
	+ Uses `rouge_scorer.RougeScorer` to score each pair of summaries.
	+ Applies Bootstrap aggregation for robustness (optional).
	+ Returns a dictionary with ROUGE scores and optional precision/recall metrics.

**Params**

* `pred_lns`: Predicted summaries (list[str])
* `tgt_lns`: Target summaries (list[str])
* `use_stemmer`: Bool to use Porter stemmer (default= True)
* `rouge_keys`: Set of ROUGE keys to compute (default: {ROUGE1, ROUGEU, ROUGEL, ROUGESUM})
* `return_precision_and_recall`: Bool to return precision and recall metrics (default= False)
* `bootstrap_aggregation`: Bool to use Bootstrap aggregation (default= True)

**Return**
Dict[score]: Dictionary with scores if aggregate; otherwise defaultdict(list) keyed by ROUGE keys.


key: chroma:document
string_value:
CODE:

def freeze_params(model: nn.Module):
    """Set requires_grad=False for each of model.parameters()"""
    for par in model.parameters():


SUMMARY:
**Summary: Freeze Model Parameters**

Freezes the trainable parameters of a neural network model by setting `requires_grad` to `False` for all model parameters.

Key Implementation Details:

* Utilizes `model.parameters()` to access all learnable parameters.
* Sets `par.requires_grad = False` for each parameter to prevent backpropagation.
* Applies operation to all model parameters with no loop iterations specified.

Notable Dependencies/Requirements:

* `nn.Module`: Parent class of the neural network model.
* No external dependencies are required or recommended.


key: chroma:document
string_value:
CODE:

def freeze_embeds(model):
    """Freeze token embeddings and positional embeddings for bart, just token embeddings for t5."""
    model_type = model.config.model_type

    if model_type == "t5":
        freeze_params(model.shared)
        for d in [model.encoder, model.decoder]:
            freeze_params(d.embed_tokens)
    elif model_type == "fsmt":
        for d in [model.model.encoder, model.model.decoder]:
            freeze_params(d.embed_positions)
            freeze_params(d.embed_tokens)
    else:
        freeze_params(model.model.shared)
        for d in [model.model.encoder, model.model.decoder]:
            freeze_params(d.embed_positions)


SUMMARY:
**Concise Summary**

`freeze_embeds` function freezes specific embeddings (positional and token embeddings) for different T5 models.

### Key Implementation Details

*   Freezes:
    *   Token embeddings (and positional embeddings if not using T5)
    *   Embed positions for FSMT (full span mask type) versions
*   Model types supported: `t5`, `fsmt`
*   Method signatures used: 
    +   `freeze_params(d.embed_positions)`
        -   Used to freeze positional embedding sub-module
    +   `freeze_params(d.embed_tokens)`
        -   Used to freeze token embedding sub-module

### Notable Dependencies/Requirements

*   External dependencies:
    +   `t5config`, `tokenization_utils`
*   Used modules and functions:
    +   ModelConfig class
    +   Shared parameters (`model.shared`)

### Primary Purpose

The primary purpose of the `freeze_embeds` function is to freeze specific embeddings (positional and token embedding) for optimal model stability, as necessary in some T5 models like FSMT.


key: chroma:document
string_value:
CODE:

def grad_status(model: nn.Module) -> Iterable:


SUMMARY:
**Summary**

The `grad_status` function takes a neural network model (`nn.Module`) as input and returns an iterable containing boolean values indicating whether each parameter in the model requires automatic gradient computation (`requires_grad=True`) or not.

**Implementation Details**

* The function iterates over all parameters (`par`) in the input model.
* Each parameter's `requires_grad` status is obtained using `model.parameters()`.
* The result is wrapped in an iterable (generator) using `(...)`.

**Dependencies/Requirements**

* None
* Assumes familiarity with PyTorch (`nn.Module`) and its associated functionality.


key: chroma:document
string_value:
CODE:

def any_requires_grad(model: nn.Module) -> bool:


SUMMARY:
**Function Summary**

`any_requires_grad` checks if any layer in the input neural network (`model`) has its gradients enabled.

**Key Implementation Details**

* Repeatedly applies `grad_status` to individual layers of the network using `any`.
* The `grad_status` function is not shown here, but it likely determines whether gradients are required for a given layer.

**Notable Dependencies or Requirements**

 None

This summary focuses on the primary purpose and key implementation details of the function, while omitting unnecessary information.


key: chroma:document
string_value:
CODE:

def assert_all_frozen(model):
    model_grads: List[bool] = list(grad_status(model))
    n_require_grad = sum(lmap(int, model_grads))
    npars = len(model_grads)


SUMMARY:
**Function Summary**

The `assert_all_frozen` function checks if all layers in a given model are frozen by verifying that none of the gradients (i.e., weights) need to be backpropagated.

**Key Details**

* The function uses the `grad_status` function to obtain a list of boolean values representing whether each weight requires gradients.
* It then calculates the proportion of weights requiring gradients using the formula: `(n_require_grad / npars) * 100`, where `n_require_grad` is the total number of weights that need gradients and `npars` is the total number of parameters in the model.
* The function asserts that no weights require gradients, raising a ValueError if this condition is not met.

**Notable Dependencies**

* `grad_status`: a function that returns a list of boolean values representing weight gradient requirements.


key: chroma:document
string_value:
CODE:

def assert_not_all_frozen(model):
    model_grads: List[bool] = list(grad_status(model))
    npars = len(model_grads)


SUMMARY:
**Code Summary**

The `assert_not_all_frozen` function validates that at least one weight in a neural network model requires gradient computation. It:

* Collects the status of each weight (frozen or non-frozen) into a list (`model_grads`)
* Asserts that any element in this list is `True` to indicate at least one non-frozen weight
* If no weights require gradients, an assertion error is raised with a message indicating the number of frozen weights.

**Technical Details**

* Uses NumPy's `grad_status` function to determine weight updates (`List[bool]` returns)
* Employs Boolean operations and assertions for validation

**No notable dependencies or requirements. The code relies solely on the provided functions and variables.**


key: chroma:document
string_value:
CODE:

def parse_numeric_n_bool_cl_kwargs(unparsed_args: List[str]) -> Dict[str, Union[int, float, bool]]:
    """
    Parse an argv list of unspecified command line args to a dict.
    Assumes all values are either numeric or boolean in the form of true/false.
    """
    result = {}
    assert len(unparsed_args) % 2 == 0, f"got odd number of unparsed args: {unparsed_args}"
    num_pairs = len(unparsed_args) // 2
    for pair_num in range(num_pairs):
        i = 2 * pair_num
        assert unparsed_args[i].startswith("--")
        if unparsed_args[i + 1].lower() == "true":
            value = True
        elif unparsed_args[i + 1].lower() == "false":
            value = False
        else:
            try:
                value = int(unparsed_args[i + 1])
            except ValueError:
                value = float(unparsed_args[i + 1])  # this can raise another informative ValueError

        result[unparsed_args[i][2:]] = value


SUMMARY:
**Summary**

`parse_numeric_n_bool_cl_kwargs` function parses an argv list of unspecified command line args to a dictionary.

**Purpose and Functionality**
 
 - Parses argv list into a dictionary with numeric/boolean values.
 - Assumes input values are either numeric or boolean in true/false format.
 

**Key Implementation Details**

- Iterates through pairs of arguments (assuming even-length argv list).
- Extracts value type from the second argument:
    * True/False for bool
    * Integer/Float for numeric
- Uses `int()`/`float()` to parse value based on its type.

**Dependencies and Requirements**

- None explicitly stated, but implies `argparse` library due to argv list structure.


key: chroma:document
string_value:
CODE:

def write_txt_file(ordered_tgt, path):
    f = Path(path).open("w")
    for ln in ordered_tgt:
        f.write(ln + "\n")


SUMMARY:
**Summary:**

`write_txt_file` function writes a list of strings (`ordered_tgt`) to a file at the specified path, appending each line followed by a newline character and flushing the buffer after each write.

**Key Implementation Details:**

* Appends lines to a file using `write()` method
* Flushes the buffer after each write using `flush()`
* Writes list elements in sequence (`for ln in ordered_tgt:`)

**Dependencies/Requirements:**

* Uses the `Path` class from the `pathlib` library to manipulate file paths.
	+ `Path(path).open("w")`: Opens a file at the specified path in write mode.
	+ No explicit external dependencies are listed, but `pathlib` and Python's built-in `io` module are assumed.


key: chroma:document
string_value:
CODE:

def chunks(lst, n):
    """Yield successive n-sized chunks from lst."""
    for i in range(0, len(lst), n):


SUMMARY:
**Summary:**

The `chunks` function takes a list `lst` and an integer `n` as input. It yields successive chunks of size `n` from the input list, starting from index 0.

Key Details:

* Implemented using a generator (`yield` statement)
* Slices the original list in increments of `n`

Notable Dependencies/Requirements:

* None mentioned, assuming `lst` is a list and `n` is an integer.


key: chroma:document
string_value:
CODE:

def check_output_dir(args, expected_items=0):
    """
    Checks whether to bail out if output_dir already exists and has more than expected_items in it

    `args`: needs to have the following attributes of `args`:
      - output_dir
      - do_train
      - overwrite_output_dir

    `expected_items`: normally 0 (default) - i.e. empty dir, but in some cases a few files are expected (e.g. recovery from OOM)
    """
    if (
        os.path.exists(args.output_dir)
        and len(os.listdir(args.output_dir)) > expected_items
        and args.do_train
        and not args.overwrite_output_dir
    ):
        raise ValueError(
            f"Output directory ({args.output_dir}) already exists and "
            f"has {len(os.listdir(args.output_dir))} items in it (expected {expected_items} items). "
            "Use --overwrite_output_dir to overcome."


SUMMARY:
**Function Summary: `check_output_dir`**

**Purpose:** Checks if an output directory exists with more expected items than its default threshold and raises an error if so.

**Key Implementation Details:**

* Uses `os.path.exists` to check for directory existence
* Verifies directory length against a predefined threshold (`expected_items`)
* Implements conditional control flow using bitwise operators
* Raises a `ValueError` exception with informative error message

**Notable Dependencies/Requirements:**

* `args`: Object with attributes `output_dir`, `do_train`, and `overwrite_outputDir`
* `os` (Module for directory existence and listing)


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/seq2seq-distillation/_test_seq2seq_examples_multi_gpu.py


key: chroma:document
string_value:
CODE:

def _dump_articles(path: Path, articles: list):
    content = "\n".join(articles)


SUMMARY:
**Summary**

`_dump_articles`: Dumps a list of article contents to a file at the specified path.

**Key Details**

* Writes concatenated articles to a file using `Path(path).open("w")`
* Utilizes `list.join` to concatenate articles with newline separators

**Notable Dependencies**

* (`os.path.Path`) for file system handling


key: chroma:document
string_value:
CODE:

def make_test_data_dir(tmp_dir):
    for split in ["train", "val", "test"]:
        _dump_articles(os.path.join(tmp_dir, f"{split}.source"), ARTICLES)
        _dump_articles(os.path.join(tmp_dir, f"{split}.target"), SUMMARIES)


SUMMARY:
**Summary**

`make_test_data_dir` function:

- Creates directories for training, validation, and testing data.
- Dumps source articles (`ARTICLES`) and target summaries (`SUMMARIES`) into these directories by split.

**Key Implementation Details**

- `split` loop iterates over ["train", "val", "test"] splits.
- `_dump_articles` function is called with different directory paths for each split, utilizing `ARTICLES` and `SUMMARIES`.
  
**Function Signature**

- `make_test_data_dir(tmp_dir)`: takes a temporary directory path (`tmp_dir`) as input.

**Notable Dependencies and Requirements**

- `_dump_articles` external function, dependency not specified.
- External variables/objects: `ARTICLES`, `SUMMARIES`.
- No inherent or implemented interfaces indicated.


key: chroma:document
string_value:
CODE:

class TestSummarizationDistillerMultiGPU(TestCasePlus):
    @classmethod
    def setUpClass(cls):
        return cls

    @require_torch_multi_gpu
    def test_multi_gpu(self):
        updates = {
            "no_teacher": True,
            "freeze_encoder": True,
            "gpus": 2,
            "overwrite_output_dir": True,
            "sortish_sampler": True,
        }
        self._test_distiller_cli_fork(updates, check_contents=False)

    def _test_distiller_cli_fork(self, updates, check_contents=True):
        default_updates = {
            "label_smoothing": 0.0,
            "early_stopping_patience": -1,
            "train_batch_size": 1,
            "eval_batch_size": 2,
            "max_epochs": 2,
            "alpha_mlm": 0.2,
            "alpha_ce": 0.8,
            "do_predict": True,
            "model_name_or_path": "sshleifer/tinier_bart",
            "teacher": CHEAP_ARGS["model_name_or_path"],
            "val_check_interval": 0.5,
        }
        default_updates.update(updates)
        args_d: dict = CHEAP_ARGS.copy()
        tmp_dir = make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir())
        output_dir = self.get_auto_remove_tmp_dir()
        args_d.update(data_dir=tmp_dir, output_dir=output_dir, **default_updates)

        def convert(k, v):
            if k in ["tgt_suffix", "server_ip", "server_port", "out", "n_tpu_cores"]:
                return ""
            if v is False or v is None:
                return ""
            if v is True:  # or len(str(v))==0:
                return f"--{k}"
            return f"--{k}={v}"

        cli_args = [x for x in (convert(k, v) for k, v in args_d.items()) if len(x)]
        cmd = [sys.executable, f"{self.test_file_dir}/distillation.py"] + cli_args
        execute_subprocess_async(cmd, env=self.get_env())

        contents = os.listdir(output_dir)
        contents = {os.path.basename(p) for p in contents}
        ckpt_files = [p for p in contents if p.endswith("ckpt")]
        assert len(ckpt_files) > 0

        self.assertIn("test_generations.txt", contents)
        self.assertIn("test_results.txt", contents)

        # get the following from the module, (we don't have access to `model` here)
        metrics_save_path = os.path.join(output_dir, "metrics.json")
        val_metric = "rouge2"

        metrics = load_json(metrics_save_path)
        # {'test': [{'test_avg_loss': 10.63731575012207, 'test_avg_rouge1': 0.0, 'test_avg_rouge2': 0.0, 'test_avg_rougeL': 0.0, 'test_avg_gen_time': 0.1822289228439331, 'test_avg_gen_len': 142.0, 'step_count': 1}]}
        print(metrics)
        last_step_stats = metrics["val"][-1]
        self.assertGreaterEqual(last_step_stats["val_avg_gen_time"], 0.01)
        self.assertIsInstance(last_step_stats[f"val_avg_{val_metric}"], float)
        self.assertEqual(len(metrics["test"]), 1)
        desired_n_evals = int(args_d["max_epochs"] * (1 / args_d["val_check_interval"]) / 2 + 1)


SUMMARY:
**Class Summary: TestSummarizationDistillerMultiGPU**

**Main Purpose and Functionality:** This class tests the multi-GPU support for summarization distillation using PyTorch.

*   Tests the command-line interface (CLI) execution with multiple GPUs enabled.
*   Verifies the output directories, generated files, and metrics.

**Key Implementation Details:**

*   The `_test_distiller_cli_fork` method executes the `distillation.py` CLI script with customized arguments.
*   It generates a temporary directory for test data and updates the default CLI arguments with user-provided values.
*   The method converts certain key-value pairs to CLI format and constructs the command.
*   It executes the command asynchronously, awaiting its completion and verifying the output directories.

**Notable Dependencies or Requirements:**

*   PyTorch
*   `distillation.py` CLI script
*   `CHEAP_ARGS` dictionary (not explicitly shown in the code snippet)
*   External libraries such as `os`, `sys`, `execute_subprocess_async`, and `get_env`

**Classes, Functions, or Variables Mentioned:**

*   `TestSummarizationDistillerMultiGPU`
*   `_test_distiller_cli_fork` method
*   `CHEAP_ARGS` dictionary


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/seq2seq-distillation/distillation.py


key: chroma:document
string_value:
CODE:

class SummarizationDistiller(SummarizationModule):
    """Supports T5, Bart, Pegasus and other models that inherit from Bart."""

    loss_names = ["loss", "ce_loss", "mlm_loss", "hid_loss_enc", "hid_loss_dec"]

    def __init__(self, hparams):
        assert Path(hparams.data_dir).exists()
        self.output_dir = Path(hparams.output_dir)
        self.output_dir.mkdir(exist_ok=True)

        save_dir = self.output_dir.joinpath("student")

        hparams.model_name_or_path = str(save_dir)  # Tell lightning we are training the student
        teacher = AutoModelForSeq2SeqLM.from_pretrained(hparams.teacher).eval()
        use_task_specific_params(teacher, hparams.task)  # We copy good generation parameters to student by default
        if hparams.student is not None:
            student = AutoModelForSeq2SeqLM.from_pretrained(hparams.student)
            use_task_specific_params(student, hparams.task)
            e_layer_ids, d_layer_ids = None, None
        else:
            student, e_layer_ids, d_layer_ids = create_student_by_copying_alternating_layers(
                teacher, e=hparams.student_encoder_layers, d=hparams.student_decoder_layers, save_path=save_dir
            )

        if hparams.length_penalty != -1:
            student.config.length_penalty = hparams.length_penalty
        hparams.tokenizer_name = hparams.teacher  # Use teacher's tokenizer
        super().__init__(hparams, model=student, config=student.config)
        assert student.config.model_type == teacher.config.model_type, (
            f"teacher, student model types should be the same, got {student.config.model_type} !="
            f" {teacher.config.model_type}"
        )

        if student.config.model_type == "t5":
            student_encoder_layers = len(student.get_encoder().block)
            student_decoder_layers = len(student.get_decoder().block)
            teacher_encoder_layers = len(teacher.get_encoder().block)
            teacher_decoder_layers = len(teacher.get_decoder().block)
        else:
            student_encoder_layers = student.config.encoder_layers
            student_decoder_layers = student.config.decoder_layers
            teacher_encoder_layers = teacher.config.encoder_layers
            teacher_decoder_layers = teacher.config.decoder_layers

        self.different_base_models = not (hparams.student is None or hparams.teacher == hparams.student)
        self.do_calc_hidden_loss = (not self.different_base_models) and hparams.alpha_hid > 0
        self.different_encoder = self.different_base_models or (student_encoder_layers != teacher_encoder_layers)
        # self.different_encoder determines whether we need to run the teacher encoder
        self.teacher = teacher
        freeze_params(self.teacher)

        if not self.different_encoder:  # To save RAM, delete teacher encoder and freeze student encoder.
            try:
                del self.teacher.model.encoder
            except AttributeError:  # T5
                del self.teacher.encoder

        if e_layer_ids is None:
            e_layer_ids = list(range(student_encoder_layers))
        if d_layer_ids is None:
            d_layer_ids = list(range(student_decoder_layers))

        self.e_layer_ids, self.d_layer_ids = e_layer_ids, d_layer_ids  # type: List[int], List[int]

        if self.do_calc_hidden_loss:  # Intermediate supervision: Decide which layers to supervise
            if hparams.supervise_forward:
                self.e_matches = get_layers_to_supervise(
                    n_student=len(self.e_layer_ids), n_teacher=teacher_encoder_layers
                )
                self.d_matches = get_layers_to_supervise(
                    n_student=len(self.d_layer_ids), n_teacher=teacher_decoder_layers
                )
            else:  # student layer should emulate hidden states of the teacher layer it was copied from
                self.e_matches = self.e_layer_ids
                self.d_matches = self.d_layer_ids
        else:
            self.e_matches = None
            self.d_matches = None

        self.ce_loss_fct = nn.KLDivLoss(reduction="batchmean")
        self.temperature = 2.0
        self.alpha_mlm = hparams.alpha_mlm
        self.alpha_ce = hparams.alpha_ce
        self.alpha_hid = hparams.alpha_hid
        gc.collect()
        torch.cuda.empty_cache()

    def calc_ce_loss(self, mask, s_logits, t_logits):
        """Copy pasted from distillbert (transformers/examples/distillation/)"""
        # mask has False at padding_idx
        sel_mask = mask[:, :, None].expand_as(s_logits)
        vocab_size = s_logits.size(-1)
        s_logits_slct = torch.masked_select(s_logits, sel_mask)  # (bs * seq_length * voc_size) modulo the 1s in mask
        t_logits_slct = torch.masked_select(t_logits, sel_mask)  # (bs * seq_length * voc_size) modulo the 1s in mask
        s_logits_slct = s_logits_slct.view(-1, vocab_size)  # (bs * seq_length, voc_size) modulo the 1s in mask
        t_logits_slct = t_logits_slct.view(-1, vocab_size)  # (bs * seq_length, voc_size) modulo the 1s in mask
        assert t_logits_slct.size() == s_logits_slct.size()
        loss_ce = (
            self.ce_loss_fct(
                nn.functional.log_softmax(s_logits_slct / self.temperature, dim=-1),
                nn.functional.softmax(t_logits_slct / self.temperature, dim=-1),
            )
            * (self.temperature) ** 2
        )
        return loss_ce

    @staticmethod
    def add_model_specific_args(parser, root_dir):
        SummarizationModule.add_model_specific_args(parser, root_dir)
        add_distill_args(parser)
        return parser

    def _step(self, batch: dict) -> tuple:
        """Compute the loss for a batch"""
        pad_token_id = self.tokenizer.pad_token_id
        input_ids, src_mask, labels = batch["input_ids"], batch["attention_mask"], batch["labels"]
        if isinstance(self.model, T5ForConditionalGeneration):
            decoder_input_ids = self.model._shift_right(labels)
        else:
            decoder_input_ids = shift_tokens_right(labels, pad_token_id)

        # noinspection PyCallingNonCallable
        student_outputs = self(
            input_ids,
            attention_mask=src_mask,
            decoder_input_ids=decoder_input_ids,
            output_hidden_states=self.do_calc_hidden_loss,
            output_attentions=False,
            use_cache=False,
        )
        lm_logits = student_outputs["logits"]

        # Same cross entropy vs. label smoothing logic as finetune.py
        assert lm_logits.shape[-1] == self.model.config.vocab_size
        if self.hparams.label_smoothing == 0:
            # Same behavior as modeling_bart.py, besides ignoring pad_token_id
            loss_fct = nn.CrossEntropyLoss(ignore_index=pad_token_id)
            student_lm_loss = loss_fct(lm_logits.view(-1, lm_logits.shape[-1]), labels.view(-1))
        else:
            lprobs = nn.functional.log_softmax(lm_logits, dim=-1)
            student_lm_loss, _ = label_smoothed_nll_loss(
                lprobs, labels, self.hparams.label_smoothing, ignore_index=pad_token_id
            )

        def zero_tensor():
            return torch.tensor(0.0).type_as(student_lm_loss)

        teacher_enc_outputs = student_outputs[
            "encoder_last_hidden_state"
        ]  # use this unless self.different_base_models
        hid_loss_enc, hid_loss_dec = zero_tensor(), zero_tensor()
        if self.different_encoder:  # compute encoder hidden state loss
            all_teacher_encoder_outputs = self.teacher.get_encoder()(
                input_ids,
                attention_mask=src_mask,
                output_hidden_states=self.do_calc_hidden_loss,
            )
            if self.different_base_models:
                teacher_enc_outputs = all_teacher_encoder_outputs["last_hidden_state"]
            elif self.do_calc_hidden_loss:
                hid_loss_enc = self.calc_hidden_loss(
                    src_mask,
                    student_outputs["encoder_hidden_states"],
                    all_teacher_encoder_outputs["hidden_states"],
                    self.e_matches,
                    normalize_hidden=self.hparams.normalize_hidden,
                )

        teacher_outputs = self.teacher(
            input_ids,
            attention_mask=src_mask,
            encoder_outputs=(teacher_enc_outputs,),
            decoder_input_ids=decoder_input_ids,
            output_hidden_states=self.do_calc_hidden_loss,
            use_cache=False,  # since we are not passing labels, never let this default to True
        )
        dec_mask = decoder_input_ids.ne(pad_token_id)
        loss_ce = self.calc_ce_loss(dec_mask, lm_logits, teacher_outputs["logits"])
        if self.do_calc_hidden_loss:  # Intermediate supervision of decoder hidden states
            hid_loss_dec = self.calc_hidden_loss(
                dec_mask,
                student_outputs["decoder_hidden_states"],
                teacher_outputs["decoder_hidden_states"],
                self.d_matches,
                normalize_hidden=self.hparams.normalize_hidden,
            )

        blended_loss = (
            self.alpha_ce * loss_ce
            + self.alpha_mlm * student_lm_loss
            + self.hparams.alpha_hid * (hid_loss_enc + hid_loss_dec)
        )
        return blended_loss, loss_ce, student_lm_loss, hid_loss_enc, hid_loss_dec

    @staticmethod
    def calc_hidden_loss(attention_mask, hidden_states, hidden_states_T, matches, normalize_hidden):
        """MSE(student_hid, teacher_hid[matches]). Called "Intermediate supervision" in paper. Inspired by TinyBERT."""
        msg = "expected list or tuple for hidden_states, got tensor of shape: "
        assert not isinstance(hidden_states, torch.Tensor), f"{msg}{hidden_states.shape}"
        assert not isinstance(hidden_states_T, torch.Tensor), f"{msg}{hidden_states_T.shape}"
        mask = attention_mask.to(hidden_states[0])
        valid_count = mask.sum() * hidden_states[0].size(-1)
        student_states = torch.stack([hidden_states[i] for i in range(len(matches))])
        teacher_states = torch.stack([hidden_states_T[j] for j in matches])
        assert student_states.shape == teacher_states.shape, f"{student_states.shape} != {teacher_states.shape}"
        if normalize_hidden:
            student_states = nn.functional.layer_norm(student_states, student_states.shape[1:])
            teacher_states = nn.functional.layer_norm(teacher_states, teacher_states.shape[1:])
        mse = nn.functional.mse_loss(student_states, teacher_states, reduction="none")
        masked_mse = (mse * mask.unsqueeze(0).unsqueeze(-1)).sum() / valid_count


SUMMARY:
This code is implementing an unsupervised learning approach called "Intermediate Supervision" by inspired by TinyBERT. The framework appears to be designed for pre-training language models, particularly for text classification and question answering tasks.

**Key Points:**

1. **Model Architecture**: The model uses a multi-layer architecture similar to BERT, which suggests that it can handle long-range dependencies in the input sequence.
2. **Pre-Training Objective**: The pre-training objective is to learn the representation of the input tokens by minimizing the loss between the student's hidden states and the teacher's hidden states, where only certain parts of the teacher's hidden state space (i.e., `matches`) are used to compute the loss.
3. **Label Smoothing**: To avoid gradient exploding due to vanishing gradients, label smoothing is applied during training.

**Functions:**

1.  `T5ForConditionalGeneration` (`self.model`): The pre-trained T5 for conditional generation.
2.  `calc_hidden_loss` (`self.calc_hidden_loss(attention_mask, hidden_states, hidden_states_T, matches, normalize_hidden)`): This function computes the intermediate supervision loss by taking the mean squared error (MSE) between the student's and teacher's hidden states on subsets of their respective feature spaces.

Overall, this code implements a flexible unsupervised learning framework for text classification, question answering, or other downstream NLP tasks.


key: chroma:document
string_value:
CODE:

def add_distill_args(parser):
    # NOTE: if --student argument was specified and the teacher and student base models
    # are different, the models still have to have the same tokenizer, specified by
    # --tokenizer_name. So, for example, you can distill from t5_large to t5_small but not
    # from bart to t5. This s because if the tokenizers are different, the output space
    # for the two models is also different and their logits are not comparable.
    parser.add_argument("--teacher", type=str)
    parser.add_argument("--alpha_ce", default=0.8, type=float)
    parser.add_argument("--alpha_mlm", default=0.2, type=float)
    parser.add_argument("--alpha_hid", default=0.0, type=float, required=False)
    parser.add_argument("--student", type=str, required=False)
    parser.add_argument("--student_decoder_layers", default=12, type=int, required=False)
    parser.add_argument("--student_encoder_layers", default=12, type=int, required=False)
    parser.add_argument("--no_teacher", action="store_true", default=False)
    parser.add_argument("--length_penalty", type=float, default=-1)
    parser.add_argument("--supervise_forward", action="store_true", default=False)


SUMMARY:
**Summary**

The `add_distill_args` function defines and adds arguments to a parser using the Hugging Face `ArgumentParser`. The purpose of this function is to collect arguments for training models in a distillation setup. Key implementation details include:

* Argument parsing for model tokenizer (`--tokenizer_name`)
* Argument parsing for alpha coefficients (`alpha_ce`, `alpha_mlm`, `alpha_hid`) with default values
* Optional arguments for student model configurations (`student`, `student_decoder_layers`, `student_encoder_layers`)
* Flag arguments for supervision and normalization of hidden states
* Required flag argument to indicate presence or absence of a teacher model

Notable dependencies:
* `ArgumentParser` class from Hugging Face library

These arguments can be used to configure the distillation process, including specifying the tokenizer, hyperparameters, and student model configuration.


key: chroma:document
string_value:
CODE:

class TranslationDistiller(SummarizationDistiller):
    """Supports T5, mBART, Marian, other models that inherit from Bart."""

    mode = "translation"
    metric_names = ["bleu"]
    default_val_metric = "bleu"

    def __init__(self, hparams, **kwargs):
        super().__init__(hparams, **kwargs)
        assert hparams.src_lang is not None
        assert hparams.tgt_lang is not None
        self.dataset_kwargs["src_lang"] = hparams.src_lang
        self.dataset_kwargs["tgt_lang"] = hparams.tgt_lang
        if self.model.config.decoder_start_token_id is None and isinstance(self.tokenizer, MBartTokenizer):
            self.decoder_start_token_id = self.tokenizer.lang_code_to_id[hparams.tgt_lang]

    def calc_generative_metrics(self, preds, target) -> dict:
        return calculate_bleu(preds, target)

    @staticmethod
    def add_model_specific_args(parser, root_dir):
        TranslationModule.add_model_specific_args(parser, root_dir)
        add_distill_args(parser)


SUMMARY:
**Class Summary: TranslationDistiller**

The `TranslationDistiller` class is a subclass of `SummarizationDistiller`, supporting translation models that inherit from Bart. Its primary purpose is to facilitate model-based machine translation tasks.

**Key Functionality:**

* Initializes with source and target language IDs
* Uses Hugging Face's `MBartTokenizer` for decoding if necessary
* Calculates generative metrics, specifically BLEU score

**Notable Implementation Details:**

* The `calc_generative_metrics` method leverages the `calculate_bleu` function to compute BLEU scores.

**Dependencies and Requirements:**

* Inherited from `SummarizationDistiller`
* Uses `hparams.src_lang`, `hparams.tgt_lang`, and `MBartTokenizer`
* Expects external dependencies on `translate` modules (e.g. Hugging Face's translation models)

**Interfaces and Dependencies:**

* Implemented `model.config.decoder_start_token_id` for manual token ID assignment
* Inherently depends on the following classes/functions:
	+ `calculate_bleu`
	+ `MBartTokenizer`
	+ `TranslateModule` (base class)


key: chroma:document
string_value:
CODE:

def create_module(args):
    if args.no_teacher:
        module_cls = TranslationModule if "translation" in args.task else SummarizationModule
    else:  # DISTILL WITH TEACHER
        module_cls = TranslationDistiller if "translation" in args.task else SummarizationDistiller
    args.setup_cls: str = module_cls.__name__
    print(f"using module {args.setup_cls}")
    model = module_cls(args)


SUMMARY:
**Summary**

`create_module` function creates and initializes a machine learning module based on the `args.no_teacher` flag. It takes an `args` object with task-specific flags as input.

*   **Primary Purpose**: Dynamically selects and instantiates a module class (`TranslationModule` or `SummarizationModule`) based on whether to use a teacher-force mode.
*   **Key Implementation Details**:
    *   Uses string comparison (``"translation"` in args.task) to determine the module class.
    *   Assigns the selected module class's name to `args.setup_cls`.
    *   Instantiates an instance of the selected module class with the provided `args`.
        *   The `module_cls` class takes the input arguments (`args`) during its initialization.
*   **Notable Dependencies or Requirements**:
    *   `module_cls.__name__` is used to set `args.setup_cls`, implying that `module_cls` is accessible and has a `__name__` attribute.

Here's a shorter version:

```python
create_module: Function 

Dynamically creates and initializes a machine learning module based on the no_teacher flag.
 
    Args:
        args (object): Task-specific flags.
    
    Returns:
        Model instance
    
        Key implementation details:
            - module selection using task flag comparison
            - instantiates the selected class with provided arguments
        
    Requires/imports:
        - module_cls with __name__ accessible attribute
```


key: chroma:document
string_value:
CODE:

def distill_main(args):
    Path(args.output_dir).mkdir(exist_ok=True)
    check_output_dir(args, expected_items=3)

    model = create_module(args)


SUMMARY:
**Summary:**

`distill_main` Function Summarization

* Primary purpose: Creates and executes the main distillation process, ensuring necessary output directories exist and importing the required module.

**Key Implementation Details:**

* Calls `check_output_dir` to validate the specified output directory matches expectations (3 items expected).
* Imports a custom `create_module` function from the provided arguments (`model = create_module(args)`).

**Notable Dependencies or Requirements:**

* External dependencies:
	+ None explicitly mentioned, but could include dependencies required by `check_output_dir` and `create_module`.
* Classes/functions referenced by exact name:
	+ `Path`, `mkdir`, `exist_ok`
* Interfaces:
	+ Expected to implement `check_output_dir` interface with `expected_items` parameter.


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/performer/run_mlm_performer.py


key: chroma:document
string_value:
CODE:

def create_learning_rate_scheduler(
    factors="constant * linear_warmup * rsqrt_decay",
    base_learning_rate=0.5,
    warmup_steps=1000,
    decay_factor=0.5,
    steps_per_decay=20000,
    steps_per_cycle=100000,
):
    """Creates learning rate schedule.
    Interprets factors in the factors string which can consist of:
    * constant: interpreted as the constant value,
    * linear_warmup: interpreted as linear warmup until warmup_steps,
    * rsqrt_decay: divide by square root of max(step, warmup_steps)
    * rsqrt_normalized_decay: divide by square root of max(step/warmup_steps, 1)
    * decay_every: Every k steps decay the learning rate by decay_factor.
    * cosine_decay: Cyclic cosine decay, uses steps_per_cycle parameter.
    Args:
      factors: string, factors separated by "*" that defines the schedule.
      base_learning_rate: float, the starting constant for the lr schedule.
      warmup_steps: int, how many steps to warm up for in the warmup schedule.
      decay_factor: float, the amount to decay the learning rate by.
      steps_per_decay: int, how often to decay the learning rate.
      steps_per_cycle: int, steps per cycle when using cosine decay.
    Returns:
      a function learning_rate(step): float -> {"learning_rate": float}, the
      step-dependent lr.
    """
    factors = [n.strip() for n in factors.split("*")]

    def step_fn(step):
        """Step to learning rate function."""
        ret = 1.0
        for name in factors:
            if name == "constant":
                ret *= base_learning_rate
            elif name == "linear_warmup":
                ret *= jnp.minimum(1.0, step / warmup_steps)
            elif name == "rsqrt_decay":
                ret /= jnp.sqrt(jnp.maximum(step, warmup_steps))
            elif name == "rsqrt_normalized_decay":
                ret *= jnp.sqrt(warmup_steps)
                ret /= jnp.sqrt(jnp.maximum(step, warmup_steps))
            elif name == "decay_every":
                ret *= decay_factor ** (step // steps_per_decay)
            elif name == "cosine_decay":
                progress = jnp.maximum(0.0, (step - warmup_steps) / float(steps_per_cycle))
                ret *= jnp.maximum(0.0, 0.5 * (1.0 + jnp.cos(jnp.pi * (progress % 1.0))))
            else:
                raise ValueError("Unknown factor %s." % name)
        return jnp.asarray(ret, dtype=jnp.float32)



SUMMARY:
**Summary:**

The `create_learning_rate_scheduler` function creates a learning rate schedule based on the provided factors and parameters.

**Purpose:**
Creates a stepping function (`step_fn`) that returns the learning rate at a given step, taking into account various scheduling factors such as warm-up, decay, and cosine decay.

**Key Implementation Details:**

* The `factors` string is parsed to determine the schedule type.
* For each factor:
	+ "constant" multiplies the base learning rate.
	+ "linear_warmup" divides by 1 - (step / warmup_steps).
	+ "rsqrt_decay" divides by sqrt(step / warmup_steps).
	+ "decay_every" decays the learning rate every steps_per_decay steps.
	+ "cosine_decay" uses cyclic cosine decay with steps_per_cycle.

**Notable Dependencies and Requirements:**

* Uses `jnp` library for numerical computations.
* Inherits from `numpy`, but since it's jax, it should work correctly.


key: chroma:document
string_value:
CODE:

def compute_metrics(logits, labels, weights, label_smoothing=0.0):
    """Compute summary metrics."""
    loss, normalizer = cross_entropy(logits, labels, weights, label_smoothing)
    acc, _ = accuracy(logits, labels, weights)
    metrics = {"loss": loss, "accuracy": acc, "normalizer": normalizer}
    metrics = jax.lax.psum(metrics, axis_name="batch")


SUMMARY:
**Function Summary: `compute_metrics`**

Computes summary metrics for classification loss and accuracy.

**Purpose**: Compute cross-entropy loss, accuracy, and normalizer for logging purposes.

**Implementation Details**:

* Computes loss using `cross_entropy logits, labels, weights, label_smoothing` function.
* Calculates accuracy using `accuracy logits, labels, weights` function.
* Returns a dictionary containing loss, accuracy, and normalizer values (`metrics`).
* Applies parallel reduction along the batch dimension of `metrics` with `jax.lax.psum`.

**Dependencies/Requirements**:

* `cross_entropy`
* `accessibility`
* `weights`
* JAX library for parallelization (`jax`)


key: chroma:document
string_value:
CODE:

def accuracy(logits, targets, weights=None):
    """Compute weighted accuracy for log probs and targets.
    Args:
     logits: [batch, length, num_classes] float array.
     targets: categorical targets [batch, length] int array.
     weights: None or array of shape [batch, length]
    Returns:
      Tuple of scalar loss and batch normalizing factor.
    """
    if logits.ndim != targets.ndim + 1:
        raise ValueError(
            "Incorrect shapes. Got shape %s logits and %s targets" % (str(logits.shape), str(targets.shape))
        )

    loss = jnp.equal(jnp.argmax(logits, axis=-1), targets)
    loss *= weights



SUMMARY:
**Function Summary:**

`accuracy` computes weighted accuracy for log probabilities and targets. It takes in `logits`, `targets`, and optional `weights`. The function returns the sum of accurate predictions (`loss`) and the sum of input weights.

**Key Implementation Details:**

* Computes accuracy via `jnp.equal(jnp.argmax(logits, axis=-1), targets)`
* Applies input weights to losses using element-wise multiplication (`loss *= weights`)
* Returns weighted sum of accurate predictions (`loss.sum()`) and weighted average of input weights (`weights.sum()`)

**Notable Dependencies/Requirements:**

* Uses JAX (Jupyter API eXtension) NumPy equivalent library, `jnp`
* Assumes input data is a batch of token sequences (`logits` and `targets`) with shape `[batch, length, num_classes]`.
* Optional input for weights with shape `[batch, length]`.


key: chroma:document
string_value:
CODE:

def cross_entropy(logits, targets, weights=None, label_smoothing=0.0):
    """Compute cross entropy and entropy for log probs and targets.
    Args:
     logits: [batch, length, num_classes] float array.
     targets: categorical targets [batch, length] int array.
     weights: None or array of shape [batch, length]
     label_smoothing: label smoothing constant, used to determine the on and off values.
    Returns:
      Tuple of scalar loss and batch normalizing factor.
    """
    if logits.ndim != targets.ndim + 1:
        raise ValueError(
            "Incorrect shapes. Got shape %s logits and %s targets" % (str(logits.shape), str(targets.shape))
        )

    vocab_size = logits.shape[-1]
    confidence = 1.0 - label_smoothing
    low_confidence = (1.0 - confidence) / (vocab_size - 1)
    normalizing_constant = -(
        confidence * jnp.log(confidence) + (vocab_size - 1) * low_confidence * jnp.log(low_confidence + 1e-20)
    )
    soft_targets = common_utils.onehot(targets, vocab_size, on_value=confidence, off_value=low_confidence)

    loss = -jnp.sum(soft_targets * log_softmax(logits), axis=-1)
    loss = loss - normalizing_constant

    if weights is not None:
        loss = loss * weights
        normalizing_factor = weights.sum()
    else:
        normalizing_factor = np.prod(targets.shape)



SUMMARY:
**Cross-Entropy Function Summary**

This function computes the cross-entropy loss and batch normalization factor. It takes in:

* `logits`: [batch, length, num_classes] float array
* `targets`: categorical targets [batch, length] int array
* Optional: `weights` None or array of shape [batch, length]
* Optional: `label_smoothing` constant (0.0 by default)

**Purpose:** Compute cross-entropy loss and batch normalization factor.

**Key Implementation Details:**

* Uses one-hot encoding to create soft targets.
* Calculates a normalizing constant to avoid division by zero.
* Applies weighting to the loss if provided.

**Return Values:**

* Tuple of scalar loss (summed across batch) and batch normalization factor.
* `weights`: array of shape [batch, length] with optional ` label_smoothing` parameter.


key: chroma:document
string_value:
CODE:

def training_step(optimizer, batch, dropout_rng):
    dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)

    def loss_fn(params):
        targets = batch.pop("labels")

        # Hide away tokens which doesn't participate in the optimization
        token_mask = jnp.where(targets > 0, 1.0, 0.0)

        logits = model(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
        loss, weight_sum = cross_entropy(logits, targets, token_mask)
        return loss / weight_sum

    step = optimizer.state.step
    lr = lr_scheduler_fn(step)
    grad_fn = jax.value_and_grad(loss_fn)
    loss, grad = grad_fn(optimizer.target)
    grad = jax.lax.pmean(grad, "batch")
    optimizer = optimizer.apply_gradient(grad, learning_rate=lr)



SUMMARY:
**Training Step Function Summary**

The `training_step` function is a key component of Stochastic Gradient Descent (SGD) for training a deep learning model. It takes in:

*   `optimizer`: the SGD optimizer instance
*   `batch`: the input batch containing labels and other relevant data
*   `dropout_rng`: a random number generator

**Primary Purpose**

The primary purpose of this function is to perform one step of gradient descent on the given batch, updating the model parameters using the calculated gradients.

**Key Implementation Details**

1.  **Loss Calculation**: The loss is computed by calling the `cross_entropy` function with the predicted logits and targets, weighted by a token mask that excludes non-participating tokens.
2.  **Gradient Calculation**: The gradients of the loss are calculated using the `value_and_grad` function, which returns both the value (loss) and the gradient.
3.  **Parameter Update**: The model parameters are updated using the optimizer's `apply_gradient` method with the calculated gradient.

**Notable Dependencies**

*   `optimizer`: an instance of a SGD optimizer algorithm
*   `cross_entropy`: a function computing cross-entropy loss between predicted logits and targets
*   `lr_scheduler_fn`: a learning rate scheduler function
*   `model`: a deep learning model instance

This summary focuses on the key aspects of the `training_step` function, highlighting its primary purpose, implementation details, and dependencies.


key: chroma:document
string_value:
CODE:

def eval_step(params, batch):
    """
    Calculate evaluation metrics on a batch.
    """
    targets = batch.pop("labels")

    # Hide away tokens which doesn't participate in the optimization
    token_mask = jnp.where(targets > 0, 1.0, 0.0)
    logits = model(**batch, params=params, train=False)[0]



SUMMARY:
**Function Summary: `eval_step`**

**Purpose:** Calculate evaluation metrics on a batch of data.

* Predicts labels using the provided model and batch data.
* Filters out tokens with no optimization participation (label > 0).
* Computes evaluation metrics between predicted logits, original targets, and token mask.

**Key Implementation Details:**

* `model(params=params, train=False)` calls model with `params` and no training mode (`train=False`)
* `compute_metrics` function computes evaluation metrics from `logits`, `targets`, and `token_mask`

**Notable Dependencies/Iferences:**

* **`jnp.where`**: Numerical element-wise comparison and assignment.
* **`model`**: Predictive model implementation (interface).
* **`batch`**: Input batch data with "labels" key.
* **`compute_metrics`**: Function to compute evaluation metrics.


key: chroma:document
string_value:
CODE:

def generate_batch_splits(samples_idx: np.ndarray, batch_size: int) -> np.ndarray:
    nb_samples = len(samples_idx)
    samples_to_remove = nb_samples % batch_size

    if samples_to_remove != 0:
        samples_idx = samples_idx[:-samples_to_remove]
    sections_split = nb_samples // batch_size
    batch_idx = np.split(samples_idx, sections_split)


SUMMARY:
**Summary:**

* `generate_batch_splits` function generates batches of identical size from input `samples_idx`.
* Primary purpose: Splits data into batches based on batch size.
* Key implementation details:
	+ Calculates the number of samples to remove from the beginning to ensure uniform batch sizes.
	+ Uses NumPy's array splitting feature (`np.split`) to divide the indices into batches.

**Key Method Signatures and Interfaces:**

* `generate_batch_splits(samples_idx, batch_size)`

**Notable Dependencies or Requirements:**

* NumPy library (`np.ndarray`, `np.split`)
* Python ( NumPy % operator for integer division) 

**Input/Output Types:**

* Input: `samples_idx` as a NumPy array of indices
* Output: List of batches from samples_idx as a list of NumPy arrays.


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/performer/modeling_flax_performer_utils.py


key: chroma:document
string_value:
CODE:

def nonnegative_softmax_kernel_feature_creator(
    data, projection_matrix, attention_dims_t, batch_dims_t, precision, is_query, normalize_data=True, eps=0.0001
):
    """
    Constructs nonnegative kernel features for fast softmax attention

    Args:
      data: input for which features are computes
      projection_matrix: random matrix used to compute features
      attention_dims_t: tuple of attention dimensions
      batch_dims_t: tuple of batch dimensions
      precision: precision parameter
      is_query: predicate indicating whether input data corresponds to queries or
        keys
      normalize_data: predicate indicating whether data should be normalized,
      eps: numerical stabilizer

    Returns:
      Random features for fast softmax attention.
    """
    del attention_dims_t
    if normalize_data:
        # We have e^{qk^T/sqrt{d}} = e^{q_norm k_norm^T}, where
        # w_norm = w * data_normalizer for w in {q,k}.
        data_normalizer = 1.0 / (jnp.sqrt(jnp.sqrt(data.shape[-1])))
    else:
        data_normalizer = 1.0
    ratio = 1.0 / jnp.sqrt(projection_matrix.shape[0])
    data_mod_shape = data.shape[0 : len(batch_dims_t)] + projection_matrix.shape
    data_thick_random_matrix = jnp.zeros(data_mod_shape) + projection_matrix

    data_dash = lax.dot_general(
        data_normalizer * data,
        data_thick_random_matrix,
        (((data.ndim - 1,), (data_thick_random_matrix.ndim - 1,)), (batch_dims_t, batch_dims_t)),
        precision=precision,
    )

    diag_data = jnp.square(data)
    diag_data = jnp.sum(diag_data, axis=data.ndim - 1)
    diag_data = (diag_data / 2.0) * data_normalizer * data_normalizer
    diag_data = jnp.expand_dims(diag_data, axis=data.ndim - 1)

    if is_query:
        last_dims_t = (len(data_dash.shape) - 1,)
        data_dash = ratio * (
            jnp.exp(data_dash - diag_data - jnp.max(data_dash, axis=last_dims_t, keepdims=True)) + eps
        )
    else:
        data_dash = ratio * (jnp.exp(data_dash - diag_data - jnp.max(data_dash)) + eps)



SUMMARY:
**Function Summary: nonnegative_softmax_kernel_feature_creator**

**Main Purpose and Functionality**
Constructs nonnegative kernel features for fast softmax attention using a projection matrix, input data, and attention dimensions.

**Key Implementation Details**

* Computes kernel features by taking the dot product of normalized input data with a random matrix in projection space.
* Applies exponentiation using Jax's `exp` function to compute softmax-like scoring.
* Normalizes output using a diagonal matrix containing normalization constants.

**Notable Dependencies or Requirements**

* `jnp`: Jax NumPy, used for numerical computations and tensor operations.
* `lax.dot_general`: A function from Jax Library for efficient generic broadcasting of dot products.
* Custom data normalizers (`data_normalizer`) that transform input data into a suitable scale.

**Return Value**
Random features for fast softmax attention, denoted as `data_dash`.


key: chroma:document
string_value:
CODE:

def sincos_softmax_kernel_feature_creator(
    data, projection_matrix, attention_dims_t, batch_dims_t, precision, normalize_data=True
):
    """
    Constructs kernel sin-cos features for fast softmax attention

    Args:
      data: input for which features are computes
      projection_matrix: random matrix used to compute features
      attention_dims_t: tuple of attention dimensions
      batch_dims_t: tuple of batch dimensions
      precision: precision parameter
      normalize_data: predicate indicating whether data should be normalized

    Returns:
      Random features for fast softmax attention.
    """
    if normalize_data:
        # We have: exp(qk^T/sqrt{d}) = exp(|q|^2/2sqrt{d}) * exp(|k|^2/2sqrt{d}) *
        # exp(-(|q*c-k*c|^2)/2), where c = 1.0 / sqrt{sqrt{d}}.
        data_normalizer = 1.0 / (jnp.sqrt(jnp.sqrt(data.shape[-1])))
    else:
        data_normalizer = 1.0
    ratio = 1.0 / jnp.sqrt(projection_matrix.shape[0])
    data_mod_shape = data.shape[0 : len(batch_dims_t)] + projection_matrix.shape
    data_thick_random_matrix = jnp.zeros(data_mod_shape) + projection_matrix

    data_dash = lax.dot_general(
        data_normalizer * data,
        data_thick_random_matrix,
        (((data.ndim - 1,), (data_thick_random_matrix.ndim - 1,)), (batch_dims_t, batch_dims_t)),
        precision=precision,
    )
    data_dash_cos = ratio * jnp.cos(data_dash)
    data_dash_sin = ratio * jnp.sin(data_dash)
    data_dash = jnp.concatenate((data_dash_cos, data_dash_sin), axis=-1)

    # Constructing D_data and data^{'}
    diag_data = jnp.square(data)
    diag_data = jnp.sum(diag_data, axis=data.ndim - 1)
    diag_data = (diag_data / 2.0) * data_normalizer * data_normalizer
    diag_data = jnp.expand_dims(diag_data, axis=data.ndim - 1)
    # Additional renormalization for numerical stability
    data_renormalizer = jnp.max(diag_data, attention_dims_t, keepdims=True)
    diag_data -= data_renormalizer
    diag_data = jnp.exp(diag_data)
    data_prime = data_dash * diag_data


SUMMARY:
**Function Summary: `sincos_softmax_kernel_feature_creator`**

Constructs kernel sin-cos features for fast softmax attention.

**Key Features:**

* Computes kernel features using a projection matrix and input data `_data`
* Employs Lanczos-style algorithms to achieve efficient computation
* Utilizes JAX libraries (e.g., `lax.dot_general`, `jnp.cos`)

**Implementation details:**

* Normalization is optional, controlled by the `normalize_data` parameter
* Features are constructed using a combination of dot-product and matrix multiplication operations

**Notable dependencies:**

* Projection matrix: `projection_matrix`
* Data input: `_data`
* Attention dimensions tuple: `attention_dims_t`
* Batch dimensions tuple: `batch_dims_t`

**Interfaces/Methods used:**

* `lax.dot_general`: for computing dot product
* `jnp.cos` and `jnp.sin`: for computing cosine and sine values, respectively

The function takes in the following parameters:
```python
data, projection_matrix, attention_dims_t, batch_dims_t, precision, normalize_data=True
```
It returns the generated kernel sin-cos features: `Random features for fast softmax attention`.


key: chroma:document
string_value:
CODE:

def generalized_kernel_feature_creator(
    data, projection_matrix, batch_dims_t, precision, kernel_fn, kernel_epsilon, normalize_data
):
    """
    Constructs kernel features for fast generalized attention

    Args:
      data: input for which features are computes
      projection_matrix: matrix used to compute features
      batch_dims_t: tuple of batch dimensions
      precision: precision parameter
      kernel_fn: kernel function used
      kernel_epsilon: additive positive term added to every feature for numerical
        stability
      normalize_data: predicate indicating whether data should be normalized

    Returns:
      Random features for fast generalized attention.
    """
    if normalize_data:
        data_normalizer = 1.0 / (jnp.sqrt(jnp.sqrt(data.shape[-1])))
    else:
        data_normalizer = 1.0
    if projection_matrix is None:
        return kernel_fn(data_normalizer * data) + kernel_epsilon
    else:
        data_mod_shape = data.shape[0 : len(batch_dims_t)] + projection_matrix.shape
        data_thick_random_matrix = jnp.zeros(data_mod_shape) + projection_matrix
        data_dash = lax.dot_general(
            data_normalizer * data,
            data_thick_random_matrix,
            (((data.ndim - 1,), (data_thick_random_matrix.ndim - 1,)), (batch_dims_t, batch_dims_t)),
            precision=precision,
        )
    data_prime = kernel_fn(data_dash) + kernel_epsilon


SUMMARY:
**Function Summary: `generalized_kernel_feature_creator`**

* **Purpose**: Constructs random features for fast generalized attention.
* **Primary Functionality**: Computes kernel features using linear transformations of input data or projection matrices.

**Key Implementation Details**:
* Uses JAX's `lax.dot_general` to perform linear transformations with precision control.
* Applies kernel function (`kernel_fn`) to transformed data and adds a small noise term (`kernel_epsilon`) for numerical stability.

**Notable Dependencies/Requirements**:
* `projection_matrix`: A matrix used to transform input data (optional).
* `jnp`: JAX numerical library, providing support for numerical computations.
* `lax.dot_general`: JAX's built-in function for generic tensor contraction.
* `kernel_fn`: A custom or pre-defined kernel function used in the computation.

**Return Value**: Random features for fast generalized attention, computed using linear transformations and a kernel function.


key: chroma:document
string_value:
CODE:

def make_fast_softmax_attention(
    qkv_dim,
    renormalize_attention=True,
    numerical_stabilizer=0.000001,
    nb_features=256,
    ortho_features=True,
    ortho_scaling=0.0,
    redraw_features=True,
    unidirectional=False,
    nonnegative_features=True,
    lax_scan_unroll=1,
):
    """Construct a fast softmax attention method."""
    logging.info(
        "Fast softmax attention: %s features and orthogonal=%s, renormalize=%s",
        nb_features,
        ortho_features,
        renormalize_attention,
    )
    if ortho_features:
        matrix_creator = functools.partial(GaussianOrthogonalRandomMatrix, nb_features, qkv_dim, scaling=ortho_scaling)
    else:
        matrix_creator = functools.partial(GaussianUnstructuredRandomMatrix, nb_features, qkv_dim)
    if nonnegative_features:

        def kernel_feature_creator(
            data, projection_matrix, attention_dims_t, batch_dims_t, precision, is_query, normalize_data=True
        ):
            return nonnegative_softmax_kernel_feature_creator(
                data,
                projection_matrix,
                attention_dims_t,
                batch_dims_t,
                precision,
                is_query,
                normalize_data,
                numerical_stabilizer,
            )

    else:

        def kernel_feature_creator(
            data, projection_matrix, attention_dims_t, batch_dims_t, precision, is_query, normalize_data=True
        ):
            del is_query
            return sincos_softmax_kernel_feature_creator(
                data, projection_matrix, attention_dims_t, batch_dims_t, precision, normalize_data
            )

    attention_fn = FastAttentionviaLowRankDecomposition(
        matrix_creator,
        kernel_feature_creator,
        renormalize_attention=renormalize_attention,
        numerical_stabilizer=numerical_stabilizer,
        redraw_features=redraw_features,
        unidirectional=unidirectional,
        lax_scan_unroll=lax_scan_unroll,
    ).dot_product_attention


SUMMARY:
**Fast Softmax Attention Function Summary**

**Main Purpose and Functionality:**
Constructs a fast softmax attention method with optional customizability for orthogonal features, renormalization, numerical stabilization, and more.

**Key Implementation Details:**

* Uses `GaussianOrthogonalRandomMatrix` or `GaussianUnstructuredRandomMatrix` to create matrices.
* Implements two kernel creation functions:
	+ With normalization (using `nonnegative_softmax_kernel_feature_creator`) when `nonnegative_features=True`.
	+ Without normalization (using `sincos_softmax_kernel_feature_creator`) when `nonnegative_features=False`.
* Defines the attention function using a low-rank decomposition, including configuration settings like renormalization and numerical stabilization.

**Notable Dependencies or Requirements:**

* `functools` for partial application of matrix creators.
* Custom functions:
	+ `GaussianOrthogonalRandomMatrix`
	+ `GaussianUnstructuredRandomMatrix`
	+ `nonnegative_softmax_kernel_feature_creator`
	+ `sincos_softmax_kernel_feature_creator`

**Primary Purpose:** Constructs a fast softmax attention method with customizable settings.


key: chroma:document
string_value:
CODE:

def make_fast_generalized_attention(
    qkv_dim,
    renormalize_attention=True,
    numerical_stabilizer=0.0,
    nb_features=256,
    features_type="deterministic",
    kernel_fn=jax.nn.relu,
    kernel_epsilon=0.001,
    redraw_features=False,
    unidirectional=False,
    lax_scan_unroll=1,
):
    """Construct a fast generalized attention menthod."""
    logging.info("Fast generalized attention.: %s features and renormalize=%s", nb_features, renormalize_attention)
    if features_type == "ortho":
        matrix_creator = functools.partial(GaussianOrthogonalRandomMatrix, nb_features, qkv_dim, scaling=False)
    elif features_type == "iid":
        matrix_creator = functools.partial(GaussianUnstructuredRandomMatrix, nb_features, qkv_dim)
    elif features_type == "deterministic":
        matrix_creator = None
    else:
        raise ValueError("Unknown feature value type")

    def kernel_feature_creator(
        data, projection_matrix, attention_dims_t, batch_dims_t, precision, is_query, normalize_data=False
    ):
        del attention_dims_t
        del is_query
        return generalized_kernel_feature_creator(
            data, projection_matrix, batch_dims_t, precision, kernel_fn, kernel_epsilon, normalize_data
        )

    attention_fn = FastAttentionviaLowRankDecomposition(
        matrix_creator,
        kernel_feature_creator,
        renormalize_attention=renormalize_attention,
        numerical_stabilizer=numerical_stabilizer,
        redraw_features=redraw_features,
        unidirectional=unidirectional,
        lax_scan_unroll=lax_scan_unroll,
    ).dot_product_attention


SUMMARY:
**Concise Summary**

Constructs a fast generalized attention method (`FastAttentionviaLowRankDecomposition`) with customizable parameters:

* `make_fast_generalized_attention`:
	+ Takes input dimensions (`qkv_dim`) and hyperparameters (`RenormalizeAttention`, `NumericalStabilizer`, etc.)
	+ Chooses between different matrix creators based on `features_type` (ortho, iid, deterministic)
	+ Creates a kernel feature creator function (`kernel_feature_creator`) using `generalized_kernel_feature_creator`
	+ Instantiates the fast generalized attention method with chosen creators and hyperparameters
	+ Returns composed dot product attention function

**Key Implementation Details**

* Uses `FastAttentionviaLowRankDecomposition` with custom matrix creators and kernel features
* Applies `GaussianOrthogonalRandomMatrix`, `GaussianUnstructuredRandomMatrix`, or `none` for different `features_type` values
* Utilizes `generalized_kernel_feature_creator` to create kernel features

**Notable Dependencies and Requirements**

* JAX (`jax.nn.relu`)
* Custom functions:
	+ `GaussianOrthogonalRandomMatrix`
	+ `GaussianUnstructuredRandomMatrix`
	+ `generalizedKernelFeatureCreator`


key: chroma:document
string_value:
CODE:

class RandomMatrix:
    r"""
    Abstract class providing a method for constructing 2D random arrays. Class is responsible for constructing 2D
    random arrays.
    """

    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def get_2d_array(self):


SUMMARY:
**Summary: Abstract RandomMatrix Class**

### Purpose and Functionality:

The `RandomMatrix` class is an abstract base class that provides a method for constructing 2D random arrays. Its primary purpose is to serve as a template for concrete implementations.

### Implementation Details:

* The class uses the Metaclass `abc.ABCMeta` to enable abstract methods.
* The `get_2d_array` method is declared as an abstract interface, required to be implemented by any subclass.

### Dependencies/Requirements:

* None explicitly mentioned
* Any subclass of `RandomMatrix` must implement the `get_2d_array` interface.


key: chroma:document
string_value:
CODE:

class GaussianUnstructuredRandomMatrix(RandomMatrix):
    def __init__(self, nb_rows, nb_columns, key):
        self.nb_rows = nb_rows
        self.nb_columns = nb_columns
        self.key = key

    def get_2d_array(self):


SUMMARY:
**Summary:**

Class `GaussianUnstructuredRandomMatrix` extends `RandomMatrix` with functionality to generate dense, unstructured matrices using Gaussian distributions.

**Key details:**

- Initializes matrix dimensions (`nb_rows`, `nb_columns`) and distribution key (`key`).
- Generated 2D array using `random.normal()` with specified key and shape `(nb_rows, nb_columns)`.

**Notable Interface:**

- Extends `RandomMatrix`
- Depends on `random`

**External Dependencies:**

- `numpy`: implied by `random.normal()`
- No external libraries required


key: chroma:document
string_value:
CODE:

class GaussianOrthogonalRandomMatrix(RandomMatrix):
    r"""
    Class providing a method to create Gaussian orthogonal matrix. Class is responsible for constructing 2D Gaussian
    orthogonal arrays.
    """

    def __init__(self, nb_rows, nb_columns, key, scaling=0):
        self.nb_rows = nb_rows
        self.nb_columns = nb_columns
        self.key = key
        self.scaling = scaling

    def get_2d_array(self):
        nb_full_blocks = int(self.nb_rows / self.nb_columns)
        block_list = []
        rng = self.key
        for _ in range(nb_full_blocks):
            rng, rng_input = jax.random.split(rng)
            unstructured_block = random.normal(rng_input, (self.nb_columns, self.nb_columns))
            q, _ = jnp.linalg.qr(unstructured_block)
            q = jnp.transpose(q)
            block_list.append(q)
        remaining_rows = self.nb_rows - nb_full_blocks * self.nb_columns
        if remaining_rows > 0:
            rng, rng_input = jax.random.split(rng)
            unstructured_block = random.normal(rng_input, (self.nb_columns, self.nb_columns))
            q, _ = jnp.linalg.qr(unstructured_block)
            q = jnp.transpose(q)
            block_list.append(q[0:remaining_rows])
        final_matrix = jnp.vstack(block_list)

        if self.scaling == 0:
            multiplier = jnp.linalg.norm(random.normal(self.key, (self.nb_rows, self.nb_columns)), axis=1)
        elif self.scaling == 1:
            multiplier = jnp.sqrt(float(self.nb_columns)) * jnp.ones((self.nb_rows))
        else:
            raise ValueError("Scaling must be one of {0, 1}. Was %s" % self._scaling)



SUMMARY:
### GaussianOrthogonalRandomMatrix Class Summary

#### Purpose and Functionality:
The `GaussianOrthogonalRandomMatrix` class generates 2D Gaussian orthogonal arrays. It takes as input the number of rows (`nb_rows`) and columns (`nb_columns`) and optionally a scaling factor (`scaling`). The `get_2d_array` method returns the constructed Gaussian orthogonal matrix.

#### Key Implementation Details:

*   Constructs Gaussian orthogonal matrices in blocks, combining them to form the final matrix.
*   Uses `jax.random.split` for sampling and `jnp.linalg.qr` for QR decomposition.
*   Applies scaling factor (`scaling`) using diagonal elements of a multiplier matrix.

#### Notable Dependencies and Requirements:

*   `RandomMatrix`: Base class providing foundation for Gaussian orthogonal matrices.
*   `jax.random` and `jax.numpy` (as `jnp`) for random sampling and numerical computations.
*   `numpy.linalg` for QR decomposition.

#### Example Method Signatures:
```python
def __init__(self, nb_rows, nb_columns, key, scaling=0):
    ...

def get_2d_array(self):
    ...
```
Inheritance and Interfaces:
The class inherits from `RandomMatrix` and provides a unique method (`get_2d_array`) to generate 2D Gaussian orthogonal arrays.


key: chroma:document
string_value:
CODE:

class FastAttention:
    r"""
    Abstract class providing a method for fast attention. Class is responsible for providing a method
    <dot_product_attention> for fast approximate attention.
    """

    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def dot_product_attention(
        self,
        query,
        key,
        value,
        dtype=jnp.float32,
        bias=None,
        axis=None,
        broadcast_dropout=True,
        dropout_rng=None,
        dropout_rate=0.0,
        deterministic=False,
        precision=None,
    ):
        """
        Computes dot-product attention given query, key, and value. This is the core function for applying fast
        approximate dot-product attention. It calculates the attention weights given query and key and combines the
        values using the attention weights. This function supports multi-dimensional inputs

        Args:
          query: queries for calculating attention with shape of [batch_size, dim1,
            dim2, ..., dimN, num_heads, mem_channels].
          key: keys for calculating attention with shape of [batch_size, dim1, dim2,
            ..., dimN, num_heads, mem_channels].
          value: values to be used in attention with shape of [batch_size, dim1,
            dim2,..., dimN, num_heads, value_channels].
          dtype: the dtype of the computation (default: float32)
          bias: bias for the attention weights. This can be used for incorporating
            autoregressive mask, padding mask, proximity bias.
          axis: axises over which the attention is applied.
          broadcast_dropout: bool: use a broadcasted dropout along batch dims.
          dropout_rng: JAX PRNGKey: to be used for dropout.
          dropout_rate: dropout rate.
          deterministic: bool, deterministic or not (to apply dropout).
          precision: numerical precision of the computation see `jax.lax.Precision`
            for details

        Returns:
          Output of shape [bs, dim1, dim2, ..., dimN,, num_heads, value_channels].
        """


SUMMARY:
**FastAttention Class Summary**

Provides an abstract implementation for fast approximate dot-product attention.

**Key Methods:**

* `dot_product_attention`: computes dot-product attention given query, key, and value. It calculates attention weights based on the input shapes and combines values using these weights.

**Implementation Details:**

* Supports multi-dimensional inputs
* Takes into account bias for autoregressive mask, padding mask, and proximity bias
* Applies dropout with optional broadcasted dropout along batch dims

**Parameters:**

* `query`, `key`, `value`: input tensors for attention calculation
* `dtype`: computation dtype (default: float32)
* `bias`: bias for attention weights
* `axis`: axes over which attention is applied
* `broadcast_dropout`: bool for broadcasted dropout along batch dims
* `dropout_rng`: JAX PRNGKey for dropout
* `dropout_rate`:
* `deterministic`: bool for deterministic or non-deterministic attention

**Returns:**

* Output tensor with shape [batch_size, ..., num_heads, value_channels]


key: chroma:document
string_value:
CODE:

def _numerator(z_slice_shape, precision, unroll=1):
    def fwd(qs, ks, vs):
        def body(p, qkv):
            (q, k, v) = qkv
            p += jnp.einsum("...m,...d->...md", k, v, precision=precision)
            X_slice = jnp.einsum("...m,...md->...d", q, p, precision=precision)
            return p, X_slice

        init_value = jnp.zeros(z_slice_shape)
        p, W = lax.scan(body, init_value, (qs, ks, vs), unroll=unroll)
        return W, (p, qs, ks, vs)

    def bwd(pqkv, W_ct):
        def body(carry, qkv_xct):
            p, p_ct = carry
            q, k, v, x_ct = qkv_xct
            q_ct = jnp.einsum("...d,...md->...m", x_ct, p, precision=precision)
            p_ct += jnp.einsum("...d,...m->...md", x_ct, q, precision=precision)
            k_ct = jnp.einsum("...md,...d->...m", p_ct, v, precision=precision)
            v_ct = jnp.einsum("...md,...m->...d", p_ct, k, precision=precision)
            p -= jnp.einsum("...m,...d->...md", k, v, precision=precision)
            return (p, p_ct), (q_ct, k_ct, v_ct)

        p, qs, ks, vs = pqkv
        _, (qs_ct, ks_ct, vs_ct) = lax.scan(
            body, (p, jnp.zeros_like(p)), (qs, ks, vs, W_ct), reverse=True, unroll=unroll
        )
        return qs_ct, ks_ct, vs_ct

    @jax.custom_vjp
    def _numerator_impl(qs, ks, vs):
        W, _ = fwd(qs, ks, vs)
        return W

    _numerator_impl.defvjp(fwd, bwd)



SUMMARY:
**Summary:**

This code defines a custom JAX `jax.custom_vjp` implementation for computing the numerator of a neural network's forward pass. The function `_numerator` takes three arguments:

- `z_slice_shape`: Shape of the input tensor
- `precision`: Precision parameter for numerical computations
- `unroll`: Unrolling factor for memory efficiency (default=1)

The function consists of two main components:

**`fwd(qs, ks, vs)`**

This inner function performs the forward pass using `jnp.einsum`. It takes three tensor arguments: `qs`, `ks`, and `vs`.

- Compute `p` and `X_slice` by recursively applying a body function to the input tensors
- Return the updated tensors

**`bwd(pqkv, W_ct)`**

This inner function computes the gradients of the loss with respect to each input tensor using reverse mode automatic differentiation.

- Take the current carry values (`p`, `qs`, `ks`, `vs`) and gradients of previous states (`X_slice`) as inputs
- Update the gradients in reverse order using another recursive body function

The `@jax.custom_vjp` decorator defines a custom VPJ (Variance Reduced Jacobian) implementation between the `_numerator_impl` forward pass and its reverse pass `bwd`.

**Key dependencies:**

* JAX library for numerical computations
* `jnp` namespace for vectorized operations

**Notable functions:**

- `fwd`
- `bwd`
- `_numerator_impl`
- `jax.custom_vjp`

Note that the `Z_slice_shape`, `precision`, and `unroll` arguments are not directly used in the implementation, but rather serve as input parameters for the function.


key: chroma:document
string_value:
CODE:

def _denominator(t_slice_shape, precision, unroll=1):
    def fwd(qs, ks):
        def body(p, qk):
            q, k = qk
            p += k
            x = jnp.einsum("...m,...m->...", q, p, precision=precision)
            return p, x

        p = jnp.zeros(t_slice_shape)
        p, R = lax.scan(body, p, (qs, ks), unroll=unroll)
        return R, (qs, ks, p)

    def bwd(qkp, R_ct):
        def body(carry, qkx):
            p, p_ct = carry
            q, k, x_ct = qkx
            q_ct = jnp.einsum("...,...m->...m", x_ct, p, precision=precision)
            p_ct += jnp.einsum("...,...m->...m", x_ct, q, precision=precision)
            k_ct = p_ct
            p -= k
            return (p, p_ct), (q_ct, k_ct)

        qs, ks, p = qkp
        _, (qs_ct, ks_ct) = lax.scan(body, (p, jnp.zeros_like(p)), (qs, ks, R_ct), reverse=True, unroll=unroll)
        return (qs_ct, ks_ct)

    @jax.custom_vjp
    def _denominator_impl(qs, ks):
        R, _ = fwd(qs, ks)
        return R

    _denominator_impl.defvjp(fwd, bwd)



SUMMARY:
**Summary: Denominator Function in JAX**

The `_denominator` function computes the convolution inverse in the frequency domain using JAX's `lunroll` and `lax.scan` functions.

**Main Purpose and Functionality:**
Computes forward and backward convolution inverses for a 1D convolutional layer with a custom Jacobian computation (`_denominator_impl`).

**Key Implementation Details:**

*   **Forward Convolution Inverse**: Computes the inverse of a 1D convolution with shape `(N, M)` given input `qs`, kernel `ks`, and precision `precision`. The output is returned along with intermediate results `p`.
    ```python
fwd(qs, ks) = lax.scan(body, p=0, (qs, ks))
```
*   **Backward Convolution Inverse**: Computes the inverse of a 1D convolution in reversed order. Given input `qkp` and output results `R_ct`. The output is returned along with intermediate results `p_ct`.
    ```python
bwd(qkp) = lax.scan(body, (p=0), (qs, ks, R_ct))
```

**Notable Dependencies or Requirements:**

*   JAX (`lux`)
*   JAX primitives (e.g., `[numpy], [linAlg])`
*   `jnp.einsum` and `lax.scan` for tensor operations
*   Custom Jacobian computation with `_denominator_impl`


key: chroma:document
string_value:
CODE:

class FastAttentionviaLowRankDecomposition(FastAttention):
    r"""
    Class providing a method for fast attention via low rank decomposition. Class is responsible for providing a method
    <dot_product_attention> for fast dot-product attention with the use of low rank decomposition (e.g. with random
    feature maps).
    """

    def __init__(
        self,
        matrix_creator,
        kernel_feature_creator,
        renormalize_attention,
        numerical_stabilizer,
        redraw_features,
        unidirectional,
        lax_scan_unroll=1,
    ):  # For optimal GPU performance, set to 16.
        rng = random.PRNGKey(0)
        self.matrix_creator = matrix_creator
        self.projection_matrix = self.draw_weights(rng)
        self.kernel_feature_creator = kernel_feature_creator
        self.renormalize_attention = renormalize_attention
        self.numerical_stabilizer = numerical_stabilizer
        self.redraw_features = redraw_features
        self.unidirectional = unidirectional
        self.lax_scan_unroll = lax_scan_unroll

    def draw_weights(self, key):
        if self.matrix_creator is None:
            return None
        matrixrng, _ = random.split(key)
        projection_matrix = self.matrix_creator(key=matrixrng).get_2d_array()
        return projection_matrix

    def dot_product_attention(
        self,
        query,
        key,
        value,
        dtype=jnp.float32,
        bias=None,
        axis=None,
        broadcast_dropout=True,
        dropout_rng=None,
        dropout_rate=0.0,
        deterministic=False,
        precision=None,
    ):
        assert key.shape[:-1] == value.shape[:-1]
        assert query.shape[0:1] == key.shape[0:1] and query.shape[-1] == key.shape[-1]
        if axis is None:
            axis = tuple(range(1, key.ndim - 2))
        if not isinstance(axis, Iterable):
            axis = (axis,)
        assert key.ndim == query.ndim
        assert key.ndim == value.ndim
        for ax in axis:
            if not (query.ndim >= 3 and 1 <= ax < query.ndim - 2):
                raise ValueError("Attention axis must be between the batch axis and the last-two axes.")
        n = key.ndim

        # Constructing projection tensor.
        if self.redraw_features:
            # TODO(kchoro): Get rid of the constant below.
            query_seed = lax.convert_element_type(jnp.ceil(jnp.sum(query) * 10000000.0), jnp.int32)
            rng = random.PRNGKey(query_seed)
            self.projection_matrix = self.draw_weights(rng)

        # batch_dims is  <bs, <non-attention dims>, num_heads>
        batch_dims = tuple(onp.delete(range(n), axis + (n - 1,)))
        # q & k -> (bs, <non-attention dims>, num_heads, <attention dims>, channels)
        qk_perm = batch_dims + axis + (n - 1,)
        k_extra_perm = axis + batch_dims + (n - 1,)
        key_extra = key.transpose(k_extra_perm)
        key = key.transpose(qk_perm)
        query = query.transpose(qk_perm)
        # v -> (bs, <non-attention dims>, num_heads, <attention dims>, channels)
        v_perm = batch_dims + axis + (n - 1,)
        value = value.transpose(v_perm)
        batch_dims_t = tuple(range(len(batch_dims)))
        attention_dims_t = tuple(range(len(batch_dims), len(batch_dims) + len(axis)))

        # Constructing tensors Q^{'} and K^{'}.
        query_prime = self.kernel_feature_creator(
            query, self.projection_matrix, attention_dims_t, batch_dims_t, precision, True
        )
        key_prime = self.kernel_feature_creator(
            key, self.projection_matrix, attention_dims_t, batch_dims_t, precision, False
        )

        if self.unidirectional:
            index = attention_dims_t[0]
            z_slice_shape = key_prime.shape[0 : len(batch_dims_t)] + (key_prime.shape[-1],) + (value.shape[-1],)

            numerator_fn = _numerator(z_slice_shape, precision, self.lax_scan_unroll)
            W = numerator_fn(
                jnp.moveaxis(query_prime, index, 0), jnp.moveaxis(key_prime, index, 0), jnp.moveaxis(value, index, 0)
            )

            # Constructing W = (Q^{'}(K^{'})^{T})_{masked}V
            W = jnp.moveaxis(W, 0, index)

            if not self.renormalize_attention:
                # Unidirectional, not-normalized attention.
                perm_inv = _invert_perm(qk_perm)
                result = W.transpose(perm_inv)
                return result
            else:
                # Unidirectional, normalized attention.
                thick_all_ones = jnp.zeros(key.shape[0:-1]) + jnp.ones(key_extra.shape[0 : len(axis)])

                index = attention_dims_t[0]
                t_slice_shape = key_prime.shape[0 : len(batch_dims_t)] + (key_prime.shape[-1],)
                denominator_fn = _denominator(t_slice_shape, precision, self.lax_scan_unroll)
                R = denominator_fn(jnp.moveaxis(query_prime, index, 0), jnp.moveaxis(key_prime, index, 0))

                R = jnp.moveaxis(R, 0, index)
        else:
            contract_query = tuple(range(len(batch_dims) + len(axis), len(batch_dims) + len(axis) + 1))
            contract_z = tuple(range(len(batch_dims), len(batch_dims) + 1))
            # Constructing Z = (K^{'})^{T}V
            # Z (bs, <non-attention dims>, num_heads, channels_m, channels_v)
            Z = lax.dot_general(
                key_prime,
                value,
                ((attention_dims_t, attention_dims_t), (batch_dims_t, batch_dims_t)),
                precision=precision,
            )
            # Constructing W = Q^{'}Z = Q^{'}(K^{'})^{T}V
            # q (bs, <non-attention dims>, num_heads, <attention dims>, channels_m)
            # Z (bs, <non-attention dims>, num_heads, channels_m, channels_v)
            # W (bs,  <non-attention dims>, num_heads, <attention dims>, channels_v)
            W = lax.dot_general(
                query_prime, Z, ((contract_query, contract_z), (batch_dims_t, batch_dims_t)), precision=precision
            )
            if not self.renormalize_attention:
                # Bidirectional, not-normalized attention.
                perm_inv = _invert_perm(qk_perm)
                result = W.transpose(perm_inv)
                return result
            else:
                # Bidirectional, normalized attention.
                thick_all_ones = jnp.zeros(key.shape[0:-1]) + jnp.ones(key_extra.shape[0 : len(axis)])
                contract_key = tuple(range(len(batch_dims), len(batch_dims) + len(axis)))
                contract_thick_all_ones = tuple(range(thick_all_ones.ndim - len(axis), thick_all_ones.ndim))
                # Construct T = (K^{'})^{T} 1_L
                # k (bs, <non-attention dims>, num_heads, <attention dims>, channels)
                T = lax.dot_general(
                    key_prime,
                    thick_all_ones,
                    ((contract_key, contract_thick_all_ones), (batch_dims_t, batch_dims_t)),
                    precision=precision,
                )

                # Construct partition function: R = Q^{'} T = Q^{'}(K^{'})^{T} 1_L
                # q_p (bs, <non-attention dims>, num_heads, <attention dims>, channs_m)
                # T   (bs, <non-attention dims>, num_heads, channels_m)
                R = lax.dot_general(
                    query_prime,
                    T,
                    (((query_prime.ndim - 1,), (T.ndim - 1,)), (batch_dims_t, range(0, len(T.shape) - 1))),
                    precision=precision,
                )

        R = R + 2 * self.numerical_stabilizer * (jnp.abs(R) <= self.numerical_stabilizer)
        R = jnp.reciprocal(R)
        R = jnp.expand_dims(R, len(R.shape))
        # W (bs, <non-attention dims>, num_heads, <attention dims>, channels_v)
        # R (bs, <non-attention dims>, num_heads, <attention dims>, extra_channel)
        result = W * R
        # back to (bs, dim1, dim2, ..., dimN, num_heads, channels)
        perm_inv = _invert_perm(qk_perm)
        result = result.transpose(perm_inv)


SUMMARY:
This is an implementation of the Fast Attention algorithm via Low Rank Decomposition in TensorFlow 2.0.

**Key Features:**

1. **Low Rank Decomposition**: The algorithm uses low rank decomposition to reduce the computational complexity of attention operations.
2. **Fast Attention**: This implementation provides a fast way to perform attention operations by using pre-computed projection matrices and kernel feature creators.
3. **Unidirectional and Bidirectional Attention**: The algorithm can be used for both unidirectional and bidirectional attention, depending on the `unidirectional` flag.
4. **Normalized Attention**: The algorithm also supports normalized attention by adding a partition function to the result.

**Main Components:**

1. `_ invert_perm` function: Inverts the permutation `perm_inv` used in the attention operation.
2. `_denominator` and `_numerator` functions: Compute the denominator and numerator of the partition function.
3. `_dot_general` function: Evaluates a matrix product with specified precision and memory requirements.
4. Main loop: Computationally intensive core of the Fast Attention algorithm, which handles low rank decomposition, attention operation, normalization, and unidirectional or bidirectional attention.

**Inputs and Outputs:**

* **Input tensors**: `query`, `key`, `value`.
* **Outputs**: The final result tensor (`result`) after performing the attention operation.


key: chroma:document
string_value:
CODE:

def _invert_perm(perm):
    perm_inv = [0] * len(perm)
    for i, j in enumerate(perm):
        perm_inv[j] = i


SUMMARY:
**Summary**

`_invert_perm` function generates the inverse permutation of input `perm`. It creates a new array `perm_inv` and assigns each element's index to its corresponding value in the original permutation.

**Key Implementation Details**

* Takes an iterable of indices as input (`perm`)
* Uses list comprehension to create `perm_inv`
* Utilizes `enumerate` for indexing and iteration

**Output**

Returns the inverse permutation as a tuple of integers.

**Notable Dependencies/Requirements**

None.


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/performer/modeling_flax_performer.py


key: chroma:document
string_value:
CODE:

class FlaxPerformerLayerNorm(nn.Module):
    """
    Layer normalization (https://arxiv.org/abs/1607.06450). Operates on the last axis of the input data.
    """

    epsilon: float = 1e-6
    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
    bias: bool = True  # If True, bias (beta) is added.
    scale: bool = True  # If True, multiply by scale (gamma). When the next layer is linear
    # (also e.g. nn.relu), this can be disabled since the scaling will be
    # done by the next layer.
    bias_init: jnp.ndarray = nn.initializers.zeros
    scale_init: jnp.ndarray = nn.initializers.ones

    @nn.compact
    def __call__(self, x):
        """
        Applies layer normalization on the input. It normalizes the activations of the layer for each given example in
        a batch independently, rather than across a batch like Batch Normalization. i.e. applies a transformation that
        maintains the mean activation within each example close to 0 and the activation standard deviation close to 1

        Args:
          x: the inputs

        Returns:
          Normalized inputs (the same shape as inputs).
        """
        features = x.shape[-1]
        mean = jnp.mean(x, axis=-1, keepdims=True)
        mean2 = jnp.mean(jax.lax.square(x), axis=-1, keepdims=True)
        var = mean2 - jax.lax.square(mean)
        mul = jax.lax.rsqrt(var + self.epsilon)
        if self.scale:
            mul = mul * jnp.asarray(self.param("gamma", self.scale_init, (features,)), self.dtype)
        y = (x - mean) * mul
        if self.bias:
            y = y + jnp.asarray(self.param("beta", self.bias_init, (features,)), self.dtype)


SUMMARY:
Here's a concise summary of the provided class code chunk:

**Class:** `FlaxPerformerLayerNorm`

**Purpose:** Layer normalization layer that normalizes activations for each example in a batch independently.

**Key Implementation:**

* Applies mean and variance calculation over last axis of input data
* Normalized inputs (x - mean) * sqrt(var + epsilon)
* Adds bias if enabled, scaled by learnable scale factor gamma

**Notable Dependencies/Requirements:**

* `jnp` for numerical computations
* `nn` for neural network functionality (e.g., `params`, `initializers`)
* `jax.lax` for lazy-style array operations
* External dependency on a linear layer or a ReLU activation function


key: chroma:document
string_value:
CODE:

class FlaxPerformerEmbedding(nn.Module):
    """
    Specify a new class for doing the embedding stuff as Flax's one use 'embedding' for the parameter name and PyTorch
    use 'weight'
    """

    vocab_size: int
    hidden_size: int
    emb_init: Callable[..., np.ndarray] = nn.initializers.normal(stddev=0.1)

    @nn.compact
    def __call__(self, inputs):
        embedding = self.param("weight", self.emb_init, (self.vocab_size, self.hidden_size))


SUMMARY:
**FlaxPerformerEmbedding Class Summary**

**Class Purpose:**
The `FlaxPerformerEmbedding` class is an implementation of an embedding layer for the Flax library in PyTorch, inspired by the performance embedding from the Flax framework.

**Key Functions and Variables:**

*   `__call__`: The main entry point of the class, which takes in input data.
*   `param("weight", self.emb_init, (self.vocab_size, self.hidden_size))`: Initializes an embedding parameter using a normal distribution with specified stddev.
*   `jnp.take(embedding, inputs, axis=0)`: Extracts the topmost dimension of the embedding into a batch.

**Dependencies and Requirements:**

*   `nn.initializers.normal(stddev=0.1)`: Used to initialize the embeddings.
*   `nn.Module`, `jnp`: Imported from the Flax library for PyTorch integration.


key: chroma:document
string_value:
CODE:

class FlaxPerformerEmbeddings(nn.Module):
    """Construct the embeddings from word, position and token_type embeddings."""

    vocab_size: int
    hidden_size: int
    type_vocab_size: int
    max_length: int

    @nn.compact
    def __call__(self, input_ids, token_type_ids, position_ids, attention_mask):
        # Embed
        w_emb = FlaxPerformerEmbedding(self.vocab_size, self.hidden_size, name="word_embeddings")(
            jnp.atleast_2d(input_ids.astype("i4"))
        )
        p_emb = FlaxPerformerEmbedding(self.max_length, self.hidden_size, name="position_embeddings")(
            jnp.atleast_2d(position_ids.astype("i4"))
        )
        t_emb = FlaxPerformerEmbedding(self.type_vocab_size, self.hidden_size, name="token_type_embeddings")(
            jnp.atleast_2d(token_type_ids.astype("i4"))
        )

        # Sum all embeddings
        summed_emb = w_emb + jnp.broadcast_to(p_emb, w_emb.shape) + t_emb

        # Layer Norm
        layer_norm = FlaxPerformerLayerNorm(name="layer_norm")(summed_emb)



SUMMARY:
**FlaxPerformerEmbeddings Class Summary**

* **Purpose:** Constructs embeddings from word, position, and token-type inputs for Transformer model embeddings.
* **Key Features:**
	+ Embeds input IDs, positions, and token types into high-dimensional spaces using `FlaxPerformerEmbedding`.
	+ Sums and normalizes embedded features with Layer Norm in `FlaxPerformerLayerNorm`.
* **Methods:**
	+ `__call__`: Takes `input_ids`, `token_type_ids`, `position_ids`, and `attention_mask` as input.
	+ Uses `FlaxPerformerEmbedding` to embed input IDs, positions, and token types with corresponding vocab sizes.
	+ Applies Layer Norm using `FlaxPerformerLayerNorm`.
* **Requirements:*
	+ `jnp` for numerical computations
	+ `nn.Module` for neural network structure definition
	+ `FlaxPerformerEmbedding` and `FlaxPerformerLayerNorm` classes for embedding and layer normalization implementations.


key: chroma:document
string_value:
CODE:

class FlaxPerformerAttention(nn.Module):
    num_heads: int
    head_size: int

    @nn.compact
    def __call__(self, hidden_state, attention_mask):
        single_head_dim = self.head_size // self.num_heads
        fast_softmax_attention = make_fast_softmax_attention(qkv_dim=single_head_dim)
        self_att = nn.attention.SelfAttention(
            num_heads=self.num_heads, qkv_features=self.head_size, name="self", attention_fn=fast_softmax_attention
        )(hidden_state, attention_mask)

        layer_norm = FlaxPerformerLayerNorm(name="layer_norm")(self_att + hidden_state)


SUMMARY:
**FlaxPerformerAttention Class Summary**

The `FlaxPerformerAttention` class is a custom neural network module implementing attention mechanism. It takes in `hidden_state` and `attention_mask`, and returns the result after self-attention and layer normalization.

**Key Implementation Details:**

* Uses `SelfAttention` from PyTorch, with:
	+ `num_heads`: 1
	+ `qkv_features`: `head_size`
	+ `attention_fn`: `fast_softmax_attention`
* Applies self-attention to `hidden_state` with `attention_mask`, and then adds it to the original `hidden_state`.
* Uses `FlaxPerformerLayerNorm` for layer normalization.

**Notable Dependencies:**

* Reifies `nn.Module` interface
* References external functions:
	+ `make_fast_softmax_attention`
* Inherets from `nn.Module`

**Primary Purpose and Functionality:**
This module computes attention weights between the query and key vectors in the input `hidden_state`, based on their similarity, as measured by the `attention_mask`. The resulting weights are then used to compute weighted sums over the input dimensions.


key: chroma:document
string_value:
CODE:

class FlaxPerformerIntermediate(nn.Module):
    output_size: int
    hidden_act: str = "gelu"

    @nn.compact
    def __call__(self, hidden_state):
        # TODO: Add ACT2FN reference to change activation function
        dense = nn.Dense(features=self.output_size, name="dense")(hidden_state)


SUMMARY:
**FlaxPerformerIntermediate Class Summary**

**Purpose and Functionality:**
The `FlaxPerformerIntermediate` class is a neural network layer that applies a dense feedforward transformation to the input `hidden_state`. It defines two main attributes:

*   `output_size`: The number of output units in the dense layer.
*   `hidden_act`: The activation function used in the dense layer, defaulting to "gelu".

**Key Implementation Details:**

*   The layer uses a compact sequence (`__call__()`) that applies:
    *   A dense feedforward transformation with `nn.Dense` (features = `output_size`, name="dense").
    *   An activation function determined by `hidden_act`.
        *   For "gelu", the activation function is retrieved from `ACT2FN` using the `self.hidden_act` value.

**Notable Dependencies and Requirements:**

*   The class inherits `nn.Module` from Hugging Face Transformers (TFX) library.
*   Repeated references:
    *   `nn.Dense()`.
        +   Parameters:
            *   `features=self.output_size`
            *   `name="dense"`
        +   Returns `dense` output tensor
    *   `ACT2FN[self.hidden_act]`.
        +   Retrieves the activation function based on `self.hidden_act`

**Code Highlight:**

```python
@nn.compact
def __call__(self, hidden_state):
    # ... (flax implementation unchanged)
    
    return ACT2FN[self.hidden_act](dense)  # Applies specified activation function
    
    # TODO: Add ACT2FN reference to change activation function
```

Overall, this class serves as a dense feedforward layer with flexible activation functions, suitable for various transformer architectures.


key: chroma:document
string_value:
CODE:

class FlaxPerformerOutput(nn.Module):
    @nn.compact
    def __call__(self, intermediate_output, attention_output):
        hidden_state = nn.Dense(attention_output.shape[-1], name="dense")(intermediate_output)
        hidden_state = FlaxPerformerLayerNorm(name="layer_norm")(hidden_state + attention_output)


SUMMARY:
**Summary:**

The `FlaxPerformerOutput` class defines a neural network module that applies a dense layer followed by layer normalization to the input `intermediate_output` and `attention_output`. The primary purpose is to produce an updated hidden state.

**Key Details:**

* Applies a dense layer (`nn.Dense`) with name "dense" to `intermediate_output`.
* Adds `attention_output` and normalizes using `FlaxPerformerLayerNorm`, both named "layer_norm".
* Returns the updated `hidden_state`.

**Notable Dependencies:**

* `nn` module (Neural Networks)
* `nn.compact` decorator ( compact neural network module definition )
* `FlaxPerformerLayerNorm` class ( custom layer normalization implementation )

Note: No external dependencies are listed in the provided code snippet. If any external libraries or dependencies were required, they would need to be included in this summary.


key: chroma:document
string_value:
CODE:

class FlaxPerformerLayer(nn.Module):
    num_heads: int
    head_size: int
    intermediate_size: int
    hidden_act: str = "gelu"

    @nn.compact
    def __call__(self, hidden_state, attention_mask):
        attention = FlaxPerformerAttention(self.num_heads, self.head_size, name="attention")(
            hidden_state, attention_mask
        )
        intermediate = FlaxPerformerIntermediate(
            self.intermediate_size, name="intermediate", hidden_act=self.hidden_act
        )(attention)
        output = FlaxPerformerOutput(name="output")(intermediate, attention)



SUMMARY:
**Summary:**

This `FlaxPerformerLayer` class implements a Self-Attention Mechanism layer from the Transformer architecture using the FLAX library. It takes in a hidden state and an attention mask as inputs.

* **Main Purpose:** Apply self-attention mechanism to transformed input sequence.
* **Key Implementation Details:**
	+ Uses `FlaxPerformerAttention` to compute self-attention weights.
	+ Applies `FlaxPerformerIntermediate` with a specified intermediate size and activation function.
	+ Combines the output of the intermediate layer with the attention scores using `FlaxPerformerOutput`.
* **Notable Dependencies/Library Functions:**
	+ Imported from `flaxnn` (FLAX library) modules:
		- `nn.Module`
		- `nn.compact`
		- `FlaxPerformerAttention`, `FlaxPerformerIntermediate`, and `FlaxPerformerOutput` classes
* **Implementation Interfaces:** conforms to PyTorch's Module API, requires attention mask input.


key: chroma:document
string_value:
CODE:

class FlaxPerformerLayerCollection(nn.Module):
    """
    Stores N BertLayer(s)
    """

    num_layers: int
    num_heads: int
    head_size: int
    intermediate_size: int
    hidden_act: str = "gelu"

    @nn.compact
    def __call__(self, inputs, attention_mask):
        assert self.num_layers > 0, f"num_layers should be >= 1, got ({self.num_layers})"

        # Initialize input / output
        input_i = inputs

        # Forward over all encoders
        for i in range(self.num_layers):
            layer = FlaxPerformerLayer(
                self.num_heads, self.head_size, self.intermediate_size, hidden_act=self.hidden_act, name=f"{i}"
            )
            input_i = layer(input_i, attention_mask)


SUMMARY:
**Summary**

The `FlaxPerformerLayerCollection` class is a PyTorch neural network module for storing and concatenating multiple instances of the `FlaxPerformerLayer` model. It takes in an input sequence `inputs` and attention mask `attention_mask`, passes it through each layer, and returns the final output.

**Key Details**

* The number of layers (`num_layers`) is validated to be greater than 0.
* The model uses the `FlaxPerformerLayer` class for each individual layer.
* The `hidden_act` parameter determines the activation function used (default: "gelu").

**Dependencies and Requirements**

* PyTorch library (`nn.Module` and `nn.compact`)
* FLAX deep learning library (`FlaxPerformerLayer`)
* No external dependencies listed.


key: chroma:document
string_value:
CODE:

class FlaxPerformerEncoder(nn.Module):
    num_layers: int
    num_heads: int
    head_size: int
    intermediate_size: int
    hidden_act: str = "gelu"

    @nn.compact
    def __call__(self, hidden_state, attention_mask):
        layer = FlaxPerformerLayerCollection(
            self.num_layers,
            self.num_heads,
            self.head_size,
            self.intermediate_size,
            name="layer",
            hidden_act=self.hidden_act,
        )(hidden_state, attention_mask)


SUMMARY:
**Summary**

The `FlaxPerformerEncoder` class is a PyTorch neural network module that implements the Transformer encoder architecture. It consists of a single layer using the `FlaxPerformerLayerCollection` mechanism.

**Key Implementation Details:**

* The main task is to encode input data (`hidden_state`) into an output sequence.
* A single `FlaxPerformerLayerCollection` instance is created with configurations:
	+ Number of layers (`num_layers`)
	+ Number of attention heads (`num_heads`)
	+ Head size (`head_size`)
	+ Intermediate size (`intermediate_size`)
	+ Activation function (`hidden_act`: `"gelu"`)

**Notable Methods and Interfaces:**

* `__call__`: The main method that takes input data (`hidden_state`) and attention mask (`attention_mask`) as arguments.
* `FlaxPerformerLayerCollection`: A custom layer class from the Flax library.

**External Dependencies:**

* PyTorch (`nn.Module`, `nn.compact`)
* Flax library (`FlaxPerformerLayerCollection`)

Note: The code snippets provided do not demonstrate the full implementation of `FlaxPerformerEncoder`. Additional code snippet or context is necessary to understand the complete functionality.


key: chroma:document
string_value:
CODE:

class FlaxPerformerPooler(nn.Module):
    @nn.compact
    def __call__(self, hidden_state):
        cls_token = hidden_state[:, 0]
        out = nn.Dense(hidden_state.shape[-1], name="dense")(cls_token)


SUMMARY:
**FlaxPerformerPooler Class Summary**

This class defines a neural network (`nn.Module`) called `FlaxPerformerPooler`, which inherits from PyTorch's `nn.Module`. The primary purpose of this class is to compute the output of a performer pool based on the input **hidden state**.

**Key Implementation Details:**

* The class applies a dense layer (with name "dense") to the first column (`cls_token`) of the hidden state, resulting in an output shape of `(batch_size, embedding_dim)`.
* The output is then passed through a hyperbolic tangent activation function (`tanh`).

**Notable Dependencies and Requirements:**

* `jax`: imported for use of JAX's `lax.tanh` activation function.
* `nn`: PyTorch's neural network module, used throughout the class.


key: chroma:document
string_value:
CODE:

class FlaxPerformerModule(nn.Module):
    vocab_size: int
    hidden_size: int
    type_vocab_size: int
    max_length: int
    num_encoder_layers: int
    num_heads: int
    head_size: int
    intermediate_size: int
    hidden_act: str = "gelu"
    add_pooling_layer: bool = True

    @nn.compact
    def __call__(self, input_ids, token_type_ids, position_ids, attention_mask):
        # Embedding
        embeddings = FlaxPerformerEmbeddings(
            self.vocab_size, self.hidden_size, self.type_vocab_size, self.max_length, name="embeddings"
        )(input_ids, token_type_ids, position_ids, attention_mask)

        # N stacked encoding layers
        encoder = FlaxPerformerEncoder(
            self.num_encoder_layers,
            self.num_heads,
            self.head_size,
            self.intermediate_size,
            hidden_act=self.hidden_act,
            name="encoder",
        )(embeddings, attention_mask)

        if not self.add_pooling_layer:
            return encoder

        pooled = FlaxPerformerPooler(name="pooler")(encoder)


SUMMARY:
**Summary:**

The `FlaxPerformerModule` is a neural network module designed for transformer-based tasks, specifically for Natural Language Processing (NLP) tasks. It consists of two key components:

1.  **Embedding Layer:** The input IDs, token type IDs, and position IDs are embedded into a high-dimensional vector space using the `FlaxPerformerEmbeddings` layer.
2.  **Encoder and Pooler Layers:** The embedded output is passed through a stack of multi-layer attention encoder layers (`FlaxPerformerEncoder`) to process the sequence-level dependencies. The final output of the encoder is transformed by a pooling layer (`FlaxPerformerPooler`) when `add_pooling_layer` is enabled.

**Implementation Details:**

*   Uses the Flax library for efficient neural network computation.
*   Employs multi-head attention and gelu activation functions.
*   Optionally includes an additional pooling layer to combine the encoder and pooled outputs.

**Notable Dependencies:**

*   `FlaxPerformerEmbeddings` and `FlaxPerformerEncoder`: custom embedding and encoding layers designed for transformer architectures.
*   `FlaxPerformerPooler`: a pooling layer class providing aggregation and normalization functions.

**Interfaces/Methods:**

*   `__call__`: The main entry point of the module, taking input IDs, token type IDs, position IDs, and attention masks as arguments.


key: chroma:document
string_value:
CODE:

class FlaxPerformerForMaskedLM(FlaxBertPreTrainedModel):
    def __init__(
        self, config: BertConfig, input_shape: Tuple = (1, 1), seed: int = 0, dtype: jnp.dtype = jnp.float32, **kwargs
    ):
        module = FlaxPerformerForMaskedLMModule(
            vocab_size=config.vocab_size,
            type_vocab_size=config.type_vocab_size,
            hidden_size=config.hidden_size,
            intermediate_size=config.intermediate_size,
            head_size=config.hidden_size,
            num_heads=config.num_attention_heads,
            num_encoder_layers=config.num_hidden_layers,
            max_length=config.max_position_embeddings,
            hidden_act=config.hidden_act,
            **kwargs,
        )

        super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype)

    def __call__(
        self,
        input_ids,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        params: dict = None,
        train: bool = False,
        dropout_rng: PRNGKey = None,
    ):
        input_ids, attention_mask, token_type_ids, position_ids = self._check_inputs(
            input_ids, attention_mask, token_type_ids, position_ids
        )

        # Handle any PRNG if needed
        rngs = {}
        if dropout_rng is not None:
            rngs["dropout"] = dropout_rng

        return self.module.apply(
            {"params": params or self.params},
            jnp.array(input_ids, dtype="i4"),
            jnp.array(attention_mask, dtype="i4"),
            jnp.array(token_type_ids, dtype="i4"),
            jnp.array(position_ids, dtype="i4"),
            not train,
            rngs=rngs,


SUMMARY:
**FlaxPerformerForMaskedLM Class Summary**

### Main Purpose and Functionality

`FlaxPerformerForMaskedLM` is a trainable BERT model variant that leverages the Flax library for its performance. It extends `FlaxBertPreTrainedModel` and inherits functionality from the underlying `FlaxPerformerForMaskedLMModule`.

**Primary Purpose:**

Implement masked language modeling (MLM) using a performer architecture.

### Key Implementation Details:

* Utilizes `FlaxPerformerForMaskedLMModule` for MLM processing.
* Applies `BertConfig`-defined parameters to the model.
* Performs input validation and data type conversions using `_check_inputs`.
* Uses a custom `apply` method to execute the model on given inputs.

### Notable Dependencies and Requirements:

* `FlaxPerformerForMaskedLMModule`
* `BertConfig`
* `jnp.dtype`
* `jnp.array`

Note: The provided class summary is concise and focused, highlighting the main purpose and key implementation details of the `FlaxPerformerForMaskedLM` class.


key: chroma:document
string_value:
CODE:

class FlaxPerformerForMaskedLMModule(nn.Module):
    vocab_size: int
    hidden_size: int
    intermediate_size: int
    head_size: int
    num_heads: int
    num_encoder_layers: int
    type_vocab_size: int
    max_length: int
    hidden_act: str
    dropout_rate: float = 0.0
    dtype: jnp.dtype = jnp.float32

    @nn.compact
    def __call__(
        self, input_ids, attention_mask=None, token_type_ids=None, position_ids=None, deterministic: bool = True
    ):
        # Model
        encoder = FlaxPerformerModule(
            vocab_size=self.vocab_size,
            hidden_size=self.hidden_size,
            type_vocab_size=self.type_vocab_size,
            max_length=self.max_length,
            num_encoder_layers=self.num_encoder_layers,
            num_heads=self.num_heads,
            head_size=self.hidden_size,
            intermediate_size=self.intermediate_size,
            hidden_act=self.hidden_act,
            add_pooling_layer=False,
            name="bert",
        )(input_ids, attention_mask, token_type_ids, position_ids)

        # Compute the prediction scores
        encoder = nn.Dropout(rate=self.dropout_rate)(encoder, deterministic=deterministic)
        logits = FlaxBertOnlyMLMHead(
            vocab_size=self.vocab_size, hidden_act=self.hidden_act, name="cls", dtype=self.dtype
        )(encoder)



SUMMARY:
Here is a concise summary of the provided class code chunk:

**Summary:**

`FlaxPerformerForMaskedLMModule` is an implementation of a masked language modeling transformer model using the Flax framework.

**Main Purpose and Functionality:**
This class performs the forward pass of the mask-learnable transformer model, taking in input IDs, attention masks, token type IDs, and position IDs as inputs.

**Key Implementation Details:**

* The model uses `FlaxPerformerModule` for encoding the input sequence.
* A dropout layer is applied to the encoder output to introduce randomness.
* The predictions are computed using a masked language modeling head implemented by `FlaxBertOnlyMLMHead`.

**Notable Dependencies and Requirements:**

* `nn.Module`
* `jnp.float32`
* `FlaxPerformerModule` (used for encoding)
* `FlaxBertOnlyMLMHead` (specific to masked language modeling)

This summary captures the main purpose, key implementation details, and notable dependencies of the provided code chunk.


key: chroma:document
string_value:
CODE:
###############
def make_es_index_snippets(es_client, passages_dset, index_name="english_wiki_kilt_snippets_100w"):
    index_config = {
        "settings": {
            "number_of_shards": 1,
            "analysis": {"analyzer": {"stop_standard": {"type": "standard", " stopwords": "_english_"}}},
        },
        "mappings": {
            "properties": {
                "article_title": {"type": "text", "analyzer": "standard", "similarity": "BM25"},
                "section_title": {"type": "text", "analyzer": "standard", "similarity": "BM25"},
                "passage_text": {"type": "text", "analyzer": "standard", "similarity": "BM25"},
            }
        },
    }
    es_client.indices.create(index=index_name, body=index_config)
    number_of_docs = passages_dset.num_rows
    progress = tqdm(unit="docs", total=number_of_docs)
    successes = 0

    def passage_generator():
        for passage in passages_dset:
            yield passage

    # create the ES index
    for ok, action in streaming_bulk(
        client=es_client,
        index=index_name,
        actions=passage_generator(),
    ):
        progress.update(1)
        successes += ok


SUMMARY:
**Function Summary**

```python
make_es_index_snippets es function creates an Elasticsearch (ES) index for building snippets.

**Primary Purpose**
The primary purpose of the `make_es_index_snippets` function is to create an ES index with specific settings and mapping configurations, suitable for harvesting wiki articles and generating snippets based on the text analysis.

**Key Implementation Details**

* Creates an ES index with `number_of_shards=1`, using the `standard` stop words and `BM25` similarity analyzer.
* Defines three properties: `article_title`, `section_title`, and `passage_text`, each typed as `text` and analyzed using the `standard` analyzer.

**Notable Dependencies or Requirements**

* References an ES client instance (`es_client`) to interact with the index.
* Uses the `streaming_bulk` function from Elasticsearch-Python library.
* Does not depend on any external libraries beyond the basic ES-Python library.
```

Note that I excluded unnecessary details, such as comments and code formatting variations, focusing on providing a concise yet accurate summary of the function's main purpose, implementation details, and dependencies or requirements.


key: chroma:document
string_value:
CODE:

def query_es_index(question, es_client, index_name="english_wiki_kilt_snippets_100w", n_results=10, min_length=20):
    q = question.lower()
    banned = ["how", "why", "what", "where", "which", "do", "does", "is", "?", "eli5", "eli5:"]
    q = " ".join([w for w in q.split() if w not in banned])
    response = es_client.search(
        index=index_name,
        body={
            "query": {
                "multi_match": {
                    "query": q,
                    "fields": ["article_title", "section_title", "passage_text^2"],
                    "type": "cross_fields",
                }
            },
            "size": 2 * n_results,
        },
    )
    hits = response["hits"]["hits"]
    support_doc = "<P> " + " <P> ".join([hit["_source"]["passage_text"] for hit in hits])
    res_list = [{k: hit["_source"][k] for k in hit["_source"] if k != "passage_text"} for hit in hits]
    for r, hit in zip(res_list, hits):
        r["passage_id"] = hit["_id"]
        r["score"] = hit["_score"]
        r["passage_text"] = hit["_source"]["passage_text"]
    res_list = [res for res in res_list if len(res["passage_text"].split()) > min_length][:n_results]


SUMMARY:
**Summary: Search English Wiki Index**

This function performs a cross-field search on the `english_wiki_kilt_snippets_100w` ES index. It takes a question and an ES client as input, and returns:

1. **Support document**: A concatenated passage text from the top 2x `n_results` hits.
2. **Search results list**: An ordered list of up to `n_results` documents with their corresponding `passage_text`, `score`, and `id`.

**Key Implementation Details:**

* Uses `elastic_search_client.search()` with a multi-match query on `article_title`, `section_title`, and `passage_text^2`.
* Applies filtering on banned words (`how`, `why`, etc.) to the question input.
* Uses cross-field search with type "cross_fields" to match across multiple fields.

**Notable Dependencies/Requirements:**

* `es_client` instance ( Elasticsearch client)
* `english_wiki_kilt_snippets_100w` index name
* External dependency on Elasticsearch (not handled in this code snippet)


key: chroma:document
string_value:
CODE:
###############
class ELI5DatasetQARetriver(Dataset):
    def __init__(self, examples_array, extra_answer_threshold=3, min_answer_length=64, training=True, n_samples=None):
        self.data = examples_array
        self.answer_thres = extra_answer_threshold
        self.min_length = min_answer_length
        self.training = training
        self.n_samples = self.data.num_rows if n_samples is None else n_samples

    def __len__(self):
        return self.n_samples

    def make_example(self, idx):
        example = self.data[idx]
        question = example["title"]
        if self.training:
            answers = [a for i, (a, sc) in enumerate(zip(example["answers"]["text"], example["answers"]["score"]))]
            answer_tab = choice(answers).split(" ")
            start_idx = randint(0, max(0, len(answer_tab) - self.min_length))
            answer_span = " ".join(answer_tab[start_idx:])
        else:
            answer_span = example["answers"]["text"][0]
        return (question, answer_span)

    def __getitem__(self, idx):


SUMMARY:
**Summary**

The `ELI5DatasetQARetriver` class is a custom dataset class designed for NLP tasks, specifically for retrieving questions and answers in an Elementary-level Informational style. The primary purpose is to provide pre-processed examples of question-answer pairs for training and testing.

* **Main Purpose**: To create a dataset for QARetriver, which generates answers based on input questions.
* **Key Implementation**:
	+ `__init__`: Initializes the dataset with example arrays, answer threshold, minimum answer length, training status, and sample size.
	+ `make_example`: Returns a pre-processed question-answer pair, either from the training data or using a random sampled response.
	+ `__getitem__`: Maps indices to examples in the dataset with wrapping around full data.
* **Notable Dependencies**:
	+ `Dataset` class (inherited)
	+ `numpy` for numerical operations
	+ `random` for random number generation

This summary should provide a clear understanding of the main purpose, implementation details, and dependencies required to use this custom dataset class.


key: chroma:document
string_value:
CODE:

class RetrievalQAEmbedder(nn.Module):
    def __init__(self, sent_encoder, dim):
        super(RetrievalQAEmbedder, self).__init__()
        self.sent_encoder = sent_encoder
        self.output_dim = 128
        self.project_q = nn.Linear(dim, self.output_dim, bias=False)
        self.project_a = nn.Linear(dim, self.output_dim, bias=False)
        self.ce_loss = nn.CrossEntropyLoss(reduction="mean")

    def embed_sentences_checkpointed(self, input_ids, attention_mask, checkpoint_batch_size=-1):
        # reproduces BERT forward pass with checkpointing
        if checkpoint_batch_size < 0 or input_ids.shape[0] < checkpoint_batch_size:
            return self.sent_encoder(input_ids, attention_mask=attention_mask)[1]
        else:
            # prepare implicit variables
            device = input_ids.device
            input_shape = input_ids.size()
            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
            head_mask = [None] * self.sent_encoder.config.num_hidden_layers
            extended_attention_mask: torch.Tensor = self.sent_encoder.get_extended_attention_mask(
                attention_mask, input_shape
            )

            # define function for checkpointing
            def partial_encode(*inputs):
                encoder_outputs = self.sent_encoder.encoder(
                    inputs[0],
                    attention_mask=inputs[1],
                    head_mask=head_mask,
                )
                sequence_output = encoder_outputs[0]
                pooled_output = self.sent_encoder.pooler(sequence_output)
                return pooled_output

            # run embedding layer on everything at once
            embedding_output = self.sent_encoder.embeddings(
                input_ids=input_ids, position_ids=None, token_type_ids=token_type_ids, inputs_embeds=None
            )
            # run encoding and pooling on one mini-batch at a time
            pooled_output_list = []
            for b in range(math.ceil(input_ids.shape[0] / checkpoint_batch_size)):
                b_embedding_output = embedding_output[b * checkpoint_batch_size : (b + 1) * checkpoint_batch_size]
                b_attention_mask = extended_attention_mask[b * checkpoint_batch_size : (b + 1) * checkpoint_batch_size]
                pooled_output = checkpoint.checkpoint(partial_encode, b_embedding_output, b_attention_mask)
                pooled_output_list.append(pooled_output)
            return torch.cat(pooled_output_list, dim=0)

    def embed_questions(self, q_ids, q_mask, checkpoint_batch_size=-1):
        q_reps = self.embed_sentences_checkpointed(q_ids, q_mask, checkpoint_batch_size)
        return self.project_q(q_reps)

    def embed_answers(self, a_ids, a_mask, checkpoint_batch_size=-1):
        a_reps = self.embed_sentences_checkpointed(a_ids, a_mask, checkpoint_batch_size)
        return self.project_a(a_reps)

    def forward(self, q_ids, q_mask, a_ids, a_mask, checkpoint_batch_size=-1):
        device = q_ids.device
        q_reps = self.embed_questions(q_ids, q_mask, checkpoint_batch_size)
        a_reps = self.embed_answers(a_ids, a_mask, checkpoint_batch_size)
        compare_scores = torch.mm(q_reps, a_reps.t())
        loss_qa = self.ce_loss(compare_scores, torch.arange(compare_scores.shape[1]).to(device))
        loss_aq = self.ce_loss(compare_scores.t(), torch.arange(compare_scores.shape[0]).to(device))
        loss = (loss_qa + loss_aq) / 2


SUMMARY:
**Summary: RetrievalQAEmbedder Class**

The `RetrievalQAEmbedder` class, derived from PyTorch's `nn.Module`, is designed for retrieving questions based on answers in a QA (Question Answer) pair. This model utilizes knowledge distillation to fine-tune a smaller sentinel model for each of the input questions.

**Primary Purpose:** Embedding and comparing question-answer pairs.

*   **Main Methods:**
    *   `embed_sentences_checkpointed`: Reproduces BERT's forward pass with checkpointing on input sentences.
    *   `embed_questions` and `embed_answers`: Project embedded sentence and answer representations to a shared space.
    *   `forward`: Compares projected question-answer embeddings using cross-entropy loss, calculating both Q-A and A-Q similarities.

**Key Implementation Details:**

*   **Sentinel Model**: Splits input sentences into mini-batches for sequential computation with checkpointing.
*   **Linear Projects**: Applies linear projections (`project_q` and `project_a`) to embedding outputs from the sentinel model.
*   **Cross-Entropy Loss**: Computes loss functions (`loss_qa` and `loss_aq`) comparing projected question-answer embeddings.

**Notable Dependencies/Requirements:**

*   `BERT-sentinel` module (imported via `self.sent_encoder`)
*   `torch`

Overall, this class integrates knowledge distillation from smaller sentiment models with cross-entropy loss for comparing QA pairs.


key: chroma:document
string_value:
CODE:

def make_qa_retriever_model(model_name="google/bert_uncased_L-8_H-512_A-8", from_file=None, device="cuda:0"):
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    bert_model = AutoModel.from_pretrained(model_name).to(device)
    # run bert_model on a dummy batch to get output dimension
    d_ids = torch.LongTensor(
        [[bert_model.config.bos_token_id if bert_model.config.bos_token_id is not None else 1]]
    ).to(device)
    d_mask = torch.LongTensor([[1]]).to(device)
    sent_dim = bert_model(d_ids, attention_mask=d_mask)[1].shape[-1]
    qa_embedder = RetrievalQAEmbedder(bert_model, sent_dim).to(device)
    if from_file is not None:
        param_dict = torch.load(from_file)  # has model weights, optimizer, and scheduler states
        qa_embedder.load_state_dict(param_dict["model"])


SUMMARY:
**Function Summary: `make_qa_retriever_model`**

 Creates a QA Retrieval Model using BERT.

### Purpose and Functionality

* Loads pre-trained BERT model (`AutoTokenizer`, `AutoModel`) and moves it to specified device (`cuda:0`).
* Defines a dummy batch to retrieve output dimension.
* Instantiates `RetrievalQAEmbedder` with prepared model, dimension, and device.
* If a saved model file is provided (`from_file`), loads pre-trained weights from it.

### Key Implementation Details

* Model loaded using `AutoTokenizer.from_pretrained` and `AutoModel.from_pretrained`.
* Output dimension retrieved by passing dummy batch to BERT model outputs.
* `RetrievalQAEmbedder` instantiated with prepared model, output dimension, and device (`device`) parameter.
* Saved model file loading via `torch.load`.

### Notable Dependencies or Requirements

* PyTorch library.
* Transformers library for AutoTokenizer and AutoModel usage.


key: chroma:document
string_value:
CODE:

def make_qa_retriever_batch(qa_list, tokenizer, max_len=64, device="cuda:0"):
    q_ls = [q for q, a in qa_list]
    a_ls = [a for q, a in qa_list]
    q_toks = tokenizer(q_ls, max_length=max_len, padding="max_length", truncation=True)
    q_ids, q_mask = (
        torch.LongTensor(q_toks["input_ids"]).to(device),
        torch.LongTensor(q_toks["attention_mask"]).to(device),
    )
    a_toks = tokenizer(a_ls, max_length=max_len, padding="max_length", truncation=True)
    a_ids, a_mask = (
        torch.LongTensor(a_toks["input_ids"]).to(device),
        torch.LongTensor(a_toks["attention_mask"]).to(device),
    )


SUMMARY:
**QA Retriever Batch Preparation Function**

Prepares batches of queries and answers for a question answering model.

### Main Purpose and Functionality

This function takes in a list of QA pairs (`qa_list`), a tokenizer (`tokenizer`), and optional parameters (`max_len`, `device`). It processes the queries and answers, tokenizing them, and returns tensors representing the input IDs and attention masks for both queries and answers on the specified device.

### Key Implementation Details

* Tokenizes queries and answers using the provided tokenizer.
* Converts tokenized inputs to LongTensor form for efficient computation.
* Returns four tensor pairs: `(q_ids, q_mask)`, `(a_ids, a_mask)`.

### Notable Dependencies or Requirements

* `tokenizer`: Hugging Face tokenizer instance
* External dependency: `torch` library for tensor computations.

### Method Signatures and Interfaces

* `make_qa_retriever_batch(qa_list, tokenizer, max_len=64, device="cuda:0")`
	+ Takes in `qa_list`, `tokenizer`, `max_len`, and `device` as inputs.
	+ Returns a tuple of four tensors: `(q_ids, q_mask, a_ids, a_mask)`.


key: chroma:document
string_value:
CODE:

def train_qa_retriever_epoch(model, dataset, tokenizer, optimizer, scheduler, args, e=0):
    model.train()
    # make iterator
    train_sampler = RandomSampler(dataset)
    model_collate_fn = functools.partial(
        make_qa_retriever_batch, tokenizer=tokenizer, max_len=args.max_length, device="cuda:0"
    )
    data_loader = DataLoader(dataset, batch_size=args.batch_size, sampler=train_sampler, collate_fn=model_collate_fn)
    epoch_iterator = tqdm(data_loader, desc="Iteration", disable=True)
    # accumulate loss since last print
    loc_steps = 0
    loc_loss = 0.0
    st_time = time()
    for step, batch in enumerate(epoch_iterator):
        q_ids, q_mask, a_ids, a_mask = batch
        pre_loss = model(q_ids, q_mask, a_ids, a_mask, checkpoint_batch_size=args.checkpoint_batch_size)
        loss = pre_loss.sum()
        # optimizer
        loss.backward()
        optimizer.step()
        scheduler.step()
        model.zero_grad()
        # some printing within the epoch
        loc_loss += loss.item()
        loc_steps += 1
        if step % args.print_freq == 0 or step == 1:
            print(
                "{:2d} {:5d} of {:5d} \t L: {:.3f} \t -- {:.3f}".format(
                    e,
                    step,
                    len(dataset) // args.batch_size,
                    loc_loss / loc_steps,
                    time() - st_time,
                )
            )
            loc_loss = 0


SUMMARY:
**Summary:**

`train_qa_retriever_epoch` function trains a QA retriever model using the given dataset, tokenizer, optimizer, scheduler, and hyperparameters.

**Key Implementation Details:**

*   The training process accumulates losses since last print (every `args.print_freq` epochs or after first epoch).
*   Batch processing and data collation are handled by `DataLoader` with customized collate function using `functools.partial`.
*   Loss calculation involves computing loss from pre-loss output, which is then passed to the optimizer for update.

**Notable Dependencies and Requirements:**

*   TensorFlow or PyTorch library
*   `Dataset` class containing sample data (not shown)
*   `RandomSampler`
*   `DataLoader` with batch size `args.batch_size`
*   `Model` class with `train` method and necessary modules for computation
*   `Optimizer` instance for gradient descent optimization
*   `Scheduler` instance for learning rate adjustment

**Method Signatures Highlighted:**

*   `model(q_ids, q_mask, a_ids, a_mask, checkpoint_batch_size)`
*   `optimizer.step()`
*   `scheduler.step()`

Note that some important implementation details like the actual computation within the `model` class (which is not shown in this snippet) and other specific hyperparameters or data handling logic are not included here for concise summary purposes.


key: chroma:document
string_value:
CODE:

def train_qa_retriever_joint_epoch(model, dataset_list, tokenizer, optimizer, scheduler, args, e=0):
    model.train()
    model_collate_fn = functools.partial(
        make_qa_retriever_batch, tokenizer=tokenizer, max_len=args.max_length, device="cuda:0"
    )
    # make iterator
    train_samplers = [RandomSampler(dataset) for dataset in dataset_list]
    data_loaders = [
        DataLoader(dataset, batch_size=args.batch_size, sampler=train_sampler, collate_fn=model_collate_fn)
        for dataset, train_sampler in zip(dataset_list, train_samplers)
    ]
    iterators = [iter(dloader) for dloader in data_loaders]
    joint_iter = zip(*iterators)
    # accumulate loss since last print
    loc_steps = 0
    loc_loss = 0.0
    st_time = time()
    for step, (batches,) in enumerate(zip(joint_iter)):
        for batch in batches:
            q_ids, q_mask, a_ids, a_mask = batch
            loss = model(q_ids, q_mask, a_ids, a_mask, checkpoint_batch_size=args.checkpoint_batch_size)
            # optimizer
            loss.backward()
            optimizer.step()
            scheduler.step()
            model.zero_grad()
            # some printing within the epoch
            loc_loss += loss.item()
            loc_steps += 1
        if step % args.print_freq == 0:
            print(
                "{:2d} {:5d} of {:5d} \t L: {:.3f} \t -- {:.3f}".format(
                    e,
                    step,
                    len(dataset_list[0]) // args.batch_size,
                    loc_loss / loc_steps,
                    time() - st_time,
                )
            )
            loc_loss = 0


SUMMARY:
**Function Summary: `train_qa_retriever_joint_epoch`**

**Main Purpose and Functionality**
This function trains the QA retriever joint model on a list of datasets. It utilizes multiple data loaders, optimizers, and schedulers to iterate over each batch of data, performing forward passes, backward passes, and parameter updates.

**Key Implementation Details**
- The model is trained using `RandomSampler` across multiple data loaders with a custom collate function (`make_qa_retriever_batch`) and a batch size determined by `args.batch_size`.
- Losses are accumulated since the last print cycle (`loc_loss += loss.item()`).
- The optimization process involves backward pass (`loss.backward()`) followed by step update (`optimizer.step()`) with scheduler updates (`scheduler.step()`).
- Printing progress every `args.print_freq` iterations, displaying training loss and time elapsed since epoch start.

**Notable Dependencies or Requirements**
- `RandomSampler` from PyTorch
- `DataLoader` from PyTorch
- Custom `make_qa_retriever_batch` function with `tokenizer`, `max_len`, and `device` arguments.
- `optimizer` and `scheduler` object from `args`.


key: chroma:document
string_value:
CODE:

def evaluate_qa_retriever(model, dataset, tokenizer, args):
    model.eval()
    # make iterator
    eval_sampler = SequentialSampler(dataset)
    model_collate_fn = functools.partial(
        make_qa_retriever_batch, tokenizer=tokenizer, max_len=args.max_length, device="cuda:0"
    )
    data_loader = DataLoader(dataset, batch_size=args.batch_size, sampler=eval_sampler, collate_fn=model_collate_fn)
    epoch_iterator = tqdm(data_loader, desc="Iteration", disable=True)
    tot_loss = 0.0
    with torch.no_grad():
        for step, batch in enumerate(epoch_iterator):
            q_ids, q_mask, a_ids, a_mask = batch
            loss = model(q_ids, q_mask, a_ids, a_mask)
            tot_loss += loss.item()


SUMMARY:
**Function Summary: `evaluate_qa_retriever`**

**(Main Purpose and Functionality)**

This function evaluates the performance of a QA retriever model on a given dataset using the Hugging Face Transformers library. It takes in a pre-trained model, dataset, tokenizer, and hyperparameters, and returns the average validation loss.

**Key Implementation Details**

* Employs a `SequentialSampler` to iterate through the dataset in a batch-wise manner.
* Utilizes a custom collate function (`make_qa_retriever_batch`) to process batches with the specified `max_len` and device.
* Uses PyTorch's DataLoader to load the data in a memory-efficient manner.

**Notable Dependencies and Requirements**

* Hugging Face Transformers library
	+ `transforms.Pipeline`: implied via `transformers`
	+ custom `make_qa_retriever_batch` function (not shown)
	+ `torchio` and `pytorch` libraries for data processing

**(Primary Purpose:)** Evaluates the performance of a QA retriever model on a specified dataset.


key: chroma:document
string_value:
CODE:

def train_qa_retriever(qar_model, qar_tokenizer, qar_train_dset, qar_valid_dset, qar_args):
    qar_optimizer = AdamW(qar_model.parameters(), lr=qar_args.learning_rate, eps=1e-8)
    qar_scheduler = get_linear_schedule_with_warmup(
        qar_optimizer,
        num_warmup_steps=100,
        num_training_steps=(qar_args.num_epochs + 1) * math.ceil(len(qar_train_dset) / qar_args.batch_size),
    )
    for e in range(qar_args.num_epochs):
        train_qa_retriever_epoch(qar_model, qar_train_dset, qar_tokenizer, qar_optimizer, qar_scheduler, qar_args, e)
        m_save_dict = {
            "model": qar_model.state_dict(),
            "optimizer": qar_optimizer.state_dict(),
            "scheduler": qar_scheduler.state_dict(),
        }
        print("Saving model {}".format(qar_args.model_save_name))
        torch.save(m_save_dict, "{}_{}.pth".format(qar_args.model_save_name, e))
        eval_loss = evaluate_qa_retriever(qar_model, qar_valid_dset, qar_tokenizer, qar_args)


SUMMARY:
**Concise Summary:**

Function:

* `train_qa_retriever`: Trains QA Retriever model using AdamW optimizer and linear scheduler.

Implementation Details:

* Uses `AdamW` optimizer with learning rate `qar_args.learning_rate`, epsilon value 1e-8.
* Utilizes linear schedule with warm-up period for epoch scaling.

Notable Operations:

* Iterates over `qar_args.num_epochs` epochs, applying training and saving process.
* Saves model state, optimizer, and scheduler at each epoch using PyTorch's `torch.save`.
* Evaluates QA Retriever model on valid dataset (`qar_valid_dset`) and prints evaluation loss.

Key Dependencies:

* `qar_model`, `qar_tokenizer`, `qar_train_dset`, `qar_valid_dset`, and `qar_args` are required input parameters.
* External dependencies include:
	+ `AdamW`
	+ `get_linear_schedule_with_warmup`
	+ `torch.save`
	+ `math.ceil`
	+ `evaluate_qa_retriever`

Primary Purpose:

The primary purpose of this function is to train the QA Retriever model on a dataset using AdamW optimizer and linear scheduler, with evaluation on a valid set at each epoch.


key: chroma:document
string_value:
CODE:
###############
class ELI5DatasetS2S(Dataset):
    def __init__(
        self, examples_array, make_doc_fun=None, extra_answer_threshold=3, document_cache=None, training=True
    ):
        self.training = training
        self.data = examples_array
        self.make_doc_function = make_doc_fun
        self.document_cache = {} if document_cache is None else document_cache
        assert not (make_doc_fun is None and document_cache is None)
        # make index of specific question-answer pairs from multi-answers
        if self.training:
            self.qa_id_list = [
                (i, j)
                for i, qa in enumerate(self.data)
                for j, (a, sc) in enumerate(zip(qa["answers"]["text"], qa["answers"]["score"]))
                if j == 0 or sc >= extra_answer_threshold
            ]
        else:
            self.qa_id_list = [(i, 0) for i in range(self.data.num_rows)]

    def __len__(self):
        return len(self.qa_id_list)

    def make_example(self, idx):
        i, j = self.qa_id_list[idx]
        example = self.data[i]
        question = example["title"] + " " + example["selftext"]
        answer = example["answers"]["text"][j]
        q_id = example["q_id"]
        if self.make_doc_function is not None:
            self.document_cache[q_id] = self.document_cache.get(q_id, self.make_doc_function(example["title"]))
        document = self.document_cache[q_id]
        in_st = "question: {} context: {}".format(
            question.lower().replace(" --t--", "").strip(),
            document.lower().strip(),
        )
        out_st = answer
        return (in_st, out_st)

    def __getitem__(self, idx):


SUMMARY:
**Class Summary**

`ELI5DatasetS2S`: A sequence-to-sequence dataset class for ELI5 (Explain Like I'm 5) models.

**Primary Purpose**
The primary purpose of this class is to create a dataset for training an ELI5 seq2seq model. It loads and processes question-answer pairs from an array, creating examples with input (`question`) and output (`answer`).

**Key Implementation Details**

* `_make_example(idx)`: Takes an index `idx` and returns the corresponding example with question and answer. If a document cache is provided, it retrieves a cached document for the question.
	+ Input: `Idx`
	+ Output: `Question-string`, `Answer-string`

* `_init(examples_array, ...)`:
	+ Loads examples array
	+ Creates `qa_id_list` based on training or testing mode (question-answer pairs with scores above threshold)
	+ Initializes document cache (if provided)

**Notable Dependencies/Requirements**

* `Dataset`: A parent class (base dataset implementation)
* `num_rows`: An attribute of the loaded data array, referring to the number of rows
* `qa_id_list`: A list of question-answer pairs (idx,i) where i is the index in the answer text (or 0 for testing)

**External Dependencies**

* None

This summary preserves key class and method names (`ELI5DatasetS2S`, `_init`, `_make_example`), focuses on primary purpose and implementation details, and highlights notable dependencies and requirements.


key: chroma:document
string_value:
CODE:

def make_qa_s2s_model(model_name="facebook/bart-large", from_file=None, device="cuda:0"):
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)
    if from_file is not None:
        param_dict = torch.load(from_file)  # has model weights, optimizer, and scheduler states
        model.load_state_dict(param_dict["model"])


SUMMARY:
**Concise Summary**

Creates a sequence-to-sequence (S2S) QA model using the Hugging Face Transformers library. The function:
*   Initializes the pre-trained `AutTokenizer` and S2S language model (`AutoModelForSeq2SeqLM`) for sequence generation.
*   Loads model weights from a specified file using PyTorch's `load_state_dict`.
*   Returns the tokenizer and loaded model on a specified device (default: CUDA 0).

**Key Implementation Details**

*   Uses pre-trained transformer models from Hugging Face via `AutoTokenizer` and `AutoModelForSeq2SeqLM`.
*   Loads model parameters using PyTorch's `load_state_dict`.

**Notable Dependencies or Requirements**

*   requires the Hugging Face Transformers library
*   uses PyTorch for loading model weights
*   assumes availability of CUDA device (if specified)


key: chroma:document
string_value:
CODE:

def make_qa_s2s_batch(qa_list, tokenizer, max_len=64, max_a_len=360, device="cuda:0"):
    q_ls = [q for q, a in qa_list]
    a_ls = [a for q, a in qa_list]
    q_toks = tokenizer(q_ls, max_length=max_len, padding="max_length", truncation=True)
    q_ids, q_mask = (
        torch.LongTensor(q_toks["input_ids"]).to(device),
        torch.LongTensor(q_toks["attention_mask"]).to(device),
    )
    a_toks = tokenizer(a_ls, max_length=min(max_len, max_a_len), padding="max_length", truncation=True)
    a_ids, a_mask = (
        torch.LongTensor(a_toks["input_ids"]).to(device),
        torch.LongTensor(a_toks["attention_mask"]).to(device),
    )
    lm_labels = a_ids[:, 1:].contiguous().clone()
    lm_labels[a_mask[:, 1:].contiguous() == 0] = -100
    model_inputs = {
        "input_ids": q_ids,
        "attention_mask": q_mask,
        "decoder_input_ids": a_ids[:, :-1].contiguous(),
        "lm_labels": lm_labels,
    }


SUMMARY:
**Function Summary: `make_qa_s2s_batch`**

Creates a batch of QA S2S model inputs.

**Purpose and Functionality**
The `make_qa_s2s_batch` function takes a list of QA pairs (`qa_list`) and tokenizer, along with optional parameters for maximum sequence lengths and device. It processes the input questions and answers, generates corresponding IDs and masks, and constructs a dictionary of model inputs for downstream usage.

**Key Implementation Details**

* Processes input questions and answers into tokenized sequences using the provided tokenizer.
* Generates `input_ids`, `attention_mask` tensors for the question and answer sequences.
* Creates `decoder_input_ids` by shifting the answer sequence to the left (except for the first token).
* Sets negative values in `lm_labels` to -100, indicating padding.
* Returns a dictionary of model inputs: `{"input_ids", "attention_mask", "decoder_input_ids", "lm_labels"}`.

**Notable Dependencies and Requirements**

* T tokenizer with `max_length`, `padding="max_length"` and `truncation=True`
* PyTorch tensor operations
* Device (device) passed to the function


key: chroma:document
string_value:
CODE:

def train_qa_s2s_epoch(model, dataset, tokenizer, optimizer, scheduler, args, e=0, curriculum=False):
    model.train()
    # make iterator
    if curriculum:
        train_sampler = SequentialSampler(dataset)
    else:
        train_sampler = RandomSampler(dataset)
    model_collate_fn = functools.partial(
        make_qa_s2s_batch, tokenizer=tokenizer, max_len=args.max_length, device="cuda:0"
    )
    data_loader = DataLoader(dataset, batch_size=args.batch_size, sampler=train_sampler, collate_fn=model_collate_fn)
    epoch_iterator = tqdm(data_loader, desc="Iteration", disable=True)
    # accumulate loss since last print
    loc_steps = 0
    loc_loss = 0.0
    st_time = time()
    for step, batch_inputs in enumerate(epoch_iterator):
        pre_loss = model(**batch_inputs)[0]
        loss = pre_loss.sum() / pre_loss.shape[0]
        loss.backward()
        # optimizer
        if step % args.backward_freq == 0:
            optimizer.step()
            scheduler.step()
            model.zero_grad()
        # some printing within the epoch
        loc_loss += loss.item()
        loc_steps += 1
        if step % args.print_freq == 0 or step == 1:
            print(
                "{:2d} {:5d} of {:5d} \t L: {:.3f} \t -- {:.3f}".format(
                    e,
                    step,
                    len(dataset) // args.batch_size,
                    loc_loss / loc_steps,
                    time() - st_time,
                )
            )
            loc_loss = 0


SUMMARY:
**Function Summary: Train QA S2S Epoch**

### Purpose and Functionality

Trains a QA S2S model on a given dataset using the provided optimizer, scheduler, and hyperparameters. The function iterates over batches of input data, computes the total loss, accumulates gradients, and updates the model parameters.

### Key Implementation Details

* Uses the `train` method on the model instance.
* Creates an iterator for training with either sequential or random sampling depending on the `curriculum` flag.
* Defines a custom collate function `model_collate_fn` to process input data.
* Accumulates loss and gradients using `loc_steps` and `loc_loss`.
* Updates model parameters at regular intervals (`backward_freq`) using `optimizer.step()`.

### Notable Dependencies and Requirements

* `torch` library for PyTorch functionality
* `data_loader` from `transformers` for processing large datasets.
* `functools.partial` to create a custom collate function.
* `DataLoader` class to manage data loading and batching.

### Interface Inheritance or Implemented

* Implements the `Dataloader` interface from the `transformers` library.

### External Dependencies

* `torch` library ( PyTorch functionality)
* `transformers` library (for large dataset processing)


key: chroma:document
string_value:
CODE:

def eval_qa_s2s_epoch(model, dataset, tokenizer, args):
    model.eval()
    # make iterator
    train_sampler = SequentialSampler(dataset)
    model_collate_fn = functools.partial(
        make_qa_s2s_batch, tokenizer=tokenizer, max_len=args.max_length, device="cuda:0"
    )
    data_loader = DataLoader(dataset, batch_size=args.batch_size, sampler=train_sampler, collate_fn=model_collate_fn)
    epoch_iterator = tqdm(data_loader, desc="Iteration", disable=True)
    # accumulate loss since last print
    loc_steps = 0
    loc_loss = 0.0
    st_time = time()
    with torch.no_grad():
        for step, batch_inputs in enumerate(epoch_iterator):
            pre_loss = model(**batch_inputs)[0]
            loss = pre_loss.sum() / pre_loss.shape[0]
            loc_loss += loss.item()
            loc_steps += 1
            if step % args.print_freq == 0:
                print(
                    "{:5d} of {:5d} \t L: {:.3f} \t -- {:.3f}".format(
                        step,
                        len(dataset) // args.batch_size,
                        loc_loss / loc_steps,
                        time() - st_time,
                    )
                )
    print(
        "Total \t L: {:.3f} \t -- {:.3f}".format(
            loc_loss / loc_steps,
            time() - st_time,
        )


SUMMARY:
Here is a concise and focused summary of the provided function code chunk:

**Function Purpose and Functionality**
The `eval_qa_s2s_epoch` function evaluates the performance of a QA S2S model on a given dataset. It processes the dataset using a data loader, calculates loss over each batch, and prints iteration progress.

**Key Implementation Details**

* The function uses asequential sampler to load the dataset in batches.
* A custom collate function (`make_qa_s2s_batch`) is applied to each batch, which involves tokenization and length limiting using the provided tokenizer.
* Loss calculation involves summing the loss tensor across all samples and dividing by the sample count.

**Notable Dependencies and Requirements**

* The `model` parameter expects a QA S2S model instance that supports automatic mixed precision training (amp).
* The `tokenizer` parameter requires a tokenization library to tokenize input data.
* External dependencies include:
	+ PyTorch for tensor operations.
	+ Torchlight/tqdm for iteration progress metering.

**Primary Methods and Variables**

* `model.eval()`: Evaluates the model under evaluation mode.
* `make_qa_s2s_batch`: A custom collate function to tokenize and length-limit input data during batch processing.
* `DataLoader` instance (`data_loader`): Loaders data into batches for processing.
* Loss scalar variable (`loc_loss`).


key: chroma:document
string_value:
CODE:

def train_qa_s2s(qa_s2s_model, qa_s2s_tokenizer, s2s_train_dset, s2s_valid_dset, s2s_args):
    s2s_optimizer = AdamW(qa_s2s_model.parameters(), lr=s2s_args.learning_rate, eps=1e-8)
    s2s_scheduler = get_linear_schedule_with_warmup(
        s2s_optimizer,
        num_warmup_steps=400,
        num_training_steps=(s2s_args.num_epochs + 1) * math.ceil(len(s2s_train_dset) / s2s_args.batch_size),
    )
    for e in range(s2s_args.num_epochs):
        train_qa_s2s_epoch(
            qa_s2s_model,
            s2s_train_dset,
            qa_s2s_tokenizer,
            s2s_optimizer,
            s2s_scheduler,
            s2s_args,
            e,
            curriculum=(e == 0),
        )
        m_save_dict = {
            "model": qa_s2s_model.state_dict(),
            "optimizer": s2s_optimizer.state_dict(),
            "scheduler": s2s_scheduler.state_dict(),
        }
        print("Saving model {}".format(s2s_args.model_save_name))
        eval_qa_s2s_epoch(qa_s2s_model, s2s_valid_dset, qa_s2s_tokenizer, s2s_args)


SUMMARY:
**Training S2S QA Model Summary**

**Main Purpose and Functionality**
The `train_qa_s2s` function trains a sequence-to-sequence (S2S) model for natural language question answering using AdamW optimizer and linear warm-up scheduler. It iterates through `num_epochs` epochs, training on the `s2s_train_dset`, evaluating on the `s2s_valid_dset`, and saving the model at each epoch.

**Key Implementation Details**

* Uses `AdamW` optimizer with learning rate `s2s_args.learning_rate`
* Linear warm-up scheduler with 400 warm-up steps
* Trains for `(num_epochs + 1) * math.ceil(len(s2s_train_dset) / batch_size)` training steps
* Saves model with state dictionaries of optimizer, scheduler, and the current model after each epoch

**Notable Dependencies or Requirements**

* `qa_s2s_model` (S2S model)
* `qa_s2s_tokenizer` (tokenizer for S2S model)
* `s2s_train_dset` and `s2s_valid_dset` (training and validation datasets)
* External dependency:
	+ `math.ceil(len(s2s_train_dset) / batch_size)` to calculate training steps (batch size not defined in snippet)

**Interfaces**

* Trains on: `s2s_train_dset`
* Evaluates on: `s2s_valid_dset`
* Saves model at each epoch with `m_save_dict` (contains optimizer, scheduler, and model state dictionaries)


key: chroma:document
string_value:
CODE:
# generate answer from input "question: ... context: <p> ..."
def qa_s2s_generate(
    question_doc,
    qa_s2s_model,
    qa_s2s_tokenizer,
    num_answers=1,
    num_beams=None,
    min_len=64,
    max_len=256,
    do_sample=False,
    temp=1.0,
    top_p=None,
    top_k=None,
    max_input_length=512,
    device="cuda:0",
):
    model_inputs = make_qa_s2s_batch(
        [(question_doc, "A")],
        qa_s2s_tokenizer,
        max_input_length,
        device=device,
    )
    n_beams = num_answers if num_beams is None else max(num_beams, num_answers)
    generated_ids = qa_s2s_model.generate(
        input_ids=model_inputs["input_ids"],
        attention_mask=model_inputs["attention_mask"],
        min_length=min_len,
        max_length=max_len,
        do_sample=do_sample,
        early_stopping=True,
        num_beams=1 if do_sample else n_beams,
        temperature=temp,
        top_k=top_k,
        top_p=top_p,
        eos_token_id=qa_s2s_tokenizer.eos_token_id,
        no_repeat_ngram_size=3,
        num_return_sequences=num_answers,
        decoder_start_token_id=qa_s2s_tokenizer.bos_token_id,
    )


SUMMARY:
**Summary**

This function generates answers to a given question using the T5-based qa-s2s model. It takes in a question document, model parameters, tokenizer, and various options for generating text.

**Main Purpose**

Generate multiple possible answers (up to `num_answers`) or samples from the input question document using the provided qa-s2s model.

**Implementation Details**

* Uses `make_qa_s2s_batch` to prepare the model inputs.
* Sets the number of beams and sampling parameters for generating text.
* Calls the `generate` method on the qa-s2s model with specified options, including minimum and maximum length constraints.
* Applies post-processing to decode and truncate input IDs using the provided tokenizer.

**Notable Dependencies**

* `qa_s2s_model`: The T5-based qa-s2s model implementation.
* `qa_s2s_tokenizer`: A tokenizer instance for both question documents and generated responses.
* `device="cuda:0"`: A device identifier (GPU) or "cpu" string.

**Notable Parameters**

* `num_answers`, `num_beams`, `max_len`, `temp` : Configuration options for answer generation.
* `do_sample`: Flag enabling sampling from the predicted output probability distribution.
* `eos_token_id` and `bos_token_id`: IDs of the start (eos) and end (bos) tokens in the input sequence.


key: chroma:document
string_value:
CODE:
###############
def embed_passages_for_retrieval(passages, tokenizer, qa_embedder, max_length=128, device="cuda:0"):
    a_toks = tokenizer(passages, max_length=max_length, padding="max_length", truncation=True)
    a_ids, a_mask = (
        torch.LongTensor(a_toks["input_ids"]).to(device),
        torch.LongTensor(a_toks["attention_mask"]).to(device),
    )
    with torch.no_grad():
        a_reps = qa_embedder.embed_answers(a_ids, a_mask).cpu().type(torch.float)


SUMMARY:
**Concise Summary**

`embed_passages_for_retrieval` Function:

* Embeds passages for retrieval using a trained QA model (`qa_embedder`) and tokenizes input passages (`passages`) with maximum length `max_length`.

**Main Purpose and Functionality:**
This function prepares passage inputs for retrieval queries by:
1. Tokenizing passages.
2. Converting tokens to IDs and attention masks.

**Key Implementation Details:**

* Use of `tokenizer` and `qa_embedder` objects.
* Conversion to CUDA device (`device="cuda:0"`).
* Utilization of `torch.no_grad()` mode for computation.

**Notable Dependencies and Requirements:**
1. `tokenizer`: passage tokenizer object.
2. `qa_embedder`: QA embedding model object.
3. `torch`: library for tensor operations.
4. External dependency on a trained QA model (implied, not explicitly listed).


key: chroma:document
string_value:
CODE:

def embed_questions_for_retrieval(q_ls, tokenizer, qa_embedder, device="cuda:0"):
    q_toks = tokenizer(q_ls, max_length=128, padding="max_length", truncation=True)
    q_ids, q_mask = (
        torch.LongTensor(q_toks["input_ids"]).to(device),
        torch.LongTensor(q_toks["attention_mask"]).to(device),
    )
    with torch.no_grad():
        q_reps = qa_embedder.embed_questions(q_ids, q_mask).cpu().type(torch.float)


SUMMARY:
**Function Summary: Embeddings for Question Retrieval**

This function `embed_questions_for_retrieval` generates embeddings for a list of questions (`q_ls`) using the provided question-answering embedder (`qa_embedder`). The resulting embeddings are computed on a CUDA device by default.

**Key Implementation Details:**

* Uses Hugging Face tokenizer to convert questions into tokenized inputs
* Embeds questions using `qa_embedder` with attention masks and batch dimension 1

**Interfaces/Optimizations:**

* Torch tensor operations for input, attention mask, and output processing
* Uses CUDA device (`device="cuda:0"`) by default for computation

**External Dependencies:**

* Hugging Face library (tokenizer)
* PyTorch library (torch)


key: chroma:document
string_value:
CODE:

def make_qa_dense_index(
    qa_embedder,
    tokenizer,
    passages_dset,
    batch_size=512,
    max_length=128,
    index_name="kilt_passages_reps.dat",
    dtype="float32",
    device="cuda:0",
):
    st_time = time()
    fp = np.memmap(index_name, dtype=dtype, mode="w+", shape=(passages_dset.num_rows, 128))
    n_batches = math.ceil(passages_dset.num_rows / batch_size)
    for i in range(n_batches):
        passages = list(passages_dset[i * batch_size : (i + 1) * batch_size]["passage_text"])
        reps = embed_passages_for_retrieval(passages, tokenizer, qa_embedder, max_length, device)
        fp[i * batch_size : (i + 1) * batch_size] = reps
        if i % 50 == 0:


SUMMARY:
Here is a concise and focused summary of the `make_qa_dense_index` function:

**Summary:**

This function creates a dense index for question-answering (QA) retrieval using passage embeddings. It takes in:

*   Passage dataset (`passages_dset`) containing text data
*   Embedder model (`qa_embedder`) and tokenizer (`tokenizer`)
*   Batch size and maximum length parameters
*   Output file name, data type, and device specification

The function:

1.  Creates a memmap file for storing the index using `np.memmap`.
2.  Computes passage embeddings using `embed_passages_for_retrieval` for each batch of passages.
3.  Writes the embedding representations to the memmap file in chunks.

**Key Algorithm/Purpose:**

*   **Passage Embeddings:** Compute dense vector representations of passages using a learned embedder model (`qa_embedder`) and tokenizer (`tokenizer`).

**Implementation Details:**

*   **Batching:** Process passages in batches of size `batch_size` to optimize computation.
*   **Memmap File:** Use `np.memmap` to store the index file, allowing for efficient random access.

**Dependencies/Requirements:**

*   `passages_dset`: Passage dataset
*   `embed_passages_for_retrieval`: A function taking passage text, tokenizer, embedder model, batch size, and device.
*   `np.memmap` (NumPy memory map)
*   `math` for calculating `n_batches`
*   `time` for measuring elapsed time (optional)


key: chroma:document
string_value:
CODE:

def evaluate_retriever(qa_list, retriever_func, scoring_func, n_ret=10, verbose=False):
    total_retriever_time = 0.0
    total_retriever_score = 0.0
    st_time = time()
    for i, (question, answer) in enumerate(qa_list):
        r_time = time()
        retrieved_passages = retriever_func(question, n_ret)
        total_retriever_time += time() - r_time
        total_retriever_score += scoring_func(retrieved_passages, answer)
        if verbose and ((i + 1) % 500 == 0 or i <= 1):
            print(
                "{:03d}: S-{:.4f} T-{:.4f} | {:.2f}".format(
                    i + 1, total_retriever_score / (i + 1), total_retriever_time / (i + 1), time() - st_time
                )
            )


SUMMARY:
**Function Summary: EvaluateRetriever**

The `EvaluateRetriever` function assesses the performance of a natural language retrieval model. It takes in:

* `qa_list`: a list of questions and their corresponding answers
* `retriever_func`: a function that retrieves passages given a question
* `scoring_func`: a function that scores retrieved passages against the correct answer

The function outputs two metrics:

* `idf_recall`: the average recall of the retrieval model based on IDF (Invertible Document Frequency)
* `retrieval_time`: the average time taken by the retrieval model to generate answers

**Key Implementation Details:**

* Uses a `for` loop to iterate over the `qa_list`, computing the score and time for each question-answer pair
* Calls `retriever_func` with `n_ret=10` times to retrieve passages, accumulating total retrieval time and score
* Prints progress updates every 500 iterations or at the first iteration

**Notable Dependencies/Requirements:**

* Assumed use of the `time` module for timing computations


key: chroma:document
string_value:
CODE:
# build a support document for the question out of Wikipedia snippets
def query_qa_dense_index(
    question, qa_embedder, tokenizer, wiki_passages, wiki_index, n_results=10, min_length=20, device="cuda:0"
):
    q_rep = embed_questions_for_retrieval([question], tokenizer, qa_embedder, device=device)
    D, I = wiki_index.search(q_rep, 2 * n_results)
    res_passages = [wiki_passages[int(i)] for i in I[0]]
    support_doc = "<P> " + " <P> ".join([p["passage_text"] for p in res_passages])
    res_list = [{k: p[k] for k in wiki_passages.column_names} for p in res_passages]
    res_list = [res for res in res_list if len(res["passage_text"].split()) > min_length][:n_results]
    for r, sc in zip(res_list, D[0]):
        r["score"] = float(sc)


SUMMARY:
**Function Summary**

`query_qa_dense_index`
------------------------

### Purpose and Functionality

Performs a retrieval-based query on a dense index to find relevant passages from Wikipedia for a given question.

### Implementation Details

*   Embeds the question using `tokenizer` and `qa_embedder`
*   Searches the wiki index with the generated representation
*   Filters search results by minimum passage length
*   Returns the support document (wiki snippets) and the top-ranked passage results

```python
def query_qa_dense_index(
    question: str, 
    qa_embedder: Callable[[str], torch.Tensor], 
    tokenizer: HugginoidTokenizer,
    wiki_passages: pandas.DataFrame, 
    wiki_index: Index,
    n_results=10,
    min_length=20,
    device="cuda:0"
) -> Tuple[str, List[Dict]]:
    # 1. Embed question using tokenizer and qa_embedder
    q_rep = embed_questions_for_retrieval([question].tokenizer, qa_embedder, device=device)
    
    # 2. Search wiki index for top answers
    D, I = wiki_index.search(q_rep, 2 * n_results)
    
    # 3. Filter search results by minimum passage length
    res_passages = [wiki_passages[int(i)] for i in I[0]]
    res_list = [{k: p[k] for k in wiki_passages.column_names} for p in res_passages]
    res_list = [res for res in res_list if len(res["passage_text"].split()) > min_length][:n_results]
    
    # 4. Calculate scores and return top results
    for r, sc in zip(res_list, D[0]):
        r["score"] = float(sc)
        
    support_doc = "<P> " + " <P> ".join([p["passage_text"].split() for p in res_passages])
    
    return support_doc, res_list
```

### Dependencies and Requirements

*   `tokenizer`: HugginoidTokenizer instance
*   `qa_embedder`: callable function embedding questions
*   `wiki_passages` : pandas DataFrame storing Wikipedia passages
*   `wiki_index` : Index type (e.g., FaissIndex)
*   `n_results`: max number of top results to return (default: 10)
*   `min_length`: minimum length for a passage filter (default: 20)


key: chroma:document
string_value:
CODE:

def batch_query_qa_dense_index(questions, qa_embedder, tokenizer, wiki_passages, wiki_index, n_results=10):
    q_rep = embed_questions_for_retrieval(questions, tokenizer, qa_embedder)
    D, I = wiki_index.search(q_rep, n_results)
    res_passages_lst = [[wiki_passages[int(i)] for i in i_lst] for i_lst in I]
    support_doc_lst = [
        "<P> " + " <P> ".join([p["passage_text"] for p in res_passages]) for res_passages in res_passages_lst
    ]
    all_res_lists = []
    for res_passages, dl in zip(res_passages_lst, D):
        res_list = [{k: p[k] for k in wiki_passages.column_names} for p in res_passages]
        for r, sc in zip(res_list, dl):
            r["score"] = float(sc)
        all_res_lists += [res_list[:]]


SUMMARY:
**Summary:**

`batch_query_qa_dense_index` Function:
---------------------------

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

Batch indexing of QA pairs for dense retrieval.

* Takes in embeddings of questions (`q_rep`) using `embed_questions_for_retrieval`
* Serves up support documents (`support_doc_lst`) with corresponding scores

### Key Implementation Details:

* Searches `wiki_index` for top `$n_results$` matching passage IDs (`I`)
* Retrieves and ranks relevant passages from `wiki_passages`

### Dependencies/Requirements
-------------------------

* `embed_questions_for_retrieval`
* `wiki_index.search`
* `wiki_passages.column_names`
* External dependencies:
	+ `DenseIndex` (not implemented in this snippet)

**Function Signatures and Interfaces**
-----------------------------------

* `batch_query_qa_dense_index(questions, qa_embedder, tokenizer, wiki_passages, wiki_index, n_results=10)` 

### Notes:

- This function does not explicitly inherit or implement any external interfaces.
- Specific implementation details of `embed_questions_for_retrieval` are not covered in this summary.


key: chroma:document
string_value:
CODE:
# find nearest neighbors of an answer or declarative text in Wikipedia snippets
def query_qa_dense_index_nn(passage, qa_embedder, tokenizer, wiki_passages, wiki_index, n_results=10, min_length=20):
    a_rep = embed_passages_for_retrieval([passage], tokenizer, qa_embedder)
    D, I = wiki_index.search(a_rep, 2 * n_results)
    res_passages = [wiki_passages[int(i)] for i in I[0]]
    support_doc = "<P> " + " <P> ".join([p["passage_text"] for p in res_passages])
    res_list = [{k: p[k] for k in wiki_passages.column_names} for p in res_passages]
    res_list = [res for res in res_list if len(res["passage_text"].split()) > min_length][:n_results]
    for r, sc, i in zip(res_list, D[0], I[0]):
        r["passage_id"] = int(i)
        r["score"] = float(sc)


SUMMARY:
**Concise Summary**

### Query QA Dense Index NN Function

**Purpose**: Finds the nearest neighbors (Q&A passages) for a given passage in a Wikipedia database.

**Functionality**: Embeds the input passage and searches for nearby Q&A passages using the `wiki_index`. The function returns:

1. A support document snippet from the top N matches.
2. A list of up to `n_results` Q&A passages with their scores.

**Key Details**:

* Embarasses an input passage using `embed_passages_for_retrieval`.
* Searches for nearby Q&A passages in the Wikipedia index (`wiki_index`) using embeddings.
* Filters results by minimum text length and score quality (`min_length`).
* Assigns each passage a corresponding score from the search results.

**Notable Dependencies**:

* `tokenizer`: Tokenizer instance for embedding passages.
* `qa_embedder`: Passage embedder instance (e.g., BERT, RoBERTa).
* `wiki_passages`: Dictionary of Wikipedia passage data.
* `wiki_index`: Index instance for efficient querying.


key: chroma:document
string_value:
CODE:

def batch_query_qa_dense_index_nn(passages, qa_embedder, tokenizer, wiki_passages, wiki_index, n_results=10):
    a_reps = embed_passages_for_retrieval(passages, tokenizer, qa_embedder)
    D, I = wiki_index.search(a_reps, n_results)
    res_passages_lst = [[wiki_passages[int(i)] for i in i_lst] for i_lst in I]
    support_doc_lst = [
        "<P> " + " <P> ".join([p["passage_text"] for p in res_passages]) for res_passages in res_passages_lst
    ]
    all_res_lists = []
    for res_passages, dl, il in zip(res_passages_lst, D, I):
        res_list = [{k: p[k] for k in wiki_passages.column_names} for p in res_passages]
        for r, sc, i in zip(res_list, dl, il):
            r["passage_id"] = int(i)
            r["score"] = float(sc)
        all_res_lists += [res_list[:]]


SUMMARY:
**Function Summary: batch_query_qa_dense_index_nn**

* **Purpose:** Returns supportive documents and lists of retrieved passage IDs with scores for each query from dense indexing the query embeddings.

### Main Purpose and Functionality
- Embeds passages using `embed_passages_for_retrieval`.
- Searches wiki index for top-N results (`n_results=10` default) based on query embeddings.
- Constructs supporting documents by joining retrieved passage texts with `<P>` tags.
- Returns supportive document list, which contains lists of retrieved passage IDs with scores.

### Key Implementation Details

* Utilizes `wiki_index.search` to retrieve top-N results, taking advantage of dense indexing for efficient querying.
* Leverages `embed_passages_for_retrieval` to embed passages using input query embedder (`qa_embedder`) and tokenizer (`tokenizer`).


key: chroma:document
string_value:
CODE:

def find_nearest_training(question, n_results=10):
    q_rep = embed_questions_for_retrieval([question], qar_tokenizer, qar_model)
    D, I = eli5_train_q_index.search(q_rep, n_results)
    nn_examples = [eli5_train[int(i)] for i in I[0]]


SUMMARY:
**Summary:**

### Function Name: find_nearest_training
#### Purpose:
Find the nearest training examples to a given question using Index-based retrieval.

#### Functionality:

- Embeds the input question into a representation space using `qar_tokenizer` and `qar_model`.
- Searches for nearest neighbors in the retrieval index (`eli5_train_q_index`) with the top `n_results=10` relevant documents.
- Retrieves and extracts the corresponding training examples from `eli5_train`.

#### Key Implementation Details:

* Embedding question into a representation space using `embed_questions_for_retrieval`.
* Index search using `eli5_train_q_index.search(q_rep, n_results)`.
* Retrieving training examples with indexing information from `I[0]`.

#### Notable Dependencies:

* `qar_tokenizer`, `qar_model`: Question Answering model and tokenizer.
* `eli5_train_q_index`: Retrieval index for training data.
* `eli5_train`: Training dataset.


key: chroma:document
string_value:
CODE:

def make_support(question, source="wiki40b", method="dense", n_results=10):
    if source == "none":
        support_doc, hit_lst = (" <P> ".join(["" for _ in range(11)]).strip(), [])
    else:
        if method == "dense":
            support_doc, hit_lst = query_qa_dense_index(
                question, qar_model, qar_tokenizer, passages, gpu_dense_index, n_results
            )
        else:
            support_doc, hit_lst = query_es_index(
                question,
                es_client,
                index_name="english_wiki40b_snippets_100w",
                n_results=n_results,
            )
    support_list = [
        (res["article_title"], res["section_title"].strip(), res["score"], res["passage_text"]) for res in hit_lst
    ]
    question_doc = "question: {} context: {}".format(question, support_doc)


SUMMARY:
Here is a concise and focused summary of the provided function code chunk:

**Function Purpose:** 
`make_support` generates a question document and a list of supporting documents for a given question.

**Key Implementation Details:**

* Uses two query methods (`query_qa_dense_index` and `query_es_index`) based on the input `method`.
* Retrieves `n_results` number of hits from either an Elasticsearch index or a dense QA index.
* Returns the original question, a formatted context document (`question_doc`), and a list of supporting documents (`support_list`).

**Notable Dependencies/Requirements:**

* `qar_model`, `qar_tokenizer`
* Elasticsearch client (`es_client`)
* Dense QA index (`gpu_dense_index`)


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/layoutlmv3/run_funsd_cord.py


key: chroma:document
string_value:
CODE:

def main():
    # See all possible arguments in src/transformers/training_args.py
    # or by passing the --help flag to this script.
    # We now keep distinct sets of args, for a cleaner separation of concerns.

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )

    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)
    datasets.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
    )
    logger.info(f"Training/evaluation parameters {training_args}")

    # Detecting last checkpoint.
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    # Set seed before initializing model.
    set_seed(training_args.seed)

    # Get the datasets
    # In distributed training, the load_dataset function guarantee that only one local process can concurrently
    # download the dataset.
    if data_args.dataset_name == "funsd":
        # Downloading and loading a dataset from the hub.
        dataset = load_dataset(
            "nielsr/funsd-layoutlmv3",
            data_args.dataset_config_name,
            cache_dir=model_args.cache_dir,
            token=True if model_args.use_auth_token else None,
        )
    elif data_args.dataset_name == "cord":
        # Downloading and loading a dataset from the hub.
        dataset = load_dataset(
            "nielsr/cord-layoutlmv3",
            data_args.dataset_config_name,
            cache_dir=model_args.cache_dir,
            token=True if model_args.use_auth_token else None,
        )
    else:
        raise ValueError("This script only supports either FUNSD or CORD out-of-the-box.")

    if training_args.do_train:
        column_names = dataset["train"].column_names
        features = dataset["train"].features
    else:
        column_names = dataset["test"].column_names
        features = dataset["test"].features

    image_column_name = "image"
    text_column_name = "words" if "words" in column_names else "tokens"
    boxes_column_name = "bboxes"
    label_column_name = (
        f"{data_args.task_name}_tags" if f"{data_args.task_name}_tags" in column_names else column_names[1]
    )

    remove_columns = column_names

    # In the event the labels are not a `Sequence[ClassLabel]`, we will need to go through the dataset to get the
    # unique labels.
    def get_label_list(labels):
        unique_labels = set()
        for label in labels:
            unique_labels = unique_labels | set(label)
        label_list = list(unique_labels)
        label_list.sort()
        return label_list

    # If the labels are of type ClassLabel, they are already integers and we have the map stored somewhere.
    # Otherwise, we have to get the list of labels manually.
    if isinstance(features[label_column_name].feature, ClassLabel):
        label_list = features[label_column_name].feature.names
        # No need to convert the labels since they are already ints.
        id2label = dict(enumerate(label_list))
        label2id = {v: k for k, v in enumerate(label_list)}
    else:
        label_list = get_label_list(datasets["train"][label_column_name])
        id2label = dict(enumerate(label_list))
        label2id = {v: k for k, v in enumerate(label_list)}
    num_labels = len(label_list)

    # Load pretrained model and processor
    #
    # Distributed training:
    # The .from_pretrained methods guarantee that only one local process can concurrently
    # download model & vocab.
    config = AutoConfig.from_pretrained(
        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
        num_labels=num_labels,
        finetuning_task=data_args.task_name,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        token=True if model_args.use_auth_token else None,
    )

    processor = AutoProcessor.from_pretrained(
        model_args.processor_name if model_args.processor_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        use_fast=True,
        revision=model_args.model_revision,
        token=True if model_args.use_auth_token else None,
        add_prefix_space=True,
        apply_ocr=False,
    )

    model = AutoModelForTokenClassification.from_pretrained(
        model_args.model_name_or_path,
        from_tf=bool(".ckpt" in model_args.model_name_or_path),
        config=config,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        token=True if model_args.use_auth_token else None,
    )

    # Set the correspondences label/ID inside the model config
    model.config.label2id = label2id
    model.config.id2label = id2label

    # Preprocessing the dataset
    # The processor does everything for us (prepare the image using LayoutLMv3ImageProcessor
    # and prepare the words, boxes and word-level labels using LayoutLMv3TokenizerFast)
    def prepare_examples(examples):
        images = examples[image_column_name]
        words = examples[text_column_name]
        boxes = examples[boxes_column_name]
        word_labels = examples[label_column_name]

        encoding = processor(
            images,
            words,
            boxes=boxes,
            word_labels=word_labels,
            truncation=True,
            padding="max_length",
            max_length=data_args.max_seq_length,
        )

        return encoding

    if training_args.do_train:
        if "train" not in dataset:
            raise ValueError("--do_train requires a train dataset")
        train_dataset = dataset["train"]
        if data_args.max_train_samples is not None:
            train_dataset = train_dataset.select(range(data_args.max_train_samples))
        with training_args.main_process_first(desc="train dataset map pre-processing"):
            train_dataset = train_dataset.map(
                prepare_examples,
                batched=True,
                remove_columns=remove_columns,
                num_proc=data_args.preprocessing_num_workers,
                load_from_cache_file=not data_args.overwrite_cache,
            )

    if training_args.do_eval:
        validation_name = "test"
        if validation_name not in dataset:
            raise ValueError("--do_eval requires a validation dataset")
        eval_dataset = dataset[validation_name]
        if data_args.max_eval_samples is not None:
            eval_dataset = eval_dataset.select(range(data_args.max_eval_samples))
        with training_args.main_process_first(desc="validation dataset map pre-processing"):
            eval_dataset = eval_dataset.map(
                prepare_examples,
                batched=True,
                remove_columns=remove_columns,
                num_proc=data_args.preprocessing_num_workers,
                load_from_cache_file=not data_args.overwrite_cache,
            )

    if training_args.do_predict:
        if "test" not in datasets:
            raise ValueError("--do_predict requires a test dataset")
        predict_dataset = datasets["test"]
        if data_args.max_predict_samples is not None:
            max_predict_samples = min(len(predict_dataset), data_args.max_predict_samples)
            predict_dataset = predict_dataset.select(range(max_predict_samples))
        with training_args.main_process_first(desc="prediction dataset map pre-processing"):
            predict_dataset = predict_dataset.map(
                prepare_examples,
                batched=True,
                remove_columns=remove_columns,
                num_proc=data_args.preprocessing_num_workers,
                load_from_cache_file=not data_args.overwrite_cache,
            )

    # Metrics
    metric = load_metric("seqeval")

    def compute_metrics(p):
        predictions, labels = p
        predictions = np.argmax(predictions, axis=2)

        # Remove ignored index (special tokens)
        true_predictions = [
            [label_list[p] for (p, l) in zip(prediction, label) if l != -100]
            for prediction, label in zip(predictions, labels)
        ]
        true_labels = [
            [label_list[l] for (p, l) in zip(prediction, label) if l != -100]
            for prediction, label in zip(predictions, labels)
        ]

        results = metric.compute(predictions=true_predictions, references=true_labels)
        if data_args.return_entity_level_metrics:
            # Unpack nested dictionaries
            final_results = {}
            for key, value in results.items():
                if isinstance(value, dict):
                    for n, v in value.items():
                        final_results[f"{key}_{n}"] = v
                else:
                    final_results[key] = value
            return final_results
        else:
            return {
                "precision": results["overall_precision"],
                "recall": results["overall_recall"],
                "f1": results["overall_f1"],
                "accuracy": results["overall_accuracy"],
            }

    # Initialize our Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset if training_args.do_train else None,
        eval_dataset=eval_dataset if training_args.do_eval else None,
        tokenizer=processor,
        data_collator=default_data_collator,
        compute_metrics=compute_metrics,
    )

    # Training
    if training_args.do_train:
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
            checkpoint = last_checkpoint
        train_result = trainer.train(resume_from_checkpoint=checkpoint)
        metrics = train_result.metrics
        trainer.save_model()  # Saves the tokenizer too for easy upload

        max_train_samples = (
            data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset)
        )
        metrics["train_samples"] = min(max_train_samples, len(train_dataset))

        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()

    # Evaluation
    if training_args.do_eval:
        logger.info("*** Evaluate ***")

        metrics = trainer.evaluate()

        max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset)
        metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset))

        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)

    # Predict
    if training_args.do_predict:
        logger.info("*** Predict ***")

        predictions, labels, metrics = trainer.predict(predict_dataset, metric_key_prefix="predict")
        predictions = np.argmax(predictions, axis=2)

        # Remove ignored index (special tokens)
        true_predictions = [
            [label_list[p] for (p, l) in zip(prediction, label) if l != -100]
            for prediction, label in zip(predictions, labels)
        ]

        trainer.log_metrics("predict", metrics)
        trainer.save_metrics("predict", metrics)

        # Save predictions
        output_predictions_file = os.path.join(training_args.output_dir, "predictions.txt")
        if trainer.is_world_process_zero():
            with open(output_predictions_file, "w") as writer:
                for prediction in true_predictions:
                    writer.write(" ".join(prediction) + "\n")

    kwargs = {"finetuned_from": model_args.model_name_or_path, "tasks": "token-classification"}
    if data_args.dataset_name is not None:
        kwargs["dataset_tags"] = data_args.dataset_name
        if data_args.dataset_config_name is not None:
            kwargs["dataset_args"] = data_args.dataset_config_name
            kwargs["dataset"] = f"{data_args.dataset_name} {data_args.dataset_config_name}"
        else:
            kwargs["dataset"] = data_args.dataset_name

    if training_args.push_to_hub:
        trainer.push_to_hub(**kwargs)
    else:


SUMMARY:
Here is a concise summary of the code:

**Overview**

This code is a PyTorch-based implementation of a sequence labeling task using Token Classification. The script trains, evaluates, and predicts the model on datasets such as Hugging Face's token classification benchmarks.

**Key Components**

1. **Preprocessing**: The script loads the dataset, preprocesses the text data, and creates tokenized representations.
2. **Model Training**: The model is trained using a Trainer class from the Hugging Face library, which handles optimization, evaluation, and checkpointing.
3. **Evaluation and Prediction**: The model's performance is evaluated during training, and predictions are made on test datasets.

**Additional Details**

* The script uses the `TokenClassification` dataset from Hugging Face, which provides benchmark models for sequence labeling tasks.
* The `Trainer` class from Hugging Face takes care of optimization, evaluation, checkpointing, and model saving, making it easy to train and deploy the model.
* The script uses the `default_data_collator` function from Hugging Face, which is a standard way to data-parallelize training on multiple GPUs.

**Note**: The code assumes that the necessary libraries (e.g., PyTorch, Hugging Face) have been installed, and the dataset files are available in the correct location.


key: chroma:document
string_value:
CODE:

def _mp_fn(index):
    # For xla_spawn (TPUs)


SUMMARY:
**Summary:**

The `_mp_fn` function is a wrapper for the `main()` function, likely used in distributed computing or job parallelization. When `xla_spawn` (a TPUs) is specified as an index, this function calls `main()` without any additional arguments.

**Key Details:**

* Implementation of function main is currently unprovided.
* Function wraps `main()` with unspecified argument handling.
* Index `index` passed to `_mp_fn` controls which backend to use.

**Notable Dependencies/Requirements:**

* xla_spawn (TPUs)
* `main()` function (implementation not provided)


key: chroma:document
string_value:
CODE:

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


SUMMARY:
**Summary**

This function `list_field` creates and returns a Django form field. It accepts optional parameters `default` (default value) and `metadata`. The primary purpose is to provide a default implementation for the form field.

**Key Implementation Details**

* Uses `field` factory method to create a new field instance.
* Passes a lambda function as the `default_factory` parameter, which returns the provided `default` value.
* Sets `metadata` attribute when calling the `list_field` function.

**Notable Dependencies or Requirements**

None mentioned in this snippet. However, it's assumed that other Django form field-related components are imported and configured elsewhere in the codebase.


key: chroma:document
string_value:
CODE:

def create_vocabulary_from_data(
    datasets: DatasetDict,
    word_delimiter_token: Optional[str] = None,
    unk_token: Optional[str] = None,
    pad_token: Optional[str] = None,
):
    # Given training and test labels create vocabulary
    def extract_all_chars(batch):
        all_text = " ".join(batch["target_text"])
        vocab = list(set(all_text))
        return {"vocab": [vocab], "all_text": [all_text]}

    vocabs = datasets.map(
        extract_all_chars,
        batched=True,
        batch_size=-1,
        keep_in_memory=True,
        remove_columns=datasets["train"].column_names,
    )

    # take union of all unique characters in each dataset
    vocab_set = (
        (set(vocabs["train"]["vocab"][0]) if "train" in vocabs else set())
        | (set(vocabs["eval"]["vocab"][0]) if "eval" in vocabs else set())
        | (set(vocabs["predict"]["vocab"][0]) if "predict" in vocabs else set())
    )

    vocab_dict = {v: k for k, v in enumerate(sorted(vocab_set))}

    # replace white space with delimiter token
    if word_delimiter_token is not None:
        vocab_dict[word_delimiter_token] = vocab_dict[" "]
        del vocab_dict[" "]

    # add unk and pad token
    if unk_token is not None:
        vocab_dict[unk_token] = len(vocab_dict)

    if pad_token is not None:
        vocab_dict[pad_token] = len(vocab_dict)



SUMMARY:
**Concise Summary**

Creates a vocabulary dictionary from training, evaluation, and prediction data.

### Main Purpose and Functionality

Creates a vocabulary dictionary (`vocab_dict`) that maps unique characters to indices. The function extracts all unique characters from the given datasets, sorts them, and adds special tokens (unk_token, pad_token) if specified.

### Key Implementation Details

* Extracts all unique characters using `extract_all_chars` function
* Takes union of unique characters from training, evaluation, and prediction datasets
* Adds unk_token, pad_token with corresponding indices

### Notable Dependencies or Requirements

* `datasets`: a dictionary containing training, evaluation, and prediction data
* `word_delimiter_token`, `unk_token`, `pad_token`: optional tokens for whitespace, unknown, and padding purposes

### Method Signatures and Interfaces

* `extract_all_chars(batch) -> dict`: extracts unique characters from batch data
* `datasets.map()` with parameters `batched=True`, `keep_in_memory=True`

### External Dependencies:

None mentioned.


key: chroma:document
string_value:
CODE:

def main():
    # See all possible arguments in src/transformers/training_args.py
    # or by passing the --help flag to this script.
    # We now keep distinct sets of args, for a cleaner separation of concerns.

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    # Detecting last checkpoint.
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif last_checkpoint is not None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )
    logger.setLevel(logging.INFO if is_main_process(training_args.local_rank) else logging.WARN)

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, "
        f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
    )
    # Set the verbosity to info of the Transformers logger (on main process only):
    if is_main_process(training_args.local_rank):
        transformers.utils.logging.set_verbosity_info()
    logger.info("Training/evaluation parameters %s", training_args)

    # Set seed before initializing model.
    set_seed(training_args.seed)

    # 1. First, let's load the dataset
    raw_datasets = DatasetDict()
    task_name = data_args.task
    lang_id = data_args.language

    if task_name is None:
        raise ValueError(
            "Set --task should be set to '<xtreme_s_task>' (e.g. 'fleurs-asr', 'mls', 'covost2', 'minds14') "
        )
    if lang_id is None:
        raise ValueError(
            "Set --language should be set to the language id of the sub dataset "
            "config to be used (e.g. 'pl', 'en.tr', 'fr-FR') or 'all'"
            " for multi-lingual fine-tuning."
        )
    if data_args.language_group is not None:
        if data_args.task != "fleurs-asr":
            raise ValueError("--language_group should only be used with --task=fleurs-asr")
        if data_args.language != "all":
            raise ValueError("--language_group should only be used with --language=all")

    if data_args.target_column_name is None:
        target_column_name = TASK_TO_TARGET_COLUMN_NAME[task_name]
    else:
        target_column_name = data_args.target_column_name

    # here we differentiate between tasks with text as the target and classification tasks
    is_text_target = target_column_name in ("transcription", "translation")

    config_name = ".".join([task_name.split("-")[0], lang_id])

    if training_args.do_train:
        raw_datasets["train"] = load_dataset(
            data_args.dataset_name,
            config_name,
            split=data_args.train_split_name,
            token=data_args.use_auth_token,
            cache_dir=model_args.cache_dir,
        )

        if data_args.audio_column_name not in raw_datasets["train"].column_names:
            raise ValueError(
                f"--audio_column_name '{data_args.audio_column_name}' not found in dataset '{data_args.dataset_name}'."
                " Make sure to set `--audio_column_name` to the correct audio column - one of"
                f" {', '.join(raw_datasets['train'].column_names)}."
            )

        if target_column_name not in raw_datasets["train"].column_names:
            raise ValueError(
                f"--target_column_name {target_column_name} not found in dataset '{data_args.dataset_name}'. "
                "Make sure to set `--target_column_name` to the correct text column - one of "
                f"{', '.join(raw_datasets['train'].column_names)}."
            )

        if data_args.max_train_samples is not None:
            raw_datasets["train"] = raw_datasets["train"].select(range(data_args.max_train_samples))

    if training_args.do_eval:
        raw_datasets["eval"] = load_dataset(
            data_args.dataset_name,
            config_name,
            split=data_args.eval_split_name,
            token=data_args.use_auth_token,
            cache_dir=model_args.cache_dir,
        )

        if data_args.max_eval_samples is not None:
            raw_datasets["eval"] = raw_datasets["eval"].select(range(data_args.max_eval_samples))

    if training_args.do_predict:
        raw_datasets["predict"] = load_dataset(
            data_args.dataset_name,
            config_name,
            split=data_args.predict_split_name,
            token=data_args.use_auth_token,
            cache_dir=model_args.cache_dir,
        )

        if data_args.max_predict_samples is not None:
            raw_datasets["predict"] = raw_datasets["predict"].select(range(data_args.max_predict_samples))

    lang_list = next(iter(raw_datasets.values())).features["lang_id"].names
    if not is_text_target:
        label_list = next(iter(raw_datasets.values())).features[target_column_name].names
        num_labels = len(label_list)

    num_workers = data_args.preprocessing_num_workers

    lang_group = data_args.language_group
    if lang_group is not None:
        with training_args.main_process_first(desc="language group filter"):
            lang_group_id = next(iter(raw_datasets.values())).features["lang_group_id"].str2int(lang_group)
            raw_datasets = raw_datasets.filter(
                lambda lang_group: lang_group == lang_group_id,
                num_proc=num_workers,
                input_columns=["lang_group_id"],
            )

    # 2. We remove some special characters from the datasets
    # that make training complicated and do not help in transcribing the speech
    # E.g. characters, such as `,` and `.` do not really have an acoustic characteristic
    # that could be easily picked up by the model
    chars_to_ignore_regex = (
        f'[{"".join(data_args.chars_to_ignore)}]' if data_args.chars_to_ignore is not None else None
    )

    def remove_special_characters(batch):
        if chars_to_ignore_regex is not None:
            batch["target_text"] = re.sub(chars_to_ignore_regex, "", batch[target_column_name]).lower() + " "
        else:
            batch["target_text"] = batch[target_column_name].lower() + " "
        return batch

    if is_text_target:
        with training_args.main_process_first(desc="dataset map special characters removal"):
            raw_datasets = raw_datasets.map(
                remove_special_characters,
                remove_columns=[target_column_name],
                desc="remove special characters from datasets",
            )

        # save special tokens for tokenizer
        word_delimiter_token = data_args.word_delimiter_token
        unk_token = data_args.unk_token
        pad_token = data_args.pad_token

    # 3. Next, let's load the config as we might need it to create
    # the tokenizer
    config = AutoConfig.from_pretrained(
        model_args.model_name_or_path, cache_dir=model_args.cache_dir, token=data_args.use_auth_token
    )

    if is_text_target:
        # 4. (Optional, for ASR and translation) If no tokenizer file is defined,
        # we create the vocabulary of the model by extracting all unique characters from
        # the training and evaluation datasets
        # We need to make sure that only first rank saves vocabulary
        # make sure all processes wait until vocab is created
        tokenizer_name_or_path = model_args.tokenizer_name_or_path
        tokenizer_kwargs = {}
        if tokenizer_name_or_path is None:
            # save vocab in training output dir
            tokenizer_name_or_path = training_args.output_dir

            vocab_file = os.path.join(tokenizer_name_or_path, "vocab.json")

            with training_args.main_process_first():
                if training_args.overwrite_output_dir and os.path.isfile(vocab_file):
                    os.remove(vocab_file)

            with training_args.main_process_first(desc="dataset map vocabulary creation"):
                if not os.path.isfile(vocab_file):
                    os.makedirs(tokenizer_name_or_path, exist_ok=True)
                    vocab_dict = create_vocabulary_from_data(
                        raw_datasets,
                        word_delimiter_token=word_delimiter_token,
                        unk_token=unk_token,
                        pad_token=pad_token,
                    )

                    # save vocab dict to be loaded into tokenizer
                    with open(vocab_file, "w") as file:
                        json.dump(vocab_dict, file)

            # if tokenizer has just been created
            # it is defined by `tokenizer_class` if present in config else by `model_type`
            if not config.is_encoder_decoder:
                tokenizer_kwargs = {
                    "config": config if config.tokenizer_class is not None else None,
                    "tokenizer_type": config.model_type if config.tokenizer_class is None else None,
                    "unk_token": unk_token,
                    "pad_token": pad_token,
                    "word_delimiter_token": word_delimiter_token,
                }
            else:
                tokenizer_kwargs = {}

    # 5. Now we can instantiate the feature extractor, tokenizer and model
    # Note for distributed training, the .from_pretrained methods guarantee that only
    # one local process can concurrently download model & vocab.

    # load feature_extractor and tokenizer
    if is_text_target:
        tokenizer = AutoTokenizer.from_pretrained(
            tokenizer_name_or_path,
            token=data_args.use_auth_token,
            **tokenizer_kwargs,
        )
    feature_extractor = AutoFeatureExtractor.from_pretrained(
        model_args.model_name_or_path, cache_dir=model_args.cache_dir, token=data_args.use_auth_token
    )

    # adapt config
    # (speech translation requires pre-configured seq2seq models)
    if task_name != "covost2":
        config.update(
            {
                "feat_proj_dropout": model_args.feat_proj_dropout,
                "attention_dropout": model_args.attention_dropout,
                "hidden_dropout": model_args.hidden_dropout,
                "final_dropout": model_args.final_dropout,
                "mask_time_prob": model_args.mask_time_prob,
                "mask_time_length": model_args.mask_time_length,
                "mask_feature_prob": model_args.mask_feature_prob,
                "mask_feature_length": model_args.mask_feature_length,
                "gradient_checkpointing": training_args.gradient_checkpointing,
                "layerdrop": model_args.layerdrop,
                "ctc_zero_infinity": model_args.ctc_zero_infinity,
                "ctc_loss_reduction": model_args.ctc_loss_reduction,
                "activation_dropout": model_args.activation_dropout,
            }
        )
        if training_args.do_train:
            if is_text_target:
                config.pad_token_id = tokenizer.pad_token_id
                config.vocab_size = len(tokenizer)
            else:
                label_to_id = {v: i for i, v in enumerate(label_list)}
                config.label2id = label_to_id
                config.id2label = {id: label for label, id in label_to_id.items()}
                config.num_labels = num_labels

    # create model
    if target_column_name == "transcription":
        model = AutoModelForCTC.from_pretrained(
            model_args.model_name_or_path,
            cache_dir=model_args.cache_dir,
            config=config,
            token=data_args.use_auth_token,
        )
    elif config.is_encoder_decoder:
        model = AutoModelForSpeechSeq2Seq.from_pretrained(
            model_args.model_name_or_path,
            cache_dir=model_args.cache_dir,
            config=config,
            token=data_args.use_auth_token,
        )
        if model.config.decoder_start_token_id is None:
            raise ValueError("Make sure that `config.decoder_start_token_id` is correctly defined")
    else:
        model = AutoModelForAudioClassification.from_pretrained(
            model_args.model_name_or_path,
            cache_dir=model_args.cache_dir,
            config=config,
            token=data_args.use_auth_token,
        )

    # freeze encoder
    if model_args.freeze_feature_encoder:
        model.freeze_feature_encoder()

    # 6. Now we preprocess the datasets including loading the audio, resampling and normalization
    # Thankfully, `datasets` takes care of automatically loading and resampling the audio,
    # so that we just need to set the correct target sampling rate and normalize the input
    # via the `feature_extractor`

    # make sure that dataset decodes audio with correct sampling rate
    dataset_sampling_rate = next(iter(raw_datasets.values())).features[data_args.audio_column_name].sampling_rate
    if dataset_sampling_rate != feature_extractor.sampling_rate:
        raw_datasets = raw_datasets.cast_column(
            data_args.audio_column_name, datasets.features.Audio(sampling_rate=feature_extractor.sampling_rate)
        )

    # derive max & min input length for sample rate & max duration
    max_input_length = data_args.max_duration_in_seconds * feature_extractor.sampling_rate
    min_input_length = data_args.min_duration_in_seconds * feature_extractor.sampling_rate
    audio_column_name = data_args.audio_column_name

    # `phoneme_language` is only relevant if the model is fine-tuned on phoneme classification
    phoneme_language = data_args.phoneme_language

    # Preprocessing the datasets.
    # We need to read the audio files as arrays and tokenize the targets.
    def prepare_dataset(batch):
        # load audio
        sample = batch[audio_column_name]

        inputs = feature_extractor(sample["array"], sampling_rate=sample["sampling_rate"])
        batch["input_values"] = inputs.input_values[0]
        batch["length"] = len(batch["input_values"])

        # encode targets
        additional_kwargs = {}
        if phoneme_language is not None:
            additional_kwargs["phonemizer_lang"] = phoneme_language

        if is_text_target:
            batch["labels"] = tokenizer(batch["target_text"], **additional_kwargs).input_ids
        else:
            batch["labels"] = batch[target_column_name]

        batch["lang"] = batch["lang_id"]

        return batch

    with training_args.main_process_first(desc="dataset map preprocessing"):
        vectorized_datasets = raw_datasets.map(
            prepare_dataset,
            remove_columns=next(iter(raw_datasets.values())).column_names,
            num_proc=num_workers,
            desc="preprocess datasets",
        )

        if training_args.do_train:

            def is_audio_in_length_range(length):
                return length > min_input_length and length < max_input_length

            # filter data that is shorter than min_input_length
            vectorized_datasets["train"] = vectorized_datasets["train"].filter(
                is_audio_in_length_range,
                num_proc=num_workers,
                input_columns=["length"],
            )

    # 7. Next, we can prepare for the training step.
    # Let's use the appropriate XTREME-S evaluation metric,
    # instantiate a data collator and the trainer

    # Define evaluation metrics during training, *i.e.* word error rate, character error rate
    eval_metric = load_metric("xtreme_s", task_name)

    # for large datasets it is advised to run the preprocessing on a
    # single machine first with ``args.preprocessing_only`` since there will mostly likely
    # be a timeout when running the script in distributed mode.
    # In a second step ``args.preprocessing_only`` can then be set to `False` to load the
    # cached dataset
    if data_args.preprocessing_only:
        logger.info(f"Data preprocessing finished. Files cached at {vectorized_datasets.cache_files}")
        return

    def asr_logits_argmax(logits, labels):
        return logits.argmax(dim=-1)

    def compute_asr_metric(pred):
        pred.label_ids[pred.label_ids == -100] = tokenizer.pad_token_id

        pred_str = tokenizer.batch_decode(pred.predictions)
        # we do not want to group tokens when computing the metrics
        label_str = tokenizer.batch_decode(pred.label_ids, group_tokens=False)

        metric = eval_metric.compute(predictions=pred_str, references=label_str)
        return metric

    def compute_classification_metric(pred):
        pred_ids = np.argmax(pred.predictions, axis=1)
        metric = eval_metric.compute(predictions=pred_ids, references=pred.label_ids)
        return metric

    # Now save everything to be able to create a single processor later
    if is_main_process(training_args.local_rank):
        # save feature extractor, tokenizer and config
        feature_extractor.save_pretrained(training_args.output_dir)
        if is_text_target:
            tokenizer.save_pretrained(training_args.output_dir)
        config.save_pretrained(training_args.output_dir)
    # wait until configs are saved in the main process before loading the processor
    if training_args.local_rank != -1:
        torch.distributed.barrier()

    if is_text_target:
        processor = AutoProcessor.from_pretrained(training_args.output_dir)
    else:
        processor = AutoFeatureExtractor.from_pretrained(training_args.output_dir)

    # Instantiate custom data collator
    data_collator = SpeechDataCollatorWithPadding(processor=processor, pad_labels=is_text_target)

    # Initialize Trainer
    if target_column_name == "translation":
        trainer = Seq2SeqTrainer(
            model=model,
            data_collator=data_collator,
            args=training_args,
            preprocess_logits_for_metrics=asr_logits_argmax if training_args.predict_with_generate else None,
            compute_metrics=compute_asr_metric if training_args.predict_with_generate else None,
            train_dataset=vectorized_datasets["train"] if training_args.do_train else None,
            eval_dataset=vectorized_datasets["eval"] if training_args.do_eval else None,
            tokenizer=feature_extractor,
        )
    else:
        trainer = Trainer(
            model=model,
            data_collator=data_collator,
            args=training_args,
            preprocess_logits_for_metrics=asr_logits_argmax if is_text_target else None,
            compute_metrics=compute_asr_metric if is_text_target else compute_classification_metric,
            train_dataset=vectorized_datasets["train"] if training_args.do_train else None,
            eval_dataset=vectorized_datasets["eval"] if training_args.do_eval else None,
            tokenizer=feature_extractor,
        )

    # 8. Finally, we can start training

    # Training
    if training_args.do_train:
        # use last checkpoint if exist
        if last_checkpoint is not None:
            checkpoint = last_checkpoint
        elif os.path.isdir(model_args.model_name_or_path):
            checkpoint = model_args.model_name_or_path
        else:
            checkpoint = None

        train_result = trainer.train(resume_from_checkpoint=checkpoint)
        trainer.save_model()

        metrics = train_result.metrics
        max_train_samples = (
            data_args.max_train_samples
            if data_args.max_train_samples is not None
            else len(vectorized_datasets["train"])
        )
        metrics["train_samples"] = min(max_train_samples, len(vectorized_datasets["train"]))

        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()

    # Evaluation on the test set
    results = {}
    if training_args.do_predict:
        logger.info(f"*** Evaluating on the `{data_args.predict_split_name}` set ***")
        if data_args.per_lang_metrics:
            # separate the `test` dataset into language-specific subsets and compute metrics for each of them
            metrics = {}
            average_metrics = defaultdict(list)
            for lang_id in range(len(lang_list)):
                lang_name = lang_list[lang_id]
                with training_args.main_process_first(desc="per-language dataset filter"):
                    lang_dataset = vectorized_datasets["predict"].filter(
                        lambda lang: lang == lang_id,
                        num_proc=num_workers,
                        input_columns=["lang"],
                    )
                lang_metrics = trainer.evaluate(lang_dataset)
                redundant_metrics = ["eval_runtime", "eval_samples_per_second", "eval_steps_per_second", "eval_epoch"]
                for metric_name, value in lang_metrics.items():
                    average_metrics[metric_name].append(value)
                    if metric_name not in redundant_metrics:
                        metrics[f"{metric_name}_{lang_name}"] = value
            for metric_name, value in average_metrics.items():
                metrics[metric_name] = np.mean(value)
        else:
            metrics = trainer.evaluate(vectorized_datasets["predict"])
        max_predict_samples = (
            data_args.max_predict_samples
            if data_args.max_predict_samples is not None
            else len(vectorized_datasets["predict"])
        )
        metrics["predict_samples"] = min(max_predict_samples, len(vectorized_datasets["predict"]))

        # make sure that the `predict` metrics end up in the log history for the model card
        trainer.log(OrderedDict(sorted(metrics.items())))

        trainer.log_metrics("predict", metrics)
        trainer.save_metrics("predict", metrics)

    # Write model card and (optionally) push to hub
    kwargs = {
        "finetuned_from": model_args.model_name_or_path,
        "tasks": task_name,
        "tags": [task_name, data_args.dataset_name],
        "dataset_args": (
            f"Config: {config_name}, Training split: {data_args.train_split_name}, Eval split:"
            f" {data_args.eval_split_name}, Predict split: {data_args.predict_split_name}"
        ),
        "dataset": f"{data_args.dataset_name.upper()} - {config_name.upper()}",
        "language": data_args.language,
    }

    if training_args.push_to_hub:
        trainer.push_to_hub(**kwargs)
    else:
        trainer.create_model_card(**kwargs)



SUMMARY:
Here is a concise summary of the provided code:

The code trains a speech recognition model using Transformers for sequential tasks with a few key configuration options. It takes input from a dataset, preprocesses it, and then trains the model using a Trainer object that has customizable options for data collator, metrics, and evaluation.

Key points include:

* Option to train with generating (e.g., translation) or not
* Customizable metrics and evaluation settings
* Train dataset and eval dataset handling
* Evaluation on test set after training
* Model saving and pushing to HUB
* Training configuration options (e.g., fine-tuning, data splitting)
* Automatic generation of model card with provided metadata.

Overall, the code is well-structured and allows for customizable training parameters.


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/visual_bert/extracting_data.py


key: chroma:document
string_value:
CODE:

class Extract:
    def __init__(self, argv=sys.argv[1:]):
        inputdir = None
        outputfile = None
        subset_list = None
        batch_size = 1
        opts, args = getopt.getopt(argv, "i:o:b:s", ["inputdir=", "outfile=", "batch_size=", "subset_list="])
        for opt, arg in opts:
            if opt in ("-i", "--inputdir"):
                inputdir = arg
            elif opt in ("-o", "--outfile"):
                outputfile = arg
            elif opt in ("-b", "--batch_size"):
                batch_size = int(arg)
            elif opt in ("-s", "--subset_list"):
                subset_list = arg

        assert inputdir is not None  # and os.path.isdir(inputdir), f"{inputdir}"
        assert outputfile is not None and not os.path.isfile(outputfile), f"{outputfile}"
        if subset_list is not None:
            with open(os.path.realpath(subset_list)) as f:
                self.subset_list = {self._vqa_file_split()[0] for x in tryload(f)}
        else:
            self.subset_list = None

        self.config = CONFIG
        if torch.cuda.is_available():
            self.config.model.device = "cuda"
        self.inputdir = os.path.realpath(inputdir)
        self.outputfile = os.path.realpath(outputfile)
        self.preprocess = Preprocess(self.config)
        self.model = GeneralizedRCNN.from_pretrained("unc-nlp/frcnn-vg-finetuned", config=self.config)
        self.batch = batch_size if batch_size != 0 else 1
        self.schema = DEFAULT_SCHEMA

    def _vqa_file_split(self, file):
        img_id = int(file.split(".")[0].split("_")[-1])
        filepath = os.path.join(self.inputdir, file)
        return (img_id, filepath)

    @property
    def file_generator(self):
        batch = []
        for i, file in enumerate(os.listdir(self.inputdir)):
            if self.subset_list is not None and i not in self.subset_list:
                continue
            batch.append(self._vqa_file_split(file))
            if len(batch) == self.batch:
                temp = batch
                batch = []
                yield list(map(list, zip(*temp)))

        for i in range(1):
            yield list(map(list, zip(*batch)))

    def __call__(self):
        # make writer
        if not TEST:
            writer = datasets.ArrowWriter(features=self.schema, path=self.outputfile)
        # do file generator
        for i, (img_ids, filepaths) in enumerate(self.file_generator):
            images, sizes, scales_yx = self.preprocess(filepaths)
            output_dict = self.model(
                images,
                sizes,
                scales_yx=scales_yx,
                padding="max_detections",
                max_detections=self.config.MAX_DETECTIONS,
                pad_value=0,
                return_tensors="np",
                location="cpu",
            )
            output_dict["boxes"] = output_dict.pop("normalized_boxes")
            if not TEST:
                output_dict["img_id"] = np.array(img_ids)
                batch = self.schema.encode_batch(output_dict)
                writer.write_batch(batch)
            if TEST:
                break
            # finalizer the writer
        if not TEST:
            num_examples, num_bytes = writer.finalize()


SUMMARY:
**Class Extract:**

* **Purpose:** Defines a class for extract data from directory, process images and perform VQA tasks using pre-trained model.

* **Key Features:**
	+ Accepts command-line arguments to specify input directory, output file, batch size, and subset list.
	+ Uses `getopt` to parse command-line arguments.
	+ Checks if input directory exists and is a directory, ensuring consistency across configurations.
	+ Optionally loads subset list from file and generates corresponding image IDs if set.

* **Notable Functions:**
	+ `_vqa_file_split()`: Extracts img ID from VQA file name.
	+ `file_generator`: Iteratively yields batches of images with corresponding file paths.
	+ `__call__`: The main entry point, responsible for generating writer and finalizing output.

* **Dependencies:**
	+ `CONFIG` (configuration) class
	+ `Preprocess`, `GeneralizedRCNN`, and `datasets.ArrowWriter` classes (imported but not shown in this snippet)

* **Interfaces:**
	+ Implements custom `file_generator` functionality and uses pre-trained model's architecture.

**Example Usage:**

```python
extract = Extract()
# ... set command-line arguments ...
extract()
```

This class provides a structured framework for processing images and performing VQA tasks, with an emphasis on flexibility and customization through command-line arguments.


key: chroma:document
string_value:
CODE:

def tryload(stream):
    try:
        data = json.load(stream)
        try:
            data = list(data.keys())
        except Exception:
            data = [d["img_id"] for d in data]
    except Exception:
        try:
            data = eval(stream.read())
        except Exception:
            data = stream.read().split("\n")


SUMMARY:
**Summary: `tryload` Function**

```python
def tryload(stream):
    """
    Attempt to load data from the given stream.

    Returns:
        list[data]: List of loaded data (image IDs or strings)
    """

    # Try JSON loading, fall back on eval if fails
    try:
        data = json.load(stream)  # Load JSON data
        # If data is a dictionary, extract image IDs; otherwise, return entire dict
        try:
            return list(data.keys())
        except Exception:  
            return [d["img_id"] for d in data]
    except Exception:
        try:
            # Evaluate Python expression from stream content
            data = eval(stream.read())
        except Exception:
            # Fall back to simple string splitting if eval fails
            data = stream.read().split("\n")
    return data
```

**Key Implementation Details:**

- Tries loading JSON data from the input stream.
- Falls back on evaluating a Python expression and/or simple string splitting if JSON loading fails.

**Dependencies/Requirements:**

- `json` module for JSON loading.
- `eval` function for evaluating Python expressions.
- No external dependencies explicitly listed, but assumes `stream` object has necessary methods (`read()`, `split("\n")`).


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/visual_bert/modeling_frcnn.py


key: chroma:document
string_value:
CODE:
# other:
def norm_box(boxes, raw_sizes):
    if not isinstance(boxes, torch.Tensor):
        normalized_boxes = boxes.copy()
    else:
        normalized_boxes = boxes.clone()
    normalized_boxes[:, :, (0, 2)] /= raw_sizes[:, 1]
    normalized_boxes[:, :, (1, 3)] /= raw_sizes[:, 0]


SUMMARY:
**Summary**

This function normalizes bounding box coordinates in a provided `torch.Tensor` by their corresponding object sizes. It scales the x and y limits (upper and lower) of the boxes while preserving center point (width, height). The scaling factor is determined from the raw object sizes.

**Key Implementation Details**

* Normalizes bounding box coordinates using raw object sizes
* Scales width/height ratio (`raw_sizes[:, 1]` for width `/raw_sizes[:, 0]` for height) separately
* Preserves center point (width, height)

**Notable Dependencies**

* `torch.Tensor` for input and output data structure

**No external dependencies listed.**

The code maintains consistent formatting using Markdown-style comments. The focus is on delivering a clear summary that highlights the function's purpose, key implementation details, and any notable dependencies.


key: chroma:document
string_value:
CODE:

def pad_list_tensors(
    list_tensors,
    preds_per_image,
    max_detections=None,
    return_tensors=None,
    padding=None,
    pad_value=0,
    location=None,
):
    """
    location will always be cpu for np tensors
    """
    if location is None:
        location = "cpu"
    assert return_tensors in {"pt", "np", None}
    assert padding in {"max_detections", "max_batch", None}
    new = []
    if padding is None:
        if return_tensors is None:
            return list_tensors
        elif return_tensors == "pt":
            if not isinstance(list_tensors, torch.Tensor):
                return torch.stack(list_tensors).to(location)
            else:
                return list_tensors.to(location)
        else:
            if not isinstance(list_tensors, list):
                return np.array(list_tensors.to(location))
            else:
                return list_tensors.to(location)
    if padding == "max_detections":
        assert max_detections is not None, "specify max number of detections per batch"
    elif padding == "max_batch":
        max_detections = max(preds_per_image)
    for i in range(len(list_tensors)):
        too_small = False
        tensor_i = list_tensors.pop(0)
        if tensor_i.ndim < 2:
            too_small = True
            tensor_i = tensor_i.unsqueeze(-1)
        assert isinstance(tensor_i, torch.Tensor)
        tensor_i = nn.functional.pad(
            input=tensor_i,
            pad=(0, 0, 0, max_detections - preds_per_image[i]),
            mode="constant",
            value=pad_value,
        )
        if too_small:
            tensor_i = tensor_i.squeeze(-1)
        if return_tensors is None:
            if location == "cpu":
                tensor_i = tensor_i.cpu()
            tensor_i = tensor_i.tolist()
        if return_tensors == "np":
            if location == "cpu":
                tensor_i = tensor_i.cpu()
            tensor_i = tensor_i.numpy()
        else:
            if location == "cpu":
                tensor_i = tensor_i.cpu()
        new.append(tensor_i)
    if return_tensors == "np":
        return np.stack(new, axis=0)
    elif return_tensors == "pt" and not isinstance(new, torch.Tensor):
        return torch.stack(new, dim=0)
    else:


SUMMARY:
**pad_list_tensors Function Summary**

### Purpose and Functionality
The `pad_list_tensors` function pads input tensors to a consistent size for batching. It takes in a list of tensors, a `preds_per_image`, and optional parameters `max_detections` and `return_tensors`. The primary purpose is to add padding to tensors to facilitate batch processing.

### Implementation Details

*   Pads each tensor in the input list based on the number of predictions per image (`preds_per_image`) or the maximum batch size (`max_detections`).
*   Applies constant mode padding with a specified `pad_value`.
*   Returns a new list with padded tensors, possibly rearranging them for PyTorch (`return_tensors="pt"`) or NumPy (`return_tensors="np"`).

### Dependencies and Requirements

*   `torch`, `nn.functional`, and `numpy` are assumed to be imported.

### Notable Parameters

*   `max_detections`: The maximum number of detections per batch for padding.
*   `padding`: Can be `"max_detections"`, `"max_batch"`, or `None`.
*   `pad_value`: The constant value used for padding.

**Method Signatures:**

*   `torch.stack` and `numpy.stack` are used to stack tensors in PyTorch and NumPy, respectively.
*   `nn.functional.pad` is used for Padding input tensors.


key: chroma:document
string_value:
CODE:

def do_nms(boxes, scores, image_shape, score_thresh, nms_thresh, mind, maxd):
    scores = scores[:, :-1]
    num_bbox_reg_classes = boxes.shape[1] // 4
    # Convert to Boxes to use the `clip` function ...
    boxes = boxes.reshape(-1, 4)
    _clip_box(boxes, image_shape)
    boxes = boxes.view(-1, num_bbox_reg_classes, 4)  # R x C x 4

    # Select max scores
    max_scores, max_classes = scores.max(1)  # R x C --> R
    num_objs = boxes.size(0)
    boxes = boxes.view(-1, 4)
    idxs = torch.arange(num_objs).to(boxes.device) * num_bbox_reg_classes + max_classes
    max_boxes = boxes[idxs]  # Select max boxes according to the max scores.

    # Apply NMS
    keep = nms(max_boxes, max_scores, nms_thresh)
    keep = keep[:maxd]
    if keep.shape[-1] >= mind and keep.shape[-1] <= maxd:
        max_boxes, max_scores = max_boxes[keep], max_scores[keep]
        classes = max_classes[keep]
        return max_boxes, max_scores, classes, keep
    else:


SUMMARY:
**Function Summary: Non-Maximum Suppression (NMS)**

### Purpose and Functionality

The `do_nms` function applies non-maximum suppression to detected bounding boxes to reduce overlaps. It takes in:

* `boxes`: detection scores and box coordinates
* `scores`: detection scores
* `image_shape`: image dimensions (not used)
* `score_thresh`, `nms_thresh`, `mind`, and `maxd`: threshold values for filtering results

The function outputs the top-scoring boxes, scores, and class labels after applying NMS.

### Key Implementation Details

* Reshapes input boxes to `R x C x 4` format for clipping
* Uses `clip_box` to restrict box coordinates within image bounds
* Selects maximum scores, classes, and corresponding boxes using indexing `(idxs := num_objs * max_classes)`
* Applies NMS using `nms` function with given threshold values

### Notable Dependencies or Requirements

* `torch`: PyTorch library used for tensor operations
* `boxes`, `scores`, `image_shape`: input data structures, assumed to be tensors


key: chroma:document
string_value:
CODE:
# Helper Functions
def _clip_box(tensor, box_size: Tuple[int, int]):
    assert torch.isfinite(tensor).all(), "Box tensor contains infinite or NaN!"
    h, w = box_size
    tensor[:, 0].clamp_(min=0, max=w)
    tensor[:, 1].clamp_(min=0, max=h)
    tensor[:, 2].clamp_(min=0, max=w)


SUMMARY:
**Summary**

 `_clip_box` function:
 Clips a `tensor` to a bounding box defined by the `box_size` parameter.
 
**Key Implementation Details**
 
* Assumes input tensor is 4D and evenly divided into channels (height, width, color, alpha)
* Clamps each channel of the input tensor against the corresponding boundary value (h x w) using `clamp_`
* Checks if all elements in the input tensor are finite and non-NaN

**Dependencies**

* `torch`: used for tensor operations (`isfinite`, `torch.isfinite`, `clamp_`)


key: chroma:document
string_value:
CODE:

def _nonempty_boxes(box, threshold: float = 0.0) -> torch.Tensor:
    widths = box[:, 2] - box[:, 0]
    heights = box[:, 3] - box[:, 1]
    keep = (widths > threshold) & (heights > threshold)


SUMMARY:
**Function Summarization:**

### `nonempty_boxes` Function
#### Purpose:
Filters non-empty bounding boxes based on a given threshold.

#### Implementation Details:

* Calculates box widths and heights by subtracting origin coordinates from each point.
* Filters keeping only boxes with both width and height above the threshold using bitwise AND operation.
* Returns boolean tensor indicating presence of non-empty boxes.

#### Parameters:
- `box` (torch.Tensor): Input bounding box tensor (`x_min`, `y_min`, `x_max`, `y_max`)
- `threshold` (float, optional): Threshold value for determining empty boxes (default: 0.0)

#### Dependencies:
- torch.Tensor

#### Notable Return:
- torch.Tensor: Filtered boolean tensor (`keep`)


key: chroma:document
string_value:
CODE:

def get_norm(norm, out_channels):
    if isinstance(norm, str):
        if len(norm) == 0:
            return None
        norm = {
            "BN": BatchNorm2d,
            "GN": lambda channels: nn.GroupNorm(32, channels),
            "nnSyncBN": nn.SyncBatchNorm,  # keep for debugging
            "": lambda x: x,
        }[norm]


SUMMARY:
Here is a concise and focused summary of the given function:

**Summary:** `get_norm` returns a normalized normalization layer based on the input string parameter `norm`.

**Key Details:**

* It uses a dictionary to map string norms to corresponding normalization classes or lambda expressions.
* Available normalization options include:
	+ BatchNorm2d
	+ GroupNorm with a fixed channel grouping factor of 32
	+ SyncBatchNorm (with a comment indicating it's kept for debugging purposes)
	+ An identity lambda expression that leaves the input unchanged
* The function takes two parameters: `norm` and `out_channels`, where `norm` is used to determine which normalization layer to return.

**Primary Purpose:**

This function allows users to dynamically select and configure normalization layers based on a string parameter, making it convenient for experimentation or prototyping.


key: chroma:document
string_value:
CODE:

def _create_grid_offsets(size: List[int], stride: int, offset: float, device):
    grid_height, grid_width = size
    shifts_x = torch.arange(
        offset * stride,
        grid_width * stride,
        step=stride,
        dtype=torch.float32,
        device=device,
    )
    shifts_y = torch.arange(
        offset * stride,
        grid_height * stride,
        step=stride,
        dtype=torch.float32,
        device=device,
    )

    shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x)
    shift_x = shift_x.reshape(-1)
    shift_y = shift_y.reshape(-1)


SUMMARY:
**Function Summary**

Creates grid offsets for 2D indexing.

**Purpose**: Calculates horizontal and vertical displacement values (`shift_x` and `shift_y`) at regular intervals, enabling uniform 2D sampling.

**Key Details**:

- **Shift Calculation**: Computes shifts using PyTorch's tensor meshgrid with stride-based steps.
  
**Method Signatures**:
- `_create_grid_offsets(size: List[int], stride: int, offset: float, device) -> Tuple(torch.Tensor, torch.Tensor)`

**Notable Dependencies/Requirements**:

- PyTorch library
- `torch.arange` and `torch.meshgrid` functions.


key: chroma:document
string_value:
CODE:

def build_backbone(cfg):
    input_shape = ShapeSpec(channels=len(cfg.MODEL.PIXEL_MEAN))
    norm = cfg.RESNETS.NORM
    stem = BasicStem(
        in_channels=input_shape.channels,
        out_channels=cfg.RESNETS.STEM_OUT_CHANNELS,
        norm=norm,
        caffe_maxpool=cfg.MODEL.MAX_POOL,
    )
    freeze_at = cfg.BACKBONE.FREEZE_AT

    if freeze_at >= 1:
        for p in stem.parameters():
            p.requires_grad = False

    out_features = cfg.RESNETS.OUT_FEATURES
    depth = cfg.RESNETS.DEPTH
    num_groups = cfg.RESNETS.NUM_GROUPS
    width_per_group = cfg.RESNETS.WIDTH_PER_GROUP
    bottleneck_channels = num_groups * width_per_group
    in_channels = cfg.RESNETS.STEM_OUT_CHANNELS
    out_channels = cfg.RESNETS.RES2_OUT_CHANNELS
    stride_in_1x1 = cfg.RESNETS.STRIDE_IN_1X1
    res5_dilation = cfg.RESNETS.RES5_DILATION
    assert res5_dilation in {1, 2}, "res5_dilation cannot be {}.".format(res5_dilation)

    num_blocks_per_stage = {50: [3, 4, 6, 3], 101: [3, 4, 23, 3], 152: [3, 8, 36, 3]}[depth]

    stages = []
    out_stage_idx = [{"res2": 2, "res3": 3, "res4": 4, "res5": 5}[f] for f in out_features]
    max_stage_idx = max(out_stage_idx)
    for idx, stage_idx in enumerate(range(2, max_stage_idx + 1)):
        dilation = res5_dilation if stage_idx == 5 else 1
        first_stride = 1 if idx == 0 or (stage_idx == 5 and dilation == 2) else 2
        stage_kargs = {
            "num_blocks": num_blocks_per_stage[idx],
            "first_stride": first_stride,
            "in_channels": in_channels,
            "bottleneck_channels": bottleneck_channels,
            "out_channels": out_channels,
            "num_groups": num_groups,
            "norm": norm,
            "stride_in_1x1": stride_in_1x1,
            "dilation": dilation,
        }

        stage_kargs["block_class"] = BottleneckBlock
        blocks = ResNet.make_stage(**stage_kargs)
        in_channels = out_channels
        out_channels *= 2
        bottleneck_channels *= 2

        if freeze_at >= stage_idx:
            for block in blocks:
                block.freeze()
        stages.append(blocks)



SUMMARY:
**Summary**

This function `build_backbone` constructs a ResNet model as per the given configuration `cfg`. The primary purpose is to create a pyramid-shaped backbone network.

**Key Implementation Details**

*   **Basic Stem**: A basic stem block with `input_shape.channels`, `out_channels=cfg.RESNETS.STEM_OUT_CHANNELS`, and norm layer configured from `cfg.RESNETS.NORM`.
*   **ResNet Stages**: Stages are created using the `num_blocks_per_stage` defined in the model architecture. Each stage consists of a series of ResBlock instances, with specific parameters such as `in_channels`, `bottleneck_channels`, and normalization layers.
*   **Freeze At**: The function takes into account the `freeze_at` parameter and freezes parameters at or below that index in the backbone.

**Notable Dependencies and Requirements**

*   `ResNet.make_stage`: A factory method to create ResBlock instances for each stage.
*   `BottleneckBlock`: A class representing the bottleneck block used in the ResNet architecture.
*   External dependencies include:

    *   **Shapespec**: a module providing utilities for working with shapes and channel numbers.

**Functions and Classes Used**

*   `build_backbone`
*   `BasicStem`
*   `BottleneckBlock`


key: chroma:document
string_value:
CODE:

def find_top_rpn_proposals(
    proposals,
    pred_objectness_logits,
    images,
    image_sizes,
    nms_thresh,
    pre_nms_topk,
    post_nms_topk,
    min_box_side_len,
    training,
):
    """Args:
        proposals (list[Tensor]): (L, N, Hi*Wi*A, 4).
        pred_objectness_logits: tensors of length L.
        nms_thresh (float): IoU threshold to use for NMS
        pre_nms_topk (int): before nms
        post_nms_topk (int): after nms
        min_box_side_len (float): minimum proposal box side
        training (bool): True if proposals are to be used in training,
    Returns:
        results (List[Dict]): stores post_nms_topk object proposals for image i.
    """
    num_images = len(images)
    device = proposals[0].device

    # 1. Select top-k anchor for every level and every image
    topk_scores = []  # #lvl Tensor, each of shape N x topk
    topk_proposals = []
    level_ids = []  # #lvl Tensor, each of shape (topk,)
    batch_idx = torch.arange(num_images, device=device)
    for level_id, proposals_i, logits_i in zip(itertools.count(), proposals, pred_objectness_logits):
        Hi_Wi_A = logits_i.shape[1]
        num_proposals_i = min(pre_nms_topk, Hi_Wi_A)

        # sort is faster than topk (https://github.com/pytorch/pytorch/issues/22812)
        # topk_scores_i, topk_idx = logits_i.topk(num_proposals_i, dim=1)
        logits_i, idx = logits_i.sort(descending=True, dim=1)
        topk_scores_i = logits_i[batch_idx, :num_proposals_i]
        topk_idx = idx[batch_idx, :num_proposals_i]

        # each is N x topk
        topk_proposals_i = proposals_i[batch_idx[:, None], topk_idx]  # N x topk x 4

        topk_proposals.append(topk_proposals_i)
        topk_scores.append(topk_scores_i)
        level_ids.append(torch.full((num_proposals_i,), level_id, dtype=torch.int64, device=device))

    # 2. Concat all levels together
    topk_scores = torch.cat(topk_scores, dim=1)
    topk_proposals = torch.cat(topk_proposals, dim=1)
    level_ids = torch.cat(level_ids, dim=0)

    # if I change to batched_nms, I wonder if this will make a difference
    # 3. For each image, run a per-level NMS, and choose topk results.
    results = []
    for n, image_size in enumerate(image_sizes):
        boxes = topk_proposals[n]
        scores_per_img = topk_scores[n]
        # I will have to take a look at the boxes clip method
        _clip_box(boxes, image_size)
        # filter empty boxes
        keep = _nonempty_boxes(boxes, threshold=min_box_side_len)
        lvl = level_ids
        if keep.sum().item() != len(boxes):
            boxes, scores_per_img, lvl = (
                boxes[keep],
                scores_per_img[keep],
                level_ids[keep],
            )

        keep = batched_nms(boxes, scores_per_img, lvl, nms_thresh)
        keep = keep[:post_nms_topk]

        res = (boxes[keep], scores_per_img[keep])
        results.append(res)

    # I wonder if it would be possible for me to pad all these things.


SUMMARY:
**Summary**

`find_top_rpn_proposals` function:

**Functionality and Purpose**
The function takes in a list of proposal tensors `proposals`, objectness logits `pred_objectness_logits`, image information `images` and sizes, NMS thresholds, and training flags. It returns post-NMS top-k object proposals for each image.

**Main Implementation**

1. Selects top-K anchor predictions for every level and image using a sorting-based approach to outperform the traditional `topk` method.
2. Concatenates all levels together on the last dimension.
3. Runs per-level NMS for each image, filters empty boxes, and selects the top-k results.

**Key Features**

* Uses sorting-based selection of top-K predictions
* Incorporates box clipping and filtering
* Employs batched NMS for efficient computation

**Dependencies and Requirements**

* `proposals` (list[Tensor])
* `pred_objectness_logits` (tensor)
* `images` and `image_sizes`
* NMS thresholds
* Training flags


key: chroma:document
string_value:
CODE:

def subsample_labels(labels, num_samples, positive_fraction, bg_label):
    """
    Returns:
        pos_idx, neg_idx (Tensor):
            1D vector of indices. The total length of both is `num_samples` or fewer.
    """
    positive = torch.nonzero((labels != -1) & (labels != bg_label)).squeeze(1)
    negative = torch.nonzero(labels == bg_label).squeeze(1)

    num_pos = int(num_samples * positive_fraction)
    # protect against not enough positive examples
    num_pos = min(positive.numel(), num_pos)
    num_neg = num_samples - num_pos
    # protect against not enough negative examples
    num_neg = min(negative.numel(), num_neg)

    # randomly select positive and negative examples
    perm1 = torch.randperm(positive.numel(), device=positive.device)[:num_pos]
    perm2 = torch.randperm(negative.numel(), device=negative.device)[:num_neg]

    pos_idx = positive[perm1]
    neg_idx = negative[perm2]


SUMMARY:
**Function Summary: subsample_labels**

Subsamples `labels` to generate pairs of indexed positive and negative samples.

### Main Purpose
Randomly selects a specified fraction (`positive_fraction`) of positive examples from the labeled data (excluding background label `bg_label`) and corresponding negative examples.

### Key Implementation Details

* Uses `torch.nonzero()` to extract indices of non-background labels.
* Randomly permutes these indices using `torch.randperm()` to create subsamples.
* Ensures sufficient positive and negative samples are selected by limiting the size of the subsample.

### Notable Dependencies
* PyTorch tensor operations (`torch.nonzero`, `torch.randperm`)

### Return Value
- Two tensors (1D vector) containing indexes of positive and negative samples, respectively.


key: chroma:document
string_value:
CODE:

def add_ground_truth_to_proposals(gt_boxes, proposals):


SUMMARY:
**Summary:**

`add_ground_truth_to_proposals`: A missing method that takes in `gt_boxes` and `proposals`. Its purpose is unclear as it raises a `NotImplementedError`. 

The main functionality remains undefined, but it is expected to add ground truth information (gt_boxes) to proposal boxes (proposals).


key: chroma:document
string_value:
CODE:

def add_ground_truth_to_proposals_single_image(gt_boxes, proposals):


SUMMARY:
**Summary**

`add_ground_truth_to_proposals_single_image`:

**Purpose**: Modifies image proposals by adding ground truth bounding box information.

**Functionality**: This function expects two inputs: `gt_boxes` (ground truth boxes) and `proposals` (image proposal boxes). It raises a `NotImplementedError` as the actual implementation is not provided, indicating that it is intended to be overridden or implemented by child classes.

**Implementation Details**: N/A

**Dependencies**: N/A


key: chroma:document
string_value:
CODE:

def _fmt_box_list(box_tensor, batch_index: int):
    repeated_index = torch.full(
        (len(box_tensor), 1),
        batch_index,
        dtype=box_tensor.dtype,
        device=box_tensor.device,
    )


SUMMARY:
**Summary:**

`_fmt_box_list` function takes `box_tensor` and `batch_index` as input. It repeats the `batch_index` along the last dimension of `box_tensor` using `torch.full`, then concatenates it with `box_tensor` along the last dimension, resulting in a tensor where each row contains `[batch_index, ...]`.

**Key implementation details:**

* Repeats `batch_index` using `torch.full`
* Concatenates repeated `batch_index` with `box_tensor` along the last dimension

**Notable dependencies or requirements:**

* Uses PyTorch library
* Assumes that `box_tensor` is a tensor of shape `(n, d)` (not specified in the provided code snippet)


key: chroma:document
string_value:
CODE:

def convert_boxes_to_pooler_format(box_lists: List[torch.Tensor]):
    pooler_fmt_boxes = torch.cat(
        [_fmt_box_list(box_list, i) for i, box_list in enumerate(box_lists)],
        dim=0,
    )


SUMMARY:
**Summary**

`convert_boxes_to_pooler_format` function:

* Takes a list of 3D `torch.Tensor` objects representing bounding boxes (`box_lists`)
* Converts each box list to the Pooler format using `_fmt_box_list(box_list, i)` and concatenates along dimension 0
* Returns the merged, Pooler-formatted tensor (`pooler_fmt_boxes`)

**Key Implementation Details**

* Uses list comprehension with `enumerate` to iterate over `box_lists`
* Utilizes `torch.cat` to concatenate box lists at dim=0

**Notable Dependencies and Requirements**

* Lists of 3D `torch.Tensor` objects
* `_fmt_box_list` function (not shown) for converting individual box lists to Pooler format


key: chroma:document
string_value:
CODE:

def assign_boxes_to_levels(
    box_lists: List[torch.Tensor],
    min_level: int,
    max_level: int,
    canonical_box_size: int,
    canonical_level: int,
):
    box_sizes = torch.sqrt(torch.cat([boxes.area() for boxes in box_lists]))
    # Eqn.(1) in FPN paper
    level_assignments = torch.floor(canonical_level + torch.log2(box_sizes / canonical_box_size + 1e-8))
    # clamp level to (min, max), in case the box size is too large or too small
    # for the available feature maps
    level_assignments = torch.clamp(level_assignments, min=min_level, max=max_level)


SUMMARY:
**Summary: `assign_boxes_to_levels` Function**

**Purpose and Functionality**
The `assign_boxes_to_levels` function assigns boxes to levels based on their area, scaling the area by the canonical box size, and then mapping it to levels using the FPN paper's logarithmic scale.

**Key Implementation Details**

* Calculates box sizes as square roots of areas from `box_lists`.
* Applies the FPN paper's equation (`log2(box_size / canon_box_size + 1e-8)`) to assign boxes to levels.
* Clips level assignments within the range `[min_level, max_level]` using `torch.clamp`.
* Returns the resulting level assignments as integers.

**Notable Dependencies**

* `box_lists`: list of tensors representing box containers (e.g., bounding boxes).
* `torch`: PyTorch library for tensor computations and logarithmic calculations.
* No external dependencies or interfaces are required.

Note: This summary focuses on the main purpose, key implementation details, and notable dependencies. It aims to provide a concise overview of the function's functionality while preserving technical accuracy and precision.


key: chroma:document
string_value:
CODE:
# Helper Classes
class _NewEmptyTensorOp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, new_shape):
        ctx.shape = x.shape
        return x.new_empty(new_shape)

    @staticmethod
    def backward(ctx, grad):
        shape = ctx.shape


SUMMARY:
**Summary**

The `_NewEmptyTensorOp` class is a PyTorch autograd function that creates a new tensor with the specified shape. It takes two inputs:

- `x`: The original tensor
- `new_shape`: The desired shape of the new tensor

The class has two methods:

* `forward(ctx, x, new_shape)`: Creates a new tensor with the specified shape and assigns it to `ctx.shape`.
* `backward(ctx, grad)`: Returns the backward pass (i.e., creating a new tensor with the same shape as the original input).

**Key Implementation Details**

* Uses `torch.autograd.Function` to enable autograd computation graph.
* Assigns the original shape of `x` to `ctx.shape` for use in the backward pass.

**Notable Dependencies and Requirements**

- Requires PyTorch (` torch.autograd.Function` )
- Inherits from no specific classes, but uses `torch.nn.Tensor` (implied by `x.new_empty(new_shape)`).


key: chroma:document
string_value:
CODE:

class ShapeSpec(namedtuple("_ShapeSpec", ["channels", "height", "width", "stride"])):
    def __new__(cls, *, channels=None, height=None, width=None, stride=None):


SUMMARY:
**Summary:**

`ShapeSpec` class defines a named tuple to represent image shape specifications. It has four attributes: `channels`, `height`, `width`, and `stride`. The constructor initializes these attributes with optional keyword parameters.

**Key Implementation Details:**

* Uses the `namedtuple` function to create a tuple with explicit attribute names.
* Has a single constructor method (`__new__`) that initializes its attributes using the `super().__new__(cls, ...)` call.

**Notable Dependencies or Requirements:**

* None explicitly stated. However, it can be assumed that this class is part of a larger module that imports the `namedtuple` function from the `collections` module or similar.


key: chroma:document
string_value:
CODE:

class Box2BoxTransform:
    """
    This R-CNN transformation scales the box's width and height
    by exp(dw), exp(dh) and shifts a box's center by the offset
    (dx * width, dy * height).
    """

    def __init__(self, weights: Tuple[float, float, float, float], scale_clamp: float = None):
        """
        Args:
            weights (4-element tuple): Scaling factors that are applied to the
                (dx, dy, dw, dh) deltas. In Fast R-CNN, these were originally set
                such that the deltas have unit variance; now they are treated as
                hyperparameters of the system.
            scale_clamp (float): When predicting deltas, the predicted box scaling
                factors (dw and dh) are clamped such that they are <= scale_clamp.
        """
        self.weights = weights
        if scale_clamp is not None:
            self.scale_clamp = scale_clamp
        else:
            """
            Value for clamping large dw and dh predictions.
            The heuristic is that we clamp such that dw and dh are no larger
            than what would transform a 16px box into a 1000px box
            (based on a small anchor, 16px, and a typical image size, 1000px).
            """
            self.scale_clamp = math.log(1000.0 / 16)

    def get_deltas(self, src_boxes, target_boxes):
        """
        Get box regression transformation deltas (dx, dy, dw, dh) that can be used
        to transform the `src_boxes` into the `target_boxes`. That is, the relation
        ``target_boxes == self.apply_deltas(deltas, src_boxes)`` is true (unless
        any delta is too large and is clamped).
        Args:
            src_boxes (Tensor): source boxes, e.g., object proposals
            target_boxes (Tensor): target of the transformation, e.g., ground-truth
                boxes.
        """
        assert isinstance(src_boxes, torch.Tensor), type(src_boxes)
        assert isinstance(target_boxes, torch.Tensor), type(target_boxes)

        src_widths = src_boxes[:, 2] - src_boxes[:, 0]
        src_heights = src_boxes[:, 3] - src_boxes[:, 1]
        src_ctr_x = src_boxes[:, 0] + 0.5 * src_widths
        src_ctr_y = src_boxes[:, 1] + 0.5 * src_heights

        target_widths = target_boxes[:, 2] - target_boxes[:, 0]
        target_heights = target_boxes[:, 3] - target_boxes[:, 1]
        target_ctr_x = target_boxes[:, 0] + 0.5 * target_widths
        target_ctr_y = target_boxes[:, 1] + 0.5 * target_heights

        wx, wy, ww, wh = self.weights
        dx = wx * (target_ctr_x - src_ctr_x) / src_widths
        dy = wy * (target_ctr_y - src_ctr_y) / src_heights
        dw = ww * torch.log(target_widths / src_widths)
        dh = wh * torch.log(target_heights / src_heights)

        deltas = torch.stack((dx, dy, dw, dh), dim=1)
        assert (src_widths > 0).all().item(), "Input boxes to Box2BoxTransform are not valid!"
        return deltas

    def apply_deltas(self, deltas, boxes):
        """
        Apply transformation `deltas` (dx, dy, dw, dh) to `boxes`.
        Args:
            deltas (Tensor): transformation deltas of shape (N, k*4), where k >= 1.
                deltas[i] represents k potentially different class-specific
                box transformations for the single box boxes[i].
            boxes (Tensor): boxes to transform, of shape (N, 4)
        """
        boxes = boxes.to(deltas.dtype)

        widths = boxes[:, 2] - boxes[:, 0]
        heights = boxes[:, 3] - boxes[:, 1]
        ctr_x = boxes[:, 0] + 0.5 * widths
        ctr_y = boxes[:, 1] + 0.5 * heights

        wx, wy, ww, wh = self.weights
        dx = deltas[:, 0::4] / wx
        dy = deltas[:, 1::4] / wy
        dw = deltas[:, 2::4] / ww
        dh = deltas[:, 3::4] / wh

        # Prevent sending too large values into torch.exp()
        dw = torch.clamp(dw, max=self.scale_clamp)
        dh = torch.clamp(dh, max=self.scale_clamp)

        pred_ctr_x = dx * widths[:, None] + ctr_x[:, None]
        pred_ctr_y = dy * heights[:, None] + ctr_y[:, None]
        pred_w = torch.exp(dw) * widths[:, None]
        pred_h = torch.exp(dh) * heights[:, None]

        pred_boxes = torch.zeros_like(deltas)
        pred_boxes[:, 0::4] = pred_ctr_x - 0.5 * pred_w  # x1
        pred_boxes[:, 1::4] = pred_ctr_y - 0.5 * pred_h  # y1
        pred_boxes[:, 2::4] = pred_ctr_x + 0.5 * pred_w  # x2
        pred_boxes[:, 3::4] = pred_ctr_y + 0.5 * pred_h  # y2


SUMMARY:
**Summary**

The `Box2BoxTransform` class is used for transforming source boxes into target boxes through box regression transformation deltas (dx, dy, dw, dh). It applies scaling factors (weights) to these deltas and shifts the centers accordingly.

**Key Features**

* Scales source width and height by exp(dw) and exp(dh)
* Shifts center by offset (dx * width, dy * height)
* Applies weight clamping for large predictions
* Returns transformed boxes with updated coordinates

**Implementation Details**

* `get_deltas` method computes transformation deltas given source and target boxes
	+ Computes dx, dy, dw, dh using inverse scaling factors
	+ Clips values to prevent sending too large numbers into torch.exp()
* `apply_deltas` method applies transformation deltas to given boxes
	+ Scales width and height using exp(dw) and exp(dh)
	+ Shifts center by offset (dx * width, dy * height)

**Notable Dependencies**

* PyTorch for tensor operations

Note: This summary preserves the exact names of classes, functions, and variables.


key: chroma:document
string_value:
CODE:

class Matcher:
    """
    This class assigns to each predicted "element" (e.g., a box) a ground-truth
    element. Each predicted element will have exactly zero or one matches; each
    ground-truth element may be matched to zero or more predicted elements.
    The matching is determined by the MxN match_quality_matrix, that characterizes
    how well each (ground-truth, prediction)-pair match each other. For example,
    if the elements are boxes, this matrix may contain box intersection-over-union
    overlap values.
    The matcher returns (a) a vector of length N containing the index of the
    ground-truth element m in [0, M) that matches to prediction n in [0, N).
    (b) a vector of length N containing the labels for each prediction.
    """

    def __init__(
        self,
        thresholds: List[float],
        labels: List[int],
        allow_low_quality_matches: bool = False,
    ):
        """
        Args:
            thresholds (list): a list of thresholds used to stratify predictions
                into levels.
            labels (list): a list of values to label predictions belonging at
                each level. A label can be one of {-1, 0, 1} signifying
                {ignore, negative class, positive class}, respectively.
            allow_low_quality_matches (bool): if True, produce additional matches or predictions with maximum match quality lower than high_threshold.
                For example, thresholds = [0.3, 0.5] labels = [0, -1, 1] All predictions with iou < 0.3 will be marked with 0 and
                thus will be considered as false positives while training. All predictions with 0.3 <= iou < 0.5 will be marked with -1 and
                thus will be ignored. All predictions with 0.5 <= iou will be marked with 1 and thus will be considered as true positives.
        """
        thresholds = thresholds[:]
        assert thresholds[0] > 0
        thresholds.insert(0, -float("inf"))
        thresholds.append(float("inf"))
        assert all(low <= high for (low, high) in zip(thresholds[:-1], thresholds[1:]))
        assert all(label_i in [-1, 0, 1] for label_i in labels)
        assert len(labels) == len(thresholds) - 1
        self.thresholds = thresholds
        self.labels = labels
        self.allow_low_quality_matches = allow_low_quality_matches

    def __call__(self, match_quality_matrix):
        """
        Args:
            match_quality_matrix (Tensor[float]): an MxN tensor, containing the pairwise quality between M ground-truth elements and N predicted
                elements. All elements must be >= 0 (due to the us of `torch.nonzero` for selecting indices in :meth:`set_low_quality_matches_`).
        Returns:
            matches (Tensor[int64]): a vector of length N, where matches[i] is a matched ground-truth index in [0, M)
            match_labels (Tensor[int8]): a vector of length N, where pred_labels[i] indicates true or false positive or ignored
        """
        assert match_quality_matrix.dim() == 2
        if match_quality_matrix.numel() == 0:
            default_matches = match_quality_matrix.new_full((match_quality_matrix.size(1),), 0, dtype=torch.int64)
            # When no gt boxes exist, we define IOU = 0 and therefore set labels
            # to `self.labels[0]`, which usually defaults to background class 0
            # To choose to ignore instead,
            # can make labels=[-1,0,-1,1] + set appropriate thresholds
            default_match_labels = match_quality_matrix.new_full(
                (match_quality_matrix.size(1),), self.labels[0], dtype=torch.int8
            )
            return default_matches, default_match_labels

        assert torch.all(match_quality_matrix >= 0)

        # match_quality_matrix is M (gt) x N (predicted)
        # Max over gt elements (dim 0) to find best gt candidate for each prediction
        matched_vals, matches = match_quality_matrix.max(dim=0)

        match_labels = matches.new_full(matches.size(), 1, dtype=torch.int8)

        for l, low, high in zip(self.labels, self.thresholds[:-1], self.thresholds[1:]):
            low_high = (matched_vals >= low) & (matched_vals < high)
            match_labels[low_high] = l

        if self.allow_low_quality_matches:
            self.set_low_quality_matches_(match_labels, match_quality_matrix)

        return matches, match_labels

    def set_low_quality_matches_(self, match_labels, match_quality_matrix):
        """
        Produce additional matches for predictions that have only low-quality matches.
        Specifically, for each ground-truth G find the set of predictions that have
        maximum overlap with it (including ties); for each prediction in that set, if
        it is unmatched, then match it to the ground-truth G.
        This function implements the RPN assignment case (i)
        in Sec. 3.1.2 of Faster R-CNN.
        """
        # For each gt, find the prediction with which it has highest quality
        highest_quality_foreach_gt, _ = match_quality_matrix.max(dim=1)
        # Find the highest quality match available, even if it is low, including ties.
        # Note that the matches qualities must be positive due to the use of
        # `torch.nonzero`.
        of_quality_inds = match_quality_matrix == highest_quality_foreach_gt[:, None]
        if of_quality_inds.dim() == 0:
            (_, pred_inds_with_highest_quality) = of_quality_inds.unsqueeze(0).nonzero().unbind(1)
        else:
            (_, pred_inds_with_highest_quality) = of_quality_inds.nonzero().unbind(1)


SUMMARY:
**Matcher Class Summary**

The `Matcher` class assigns predicted "elements" (e.g., boxes) to ground-truth elements based on the provided `match_quality_matrix`. It determines the best match for each prediction by finding the maximum quality overlap with ground-truth elements.

**Key Functionality:**

1. `__init__`: Initializes the matcher with thresholds, labels, and options for low-quality matches.
2. `__call__`: Takes a `match_quality_matrix` as input and returns:
	* `matches`: A vector of length N containing the index of the ground-truth element that matches each prediction.
	* `match_labels`: A tensor of length N indicating whether a prediction is:
		+ True positive (1)
		+ False negative (-1)
		+ Ignore (-2, not implemented in provided code)

**Important Methods:**

1. `set_low_quality_matches_`: Implements the RPN assignment case and produces additional matches for predictions with low-quality matches.

**Dependencies:**

* `match_quality_matrix` (MxN tensor) of pairwise quality between ground-truth elements and predicted elements.
* `labels` (list) of values to label predictions belonging at each level.
* `thresholds` (list) of thresholds used to stratify predictions into levels.
* External dependencies:
	+ Torch library for tensor operations.

**Notable Interfaces:**

* `max`: Used to find the maximum quality overlap between ground-truth elements and predicted elements.
* `nonzero`: Used to select indices in `match_quality_matrix` based on high-quality matches.


key: chroma:document
string_value:
CODE:

class RPNOutputs:
    def __init__(
        self,
        box2box_transform,
        anchor_matcher,
        batch_size_per_image,
        positive_fraction,
        images,
        pred_objectness_logits,
        pred_anchor_deltas,
        anchors,
        boundary_threshold=0,
        gt_boxes=None,
        smooth_l1_beta=0.0,
    ):
        """
        Args:
            box2box_transform (Box2BoxTransform): :class:`Box2BoxTransform` instance for anchor-proposal transformations.
            anchor_matcher (Matcher): :class:`Matcher` instance for matching anchors to ground-truth boxes; used to determine training labels.
            batch_size_per_image (int): number of proposals to sample when training
            positive_fraction (float): target fraction of sampled proposals that should be positive
            images (ImageList): :class:`ImageList` instance representing N input images
            pred_objectness_logits (list[Tensor]): A list of L elements. Element i is a tensor of shape (N, A, Hi, W)
            pred_anchor_deltas (list[Tensor]): A list of L elements. Element i is a tensor of shape (N, A*4, Hi, Wi)
            anchors (list[torch.Tensor]): nested list of boxes. anchors[i][j] at (n, l) stores anchor array for feature map l
            boundary_threshold (int): if >= 0, then anchors that extend beyond the image boundary by more than boundary_thresh are not used in training.
            gt_boxes (list[Boxes], optional): A list of N elements.
            smooth_l1_beta (float): The transition point between L1 and L2 lossn. When set to 0, the loss becomes L1. When +inf, it is ignored
        """
        self.box2box_transform = box2box_transform
        self.anchor_matcher = anchor_matcher
        self.batch_size_per_image = batch_size_per_image
        self.positive_fraction = positive_fraction
        self.pred_objectness_logits = pred_objectness_logits
        self.pred_anchor_deltas = pred_anchor_deltas

        self.anchors = anchors
        self.gt_boxes = gt_boxes
        self.num_feature_maps = len(pred_objectness_logits)
        self.num_images = len(images)
        self.boundary_threshold = boundary_threshold
        self.smooth_l1_beta = smooth_l1_beta

    def _get_ground_truth(self):
        raise NotImplementedError()

    def predict_proposals(self):
        # pred_anchor_deltas: (L, N, ? Hi, Wi)
        # anchors:(N, L, -1, B)
        # here we loop over specific feature map, NOT images
        proposals = []
        anchors = self.anchors.transpose(0, 1)
        for anchors_i, pred_anchor_deltas_i in zip(anchors, self.pred_anchor_deltas):
            B = anchors_i.size(-1)
            N, _, Hi, Wi = pred_anchor_deltas_i.shape
            anchors_i = anchors_i.flatten(start_dim=0, end_dim=1)
            pred_anchor_deltas_i = pred_anchor_deltas_i.view(N, -1, B, Hi, Wi).permute(0, 3, 4, 1, 2).reshape(-1, B)
            proposals_i = self.box2box_transform.apply_deltas(pred_anchor_deltas_i, anchors_i)
            # Append feature map proposals with shape (N, Hi*Wi*A, B)
            proposals.append(proposals_i.view(N, -1, B))
        proposals = torch.stack(proposals)
        return proposals

    def predict_objectness_logits(self):
        """
        Returns:
            pred_objectness_logits (list[Tensor]) -> (N, Hi*Wi*A).
        """
        pred_objectness_logits = [
            # Reshape: (N, A, Hi, Wi) -> (N, Hi, Wi, A) -> (N, Hi*Wi*A)
            score.permute(0, 2, 3, 1).reshape(self.num_images, -1)
            for score in self.pred_objectness_logits
        ]


SUMMARY:
**RPNOutputs Class Summary**

**Primary Purpose:** The `RPNOutputs` class is responsible for processing anchor proposals and generating objectness scores, ultimately used for regression tasks within object detection frameworks.

**Key Implementation Details:**

*   The class utilizes two primary methods:
    *   `predict_proposals`: generates anchor proposals by applying the box2box transformation to predicted anchor deltas.
    *   `predict_objectness_logits`: computes objectness scores using individual logits tensors from `self.pred_objectness_logits`.

**Notable Dependencies and Requirements:**

*   `Box2BoxTransform` instance (`self.box2box_transform`)
*   `Matcher` instance (`self.anchor_matcher`)
*   `ImageList` instance (`images`)
*   Individual `torch.Tensor` instances for prediction logits (`self.pred_objectness_logits`, `self.pred_anchor_deltas`)


key: chroma:document
string_value:
CODE:
# Main Classes
class Conv2d(nn.Conv2d):
    def __init__(self, *args, **kwargs):
        norm = kwargs.pop("norm", None)
        activation = kwargs.pop("activation", None)
        super().__init__(*args, **kwargs)

        self.norm = norm
        self.activation = activation

    def forward(self, x):
        if x.numel() == 0 and self.training:
            assert not isinstance(self.norm, nn.SyncBatchNorm)
        if x.numel() == 0:
            assert not isinstance(self.norm, nn.GroupNorm)
            output_shape = [
                (i + 2 * p - (di * (k - 1) + 1)) // s + 1
                for i, p, di, k, s in zip(
                    x.shape[-2:],
                    self.padding,
                    self.dilation,
                    self.kernel_size,
                    self.stride,
                )
            ]
            output_shape = [x.shape[0], self.weight.shape[0]] + output_shape
            empty = _NewEmptyTensorOp.apply(x, output_shape)
            if self.training:
                _dummy = sum(x.view(-1)[0] for x in self.parameters()) * 0.0
                return empty + _dummy
            else:
                return empty

        x = super().forward(x)
        if self.norm is not None:
            x = self.norm(x)
        if self.activation is not None:
            x = self.activation(x)


SUMMARY:
**Class Summary: Conv2d**

This class inherits from PyTorch's `nn.Conv2d` and customizes its behavior:

### Primary Purpose

Conv2d is a 2D convolutional neural network layer that applies spatial convolutions on input images.

### Key Implementation Details

*   The layer takes optional norm (`norm`) and activation (`activation`) layers in its `__init__` method.
*   In the `forward` method:
    *   It checks if the input tensor is empty and only initializes bias-term calculations when training (`self.training`).
    *   It calculates the output shape based on padding, dilation, kernel size, and stride.
    *   If not in evaluation mode (`self.training`), it adds a dummy bias term to enable PyTorch tensor computation graph.

### Notable Dependencies or Requirements

*   `nn.Conv2d`: base 2D convolutional neural network layer class.
*   `_NewEmptyTensorOp` and `_dummy`: custom PyTorch tensor operations, likely part of the implementation using Python scripts (Python is not used directly in this class).

### Key Method Signatures and Interfaces

*   `__init__(norm, activation)`: inherits from `nn.Conv2d.__init__`.
*   `forward(x)`: custom method that expands on the PyTorch's default behavior with the addition of checking dimensions for norm layer usage and computing output shape.

This implementation prioritizes the ability to conditionally add bias to convolutional layers under training while preserving stability for inference modes.


key: chroma:document
string_value:
CODE:

class LastLevelMaxPool(nn.Module):
    """
    This module is used in the original FPN to generate a downsampled P6 feature from P5.
    """

    def __init__(self):
        super().__init__()
        self.num_levels = 1
        self.in_feature = "p5"

    def forward(self, x):


SUMMARY:
**Summary**

The `LastLevelMaxPool` class is a PyTorch neural network module used in the Focal Pyramid Network (FPN) architecture. It downsamples the P5 feature using Max Pooling.

**Key Details**

* **Implementation**: Uses `nn.functional.max_pool2d` to apply max pooling with kernel size 1, stride 2, and zero padding.
* **Input/Output**: Takes a single input `x`. Returns a list of downsampled features.

**Notable Parameters/Dependencies**

* `self.num_levels = 1`: Specifies the number of downsampling levels (Fixed at 1).
* `nn.Module` and `nn.functional.max_pool2d` are used.
* No external dependencies listed.


key: chroma:document
string_value:
CODE:

class LastLevelP6P7(nn.Module):
    """
    This module is used in RetinaNet to generate extra layers, P6 and P7 from C5 feature.
    """

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.num_levels = 2
        self.in_feature = "res5"
        self.p6 = nn.Conv2d(in_channels, out_channels, 3, 2, 1)
        self.p7 = nn.Conv2d(out_channels, out_channels, 3, 2, 1)

    def forward(self, c5):
        p6 = self.p6(c5)
        p7 = self.p7(nn.functional.relu(p6))


SUMMARY:
**Class Summary: LastLevelP6P7**

This module is used in RetinaNet to generate extra layers (P6 and P7) from the C5 feature.

### Main Purpose and Functionality

The primary purpose of this class is to create two convolutional layers (`p6` and `p7`) that take a input feature `c5` and produce two corresponding output features. The layer uses ReLU activation function from PyTorch's `nn.functional.relu()`.

### Key Implementation Details

*   **Conv2d Layers:**
    *   `self.p6`: A convolutional layer with `in_channels=in_channels`, `out_channels=out_channels`, kernel size 3, stride 2, and padding 1.
    *   `self.p7`: Another convolutional layer with the same parameters as `self.p6` but takes `out_channels=out_channels` as input.

*   **Forward Method:**
    + Takes `c5` feature as input
    + Passes it through `self.p6` to obtain P6 feature
    + Applies ReLU activation function to P6 feature and passes it through `self.p7` to obtain P7 feature
    + Returns both P6 and P7 features as a list

### Dependencies/Requirements

*   PyTorch's neural network module (`nn}`)
*   PyTorch's convolutional layer class (Conv2d)
*   ReLU activation function from PyTorch's nn.functional module


key: chroma:document
string_value:
CODE:

class BasicStem(nn.Module):
    def __init__(self, in_channels=3, out_channels=64, norm="BN", caffe_maxpool=False):
        super().__init__()
        self.conv1 = Conv2d(
            in_channels,
            out_channels,
            kernel_size=7,
            stride=2,
            padding=3,
            bias=False,
            norm=get_norm(norm, out_channels),
        )
        self.caffe_maxpool = caffe_maxpool
        # use pad 1 instead of pad zero

    def forward(self, x):
        x = self.conv1(x)
        x = nn.functional.relu_(x)
        if self.caffe_maxpool:
            x = nn.functional.max_pool2d(x, kernel_size=3, stride=2, padding=0, ceil_mode=True)
        else:
            x = nn.functional.max_pool2d(x, kernel_size=3, stride=2, padding=1)
        return x

    @property
    def out_channels(self):
        return self.conv1.out_channels

    @property
    def stride(self):


SUMMARY:
**Class Summary: BasicStem**

The `BasicStem` class is a PyTorch neural network module that applies a convolutional and max pooling operation to input data. It features:

*   A single convolution layer (`conv1`) with:
    *   3 input channels, 64 output channels
    *   kernel size of 7, stride of 2, padding of 3
    *   No bias term and norm layer (using `get_norm` function) applied to the output

*   Optional addition of Caffe-style max pooling using `caffe_maxpool`
*   Returns the output after convolution and activation.

**Key Method Signatures:**

*   `forward(x)`: Input data `x` is passed through the convolutional operation and then a rectified linear unit (ReLU) activation.
*   `out_channels`: Property getter for the output channels of `conv1`.
*   `stride`: Constant property return value 4.

**Notable Attributes/Dependencies:**

*   Conv2d layer (`self.conv1`) used in the forward pass
*   norm layer (`get_norm` function) applied to the output of `conv1`
*   External dependency on functions from torch.nn and nn.functional


key: chroma:document
string_value:
CODE:

class ResNetBlockBase(nn.Module):
    def __init__(self, in_channels, out_channels, stride):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.stride = stride

    def freeze(self):
        for p in self.parameters():
            p.requires_grad = False


SUMMARY:
**ResNetBlockBase Class Summary**

* **Main Purpose and Functionality**: A base class for ResNet blocks, which defines the architecture and parameters of a residual block.
	+ Input: `in_channels`, `out_channels`, `stride`
* **Key Implementation Details**:
	+ Uses PyTorch's built-in `nn.Module` to define the neural network layer.
	+ Defines the `freeze` method, which sets all model weights to non-trainable by setting `requires_grad=False`.
* **Notable Dependencies/Requirements**:
	+ PyTorch library (`nn.Module`)
	+ None
* **Interfaces**: Inherited from `nn.Module`

This summary focuses on the key aspects of the code, highlighting the main purpose and functionality of the class, as well as its implementation details.


key: chroma:document
string_value:
CODE:

class BottleneckBlock(ResNetBlockBase):
    def __init__(
        self,
        in_channels,
        out_channels,
        bottleneck_channels,
        stride=1,
        num_groups=1,
        norm="BN",
        stride_in_1x1=False,
        dilation=1,
    ):
        super().__init__(in_channels, out_channels, stride)

        if in_channels != out_channels:
            self.shortcut = Conv2d(
                in_channels,
                out_channels,
                kernel_size=1,
                stride=stride,
                bias=False,
                norm=get_norm(norm, out_channels),
            )
        else:
            self.shortcut = None

        # The original MSRA ResNet models have stride in the first 1x1 conv
        # The subsequent fb.torch.resnet and Caffe2 ResNe[X]t implementations have
        # stride in the 3x3 conv
        stride_1x1, stride_3x3 = (stride, 1) if stride_in_1x1 else (1, stride)

        self.conv1 = Conv2d(
            in_channels,
            bottleneck_channels,
            kernel_size=1,
            stride=stride_1x1,
            bias=False,
            norm=get_norm(norm, bottleneck_channels),
        )

        self.conv2 = Conv2d(
            bottleneck_channels,
            bottleneck_channels,
            kernel_size=3,
            stride=stride_3x3,
            padding=1 * dilation,
            bias=False,
            groups=num_groups,
            dilation=dilation,
            norm=get_norm(norm, bottleneck_channels),
        )

        self.conv3 = Conv2d(
            bottleneck_channels,
            out_channels,
            kernel_size=1,
            bias=False,
            norm=get_norm(norm, out_channels),
        )

    def forward(self, x):
        out = self.conv1(x)
        out = nn.functional.relu_(out)

        out = self.conv2(out)
        out = nn.functional.relu_(out)

        out = self.conv3(out)

        if self.shortcut is not None:
            shortcut = self.shortcut(x)
        else:
            shortcut = x

        out += shortcut
        out = nn.functional.relu_(out)


SUMMARY:
**Summary: BottleneckBlock Class**

The `BottleneckBlock` class is a subclass of `ResNetBlockBase`, implementing the bottleneck structure of ResNet models. Its primary purpose is to reduce computational cost by downsampling the input feature map while maintaining channel numbers.

**Key Implementation Details:**

* The block consists of three convolutional layers with ReLU activation:
	+ `conv1`: 1x1 convolution with stride in first layer
	+ `conv2`: 3x3 convolution with stride and dilation based on `stride_in_1x1`
	+ `conv3`: 1x1 convolution with out-channel size matching the input
* A shortcut connection is added if channel numbers are different, utilizing a shared 1x1 convolution (`self.shortcut`).

**Notable Dependencies and Requirements:**

* `Conv2d`: Shared convolutional layer used throughout the block
* `get_norm`: Function to compute normalization layers (e.g., BatchNorm)
* `nn.functional`: Import of PyTorch activation functions (ReLU)

**Key Method Signatures:**

* `__init__`: Initializes the bottleneck block with given parameters
* `forward`: Defines the forward pass through the block, applying ReLU activations and outputting the final result


key: chroma:document
string_value:
CODE:

class Backbone(nn.Module, metaclass=ABCMeta):
    def __init__(self):
        super().__init__()

    @abstractmethod
    def forward(self):
        pass

    @property
    def size_divisibility(self):
        """
        Some backbones require the input height and width to be divisible by a specific integer. This is
        typically true for encoder / decoder type networks with lateral connection (e.g., FPN) for which feature maps need to match
        dimension in the "bottom up" and "top down" paths. Set to 0 if no specific input size divisibility is required.
        """
        return 0

    def output_shape(self):
        return {
            name: ShapeSpec(
                channels=self._out_feature_channels[name],
                stride=self._out_feature_strides[name],
            )
            for name in self._out_features
        }

    @property
    def out_features(self):
        """deprecated"""
        return self._out_features

    @property
    def out_feature_strides(self):
        """deprecated"""
        return {f: self._out_feature_strides[f] for f in self._out_features}

    @property
    def out_feature_channels(self):
        """deprecated"""


SUMMARY:
**Backbone Class Summary**

### Purpose and Functionality

This class `Backbone` defines a basic backbone network architecture, acting as the foundation for further layers in the neural network.

The primary purpose of this class is to serve as an abstraction layer, providing methods to handle feature extraction from input data through convolutional operations. It serves as a parent class for other modules and encoders in the network.

### Key Implementation Details

*   **Abstract `forward` Method**: This method is intended to be overridden by child classes, where it will define how the input data flows through the network.
*   **Output Shape Calculation**: The `output_shape` method calculates the output shape and features of the backbone at each layer. It uses shapes from the `_out_feature_channels`, `_out_feature_strides`, and `name` variables to define the channel-wise specifications.
*   **Size Divisibility Property**: This class ensures that feature maps produced by its layers must have a specific size divisibility, set to 0 by default.

### Notable Dependencies or Requirements

*   It uses PyTorch's built-in module management (`nn.Module`) and provides `ABCMeta` as the metaclass for customizing abstract base classes.
*   External dependencies include:
    *   `ShapeSpec`: A class from PyTorch that describes a set of channels with specified strides (defined in another file).
        It is not explicitly referenced here but has been used for specifying output shapes.


key: chroma:document
string_value:
CODE:

class ResNet(Backbone):
    def __init__(self, stem, stages, num_classes=None, out_features=None):
        """
        Args:
            stem (nn.Module): a stem module
            stages (list[list[ResNetBlock]]): several (typically 4) stages, each contains multiple :class:`ResNetBlockBase`.
            num_classes (None or int): if None, will not perform classification.
            out_features (list[str]): name of the layers whose outputs should be returned in forward. Can be anything in:
            "stem", "linear", or "res2" ... If None, will return the output of the last layer.
        """
        super(ResNet, self).__init__()
        self.stem = stem
        self.num_classes = num_classes

        current_stride = self.stem.stride
        self._out_feature_strides = {"stem": current_stride}
        self._out_feature_channels = {"stem": self.stem.out_channels}

        self.stages_and_names = []
        for i, blocks in enumerate(stages):
            for block in blocks:
                assert isinstance(block, ResNetBlockBase), block
                curr_channels = block.out_channels
            stage = nn.Sequential(*blocks)
            name = "res" + str(i + 2)
            self.add_module(name, stage)
            self.stages_and_names.append((stage, name))
            self._out_feature_strides[name] = current_stride = int(
                current_stride * np.prod([k.stride for k in blocks])
            )
            self._out_feature_channels[name] = blocks[-1].out_channels

        if num_classes is not None:
            self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
            self.linear = nn.Linear(curr_channels, num_classes)

            # Sec 5.1 in "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour":
            # "The 1000-way fully-connected layer is initialized by
            # drawing weights from a zero-mean Gaussian with std of 0.01."
            nn.init.normal_(self.linear.weight, stddev=0.01)
            name = "linear"

        if out_features is None:
            out_features = [name]
        self._out_features = out_features
        assert len(self._out_features)
        children = [x[0] for x in self.named_children()]
        for out_feature in self._out_features:
            assert out_feature in children, "Available children: {}".format(", ".join(children))

    def forward(self, x):
        outputs = {}
        x = self.stem(x)
        if "stem" in self._out_features:
            outputs["stem"] = x
        for stage, name in self.stages_and_names:
            x = stage(x)
            if name in self._out_features:
                outputs[name] = x
        if self.num_classes is not None:
            x = self.avgpool(x)
            x = self.linear(x)
            if "linear" in self._out_features:
                outputs["linear"] = x
        return outputs

    def output_shape(self):
        return {
            name: ShapeSpec(
                channels=self._out_feature_channels[name],
                stride=self._out_feature_strides[name],
            )
            for name in self._out_features
        }

    @staticmethod
    def make_stage(
        block_class,
        num_blocks,
        first_stride=None,
        *,
        in_channels,
        out_channels,
        **kwargs,
    ):
        """
        Usually, layers that produce the same feature map spatial size
        are defined as one "stage".
        Under such definition, stride_per_block[1:] should all be 1.
        """
        if first_stride is not None:
            assert "stride" not in kwargs and "stride_per_block" not in kwargs
            kwargs["stride_per_block"] = [first_stride] + [1] * (num_blocks - 1)
        blocks = []
        for i in range(num_blocks):
            curr_kwargs = {}
            for k, v in kwargs.items():
                if k.endswith("_per_block"):
                    assert (
                        len(v) == num_blocks
                    ), f"Argument '{k}' of make_stage should have the same length as num_blocks={num_blocks}."
                    newk = k[: -len("_per_block")]
                    assert newk not in kwargs, f"Cannot call make_stage with both {k} and {newk}!"
                    curr_kwargs[newk] = v[i]
                else:
                    curr_kwargs[k] = v

            blocks.append(block_class(in_channels=in_channels, out_channels=out_channels, **curr_kwargs))
            in_channels = out_channels



SUMMARY:
**ResNet Class Summary**

### Purpose and Functionality
The ResNet class is a PyTorch implementation of the Residual Network architecture. It extends the Backbone class and provides a way to define a residual network with multiple stages.

### Key Implementation Details

* The `__init__` method initializes the stem module, stages, and number of classes.
	+ It defines the stage modules using `make_stage`, which creates a sequence of ResNetBlockBase blocks.
	+ It sets up the output feature strides and channels for each stage and block.
* The `forward` method processes input data through the network:
	+ It first passes the input through the stem module.
	+ Then, it applies each stage module in sequence, storing the output at each stage.
	+ If classification is enabled, it adds a linear layer to the output of the last stage before returning the outputs.

### Notable Dependencies and Requirements

* `nn.Module`: A PyTorch module base class.
* `ResNetBlockBase`: The base class for ResNet block modules (not shown in this code snippet).
* `nn.Sequential`: A PyTorch sequential container.
* `ResNet.make_stage`: A helper method to create stages.

### Additional Notes

The `make_stage` method is used to define each stage of the network, which typically consists of multiple blocks. The stages are defined sequentially in the `__init__` method using `self.add_module`.


key: chroma:document
string_value:
CODE:

class ROIPooler(nn.Module):
    """
    Region of interest feature map pooler that supports pooling from one or more
    feature maps.
    """

    def __init__(
        self,
        output_size,
        scales,
        sampling_ratio,
        canonical_box_size=224,
        canonical_level=4,
    ):
        super().__init__()
        # assumption that stride is a power of 2.
        min_level = -math.log2(scales[0])
        max_level = -math.log2(scales[-1])

        # a bunch of testing
        assert math.isclose(min_level, int(min_level)) and math.isclose(max_level, int(max_level))
        assert len(scales) == max_level - min_level + 1, "not pyramid"
        assert 0 < min_level and min_level <= max_level
        if isinstance(output_size, int):
            output_size = (output_size, output_size)
        assert len(output_size) == 2 and isinstance(output_size[0], int) and isinstance(output_size[1], int)
        if len(scales) > 1:
            assert min_level <= canonical_level and canonical_level <= max_level
        assert canonical_box_size > 0

        self.output_size = output_size
        self.min_level = int(min_level)
        self.max_level = int(max_level)
        self.level_poolers = nn.ModuleList(RoIPool(output_size, spatial_scale=scale) for scale in scales)
        self.canonical_level = canonical_level
        self.canonical_box_size = canonical_box_size

    def forward(self, feature_maps, boxes):
        """
        Args:
            feature_maps: List[torch.Tensor(N,C,W,H)]
            box_lists: list[torch.Tensor])
        Returns:
            A tensor of shape(N*B, Channels, output_size, output_size)
        """
        x = list(feature_maps.values())
        num_level_assignments = len(self.level_poolers)
        assert len(x) == num_level_assignments and len(boxes) == x[0].size(0)

        pooler_fmt_boxes = convert_boxes_to_pooler_format(boxes)

        if num_level_assignments == 1:
            return self.level_poolers[0](x[0], pooler_fmt_boxes)

        level_assignments = assign_boxes_to_levels(
            boxes,
            self.min_level,
            self.max_level,
            self.canonical_box_size,
            self.canonical_level,
        )

        num_boxes = len(pooler_fmt_boxes)
        num_channels = x[0].shape[1]
        output_size = self.output_size[0]

        dtype, device = x[0].dtype, x[0].device
        output = torch.zeros(
            (num_boxes, num_channels, output_size, output_size),
            dtype=dtype,
            device=device,
        )

        for level, (x_level, pooler) in enumerate(zip(x, self.level_poolers)):
            inds = torch.nonzero(level_assignments == level).squeeze(1)
            pooler_fmt_boxes_level = pooler_fmt_boxes[inds]
            output[inds] = pooler(x_level, pooler_fmt_boxes_level)



SUMMARY:
**ROIPooler Class Summary**

The `ROIPooler` class is a PyTorch module that implements a region of interest (ROI) feature map pooler. It supports pooling from one or more feature maps and returns a tensor of shape `(N*B, Channels, output_size, output_size)`.

**Main Purpose and Functionality**

The primary purpose of this class is to pool feature maps of varying scales and return a set of pooled feature maps with scaled spatial coordinates.

**Key Implementation Details**

* The `forward` method takes in a list of feature maps and boxes as input.
* It first converts the boxes to a format suitable for pooling using the `convert_boxes_to_pooler_format` function.
* If there's only one level assignment, it pools the feature map directly using the `level_poolers`.
* Otherwise, it assigns each box to its corresponding level using the `assign_boxes_to_levels` function and then pools the feature maps in each level.
* The output is a tensor of pooled feature maps with scaled spatial coordinates.

**Notable Dependencies and Requirements**

* `torch`
* `math`
* `nn.ModuleList`
* `RoIPool`
* `convert_boxes_to_pooler_format`
* `assign_boxes_to_levels`

**Inherited or Implemented Interfaces**

The class implements the PyTorch module interface, inheriting from `nn.Module`. It uses `PyTorch's nn.Module` to create a list of level poolers (`level_poolers`) and employs custom functions for converting boxes and assigning box levels.


key: chroma:document
string_value:
CODE:

class ROIOutputs:
    def __init__(self, cfg, training=False):
        self.smooth_l1_beta = cfg.ROI_BOX_HEAD.SMOOTH_L1_BETA
        self.box2box_transform = Box2BoxTransform(weights=cfg.ROI_BOX_HEAD.BBOX_REG_WEIGHTS)
        self.training = training
        self.score_thresh = cfg.ROI_HEADS.SCORE_THRESH_TEST
        self.min_detections = cfg.MIN_DETECTIONS
        self.max_detections = cfg.MAX_DETECTIONS

        nms_thresh = cfg.ROI_HEADS.NMS_THRESH_TEST
        if not isinstance(nms_thresh, list):
            nms_thresh = [nms_thresh]
        self.nms_thresh = nms_thresh

    def _predict_boxes(self, proposals, box_deltas, preds_per_image):
        num_pred = box_deltas.size(0)
        B = proposals[0].size(-1)
        K = box_deltas.size(-1) // B
        box_deltas = box_deltas.view(num_pred * K, B)
        proposals = torch.cat(proposals, dim=0).unsqueeze(-2).expand(num_pred, K, B)
        proposals = proposals.reshape(-1, B)
        boxes = self.box2box_transform.apply_deltas(box_deltas, proposals)
        return boxes.view(num_pred, K * B).split(preds_per_image, dim=0)

    def _predict_objs(self, obj_logits, preds_per_image):
        probs = nn.functional.softmax(obj_logits, dim=-1)
        probs = probs.split(preds_per_image, dim=0)
        return probs

    def _predict_attrs(self, attr_logits, preds_per_image):
        attr_logits = attr_logits[..., :-1].softmax(-1)
        attr_probs, attrs = attr_logits.max(-1)
        return attr_probs.split(preds_per_image, dim=0), attrs.split(preds_per_image, dim=0)

    @torch.no_grad()
    def inference(
        self,
        obj_logits,
        attr_logits,
        box_deltas,
        pred_boxes,
        features,
        sizes,
        scales=None,
    ):
        # only the pred boxes is the
        preds_per_image = [p.size(0) for p in pred_boxes]
        boxes_all = self._predict_boxes(pred_boxes, box_deltas, preds_per_image)
        obj_scores_all = self._predict_objs(obj_logits, preds_per_image)  # list of length N
        attr_probs_all, attrs_all = self._predict_attrs(attr_logits, preds_per_image)
        features = features.split(preds_per_image, dim=0)

        # fun for each image too, also I can experiment and do multiple images
        final_results = []
        zipped = zip(boxes_all, obj_scores_all, attr_probs_all, attrs_all, sizes)
        for i, (boxes, obj_scores, attr_probs, attrs, size) in enumerate(zipped):
            for nms_t in self.nms_thresh:
                outputs = do_nms(
                    boxes,
                    obj_scores,
                    size,
                    self.score_thresh,
                    nms_t,
                    self.min_detections,
                    self.max_detections,
                )
                if outputs is not None:
                    max_boxes, max_scores, classes, ids = outputs
                    break

            if scales is not None:
                scale_yx = scales[i]
                max_boxes[:, 0::2] *= scale_yx[1]
                max_boxes[:, 1::2] *= scale_yx[0]

            final_results.append(
                (
                    max_boxes,
                    classes,
                    max_scores,
                    attrs[ids],
                    attr_probs[ids],
                    features[i][ids],
                )
            )
        boxes, classes, class_probs, attrs, attr_probs, roi_features = map(list, zip(*final_results))
        return boxes, classes, class_probs, attrs, attr_probs, roi_features

    def training(self, obj_logits, attr_logits, box_deltas, pred_boxes, features, sizes):
        pass

    def __call__(
        self,
        obj_logits,
        attr_logits,
        box_deltas,
        pred_boxes,
        features,
        sizes,
        scales=None,
    ):
        if self.training:
            raise NotImplementedError()
        return self.inference(
            obj_logits,
            attr_logits,
            box_deltas,
            pred_boxes,
            features,
            sizes,
            scales=scales,


SUMMARY:
**Class: ROIOutputs**

**Primary Purpose:** This class is designed to process output data from an object detection model, performing non-maximum suppression and bounding box transformations.

**Key Features:**

*   Initializes the class with configuration settings for smoothing L1 beta, box-to-box transformation weights, and bounding box threshold values.
*   `Predict Boxes`: Applies the box-to-box transformation to predict bboxes given proposals and deltas using a Box2BoxTransform object.
*   `Predict Objects`: Computes softmax probabilities of objects from object logits using nn.functional.softmax.
*   `Predict Attributes`: Softmax attributes from attribute logits, separating classes and IDs.
*   `inference`: Main function to predict bounding boxes, class labels, scores, attributes, and features for a list of objects given the input data.

**Implementation Highlights:**

*   _predict_boxes_ method groups box predictions across images based on _preds_per_image_, applying non-maximum suppression with an array of thresholds.
*   The process employs different techniques to handle object attributes prediction.
*   It provides support for scales adjustments, optimizing bounding box dimensions.
*   The class supports a training mode (currently incomplete), allowing the model to train its configuration parameters from input data.


**Notable Dependencies and Interfaces:**

*   nn.functional.softmax
*   Box2BoxTransform (weights)
 

The function's architecture enables efficient object detection processing with adaptable scaling options.


key: chroma:document
string_value:
CODE:

class Res5ROIHeads(nn.Module):
    """
    ROIHeads perform all per-region computation in an R-CNN.
    It contains logic of cropping the regions, extract per-region features
    (by the res-5 block in this case), and make per-region predictions.
    """

    def __init__(self, cfg, input_shape):
        super().__init__()
        self.batch_size_per_image = cfg.RPN.BATCH_SIZE_PER_IMAGE
        self.positive_sample_fraction = cfg.ROI_HEADS.POSITIVE_FRACTION
        self.in_features = cfg.ROI_HEADS.IN_FEATURES
        self.num_classes = cfg.ROI_HEADS.NUM_CLASSES
        self.proposal_append_gt = cfg.ROI_HEADS.PROPOSAL_APPEND_GT
        self.feature_strides = {k: v.stride for k, v in input_shape.items()}
        self.feature_channels = {k: v.channels for k, v in input_shape.items()}
        self.cls_agnostic_bbox_reg = cfg.ROI_BOX_HEAD.CLS_AGNOSTIC_BBOX_REG
        self.stage_channel_factor = 2**3  # res5 is 8x res2
        self.out_channels = cfg.RESNETS.RES2_OUT_CHANNELS * self.stage_channel_factor

        # self.proposal_matcher = Matcher(
        #     cfg.ROI_HEADS.IOU_THRESHOLDS,
        #     cfg.ROI_HEADS.IOU_LABELS,
        #     allow_low_quality_matches=False,
        # )

        pooler_resolution = cfg.ROI_BOX_HEAD.POOLER_RESOLUTION
        pooler_scales = (1.0 / self.feature_strides[self.in_features[0]],)
        sampling_ratio = cfg.ROI_BOX_HEAD.POOLER_SAMPLING_RATIO
        res5_halve = cfg.ROI_BOX_HEAD.RES5HALVE
        use_attr = cfg.ROI_BOX_HEAD.ATTR
        num_attrs = cfg.ROI_BOX_HEAD.NUM_ATTRS

        self.pooler = ROIPooler(
            output_size=pooler_resolution,
            scales=pooler_scales,
            sampling_ratio=sampling_ratio,
        )

        self.res5 = self._build_res5_block(cfg)
        if not res5_halve:
            """
            Modifications for VG in RoI heads:
            1. Change the stride of conv1 and shortcut in Res5.Block1 from 2 to 1
            2. Modifying all conv2 with (padding: 1 --> 2) and (dilation: 1 --> 2)
            """
            self.res5[0].conv1.stride = (1, 1)
            self.res5[0].shortcut.stride = (1, 1)
            for i in range(3):
                self.res5[i].conv2.padding = (2, 2)
                self.res5[i].conv2.dilation = (2, 2)

        self.box_predictor = FastRCNNOutputLayers(
            self.out_channels,
            self.num_classes,
            self.cls_agnostic_bbox_reg,
            use_attr=use_attr,
            num_attrs=num_attrs,
        )

    def _build_res5_block(self, cfg):
        stage_channel_factor = self.stage_channel_factor  # res5 is 8x res2
        num_groups = cfg.RESNETS.NUM_GROUPS
        width_per_group = cfg.RESNETS.WIDTH_PER_GROUP
        bottleneck_channels = num_groups * width_per_group * stage_channel_factor
        out_channels = self.out_channels
        stride_in_1x1 = cfg.RESNETS.STRIDE_IN_1X1
        norm = cfg.RESNETS.NORM

        blocks = ResNet.make_stage(
            BottleneckBlock,
            3,
            first_stride=2,
            in_channels=out_channels // 2,
            bottleneck_channels=bottleneck_channels,
            out_channels=out_channels,
            num_groups=num_groups,
            norm=norm,
            stride_in_1x1=stride_in_1x1,
        )
        return nn.Sequential(*blocks)

    def _shared_roi_transform(self, features, boxes):
        x = self.pooler(features, boxes)
        return self.res5(x)

    def forward(self, features, proposal_boxes, gt_boxes=None):
        if self.training:
            """
            see https://github.com/airsplay/py-bottom-up-attention/\
                    blob/master/detectron2/modeling/roi_heads/roi_heads.py
            """
            raise NotImplementedError()

        assert not proposal_boxes[0].requires_grad
        box_features = self._shared_roi_transform(features, proposal_boxes)
        feature_pooled = box_features.mean(dim=[2, 3])  # pooled to 1x1
        obj_logits, attr_logits, pred_proposal_deltas = self.box_predictor(feature_pooled)


SUMMARY:
**Res5ROIHeads Class Summary**

### Purpose and Functionality

The `Res5ROIHeads` class is responsible for performing region-of-interest (ROI) computation in a Faster R-CNN architecture. It extracts features from predefined ROIs, performs per-region predictions, and predicts bounding box deltas.

The primary purpose of this class is to extract features from the res-5 block of the ResNet backbone using ROI pooling and then pass these features through a FastRCNN output layer for object detection and attribute prediction.

### Key Implementation Details

* The class uses a `ROIPooler` layer for ROI pooling and a `FastRCNNOutputLayers` layer for object detection and attribute prediction.
* The `Res5Block` is modified to change the stride of conv1 and shortcut in Res5.Block1 from 2 to 1.
* The output channels of res-5 are calculated using the stage channel factor.

### Notable Dependencies or Requirements

* `nn.Module`
* `ROIPooler`
* `FastRCNNOutputLayers`
* `ResNet` (specifically, the `BottleneckBlock` and `make_stage`)
* `cfg.ROI_HEADS` configuration parameters
* External dependencies:
	+ `Matcher` (commented out)
	+ Additional modules for object detection and attribute prediction (not shown in this code snippet)

### Interfaces Implemented

This class implements interfaces such as:

* `nn.Module`
* `ROIPooler`
* `FastRCNNOutputLayers`

Note: This summary focuses on the key purpose, implementation details, and notable dependencies of the `Res5ROIHeads` class.


key: chroma:document
string_value:
CODE:

class AnchorGenerator(nn.Module):
    """
    For a set of image sizes and feature maps, computes a set of anchors.
    """

    def __init__(self, cfg, input_shape: List[ShapeSpec]):
        super().__init__()
        sizes = cfg.ANCHOR_GENERATOR.SIZES
        aspect_ratios = cfg.ANCHOR_GENERATOR.ASPECT_RATIOS
        self.strides = [x.stride for x in input_shape]
        self.offset = cfg.ANCHOR_GENERATOR.OFFSET
        assert 0.0 <= self.offset < 1.0, self.offset

        """
        sizes (list[list[int]]): sizes[i] is the list of anchor sizes for feat map i
            1. given in absolute lengths in units of the input image;
            2. they do not dynamically scale if the input image size changes.
        aspect_ratios (list[list[float]])
        strides (list[int]): stride of each input feature.
        """

        self.num_features = len(self.strides)
        self.cell_anchors = nn.ParameterList(self._calculate_anchors(sizes, aspect_ratios))
        self._spacial_feat_dim = 4

    def _calculate_anchors(self, sizes, aspect_ratios):
        # If one size (or aspect ratio) is specified and there are multiple feature
        # maps, then we "broadcast" anchors of that single size (or aspect ratio)
        if len(sizes) == 1:
            sizes *= self.num_features
        if len(aspect_ratios) == 1:
            aspect_ratios *= self.num_features
        assert self.num_features == len(sizes)
        assert self.num_features == len(aspect_ratios)

        cell_anchors = [self.generate_cell_anchors(s, a).float() for s, a in zip(sizes, aspect_ratios)]

        return cell_anchors

    @property
    def box_dim(self):
        return self._spacial_feat_dim

    @property
    def num_cell_anchors(self):
        """
        Returns:
            list[int]: Each int is the number of anchors at every pixel location, on that feature map.
        """
        return [len(cell_anchors) for cell_anchors in self.cell_anchors]

    def grid_anchors(self, grid_sizes):
        anchors = []
        for size, stride, base_anchors in zip(grid_sizes, self.strides, self.cell_anchors):
            shift_x, shift_y = _create_grid_offsets(size, stride, self.offset, base_anchors.device)
            shifts = torch.stack((shift_x, shift_y, shift_x, shift_y), dim=1)

            anchors.append((shifts.view(-1, 1, 4) + base_anchors.view(1, -1, 4)).reshape(-1, 4))

        return anchors

    def generate_cell_anchors(self, sizes=(32, 64, 128, 256, 512), aspect_ratios=(0.5, 1, 2)):
        """
        anchors are continuous geometric rectangles
        centered on one feature map point sample.
        We can later build the set of anchors
        for the entire feature map by tiling these tensors
        """

        anchors = []
        for size in sizes:
            area = size**2.0
            for aspect_ratio in aspect_ratios:
                w = math.sqrt(area / aspect_ratio)
                h = aspect_ratio * w
                x0, y0, x1, y1 = -w / 2.0, -h / 2.0, w / 2.0, h / 2.0
                anchors.append([x0, y0, x1, y1])
        return nn.Parameter(torch.tensor(anchors))

    def forward(self, features):
        """
        Args:
            features List[torch.Tensor]: list of feature maps on which to generate anchors.
        Returns:
            torch.Tensor: a list of #image elements.
        """
        num_images = features[0].size(0)
        grid_sizes = [feature_map.shape[-2:] for feature_map in features]
        anchors_over_all_feature_maps = self.grid_anchors(grid_sizes)
        anchors_over_all_feature_maps = torch.stack(anchors_over_all_feature_maps)


SUMMARY:
**AnchorGenerator Class Summary**

**Main Purpose and Functionality:**
The `AnchorGenerator` class is responsible for generating anchors for object detection tasks. It takes in a set of image sizes and feature maps, computes the corresponding anchors using pre-defined size scales and aspect ratios.

**Key Implementation Details:**

* The `_calculate_anchors` method broadcasts anchor sizes and aspect ratios to match the number of feature maps.
* The `generate_cell_anchors` method defines the geometric rectangles that make up a cell of anchors with various aspect ratios and sizes.
* The `grid_anchors` method generates anchors for each grid size by shifting the base anchors horizontally, vertically, and bilaterally.

**Notable Dependencies and Requirements:**

* `nn.Module`: Inherit from PyTorch's neural network module.
* `cfg.ANCHOR_GENERATOR.SIZES` and `cfg.ANCHOR_GENERATOR.ASPECT_RATIOS`: Configuration settings for anchor sizes and aspect ratios.
* `_create_grid_offsets`, `math.sqrt`, `torch.stack`, `torch.view`, `torch.reshape`: PyTorch functions used in the implementation.

**Method Signatures:**

* `__init__(cfg, input_shape: List[ShapeSpec])`
* `_calculate_anchors(sizes, aspect_ratios)`
* `generate_cell_anchors(sizes=(32, 64, 128, 256, 512), aspect_ratios=(0.5, 1, 2))`
* `grid_anchors(grid_sizes)`
* `forward(features: List[torch.Tensor])`

**Inherited or Implemented Interfaces:**

* None. This class does not inherit from any PyTorch module's interfaces. However, it uses various PyTorch functions and classes (e.g., `nn.Parameter`, `torch.Tensor`).


key: chroma:document
string_value:
CODE:

class RPNHead(nn.Module):
    """
    RPN classification and regression heads. Uses a 3x3 conv to produce a shared
    hidden state from which one 1x1 conv predicts objectness logits for each anchor
    and a second 1x1 conv predicts bounding-box deltas specifying how to deform
    each anchor into an object proposal.
    """

    def __init__(self, cfg, input_shape: List[ShapeSpec]):
        super().__init__()

        # Standard RPN is shared across levels:
        in_channels = [s.channels for s in input_shape]
        assert len(set(in_channels)) == 1, "Each level must have the same channel!"
        in_channels = in_channels[0]

        anchor_generator = AnchorGenerator(cfg, input_shape)
        num_cell_anchors = anchor_generator.num_cell_anchors
        box_dim = anchor_generator.box_dim
        assert len(set(num_cell_anchors)) == 1, "Each level must have the same number of cell anchors"
        num_cell_anchors = num_cell_anchors[0]

        if cfg.PROPOSAL_GENERATOR.HIDDEN_CHANNELS == -1:
            hid_channels = in_channels
        else:
            hid_channels = cfg.PROPOSAL_GENERATOR.HIDDEN_CHANNELS
            # Modifications for VG in RPN (modeling/proposal_generator/rpn.py)
            # Use hidden dim  instead fo the same dim as Res4 (in_channels)

        # 3x3 conv for the hidden representation
        self.conv = nn.Conv2d(in_channels, hid_channels, kernel_size=3, stride=1, padding=1)
        # 1x1 conv for predicting objectness logits
        self.objectness_logits = nn.Conv2d(hid_channels, num_cell_anchors, kernel_size=1, stride=1)
        # 1x1 conv for predicting box2box transform deltas
        self.anchor_deltas = nn.Conv2d(hid_channels, num_cell_anchors * box_dim, kernel_size=1, stride=1)

        for layer in [self.conv, self.objectness_logits, self.anchor_deltas]:
            nn.init.normal_(layer.weight, std=0.01)
            nn.init.constant_(layer.bias, 0)

    def forward(self, features):
        """
        Args:
            features (list[Tensor]): list of feature maps
        """
        pred_objectness_logits = []
        pred_anchor_deltas = []
        for x in features:
            t = nn.functional.relu(self.conv(x))
            pred_objectness_logits.append(self.objectness_logits(t))
            pred_anchor_deltas.append(self.anchor_deltas(t))


SUMMARY:
**RPNHead Class Summary**

### Purpose and Functionality

The `RPNHead` class is a neural module that implements the RPN (Region Proposal Network) classification and regression heads. It takes a shared hidden state from three 3x3 convolutions and predicts objectness logits for each anchor and bounding-box deltas specifying how to deform each anchor into an object proposal.

### Key Implementation Details

* **Convolutional Layers**:
	+ Shared `conv` layer produces a 3x3 hidden representation with `hid_channels` channels.
	+ Two separate `1x1` convolutions: `objectness_logits` predicts objectness logits for each anchor, and `anchor_deltas` predicts bounding-box deltas specifying how to deform each anchor into an object proposal.
* **Weight Initialization**: Weights are initialized using a normal distribution with standard deviation 0.01.

### Notable Dependencies and Requirements

* `AnchorGenerator`: generates anchor coordinates based on image shape and scale factor
	+ Implemented interface: `num_cell_anchors` and `box_dim`
* External dependencies:
	+ PyTorch (`nn.Module`, `nn.Conv2d`, `nn.functional.relu`)
	+ Optional configuration file (`cfg.PROPOSAL_GENERATOR.HIDDEN_CHANNELS`)


key: chroma:document
string_value:
CODE:

class RPN(nn.Module):
    """
    Region Proposal Network, introduced by the Faster R-CNN paper.
    """

    def __init__(self, cfg, input_shape: Dict[str, ShapeSpec]):
        super().__init__()

        self.min_box_side_len = cfg.PROPOSAL_GENERATOR.MIN_SIZE
        self.in_features = cfg.RPN.IN_FEATURES
        self.nms_thresh = cfg.RPN.NMS_THRESH
        self.batch_size_per_image = cfg.RPN.BATCH_SIZE_PER_IMAGE
        self.positive_fraction = cfg.RPN.POSITIVE_FRACTION
        self.smooth_l1_beta = cfg.RPN.SMOOTH_L1_BETA
        self.loss_weight = cfg.RPN.LOSS_WEIGHT

        self.pre_nms_topk = {
            True: cfg.RPN.PRE_NMS_TOPK_TRAIN,
            False: cfg.RPN.PRE_NMS_TOPK_TEST,
        }
        self.post_nms_topk = {
            True: cfg.RPN.POST_NMS_TOPK_TRAIN,
            False: cfg.RPN.POST_NMS_TOPK_TEST,
        }
        self.boundary_threshold = cfg.RPN.BOUNDARY_THRESH

        self.anchor_generator = AnchorGenerator(cfg, [input_shape[f] for f in self.in_features])
        self.box2box_transform = Box2BoxTransform(weights=cfg.RPN.BBOX_REG_WEIGHTS)
        self.anchor_matcher = Matcher(
            cfg.RPN.IOU_THRESHOLDS,
            cfg.RPN.IOU_LABELS,
            allow_low_quality_matches=True,
        )
        self.rpn_head = RPNHead(cfg, [input_shape[f] for f in self.in_features])

    def training(self, images, image_shapes, features, gt_boxes):
        pass

    def inference(self, outputs, images, image_shapes, features, gt_boxes=None):
        outputs = find_top_rpn_proposals(
            outputs.predict_proposals(),
            outputs.predict_objectness_logits(),
            images,
            image_shapes,
            self.nms_thresh,
            self.pre_nms_topk[self.training],
            self.post_nms_topk[self.training],
            self.min_box_side_len,
            self.training,
        )

        results = []
        for img in outputs:
            im_boxes, img_box_logits = img
            img_box_logits, inds = img_box_logits.sort(descending=True)
            im_boxes = im_boxes[inds]
            results.append((im_boxes, img_box_logits))

        (proposal_boxes, logits) = tuple(map(list, zip(*results)))
        return proposal_boxes, logits

    def forward(self, images, image_shapes, features, gt_boxes=None):
        """
        Args:
            images (torch.Tensor): input images of length `N`
            features (dict[str: Tensor])
            gt_instances
        """
        # features is dict, key = block level, v = feature_map
        features = [features[f] for f in self.in_features]
        pred_objectness_logits, pred_anchor_deltas = self.rpn_head(features)
        anchors = self.anchor_generator(features)
        outputs = RPNOutputs(
            self.box2box_transform,
            self.anchor_matcher,
            self.batch_size_per_image,
            self.positive_fraction,
            images,
            pred_objectness_logits,
            pred_anchor_deltas,
            anchors,
            self.boundary_threshold,
            gt_boxes,
            self.smooth_l1_beta,
        )
        # For RPN-only models, the proposals are the final output

        if self.training:
            raise NotImplementedError()
            return self.training(outputs, images, image_shapes, features, gt_boxes)
        else:


SUMMARY:
**RPN (Region Proposal Network) Class Summary**

**Main Purpose and Functionality:**
The RPN class implements the Region Proposal Network, introduced by the Faster R-CNN paper. It takes an input image, feature maps, and optional ground truth boxes to generate region proposals.

**Key Implementation Details:**

* The network consists of four main components:
	1. `AnchorGenerator`: generates anchors from feature maps.
	2. `Box2BoxTransform`: transforms anchors into proposal boxes.
	3. `Matcher` (AnchorMatcher): matches anchors with ground truth boxes.
	4. `RPNHead`: processes top-level features to generate objectness logits and anchor deltas.
* The network performs two main operations:
	1. Training: computes losses and updates parameters using the `training` method.
	2. Inference: generates proposals using the `inference` method.

**Notable Dependencies and Requirements:**

* Related classes/functions:
	+ `Box2BoxTransform`
	+ `Matcher` (AnchorMatcher)
	+ `RPNHead`
	+
* External dependencies:
	+ `cfg.RPN`: configuration settings
	+ `torch.Tensor`
* Inherited or implemented interfaces:
	+ `nn.Module`

**Code Structure:**

* The class has an initializer (`__init__`) that sets up network parameters.
* Two main methods: `training` and `inference`, which perform the different operations.

Note: The code summary is based on a concise analysis of the provided class code chunk. If further information or clarification is required, please let me know.


key: chroma:document
string_value:
CODE:

class FastRCNNOutputLayers(nn.Module):
    """
    Two linear layers for predicting Fast R-CNN outputs:
      (1) proposal-to-detection box regression deltas
      (2) classification scores
    """

    def __init__(
        self,
        input_size,
        num_classes,
        cls_agnostic_bbox_reg,
        box_dim=4,
        use_attr=False,
        num_attrs=-1,
    ):
        """
        Args:
            input_size (int): channels, or (channels, height, width)
            num_classes (int)
            cls_agnostic_bbox_reg (bool)
            box_dim (int)
        """
        super().__init__()

        if not isinstance(input_size, int):
            input_size = np.prod(input_size)

        # (do + 1 for background class)
        self.cls_score = nn.Linear(input_size, num_classes + 1)
        num_bbox_reg_classes = 1 if cls_agnostic_bbox_reg else num_classes
        self.bbox_pred = nn.Linear(input_size, num_bbox_reg_classes * box_dim)

        self.use_attr = use_attr
        if use_attr:
            """
            Modifications for VG in RoI heads
            Embedding: {num_classes + 1} --> {input_size // 8}
            Linear: {input_size + input_size // 8} --> {input_size // 4}
            Linear: {input_size // 4} --> {num_attrs + 1}
            """
            self.cls_embedding = nn.Embedding(num_classes + 1, input_size // 8)
            self.fc_attr = nn.Linear(input_size + input_size // 8, input_size // 4)
            self.attr_score = nn.Linear(input_size // 4, num_attrs + 1)

        nn.init.normal_(self.cls_score.weight, std=0.01)
        nn.init.normal_(self.bbox_pred.weight, std=0.001)
        for item in [self.cls_score, self.bbox_pred]:
            nn.init.constant_(item.bias, 0)

    def forward(self, roi_features):
        if roi_features.dim() > 2:
            roi_features = torch.flatten(roi_features, start_dim=1)
        scores = self.cls_score(roi_features)
        proposal_deltas = self.bbox_pred(roi_features)
        if self.use_attr:
            _, max_class = scores.max(-1)  # [b, c] --> [b]
            cls_emb = self.cls_embedding(max_class)  # [b] --> [b, 256]
            roi_features = torch.cat([roi_features, cls_emb], -1)  # [b, 2048] + [b, 256] --> [b, 2304]
            roi_features = self.fc_attr(roi_features)
            roi_features = nn.functional.relu(roi_features)
            attr_scores = self.attr_score(roi_features)
            return scores, attr_scores, proposal_deltas
        else:


SUMMARY:
**FastRCNNOutputLayers Class Summary**

This class implements two linear layers for predicting Fast R-CNN outputs: 

1.  **Proposal-to-Detection Box Regression Deltas**: Predicted by `self.bbox_pred`.
2.  **Classification Scores**: Predicted by `self.cls_score`.

**Key Implementation Details:**

*   The `cls_agnostic_bbox_reg` flag controls whether to use per-class or class-agnostic box regression.
*   The `use_attr` parameter enables modifications for VG in RoI heads, which involves additional linear and embedding layers.

**Notable Dependencies and Requirements:**

*   (`nn.Linear`): Used for linear layers.
*   (`nn.init.normal_`, `nn.init.constant_`): Initialization functions for weights and biases.
*   (`torch.flatten`, `torch.cat`, `nn.functional.relu`): Functions used in the forward pass.

**Class Parameters:**

*   `input_size`: Channels, height, or width of input features.
*   `num_classes`: Number of classes to predict (including background class).
*   `cls_agnostic_bbox_reg`: Flag controlling box regression method.
*   `box_dim`: Dimensionality of box feature space.
*   `use_attr`: Flag enabling VG modifications for RoI heads.

**Interface and Inheritance:**

*   This module inherits from PyTorch's (`nn`) built-in module.


key: chroma:document
string_value:
CODE:

class GeneralizedRCNN(nn.Module):
    def __init__(self, cfg):
        super().__init__()

        self.device = torch.device(cfg.MODEL.DEVICE)
        self.backbone = build_backbone(cfg)
        self.proposal_generator = RPN(cfg, self.backbone.output_shape())
        self.roi_heads = Res5ROIHeads(cfg, self.backbone.output_shape())
        self.roi_outputs = ROIOutputs(cfg)
        self.to(self.device)

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        config = kwargs.pop("config", None)
        state_dict = kwargs.pop("state_dict", None)
        cache_dir = kwargs.pop("cache_dir", None)
        from_tf = kwargs.pop("from_tf", False)
        force_download = kwargs.pop("force_download", False)
        resume_download = kwargs.pop("resume_download", False)
        proxies = kwargs.pop("proxies", None)
        local_files_only = kwargs.pop("local_files_only", False)
        use_cdn = kwargs.pop("use_cdn", True)

        # Load config if we don't provide a configuration
        if not isinstance(config, Config):
            config_path = config if config is not None else pretrained_model_name_or_path
            # try:
            config = Config.from_pretrained(
                config_path,
                cache_dir=cache_dir,
                force_download=force_download,
                resume_download=resume_download,
                proxies=proxies,
                local_files_only=local_files_only,
            )

        # Load model
        if pretrained_model_name_or_path is not None:
            if os.path.isdir(pretrained_model_name_or_path):
                if os.path.isfile(os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)):
                    # Load from a PyTorch checkpoint
                    archive_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)
                else:
                    raise EnvironmentError(
                        "Error no file named {} found in directory {} ".format(
                            WEIGHTS_NAME,
                            pretrained_model_name_or_path,
                        )
                    )
            elif os.path.isfile(pretrained_model_name_or_path) or is_remote_url(pretrained_model_name_or_path):
                archive_file = pretrained_model_name_or_path
            elif os.path.isfile(pretrained_model_name_or_path + ".index"):
                assert from_tf, "We found a TensorFlow checkpoint at {}, please set from_tf to True to load from this checkpoint".format(
                    pretrained_model_name_or_path + ".index"
                )
                archive_file = pretrained_model_name_or_path + ".index"
            else:
                archive_file = hf_bucket_url(
                    pretrained_model_name_or_path,
                    filename=WEIGHTS_NAME,
                    use_cdn=use_cdn,
                )

            try:
                # Load from URL or cache if already cached
                resolved_archive_file = cached_path(
                    archive_file,
                    cache_dir=cache_dir,
                    force_download=force_download,
                    proxies=proxies,
                    resume_download=resume_download,
                    local_files_only=local_files_only,
                )
                if resolved_archive_file is None:
                    raise EnvironmentError
            except EnvironmentError:
                msg = f"Can't load weights for '{pretrained_model_name_or_path}'."
                raise EnvironmentError(msg)

            if resolved_archive_file == archive_file:
                print("loading weights file {}".format(archive_file))
            else:
                print("loading weights file {} from cache at {}".format(archive_file, resolved_archive_file))
        else:
            resolved_archive_file = None

        # Instantiate model.
        model = cls(config)

        if state_dict is None:
            try:
                try:
                    state_dict = torch.load(resolved_archive_file, map_location="cpu")
                except Exception:
                    state_dict = load_checkpoint(resolved_archive_file)

            except Exception:
                raise OSError(
                    "Unable to load weights from pytorch checkpoint file. "
                    "If you tried to load a PyTorch model from a TF 2.0 checkpoint, please set from_tf=True. "
                )

        missing_keys = []
        unexpected_keys = []
        error_msgs = []

        # Convert old format to new format if needed from a PyTorch state_dict
        old_keys = []
        new_keys = []
        for key in state_dict.keys():
            new_key = None
            if "gamma" in key:
                new_key = key.replace("gamma", "weight")
            if "beta" in key:
                new_key = key.replace("beta", "bias")
            if new_key:
                old_keys.append(key)
                new_keys.append(new_key)
        for old_key, new_key in zip(old_keys, new_keys):
            state_dict[new_key] = state_dict.pop(old_key)

        # copy state_dict so _load_from_state_dict can modify it
        metadata = getattr(state_dict, "_metadata", None)
        state_dict = state_dict.copy()
        if metadata is not None:
            state_dict._metadata = metadata

        model_to_load = model
        model_to_load.load_state_dict(state_dict)

        if model.__class__.__name__ != model_to_load.__class__.__name__:
            base_model_state_dict = model_to_load.state_dict().keys()
            head_model_state_dict_without_base_prefix = [
                key.split(cls.base_model_prefix + ".")[-1] for key in model.state_dict().keys()
            ]
            missing_keys.extend(head_model_state_dict_without_base_prefix - base_model_state_dict)

        if len(unexpected_keys) > 0:
            print(
                f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when"
                f" initializing {model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are"
                f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task or"
                " with another architecture (e.g. initializing a BertForSequenceClassification model from a"
                " BertForPreTraining model).\n- This IS NOT expected if you are initializing"
                f" {model.__class__.__name__} from the checkpoint of a model that you expect to be exactly identical"
                " (initializing a BertForSequenceClassification model from a BertForSequenceClassification model)."
            )
        else:
            print(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n")
        if len(missing_keys) > 0:
            print(
                f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at"
                f" {pretrained_model_name_or_path} and are newly initialized: {missing_keys}\nYou should probably"
                " TRAIN this model on a down-stream task to be able to use it for predictions and inference."
            )
        else:
            print(
                f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at"
                f" {pretrained_model_name_or_path}.\nIf your task is similar to the task the model of the checkpoint"
                f" was trained on, you can already use {model.__class__.__name__} for predictions without further"
                " training."
            )
        if len(error_msgs) > 0:
            raise RuntimeError(
                "Error(s) in loading state_dict for {}:\n\t{}".format(
                    model.__class__.__name__, "\n\t".join(error_msgs)
                )
            )
        # Set model in evaluation mode to deactivate DropOut modules by default
        model.eval()

        return model

    def forward(
        self,
        images,
        image_shapes,
        gt_boxes=None,
        proposals=None,
        scales_yx=None,
        **kwargs,
    ):
        """
        kwargs:
            max_detections (int), return_tensors {"np", "pt", None}, padding {None,
            "max_detections"}, pad_value (int), location = {"cuda", "cpu"}
        """
        if self.training:
            raise NotImplementedError()
        return self.inference(
            images=images,
            image_shapes=image_shapes,
            gt_boxes=gt_boxes,
            proposals=proposals,
            scales_yx=scales_yx,
            **kwargs,
        )

    @torch.no_grad()
    def inference(
        self,
        images,
        image_shapes,
        gt_boxes=None,
        proposals=None,
        scales_yx=None,
        **kwargs,
    ):
        # run images through backbone
        original_sizes = image_shapes * scales_yx
        features = self.backbone(images)

        # generate proposals if none are available
        if proposals is None:
            proposal_boxes, _ = self.proposal_generator(images, image_shapes, features, gt_boxes)
        else:
            assert proposals is not None

        # pool object features from either gt_boxes, or from proposals
        obj_logits, attr_logits, box_deltas, feature_pooled = self.roi_heads(features, proposal_boxes, gt_boxes)

        # prepare FRCNN Outputs and select top proposals
        boxes, classes, class_probs, attrs, attr_probs, roi_features = self.roi_outputs(
            obj_logits=obj_logits,
            attr_logits=attr_logits,
            box_deltas=box_deltas,
            pred_boxes=proposal_boxes,
            features=feature_pooled,
            sizes=image_shapes,
            scales=scales_yx,
        )

        # will we pad???
        subset_kwargs = {
            "max_detections": kwargs.get("max_detections", None),
            "return_tensors": kwargs.get("return_tensors", None),
            "pad_value": kwargs.get("pad_value", 0),
            "padding": kwargs.get("padding", None),
        }
        preds_per_image = torch.tensor([p.size(0) for p in boxes])
        boxes = pad_list_tensors(boxes, preds_per_image, **subset_kwargs)
        classes = pad_list_tensors(classes, preds_per_image, **subset_kwargs)
        class_probs = pad_list_tensors(class_probs, preds_per_image, **subset_kwargs)
        attrs = pad_list_tensors(attrs, preds_per_image, **subset_kwargs)
        attr_probs = pad_list_tensors(attr_probs, preds_per_image, **subset_kwargs)
        roi_features = pad_list_tensors(roi_features, preds_per_image, **subset_kwargs)
        subset_kwargs["padding"] = None
        preds_per_image = pad_list_tensors(preds_per_image, None, **subset_kwargs)
        sizes = pad_list_tensors(image_shapes, None, **subset_kwargs)
        normalized_boxes = norm_box(boxes, original_sizes)
        return OrderedDict(
            {
                "obj_ids": classes,
                "obj_probs": class_probs,
                "attr_ids": attrs,
                "attr_probs": attr_probs,
                "boxes": boxes,
                "sizes": sizes,
                "preds_per_image": preds_per_image,
                "roi_features": roi_features,
                "normalized_boxes": normalized_boxes,
            }


SUMMARY:
This code defines a custom PyTorch model, `FastRCNNModel`, which is based on the Faster R-CNN architecture. The model takes in a set of input images and uses a backbone network to extract features.

**Key Features:**

* It generates proposals using the `proposal_generator` method.
* It pools object features from either `gt_boxes` or `proposals`.
* It predicts object, attribute, and box logits using the `roi_heads` method.
* It prepares FRCNN outputs by selecting top proposals and computing class probabilities and features.

**Method Summary:**

The code consists of two main methods:

1.  `_create_model`: Creates a new instance of the FastRCNN model, setting its backbone network, proposal generator, and ROI heads.
2.  `forward` and `inference`: Defines the forward pass through the model, taking in input images, image shapes, and optional parameters such as ground truth boxes and proposals.

**Output:**

The output of the model is an ordered dictionary containing various tensors, including:

*   `obj_ids`, `obj_probs`, `attr_ids`, and `attr_probs`: Class and attribute tensor outputs.
*   `boxes` and `sizes`: Bounding box tensor outputs with image scale information.
*   `preds_per_image`, `roi_features`, and `normalized_boxes`: Other tensors containing prediction counts, feature vectors, and normalized bounding boxes.

**Methodality:**

The code utilizes several PyTorch features, such as:

*   `torch.no_grad()` for setting the computation graph to not track gradients.
*   PyTorch's built-in support for tensors using `torch.tensor`.
*   Custom modules such as `FastRCNNBackbone`, `ProposalGenerator`, and ` RooneyHead`, which define the backbone network, proposal generation, and ROI prediction.

**Key Parameters:**

The code takes in several optional parameters:

*   `max_detections` (int): The maximum number of detected objects per image.
*   `return_tensors` (`{"np", "pt", None}`): The format of the output tensor.
*   `pad_value` (int): The padded value for tensors with padding.
*   `padding` (`None` or `"max_detections"`): The padding strategy for tensors.

Overall, this code defines a detailed implementation of a Faster R-CNN architecture in PyTorch.


