Merge pull request #1 from huggingface/smangrul/add-code

add code

.gitignore

@@ -0,0 +1,141 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# VSCode
+.vscode
+
+# IntelliJ
+.idea
+
+# Mac .DS_Store
+.DS_Store
+
+# More test things
+wandb

MANIFEST.in

@@ -0,0 +1 @@
+include LICENSE

Makefile

@@ -0,0 +1,19 @@
+.PHONY: quality style test docs
+
+check_dirs := src
+
+# Check that source code meets quality standards
+
+# this target runs checks on all files
+quality:
+	black --check $(check_dirs)
+	isort --check-only $(check_dirs)
+	flake8 $(check_dirs)
+	python utils/style_doc.py src --max_len 119 --check_only
+
+# Format source code automatically and check is there are any problems left that need manual fixing
+style:
+	black $(check_dirs)
+	isort $(check_dirs)
+	python utils/style_doc.py src --max_len 119
+	

README.md

@@ -1,2 +1,45 @@
-# pets
-Parameter-Efficient Tuning at Scale
+# 🤗 PET
+Parameter-Efficient Tuning. Intergrated with 🤗 Accelerate to scale seamlessly to large models using PyTorch FSDP. 
+
+Supported methods:
+
+1. LoRA
+2. Prefix Tuning
+3. P-Tuning
+4. Prompt Tuning 
+
+## Models support matrix
+
+### Sequence Classification
+|   Model         | LoRA | Prefix Tuning  | P-Tuning | Prompt Tuning  | 
+| --------- | ---- | ---- | ---- | ----  |
+| BERT           | ✅  | ✅  | ✅  | ✅  |  
+| RoBERTa        | ✅  | ✅  | ✅  | ✅  |
+| GPT-2          | ✅  | ✅  | ✅  | ✅  | 
+| Bloom          | ✅  | ✅  | ✅  | ✅  |   
+| OPT            | ✅  | ✅  | ✅  | ✅  |
+| GPT-Neo        | ✅  | ✅  | ✅  | ✅  |
+| GPT-J          | ✅  | ✅  | ✅  | ✅  |
+| Deberta        | ✅  |     |     |     | 
+| Deberta-v2     | ✅  |     |     |     |
+
+### Causal Language Modeling
+|   Model         | LoRA | Prefix Tuning  | P-Tuning | Prompt Tuning  |
+| --------- | ---- | ---- | ---- | ----  |
+| GPT-2          | ✅  | ✅  | ✅  | ✅  |
+| Bloom          | ✅  | ✅  | ✅  | ✅  |
+| OPT            | ✅  | ✅  | ✅  | ✅  |
+| GPT-Neo        | ✅  | ✅  | ✅  | ✅  |
+| GPT-J          | ✅  | ✅  | ✅  | ✅  |
+
+### Conditional Generation
+|   Model         | LoRA | Prefix Tuning  | P-Tuning | Prompt Tuning  | 
+| --------- | ---- | ---- | ---- | ---- |
+| T5        | ✅   | ✅   | ✅   | ✅   |
+| BART      | ✅   | ✅   | ✅   | ✅   |
+
+
+## Caveats:
+1. Doesn't work currently with DeeSpeed ZeRO Stage-3. Extending support with DeeSpeed ZeRO Stage-3 is in backlog.
+
+

pyproject.toml

@@ -0,0 +1,3 @@
+[tool.black]
+line-length = 119
+target-version = ['py36']

requirements.txt

@@ -0,0 +1,4 @@
+transformers
+accelerate
+loralib
+evaluate

setup.cfg

@@ -0,0 +1,23 @@
+[isort]
+default_section = FIRSTPARTY
+ensure_newline_before_comments = True
+force_grid_wrap = 0
+include_trailing_comma = True
+known_first_party = pet
+known_third_party =
+    numpy
+    torch
+    accelerate
+    transformers
+
+line_length = 119
+lines_after_imports = 2
+multi_line_output = 3
+use_parentheses = True
+
+[flake8]
+ignore = E203, E722, E501, E741, W503, W605
+max-line-length = 119
+
+[tool:pytest]
+doctest_optionflags=NUMBER NORMALIZE_WHITESPACE ELLIPSIS

setup.py

@@ -0,0 +1,78 @@
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from setuptools import setup
+from setuptools import find_packages
+
+extras = {}
+extras["quality"] = ["black ~= 22.0", "isort >= 5.5.4", "flake8 >= 3.8.3"]
+extras["dev"] = extras["quality"]
+
+setup(
+    name="pets",
+    version="0.1.0.dev0",
+    description="Parameter-Efficient Tuning (PET)",
+    long_description=open("README.md", "r", encoding="utf-8").read(),
+    long_description_content_type="text/markdown",
+    keywords="deep learning",
+    license="Apache",
+    author="The HuggingFace team",
+    author_email="sourab@huggingface.co",
+    url="https://github.com/huggingface/pets",
+    package_dir={"": "src"},
+    packages=find_packages("src"),
+    entry_points={},
+    python_requires=">=3.7.0",
+    install_requires=[
+        "numpy>=1.17",
+        "packaging>=20.0",
+        "psutil",
+        "pyyaml",
+        "torch>=1.4.0",
+        "transformers",
+        "accelerate",
+    ],
+    extras_require=extras,
+    classifiers=[
+        "Development Status :: 5 - Production/Stable",
+        "Intended Audience :: Developers",
+        "Intended Audience :: Education",
+        "Intended Audience :: Science/Research",
+        "License :: OSI Approved :: Apache Software License",
+        "Operating System :: OS Independent",
+        "Programming Language :: Python :: 3",
+        "Programming Language :: Python :: 3.7",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    ],
+)
+
+# Release checklist
+# 1. Change the version in __init__.py and setup.py.
+# 2. Commit these changes with the message: "Release: VERSION"
+# 3. Add a tag in git to mark the release: "git tag VERSION -m 'Adds tag VERSION for pypi' "
+#    Push the tag to git: git push --tags origin main
+# 4. Run the following commands in the top-level directory:
+#      python setup.py bdist_wheel
+#      python setup.py sdist
+# 5. Upload the package to the pypi test server first:
+#      twine upload dist/* -r pypitest
+#      twine upload dist/* -r pypitest --repository-url=https://test.pypi.org/legacy/
+# 6. Check that you can install it in a virtualenv by running:
+#      pip install -i https://testpypi.python.org/pypi accelerate
+#      accelerate env
+#      accelerate test
+# 7. Upload the final version to actual pypi:
+#      twine upload dist/* -r pypi
+# 8. Add release notes to the tag in github once everything is looking hunky-dory.
+# 9. Update the version in __init__.py, setup.py to the new version "-dev" and push to master

src/pet/__init__.py

@@ -0,0 +1,30 @@
+# flake8: noqa
+# There's no way to ignore "F401 '...' imported but unused" warnings in this
+# module, but to preserve other warnings. So, don't check this module at all.
+
+__version__ = "0.1.0.dev0"
+
+from .mapping import MODEL_TYPE_TO_PET_MODEL_MAPPING, PET_TYPE_TO_CONFIG_MAPPING, get_pet_config, get_pet_model
+from .pet_model import PETModel, PETModelForCausalLM, PETModelForSeq2SeqLM, PETModelForSequenceClassification
+from .tuners import (
+    LoRAConfig,
+    LoRAModel,
+    PrefixEncoder,
+    PrefixTuningConfig,
+    PromptEmbedding,
+    PromptEncoder,
+    PromptEncoderConfig,
+    PromptEncoderReparameterizationType,
+    PromptTuningConfig,
+    PromptTuningInit,
+)
+from .utils import (
+    PETConfig,
+    PETType,
+    PromptLearningConfig,
+    TaskType,
+    bloom_model_postprocess_past_key_value,
+    get_pet_model_state_dict,
+    set_pet_model_state_dict,
+    shift_tokens_right,
+)

src/pet/mapping.py

@@ -0,0 +1,102 @@
+from .pet_model import PETModelForCausalLM, PETModelForSeq2SeqLM, PETModelForSequenceClassification
+from .tuners import LoRAConfig, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig
+from .utils import PETType
+
+
+MODEL_TYPE_TO_PET_MODEL_MAPPING = {
+    "SEQ_CLS": PETModelForSequenceClassification,
+    "SEQ_2_SEQ_LM": PETModelForSeq2SeqLM,
+    "CAUSAL_LM": PETModelForCausalLM,
+}
+
+PET_TYPE_TO_CONFIG_MAPPING = {
+    "PROMPT_TUNING": PromptTuningConfig,
+    "PREFIX_TUNING": PrefixTuningConfig,
+    "P_TUNING": PromptEncoderConfig,
+    "LORA": LoRAConfig,
+}
+
+TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = {
+    "t5": ["q", "v"],
+    "bart": ["q_proj", "v_proj"],
+    "gpt2": ["c_attn"],
+    "bloom": ["query_key_value"],
+    "opt": ["q_proj", "v_proj"],
+    "gptj": ["q_proj", "v_proj"],
+    "gpt_neox": ["query_key_value"],
+    "gpt_neo": ["q_proj", "v_proj"],
+    "bert": ["query", "value"],
+    "roberta": ["query", "value"],
+    "electra": ["query", "value"],
+    "deberta-v2": ["query_proj", "value_proj"],
+    "deberta": ["in_proj"],
+}
+
+
+def get_pet_config(config_dict):
+    return PET_TYPE_TO_CONFIG_MAPPING[config_dict["pet_type"]](**config_dict)
+
+
+def _prepare_prompt_learning_config(pet_config, model_config):
+    if pet_config.num_layers is None:
+        if "num_hidden_layers" in model_config:
+            num_layers = model_config["num_hidden_layers"]
+        elif "num_layers" in model_config:
+            num_layers = model_config["num_layers"]
+        elif "n_layer" in model_config:
+            num_layers = model_config["n_layer"]
+        else:
+            raise ValueError("Please specify `num_layers` in `pet_config`")
+        pet_config.num_layers = num_layers
+
+    if pet_config.token_dim is None:
+        if "hidden_size" in model_config:
+            token_dim = model_config["hidden_size"]
+        elif "n_embd" in model_config:
+            token_dim = model_config["n_embd"]
+        elif "d_model" in model_config:
+            token_dim = model_config["d_model"]
+        else:
+            raise ValueError("Please specify `token_dim` in `pet_config`")
+        pet_config.token_dim = token_dim
+
+    if pet_config.num_attention_heads is None:
+        if "num_attention_heads" in model_config:
+            num_attention_heads = model_config["num_attention_heads"]
+        elif "n_head" in model_config:
+            num_attention_heads = model_config["n_head"]
+        elif "num_heads" in model_config:
+            num_attention_heads = model_config["num_heads"]
+        elif "encoder_attention_heads" in model_config:
+            num_attention_heads = model_config["encoder_attention_heads"]
+        else:
+            raise ValueError("Please specify `num_attention_heads` in `pet_config`")
+        pet_config.num_attention_heads = num_attention_heads
+
+    if pet_config.encoder_hidden_size is None:
+        pet_config.encoder_hidden_size = token_dim
+
+    return pet_config
+
+
+def _prepare_lora_config(pet_config, model_config):
+    if pet_config.target_modules is None:
+        if model_config["model_type"] not in TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING:
+            raise ValueError("Please specify `target_modules` in `pet_config`")
+        pet_config.target_modules = TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING[model_config["model_type"]]
+    if len(pet_config.target_modules) == 1:
+        pet_config.fan_in_fan_out = True
+        pet_config.enable_lora = [True, False, True]
+    if pet_config.inference_mode:
+        pet_config.merge_weights = True
+    return pet_config
+
+
+def get_pet_model(model, pet_config):
+    model_config = model.config.to_dict()
+    if pet_config.pet_type != PETType.LORA:
+        pet_config = _prepare_prompt_learning_config(pet_config, model_config)
+    else:
+        pet_config = _prepare_lora_config(pet_config, model_config)
+
+    return MODEL_TYPE_TO_PET_MODEL_MAPPING[pet_config.task_type](model, pet_config)

src/pet/pet_model.py

@@ -0,0 +1,410 @@
+import inspect
+import warnings
+
+import torch
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import SequenceClassifierOutput
+
+from .tuners import LoRAModel, PrefixEncoder, PromptEmbedding, PromptEncoder
+from .utils import PETConfig, PETType, TaskType, shift_tokens_right
+
+
+class PETModel(torch.nn.Module):
+    def __init__(self, model, pet_config: PETConfig):
+        super().__init__()
+        self.pet_config = pet_config
+        self.base_model = model
+        self.modules_to_save = None
+        if pet_config.pet_type != PETType.LORA:
+            self._setup_prompt_encoder()
+        else:
+            self.base_model = LoRAModel(pet_config, model)
+
+    def _setup_prompt_encoder(self):
+        num_transformer_submodules = 0
+        transformer_backbone = None
+        for name, module in self.base_model.named_children():
+            if isinstance(module, PreTrainedModel):
+                # Make sure to freeze Tranformers model
+                for param in module.parameters():
+                    param.requires_grad = False
+                if transformer_backbone is None:
+                    transformer_backbone = module
+                    self.transformer_backbone_name = name
+                num_transformer_submodules += 1
+        self.pet_config.num_transformer_submodules = 2 if self.pet_config.task_type == TaskType.SEQ_2_SEQ_LM else 1
+
+        for named_param, value in list(transformer_backbone.named_parameters()):
+            if value.shape[0] == self.base_model.config.vocab_size:
+                self.word_embeddings = transformer_backbone.get_submodule(named_param.replace(".weight", ""))
+                break
+
+        if self.pet_config.pet_type == PETType.PROMPT_TUNING:
+            prompt_encoder = PromptEmbedding(self.pet_config, self.word_embeddings)
+        elif self.pet_config.pet_type == PETType.P_TUNING:
+            prompt_encoder = PromptEncoder(self.pet_config)
+        elif self.pet_config.pet_type == PETType.PREFIX_TUNING:
+            prompt_encoder = PrefixEncoder(self.pet_config)
+        else:
+            raise ValueError("Not supported")
+        self.prompt_encoder = prompt_encoder
+        self.prompt_tokens = torch.arange(
+            self.pet_config.num_virtual_tokens * self.pet_config.num_transformer_submodules
+        ).long()
+
+    def get_prompt_embedding_to_save(self):
+        prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(1, -1).to(self.base_model.device)
+        if self.pet_config.pet_type == PETType.PREFIX_TUNING:
+            prompt_tokens = prompt_tokens[:, : self.pet_config.num_virtual_tokens]
+        prompt_embeddings = self.prompt_encoder(prompt_tokens)
+        return prompt_embeddings[0].detach().cpu()
+
+    def get_prompt(self, batch_size):
+        prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(batch_size, -1).to(self.base_model.device)
+        if self.pet_config.pet_type == PETType.PREFIX_TUNING:
+            prompt_tokens = prompt_tokens[:, : self.pet_config.num_virtual_tokens]
+            if self.pet_config.inference_mode:
+                past_key_values = self.prompt_encoder.embedding.weight.repeat(batch_size, 1, 1)
+            else:
+                past_key_values = self.prompt_encoder(prompt_tokens)
+            past_key_values = past_key_values.view(
+                batch_size,
+                self.pet_config.num_virtual_tokens,
+                self.pet_config.num_layers * 2,
+                self.pet_config.num_attention_heads,
+                self.pet_config.token_dim // self.pet_config.num_attention_heads,
+            )
+            if self.pet_config.num_transformer_submodules == 2:
+                past_key_values = torch.cat([past_key_values, past_key_values], dim=2)
+            past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split(
+                self.pet_config.num_transformer_submodules * 2
+            )
+            if self.pet_config.postprocess_past_key_value_function is not None:
+                post_process_fn = self.pet_config.postprocess_past_key_value_function
+                past_key_values = post_process_fn(past_key_values)
+            return past_key_values
+        else:
+            if self.pet_config.inference_mode:
+                prompts = self.prompt_encoder.embedding.weight.repeat(batch_size, 1, 1)
+            else:
+                prompts = self.prompt_encoder(prompt_tokens)
+            return prompts
+
+    def print_trainable_parameters(self):
+        trainable_params = 0
+        all_param = 0
+        for _, param in self.named_parameters():
+            all_param += param.numel()
+            if param.requires_grad:
+                trainable_params += param.numel()
+        print(
+            f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}"
+        )
+
+
+class PETModelForSequenceClassification(PETModel):
+    def __init__(self, model, pet_config: PETConfig):
+        super().__init__(model, pet_config)
+        self.config = self.base_model.config
+        self.modules_to_save = ["classifier"]
+
+        for name, module in self.base_model.named_children():
+            if isinstance(module, torch.nn.Linear):
+                self.cls_layer_name = name
+                break
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        **kwargs,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if self.pet_config.pet_type == PETType.LORA:
+            return self.base_model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                inputs_embeds=inputs_embeds,
+                labels=labels,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                **kwargs,
+            )
+
+        batch_size = input_ids.shape[0]
+        if attention_mask is not None:
+            # concat prompt attention mask
+            prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(
+                self.base_model.device
+            )
+            attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+        if kwargs.get("position_ids", None) is not None:
+            warnings.warn("Position ids are not supported for parameter efficient tuning. Ignoring position ids.")
+            kwargs["position_ids"] = None
+        kwargs.update(
+            {
+                "attention_mask": attention_mask,
+                "labels": labels,
+                "output_attentions": output_attentions,
+                "output_hidden_states": output_hidden_states,
+                "return_dict": return_dict,
+            }
+        )
+
+        if self.pet_config.pet_type == PETType.PREFIX_TUNING:
+            return self.prefix_tuning_forward(input_ids=input_ids, **kwargs)
+        else:
+            if kwargs.get("token_type_ids", None) is not None:
+                kwargs["token_type_ids"] = torch.cat(
+                    (
+                        torch.zeros(batch_size, self.pet_config.num_virtual_tokens).to(self.base_model.device),
+                        kwargs["token_type_ids"],
+                    ),
+                    dim=1,
+                ).long()
+            if inputs_embeds is None:
+                inputs_embeds = self.word_embeddings(input_ids)
+            prompts = self.get_prompt(batch_size=batch_size)
+            inputs_embeds = torch.cat((prompts, inputs_embeds), dim=1)
+            return self.base_model(inputs_embeds=inputs_embeds, **kwargs)
+
+    def prefix_tuning_forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        **kwargs,
+    ):
+        batch_size = input_ids.shape[0]
+        past_key_values = self.get_prompt(batch_size)
+        fwd_params = list(inspect.signature(self.base_model.forward).parameters.keys())
+        kwargs.update(
+            {
+                "input_ids": input_ids,
+                "attention_mask": attention_mask,
+                "inputs_embeds": inputs_embeds,
+                "output_attentions": output_attentions,
+                "output_hidden_states": output_hidden_states,
+                "return_dict": return_dict,
+                "past_key_values": past_key_values,
+            }
+        )
+        if "past_key_values" in fwd_params:
+            return self.base_model(labels=labels, **kwargs)
+        else:
+            transformer_backbone_name = self.base_model.get_submodule(self.transformer_backbone_name)
+            fwd_params = list(inspect.signature(transformer_backbone_name.forward).parameters.keys())
+            if "past_key_values" not in fwd_params:
+                raise ValueError("Model does not support past key values which are required for prefix tuning.")
+            outputs = transformer_backbone_name(**kwargs)
+            pooled_output = outputs[1] if len(outputs) > 1 else outputs[0]
+            if "dropout" in [name for name, _ in list(self.base_model.named_children())]:
+                pooled_output = self.base_model.dropout(pooled_output)
+            logits = self.base_model.get_submodule(self.cls_layer_name)(pooled_output)
+
+            loss = None
+            if labels is not None:
+                if self.config.problem_type is None:
+                    if self.base_model.num_labels == 1:
+                        self.config.problem_type = "regression"
+                    elif self.base_model.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                        self.config.problem_type = "single_label_classification"
+                    else:
+                        self.config.problem_type = "multi_label_classification"
+
+                if self.config.problem_type == "regression":
+                    loss_fct = MSELoss()
+                    if self.base_model.num_labels == 1:
+                        loss = loss_fct(logits.squeeze(), labels.squeeze())
+                    else:
+                        loss = loss_fct(logits, labels)
+                elif self.config.problem_type == "single_label_classification":
+                    loss_fct = CrossEntropyLoss()
+                    loss = loss_fct(logits.view(-1, self.base_model.num_labels), labels.view(-1))
+                elif self.config.problem_type == "multi_label_classification":
+                    loss_fct = BCEWithLogitsLoss()
+                    loss = loss_fct(logits, labels)
+            if not return_dict:
+                output = (logits,) + outputs[2:]
+                return ((loss,) + output) if loss is not None else output
+
+            return SequenceClassifierOutput(
+                loss=loss,
+                logits=logits,
+                hidden_states=outputs.hidden_states,
+                attentions=outputs.attentions,
+            )
+
+
+class PETModelForCausalLM(PETModel):
+    def __init__(self, model, pet_config: PETConfig):
+        super().__init__(model, pet_config)
+        self.config = self.base_model.config
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        **kwargs,
+    ):
+        if self.pet_config.pet_type == PETType.LORA:
+            return self.base_model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                inputs_embeds=inputs_embeds,
+                labels=labels,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                **kwargs,
+            )
+
+        batch_size = input_ids.shape[0]
+        if attention_mask is not None:
+            # concat prompt attention mask
+            prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(
+                self.base_model.device
+            )
+            attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+
+        if kwargs.get("position_ids", None) is not None:
+            warnings.warn("Position ids are not supported for parameter efficient tuning. Ignoring position ids.")
+            kwargs["position_ids"] = None
+        if kwargs.get("token_type_ids", None) is not None:
+            warnings.warn("Token type ids are not supported for parameter efficient tuning. Ignoring token type ids")
+            kwargs["token_type_ids"] = None
+        kwargs.update(
+            {
+                "attention_mask": attention_mask,
+                "labels": labels,
+                "output_attentions": output_attentions,
+                "output_hidden_states": output_hidden_states,
+                "return_dict": return_dict,
+            }
+        )
+
+        if self.pet_config.pet_type == PETType.PREFIX_TUNING:
+            past_key_values = self.get_prompt(batch_size)
+            return self.base_model(input_ids=input_ids, past_key_values=past_key_values, **kwargs)
+        else:
+            if inputs_embeds is None:
+                inputs_embeds = self.word_embeddings(input_ids)
+            # concat prompt labels
+            if labels is not None:
+                prefix_labels = torch.full((batch_size, self.pet_config.num_virtual_tokens), -100).to(
+                    self.base_model.device
+                )
+                kwargs["labels"] = torch.cat((prefix_labels, labels), dim=1)
+            prompts = self.get_prompt(batch_size=batch_size)
+            inputs_embeds = torch.cat((prompts, inputs_embeds), dim=1)
+            return self.base_model(inputs_embeds=inputs_embeds, **kwargs)
+
+
+class PETModelForSeq2SeqLM(PETModel):
+    def __init__(self, model, pet_config: PETConfig):
+        super().__init__(model, pet_config)
+        self.config = self.base_model.config
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        decoder_input_ids=None,
+        decoder_attention_mask=None,
+        decoder_inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        **kwargs,
+    ):
+        if self.pet_config.pet_type == PETType.LORA:
+            return self.base_model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                inputs_embeds=inputs_embeds,
+                decoder_input_ids=decoder_input_ids,
+                decoder_attention_mask=decoder_attention_mask,
+                decoder_inputs_embeds=decoder_inputs_embeds,
+                labels=labels,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                **kwargs,
+            )
+
+        batch_size = input_ids.shape[0]
+        if decoder_attention_mask is not None:
+            # concat prompt attention mask
+            prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(
+                self.base_model.device
+            )
+            decoder_attention_mask = torch.cat((prefix_attention_mask, decoder_attention_mask), dim=1)
+
+        if kwargs.get("position_ids", None) is not None:
+            warnings.warn("Position ids are not supported for parameter efficient tuning. Ignoring position ids.")
+            kwargs["position_ids"] = None
+        if kwargs.get("token_type_ids", None) is not None:
+            warnings.warn("Token type ids are not supported for parameter efficient tuning. Ignoring token type ids")
+            kwargs["token_type_ids"] = None
+        kwargs.update(
+            {
+                "attention_mask": attention_mask,
+                "decoder_attention_mask": decoder_attention_mask,
+                "labels": labels,
+                "output_attentions": output_attentions,
+                "output_hidden_states": output_hidden_states,
+                "return_dict": return_dict,
+            }
+        )
+
+        if self.pet_config.pet_type == PETType.PREFIX_TUNING:
+            past_key_values = self.get_prompt(batch_size)
+            return self.base_model(
+                input_ids=input_ids, decoder_input_ids=decoder_input_ids, past_key_values=past_key_values, **kwargs
+            )
+        else:
+            if inputs_embeds is None:
+                inputs_embeds = self.word_embeddings(input_ids)
+            if decoder_inputs_embeds is None and decoder_input_ids is None:
+                decoder_input_ids = shift_tokens_right(
+                    labels, self.config.pad_token_id, self.config.decoder_start_token_id
+                )
+                decoder_inputs_embeds = self.word_embeddings(decoder_input_ids)
+
+            if attention_mask is not None:
+                # concat prompt attention mask
+                prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(
+                    self.base_model.device
+                )
+                kwargs["attention_mask"] = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+            # concat prompt labels
+            if labels is not None:
+                prefix_labels = torch.full((batch_size, self.pet_config.num_virtual_tokens), -100).to(
+                    self.base_model.device
+                )
+                kwargs["labels"] = torch.cat((prefix_labels, labels), dim=1)
+            prompts = self.get_prompt(batch_size=batch_size)
+            inputs_embeds = torch.cat((prompts[:, : self.pet_config.num_virtual_tokens], inputs_embeds), dim=1)
+            decoder_inputs_embeds = torch.cat(
+                (prompts[:, self.pet_config.num_virtual_tokens :], decoder_inputs_embeds), dim=1
+            )
+            return self.base_model(inputs_embeds=inputs_embeds, decoder_inputs_embeds=decoder_inputs_embeds, **kwargs)

src/pet/tuners/__init__.py

@@ -0,0 +1,8 @@
+# flake8: noqa
+# There's no way to ignore "F401 '...' imported but unused" warnings in this
+# module, but to preserve other warnings. So, don't check this module at all
+
+from .lora import LoRAConfig, LoRAModel
+from .p_tuning import PromptEncoder, PromptEncoderConfig, PromptEncoderReparameterizationType
+from .prefix_tuning import PrefixEncoder, PrefixTuningConfig
+from .prompt_tuning import PromptEmbedding, PromptTuningConfig, PromptTuningInit

src/pet/tuners/lora.py

@@ -0,0 +1,73 @@
+# todo
+from dataclasses import dataclass, field
+from typing import Optional
+
+import torch
+from transformers.pytorch_utils import Conv1D
+
+import loralib as lora
+from loralib import mark_only_lora_as_trainable
+
+from ..utils import PETConfig
+
+
+@dataclass
+class LoRAConfig(PETConfig):
+    r: int = field(default=8, metadata={"help": "LoRA attention dimension"})
+    target_modules: Optional[list] = field(default=None, metadata={"help": "List of modules to replace with LoRA"})
+    lora_alpha: int = field(default=None, metadata={"help": "LoRA alpha"})
+    lora_dropout: float = field(default=None, metadata={"help": "LoRA dropout"})
+    merge_weights: bool = field(
+        default=False, metadata={"help": "Merge weights of the original model and the LoRA model"}
+    )
+    fan_in_fan_out: bool = field(
+        default=False,
+        metadata={"help": "Set this to True if the layer to replace stores weight like (fan_in, fan_out)"},
+    )
+    enable_lora: Optional[list[bool]] = field(default=None, metadata={"help": "Used with `lora.MergedLinear`."})
+    bias: str = field(default="none", metadata={"help": "Bias type for LoRA. Can be 'none', 'all' or 'lora_only'"})
+
+
+class LoRAModel(torch.nn.Module):
+    def __init__(self, config, model):
+        super().__init__()
+        self.config = config
+        self.model = model
+        self.find_and_replace()
+        mark_only_lora_as_trainable(self.model, self.config.bias)
+
+    def find_and_replace(self):
+        kwargs = {
+            "r": self.config.r,
+            "lora_alpha": self.config.lora_alpha,
+            "lora_dropout": self.config.lora_dropout,
+            "fan_in_fan_out": self.config.fan_in_fan_out,
+            "merge_weights": self.config.merge_weights,
+        }
+        key_list = [key for key, _ in self.model.named_modules()]
+        for key in key_list:
+            if any(key.endswith(target_key) for target_key in self.config.target_modules):
+                parent, target, target_name = self.get_submodules(key)
+                # print(parent, target, target_name)
+                if isinstance(target, torch.nn.Linear):
+                    new_module = lora.Linear(target.in_features, target.out_features, **kwargs)
+                elif isinstance(target, Conv1D):
+                    kwargs.update({"enable_lora": self.config.enable_lora})
+                    in_features, out_features = target.weight.shape
+                    new_module = lora.MergedLinear(in_features, out_features, **kwargs)
+                self.replace_module(parent, target_name, new_module, target)
+
+    def get_submodules(self, key):
+        parent = self.model.get_submodule(".".join(key.split(".")[:-1]))
+        target_name = key.split(".")[-1]
+        target = self.model.get_submodule(key)
+        return parent, target, target_name
+
+    def replace_module(self, parent_module, child_name, new_module, old_module):
+        setattr(parent_module, child_name, new_module)
+        new_module.weight = old_module.weight
+        if old_module.bias is not None:
+            new_module.bias = old_module.bias
+
+    def forward(self, *args, **kwargs):
+        return self.model(*args, **kwargs)

src/pet/tuners/p_tuning.py

@@ -0,0 +1,99 @@
+import enum
+from dataclasses import dataclass, field
+from typing import Union
+
+import torch
+
+from ..utils import PromptLearningConfig
+
+
+class PromptEncoderReparameterizationType(str, enum.Enum):
+    MLP = "MLP"
+    LSTM = "LSTM"
+
+
+@dataclass
+class PromptEncoderConfig(PromptLearningConfig):
+    encoder_reparameterization_type: Union[str, PromptEncoderReparameterizationType] = field(
+        default=PromptEncoderReparameterizationType.MLP,
+        metadata={"help": "How to reparameterize the prompt encoder"},
+    )
+    encoder_hidden_size: int = field(
+        default=None,
+        metadata={"help": "The hidden size of the prompt encoder reparameterization"},
+    )
+    encoder_num_layers: int = field(
+        default=2,
+        metadata={"help": "The number of layers of the prompt encoder reparameterization"},
+    )
+    encoder_dropout: float = field(
+        default=0.0,
+        metadata={"help": "The dropout of the prompt encoder reparameterization"},
+    )
+
+
+# Based on https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/nlp/modules/common/prompt_encoder.py
+# with some refactor
+class PromptEncoder(torch.nn.Module):
+    """
+    The prompt encoder network that is used to generate the virtual token embeddings for p-tuning.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.token_dim = config.token_dim
+        self.input_size = self.token_dim
+        self.output_size = self.token_dim
+        self.hidden_size = config.encoder_hidden_size
+        self.total_virtual_tokens = config.num_virtual_tokens * config.num_transformer_submodules
+        self.encoder_type = config.encoder_reparameterization_type
+
+        # embedding
+        self.embedding = torch.nn.Embedding(self.total_virtual_tokens, self.token_dim)
+        if not config.inference_mode:
+            if self.encoder_type == PromptEncoderReparameterizationType.LSTM:
+                lstm_dropout = config.encoder_dropout
+                num_layers = config.encoder_num_layers
+                # LSTM
+                self.lstm_head = torch.nn.LSTM(
+                    input_size=self.input_size,
+                    hidden_size=self.hidden_size,
+                    num_layers=num_layers,
+                    dropout=lstm_dropout,
+                    bidirectional=True,
+                    batch_first=True,
+                )
+
+                self.mlp_head = torch.nn.Sequential(
+                    torch.nn.Linear(self.hidden_size * 2, self.hidden_size * 2),
+                    torch.nn.ReLU(),
+                    torch.nn.Linear(self.hidden_size * 2, self.output_size),
+                )
+
+            elif self.encoder_type == PromptEncoderReparameterizationType.MLP:
+                layers = [
+                    torch.nn.Linear(self.input_size, self.hidden_size),
+                    torch.nn.ReLU(),
+                ]
+                layers.extend(
+                    [
+                        torch.nn.Linear(self.hidden_size, self.hidden_size),
+                        torch.nn.ReLU(),
+                    ]
+                )
+                layers.append(torch.nn.Linear(self.hidden_size, self.output_size))
+                self.mlp_head = torch.nn.Sequential(*layers)
+
+            else:
+                raise ValueError("Prompt encoder type not recognized. Please use one of MLP (recommended) or LSTM.")
+
+    def forward(self, indices):
+        input_embeds = self.embedding(indices)
+        if self.encoder_type == PromptEncoderReparameterizationType.LSTM:
+            output_embeds = self.mlp_head(self.lstm_head(input_embeds)[0])
+        elif self.encoder_type == PromptEncoderReparameterizationType.MLP:
+            output_embeds = self.mlp_head(input_embeds)
+        else:
+            raise ValueError("Prompt encoder type not recognized. Please use one of MLP (recommended) or LSTM.")
+
+        return output_embeds

src/pet/tuners/prefix_tuning.py

@@ -0,0 +1,60 @@
+from dataclasses import dataclass, field
+from typing import Callable, Optional
+
+import torch
+
+from ..utils import PromptLearningConfig
+
+
+@dataclass
+class PrefixTuningConfig(PromptLearningConfig):
+    encoder_hidden_size: int = field(
+        default=None,
+        metadata={"help": "The hidden size of the encoder"},
+    )
+    prefix_projection: bool = field(
+        default=False,
+        metadata={"help": "Whether to project the prefix tokens"},
+    )
+    postprocess_past_key_value_function: Optional[Callable] = field(
+        default=None,
+        metadata={"help": "The function to postprocess the past key value"},
+    )
+
+
+# Based on https://github.com/THUDM/P-tuning-v2/blob/main/model/prefix_encoder.py
+# with some refactor
+class PrefixEncoder(torch.nn.Module):
+    r"""
+    The torch.nn model to encode the prefix
+
+    Input shape: (batch_size, num_virtual_tokens)
+
+    Output shape: (batch_size, num_virtual_tokens, 2*layers*hidden)
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.prefix_projection = config.prefix_projection
+        token_dim = config.token_dim
+        num_layers = config.num_layers
+        encoder_hidden_size = config.encoder_hidden_size
+        num_virtual_tokens = config.num_virtual_tokens
+        if self.prefix_projection and not config.inference_mode:
+            # Use a two-layer MLP to encode the prefix
+            self.embedding = torch.nn.Embedding(num_virtual_tokens, token_dim)
+            self.trans = torch.nn.Sequential(
+                torch.nn.Linear(token_dim, encoder_hidden_size),
+                torch.nn.Tanh(),
+                torch.nn.Linear(encoder_hidden_size, num_layers * 2 * token_dim),
+            )
+        else:
+            self.embedding = torch.nn.Embedding(num_virtual_tokens, num_layers * 2 * token_dim)
+
+    def forward(self, prefix: torch.Tensor):
+        if self.prefix_projection:
+            prefix_tokens = self.embedding(prefix)
+            past_key_values = self.trans(prefix_tokens)
+        else:
+            past_key_values = self.embedding(prefix)
+        return past_key_values

src/pet/tuners/prompt_tuning.py

@@ -0,0 +1,63 @@
+import enum
+import math
+from dataclasses import dataclass, field
+from typing import Optional, Union
+
+import torch
+
+from ..utils import PromptLearningConfig
+
+
+class PromptTuningInit(str, enum.Enum):
+    TEXT = "TEXT"
+    RANDOM = "RANDOM"
+
+
+@dataclass
+class PromptTuningConfig(PromptLearningConfig):
+    prompt_tuning_init: Union[PromptTuningInit, str] = field(
+        default=PromptTuningInit.RANDOM,
+        metadata={"help": "How to initialize the prompt tuning parameters"},
+    )
+    prompt_tuning_init_text: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "The text to use for prompt tuning initialization. Only used if prompt_tuning_init is `TEXT`"
+        },
+    )
+    tokenizer_name_or_path: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "The tokenizer to use for prompt tuning initialization. Only used if prompt_tuning_init is `TEXT`"
+        },
+    )
+
+
+class PromptEmbedding(torch.nn.Module):
+    def __init__(self, config, word_embeddings):
+        super().__init__()
+
+        total_virtual_tokens = config.num_virtual_tokens * config.num_transformer_submodules
+        self.embedding = torch.nn.Embedding(total_virtual_tokens, config["token_dim"])
+        if config.prompt_tuning_init == PromptTuningInit.TEXT:
+            from transformers import AutoTokenizer
+
+            self.tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_name_or_path)
+            self.init_text = config.prompt_tuning_init_text
+            init_token_ids = self.tokenizer(self.init_text)["input_ids"]
+            # Trim or iterate until num_text_tokens matches total_virtual_tokens
+            num_text_tokens = len(init_token_ids)
+            if num_text_tokens > total_virtual_tokens:
+                init_token_ids = init_token_ids[:total_virtual_tokens]
+            elif num_text_tokens < total_virtual_tokens:
+                num_reps = math.ceil(total_virtual_tokens / num_text_tokens)
+                init_token_ids = init_token_ids * num_reps
+            init_token_ids = init_token_ids[:total_virtual_tokens]
+
+            word_embedding_weights = word_embeddings(torch.LongTensor(init_token_ids)).detach().clone()
+            self.embedding.weight = torch.nn.Parameter(word_embedding_weights)
+
+    def forward(self, indices):
+        # Just get embeddings
+        prompt_embeddings = self.embedding(indices)
+        return prompt_embeddings

src/pet/utils/__init__.py

@@ -0,0 +1,7 @@
+# flake8: noqa
+# There's no way to ignore "F401 '...' imported but unused" warnings in this
+# module, but to preserve other warnings. So, don't check this module at all
+
+from .config import PETConfig, PETType, PromptLearningConfig, TaskType
+from .other import bloom_model_postprocess_past_key_value, shift_tokens_right
+from .save_and_load import get_pet_model_state_dict, set_pet_model_state_dict

src/pet/utils/config.py

@@ -0,0 +1,36 @@
+import enum
+from dataclasses import dataclass, field
+from typing import Optional, Union
+
+
+class PETType(str, enum.Enum):
+    PROMPT_TUNING = "PROMPT_TUNING"
+    P_TUNING = "P_TUNING"
+    PREFIX_TUNING = "PREFIX_TUNING"
+    LORA = "LORA"
+
+
+class TaskType(str, enum.Enum):
+    SEQ_CLS = "SEQ_CLS"
+    SEQ_2_SEQ_LM = "SEQ_2_SEQ_LM"
+    CAUSAL_LM = "CAUSAL_LM"
+
+
+@dataclass
+class PETConfig:
+    """
+    This is the configuration class to store the configuration of a :class:`~pet.PETModel`.
+    """
+
+    pet_type: Union[str, PETType] = field(default=None, metadata={"help": "PET type"})
+    task_type: Union[str, TaskType] = field(default=None, metadata={"help": "Task type"})
+    inference_mode: bool = field(default=False, metadata={"help": "Whether to use inference mode"})
+
+
+@dataclass
+class PromptLearningConfig(PETConfig):
+    num_virtual_tokens: int = field(default=None, metadata={"help": "Number of virtual tokens"})
+    token_dim: int = field(default=None, metadata={"help": "Dimension of virtual tokens"})
+    num_transformer_submodules: Optional[int] = field(default=1, metadata={"help": "Number of transformer submodules"})
+    num_attention_heads: Optional[int] = field(default=None, metadata={"help": "Number of attention heads"})
+    num_layers: Optional[int] = field(default=None, metadata={"help": "Number of transformer layers"})

src/pet/utils/other.py

@@ -0,0 +1,32 @@
+import torch
+
+
+# needed for prefix-tuning of bloom model
+def bloom_model_postprocess_past_key_value(past_key_values):
+    past_key_values = torch.cat(past_key_values)
+    total_layers, batch_size, num_attention_heads, num_virtual_tokens, head_dim = past_key_values.shape
+    keys = past_key_values[: total_layers // 2]
+    keys = keys.transpose(2, 3).reshape(
+        total_layers // 2, batch_size * num_attention_heads, head_dim, num_virtual_tokens
+    )
+    values = past_key_values[total_layers // 2 :]
+    values = values.reshape(total_layers // 2, batch_size * num_attention_heads, num_virtual_tokens, head_dim)
+
+    return tuple(zip(keys, values))
+
+
+# copied from transformers.models.bart.modeling_bart
+def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):
+    """
+    Shift input ids one token 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] = decoder_start_token_id
+
+    if pad_token_id is None:
+        raise ValueError("self.model.config.pad_token_id has to be defined.")
+    # replace possible -100 values in labels by `pad_token_id`
+    shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
+
+    return shifted_input_ids

src/pet/utils/save_and_load.py

@@ -0,0 +1,27 @@
+from loralib import lora_state_dict
+
+from .config import PETType
+
+
+def get_pet_model_state_dict(model):
+    if model.pet_config.pet_type == PETType.LORA:
+        return lora_state_dict(model)
+    else:
+        to_return = {}
+        state_dict = model.state_dict()
+        prompt_embeddings = model.get_prompt_embedding_to_save()
+        to_return["prompt_embeddings"] = prompt_embeddings
+        if model.modules_to_save is not None:
+            for key, value in state_dict.items():
+                if any(module_name in key for module_name in model.modules_to_save):
+                    to_return[key] = value
+        return to_return
+
+
+def set_pet_model_state_dict(model, pet_model_state_dict):
+    model.load_state_dict(pet_model_state_dict, strict=False)
+    if model.pet_config.pet_type != PETType.LORA:
+        model.prompt_encoder.embedding.load_state_dict(
+            {"weight": pet_model_state_dict["prompt_embeddings"]}, strict=True
+        )
+    return model

utils/style_doc.py

@@ -0,0 +1,556 @@
+# coding=utf-8
+# Copyright 2020 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Style utils for the .rst and the docstrings."""
+
+import argparse
+import os
+import re
+import warnings
+
+import black
+
+
+BLACK_AVOID_PATTERNS = {}
+
+
+# Regexes
+# Re pattern that catches list introduction (with potential indent)
+_re_list = re.compile(r"^(\s*-\s+|\s*\*\s+|\s*\d+\.\s+)")
+# Re pattern that catches code block introduction (with potential indent)
+_re_code = re.compile(r"^(\s*)```(.*)$")
+# Re pattern that catches rst args blocks of the form `Parameters:`.
+_re_args = re.compile("^\s*(Args?|Arguments?|Params?|Parameters?):\s*$")
+# Re pattern that catches return blocks of the form `Return:`.
+_re_returns = re.compile("^\s*Returns?:\s*$")
+# Matches the special tag to ignore some paragraphs.
+_re_doc_ignore = re.compile(r"(\.\.|#)\s*docstyle-ignore")
+# Re pattern that matches <Tip>, </Tip> and <Tip warning={true}> blocks.
+_re_tip = re.compile("^\s*</?Tip(>|\s+warning={true}>)\s*$")
+
+DOCTEST_PROMPTS = [">>>", "..."]
+
+
+def is_empty_line(line):
+    return len(line) == 0 or line.isspace()
+
+
+def find_indent(line):
+    """
+    Returns the number of spaces that start a line indent.
+    """
+    search = re.search("^(\s*)(?:\S|$)", line)
+    if search is None:
+        return 0
+    return len(search.groups()[0])
+
+
+def parse_code_example(code_lines):
+    """
+    Parses a code example
+
+    Args:
+        code_lines (`List[str]`): The code lines to parse.
+        max_len (`int`): The maximum length per line.
+
+    Returns:
+        (List[`str`], List[`str`]): The list of code samples and the list of outputs.
+    """
+    has_doctest = code_lines[0][:3] in DOCTEST_PROMPTS
+
+    code_samples = []
+    outputs = []
+    in_code = True
+    current_bit = []
+
+    for line in code_lines:
+        if in_code and has_doctest and not is_empty_line(line) and line[:3] not in DOCTEST_PROMPTS:
+            code_sample = "\n".join(current_bit)
+            code_samples.append(code_sample.strip())
+            in_code = False
+            current_bit = []
+        elif not in_code and line[:3] in DOCTEST_PROMPTS:
+            output = "\n".join(current_bit)
+            outputs.append(output.strip())
+            in_code = True
+            current_bit = []
+
+        # Add the line without doctest prompt
+        if line[:3] in DOCTEST_PROMPTS:
+            line = line[4:]
+        current_bit.append(line)
+
+    # Add last sample
+    if in_code:
+        code_sample = "\n".join(current_bit)
+        code_samples.append(code_sample.strip())
+    else:
+        output = "\n".join(current_bit)
+        outputs.append(output.strip())
+
+    return code_samples, outputs
+
+
+def format_code_example(code: str, max_len: int, in_docstring: bool = False):
+    """
+    Format a code example using black. Will take into account the doctest syntax as well as any initial indentation in
+    the code provided.
+
+    Args:
+        code (`str`): The code example to format.
+        max_len (`int`): The maximum length per line.
+        in_docstring (`bool`, *optional*, defaults to `False`): Whether or not the code example is inside a docstring.
+
+    Returns:
+        `str`: The formatted code.
+    """
+    code_lines = code.split("\n")
+
+    # Find initial indent
+    idx = 0
+    while idx < len(code_lines) and is_empty_line(code_lines[idx]):
+        idx += 1
+    if idx >= len(code_lines):
+        return "", ""
+    indent = find_indent(code_lines[idx])
+
+    # Remove the initial indent for now, we will had it back after styling.
+    # Note that l[indent:] works for empty lines
+    code_lines = [l[indent:] for l in code_lines[idx:]]
+    has_doctest = code_lines[0][:3] in DOCTEST_PROMPTS
+
+    code_samples, outputs = parse_code_example(code_lines)
+
+    # Let's blackify the code! We put everything in one big text to go faster.
+    delimiter = "\n\n### New code sample ###\n"
+    full_code = delimiter.join(code_samples)
+    line_length = max_len - indent
+    if has_doctest:
+        line_length -= 4
+
+    for k, v in BLACK_AVOID_PATTERNS.items():
+        full_code = full_code.replace(k, v)
+    try:
+        mode = black.Mode(target_versions={black.TargetVersion.PY37}, line_length=line_length)
+        formatted_code = black.format_str(full_code, mode=mode)
+        error = ""
+    except Exception as e:
+        formatted_code = full_code
+        error = f"Code sample:\n{full_code}\n\nError message:\n{e}"
+
+    # Let's get back the formatted code samples
+    for k, v in BLACK_AVOID_PATTERNS.items():
+        formatted_code = formatted_code.replace(v, k)
+    # Triple quotes will mess docstrings.
+    if in_docstring:
+        formatted_code = formatted_code.replace('"""', "'''")
+
+    code_samples = formatted_code.split(delimiter)
+    # We can have one output less than code samples
+    if len(outputs) == len(code_samples) - 1:
+        outputs.append("")
+
+    formatted_lines = []
+    for code_sample, output in zip(code_samples, outputs):
+        # black may have added some new lines, we remove them
+        code_sample = code_sample.strip()
+        in_triple_quotes = False
+        in_decorator = False
+        for line in code_sample.strip().split("\n"):
+            if has_doctest and not is_empty_line(line):
+                prefix = (
+                    "... "
+                    if line.startswith(" ") or line in [")", "]", "}"] or in_triple_quotes or in_decorator
+                    else ">>> "
+                )
+            else:
+                prefix = ""
+            indent_str = "" if is_empty_line(line) else (" " * indent)
+            formatted_lines.append(indent_str + prefix + line)
+
+            if '"""' in line:
+                in_triple_quotes = not in_triple_quotes
+            if line.startswith(" "):
+                in_decorator = False
+            if line.startswith("@"):
+                in_decorator = True
+
+        formatted_lines.extend([" " * indent + line for line in output.split("\n")])
+        if not output.endswith("===PT-TF-SPLIT==="):
+            formatted_lines.append("")
+
+    result = "\n".join(formatted_lines)
+    return result.rstrip(), error
+
+
+def format_text(text, max_len, prefix="", min_indent=None):
+    """
+    Format a text in the biggest lines possible with the constraint of a maximum length and an indentation.
+
+    Args:
+        text (`str`): The text to format
+        max_len (`int`): The maximum length per line to use
+        prefix (`str`, *optional*, defaults to `""`): A prefix that will be added to the text.
+            The prefix doesn't count toward the indent (like a - introducing a list).
+        min_indent (`int`, *optional*): The minimum indent of the text.
+            If not set, will default to the length of the `prefix`.
+
+    Returns:
+        `str`: The formatted text.
+    """
+    text = re.sub(r"\s+", " ", text)
+    if min_indent is not None:
+        if len(prefix) < min_indent:
+            prefix = " " * (min_indent - len(prefix)) + prefix
+
+    indent = " " * len(prefix)
+    new_lines = []
+    words = text.split(" ")
+    current_line = f"{prefix}{words[0]}"
+    for word in words[1:]:
+        try_line = f"{current_line} {word}"
+        if len(try_line) > max_len:
+            new_lines.append(current_line)
+            current_line = f"{indent}{word}"
+        else:
+            current_line = try_line
+    new_lines.append(current_line)
+    return "\n".join(new_lines)
+
+
+def split_line_on_first_colon(line):
+    splits = line.split(":")
+    return splits[0], ":".join(splits[1:])
+
+
+def style_docstring(docstring, max_len):
+    """
+    Style a docstring by making sure there is no useless whitespace and the maximum horizontal space is used.
+
+    Args:
+        docstring (`str`): The docstring to style.
+        max_len (`int`): The maximum length of each line.
+
+    Returns:
+        `str`: The styled docstring
+    """
+    lines = docstring.split("\n")
+    new_lines = []
+
+    # Initialization
+    current_paragraph = None
+    current_indent = -1
+    in_code = False
+    param_indent = -1
+    prefix = ""
+    black_errors = []
+
+    # Special case for docstrings that begin with continuation of Args with no Args block.
+    idx = 0
+    while idx < len(lines) and is_empty_line(lines[idx]):
+        idx += 1
+    if (
+        len(lines[idx]) > 1
+        and lines[idx].rstrip().endswith(":")
+        and find_indent(lines[idx + 1]) > find_indent(lines[idx])
+    ):
+        param_indent = find_indent(lines[idx])
+
+    for idx, line in enumerate(lines):
+        # Doing all re searches once for the one we need to repeat.
+        list_search = _re_list.search(line)
+        code_search = _re_code.search(line)
+
+        # Are we starting a new paragraph?
+        # New indentation or new line:
+        new_paragraph = find_indent(line) != current_indent or is_empty_line(line)
+        # List item
+        new_paragraph = new_paragraph or list_search is not None
+        # Code block beginning
+        new_paragraph = new_paragraph or code_search is not None
+        # Beginning/end of tip
+        new_paragraph = new_paragraph or _re_tip.search(line)
+
+        # In this case, we treat the current paragraph
+        if not in_code and new_paragraph and current_paragraph is not None and len(current_paragraph) > 0:
+            paragraph = " ".join(current_paragraph)
+            new_lines.append(format_text(paragraph, max_len, prefix=prefix, min_indent=current_indent))
+            current_paragraph = None
+
+        if code_search is not None:
+            if not in_code:
+                current_paragraph = []
+                current_indent = len(code_search.groups()[0])
+                current_code = code_search.groups()[1]
+                prefix = ""
+                if current_indent < param_indent:
+                    param_indent = -1
+            else:
+                current_indent = -1
+                code = "\n".join(current_paragraph)
+                if current_code in ["py", "python"]:
+                    formatted_code, error = format_code_example(code, max_len, in_docstring=True)
+                    new_lines.append(formatted_code)
+                    if len(error) > 0:
+                        black_errors.append(error)
+                else:
+                    new_lines.append(code)
+                current_paragraph = None
+            new_lines.append(line)
+            in_code = not in_code
+
+        elif in_code:
+            current_paragraph.append(line)
+        elif is_empty_line(line):
+            current_paragraph = None
+            current_indent = -1
+            prefix = ""
+            new_lines.append(line)
+        elif list_search is not None:
+            prefix = list_search.groups()[0]
+            current_indent = len(prefix)
+            current_paragraph = [line[current_indent:]]
+        elif _re_args.search(line):
+            new_lines.append(line)
+            param_indent = find_indent(lines[idx + 1])
+        elif _re_tip.search(line):
+            # Add a new line before if not present
+            if not is_empty_line(new_lines[-1]):
+                new_lines.append("")
+            new_lines.append(line)
+            # Add a new line after if not present
+            if idx < len(lines) - 1 and not is_empty_line(lines[idx + 1]):
+                new_lines.append("")
+        elif current_paragraph is None or find_indent(line) != current_indent:
+            indent = find_indent(line)
+            # Special behavior for parameters intros.
+            if indent == param_indent:
+                # Special rules for some docstring where the Returns blocks has the same indent as the parameters.
+                if _re_returns.search(line) is not None:
+                    param_indent = -1
+                    new_lines.append(line)
+                elif len(line) < max_len:
+                    new_lines.append(line)
+                else:
+                    intro, description = split_line_on_first_colon(line)
+                    new_lines.append(intro + ":")
+                    if len(description) != 0:
+                        if find_indent(lines[idx + 1]) > indent:
+                            current_indent = find_indent(lines[idx + 1])
+                        else:
+                            current_indent = indent + 4
+                        current_paragraph = [description.strip()]
+                        prefix = ""
+            else:
+                # Check if we have exited the parameter block
+                if indent < param_indent:
+                    param_indent = -1
+
+                current_paragraph = [line.strip()]
+                current_indent = find_indent(line)
+                prefix = ""
+        elif current_paragraph is not None:
+            current_paragraph.append(line.lstrip())
+
+    if current_paragraph is not None and len(current_paragraph) > 0:
+        paragraph = " ".join(current_paragraph)
+        new_lines.append(format_text(paragraph, max_len, prefix=prefix, min_indent=current_indent))
+
+    return "\n".join(new_lines), "\n\n".join(black_errors)
+
+
+def style_docstrings_in_code(code, max_len=119):
+    """
+    Style all docstrings in some code.
+
+    Args:
+        code (`str`): The code in which we want to style the docstrings.
+        max_len (`int`): The maximum number of characters per line.
+
+    Returns:
+        `Tuple[str, str]`: A tuple with the clean code and the black errors (if any)
+    """
+    # fmt: off
+    splits = code.split('\"\"\"')
+    splits = [
+        (s if i % 2 == 0 or _re_doc_ignore.search(splits[i - 1]) is not None else style_docstring(s, max_len=max_len))
+        for i, s in enumerate(splits)
+    ]
+    black_errors = "\n\n".join([s[1] for s in splits if isinstance(s, tuple) and len(s[1]) > 0])
+    splits = [s[0] if isinstance(s, tuple) else s for s in splits]
+    clean_code = '\"\"\"'.join(splits)
+    # fmt: on
+
+    return clean_code, black_errors
+
+
+def style_file_docstrings(code_file, max_len=119, check_only=False):
+    """
+    Style all docstrings in a given file.
+
+    Args:
+        code_file (`str` or `os.PathLike`): The file in which we want to style the docstring.
+        max_len (`int`): The maximum number of characters per line.
+        check_only (`bool`, *optional*, defaults to `False`):
+            Whether to restyle file or just check if they should be restyled.
+
+    Returns:
+        `bool`: Whether or not the file was or should be restyled.
+    """
+    with open(code_file, "r", encoding="utf-8", newline="\n") as f:
+        code = f.read()
+
+    clean_code, black_errors = style_docstrings_in_code(code, max_len=max_len)
+
+    diff = clean_code != code
+    if not check_only and diff:
+        print(f"Overwriting content of {code_file}.")
+        with open(code_file, "w", encoding="utf-8", newline="\n") as f:
+            f.write(clean_code)
+
+    return diff, black_errors
+
+
+def style_mdx_file(mdx_file, max_len=119, check_only=False):
+    """
+    Style a MDX file by formatting all Python code samples.
+
+    Args:
+        mdx_file (`str` or `os.PathLike`): The file in which we want to style the examples.
+        max_len (`int`): The maximum number of characters per line.
+        check_only (`bool`, *optional*, defaults to `False`):
+            Whether to restyle file or just check if they should be restyled.
+
+    Returns:
+        `bool`: Whether or not the file was or should be restyled.
+    """
+    with open(mdx_file, "r", encoding="utf-8", newline="\n") as f:
+        content = f.read()
+
+    lines = content.split("\n")
+    current_code = []
+    current_language = ""
+    in_code = False
+    new_lines = []
+    black_errors = []
+
+    for line in lines:
+        if _re_code.search(line) is not None:
+            in_code = not in_code
+            if in_code:
+                current_language = _re_code.search(line).groups()[1]
+                current_code = []
+            else:
+                code = "\n".join(current_code)
+                if current_language in ["py", "python"]:
+                    code, error = format_code_example(code, max_len)
+                    if len(error) > 0:
+                        black_errors.append(error)
+                new_lines.append(code)
+
+            new_lines.append(line)
+        elif in_code:
+            current_code.append(line)
+        else:
+            new_lines.append(line)
+
+    if in_code:
+        raise ValueError(f"There was a problem when styling {mdx_file}. A code block is opened without being closed.")
+
+    clean_content = "\n".join(new_lines)
+    diff = clean_content != content
+    if not check_only and diff:
+        print(f"Overwriting content of {mdx_file}.")
+        with open(mdx_file, "w", encoding="utf-8", newline="\n") as f:
+            f.write(clean_content)
+
+    return diff, "\n\n".join(black_errors)
+
+
+def style_doc_files(*files, max_len=119, check_only=False):
+    """
+    Applies doc styling or checks everything is correct in a list of files.
+
+    Args:
+        files (several `str` or `os.PathLike`): The files to treat.
+        max_len (`int`): The maximum number of characters per line.
+        check_only (`bool`, *optional*, defaults to `False`):
+            Whether to restyle file or just check if they should be restyled.
+
+    Returns:
+        List[`str`]: The list of files changed or that should be restyled.
+    """
+    changed = []
+    black_errors = []
+    for file in files:
+        # Treat folders
+        if os.path.isdir(file):
+            files = [os.path.join(file, f) for f in os.listdir(file)]
+            files = [f for f in files if os.path.isdir(f) or f.endswith(".mdx") or f.endswith(".py")]
+            changed += style_doc_files(*files, max_len=max_len, check_only=check_only)
+        # Treat mdx
+        elif file.endswith(".mdx"):
+            try:
+                diff, black_error = style_mdx_file(file, max_len=max_len, check_only=check_only)
+                if diff:
+                    changed.append(file)
+                if len(black_error) > 0:
+                    black_errors.append(
+                        f"There was a problem while formatting an example in {file} with black:\m{black_error}"
+                    )
+            except Exception:
+                print(f"There is a problem in {file}.")
+                raise
+        # Treat python files
+        elif file.endswith(".py"):
+            try:
+                diff, black_error = style_file_docstrings(file, max_len=max_len, check_only=check_only)
+                if diff:
+                    changed.append(file)
+                if len(black_error) > 0:
+                    black_errors.append(
+                        f"There was a problem while formatting an example in {file} with black:\m{black_error}"
+                    )
+            except Exception:
+                print(f"There is a problem in {file}.")
+                raise
+        else:
+            warnings.warn(f"Ignoring {file} because it's not a py or an mdx file or a folder.")
+    if len(black_errors) > 0:
+        black_message = "\n\n".join(black_errors)
+        raise ValueError(
+            "Some code examples can't be interpreted by black, which means they aren't regular python:\n\n"
+            + black_message
+            + "\n\nMake sure to fix the corresponding docstring or doc file, or remove the py/python after ``` if it "
+            + "was not supposed to be a Python code sample."
+        )
+    return changed
+
+
+def main(*files, max_len=119, check_only=False):
+    changed = style_doc_files(*files, max_len=max_len, check_only=check_only)
+    if check_only and len(changed) > 0:
+        raise ValueError(f"{len(changed)} files should be restyled!")
+    elif len(changed) > 0:
+        print(f"Cleaned {len(changed)} files!")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("files", nargs="+", help="The file(s) or folder(s) to restyle.")
+    parser.add_argument("--max_len", type=int, help="The maximum length of lines.")
+    parser.add_argument("--check_only", action="store_true", help="Whether to only check and not fix styling issues.")
+    args = parser.parse_args()
+
+    main(*args.files, max_len=args.max_len, check_only=args.check_only)