diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..da99824 --- /dev/null +++ b/.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 \ No newline at end of file diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..1aba38f --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1 @@ +include LICENSE diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..e1c15c5 --- /dev/null +++ b/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 + \ No newline at end of file diff --git a/README.md b/README.md index 9d4e877..0e1c270 100644 --- a/README.md +++ b/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. + + diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..b7465bb --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,3 @@ +[tool.black] +line-length = 119 +target-version = ['py36'] diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..10f15c0 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +transformers +accelerate +loralib +evaluate diff --git a/setup.cfg b/setup.cfg new file mode 100644 index 0000000..6b26312 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..6f74e46 --- /dev/null +++ b/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 diff --git a/src/pet/__init__.py b/src/pet/__init__.py new file mode 100644 index 0000000..734d7c2 --- /dev/null +++ b/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, +) diff --git a/src/pet/mapping.py b/src/pet/mapping.py new file mode 100644 index 0000000..0cb2210 --- /dev/null +++ b/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) diff --git a/src/pet/pet_model.py b/src/pet/pet_model.py new file mode 100644 index 0000000..1398146 --- /dev/null +++ b/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) diff --git a/src/pet/tuners/__init__.py b/src/pet/tuners/__init__.py new file mode 100644 index 0000000..22fd5c8 --- /dev/null +++ b/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 diff --git a/src/pet/tuners/lora.py b/src/pet/tuners/lora.py new file mode 100644 index 0000000..eb30bad --- /dev/null +++ b/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) diff --git a/src/pet/tuners/p_tuning.py b/src/pet/tuners/p_tuning.py new file mode 100644 index 0000000..bb5a5a8 --- /dev/null +++ b/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 diff --git a/src/pet/tuners/prefix_tuning.py b/src/pet/tuners/prefix_tuning.py new file mode 100644 index 0000000..d709baf --- /dev/null +++ b/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 diff --git a/src/pet/tuners/prompt_tuning.py b/src/pet/tuners/prompt_tuning.py new file mode 100644 index 0000000..fdad421 --- /dev/null +++ b/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 diff --git a/src/pet/utils/__init__.py b/src/pet/utils/__init__.py new file mode 100644 index 0000000..aa359db --- /dev/null +++ b/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 diff --git a/src/pet/utils/config.py b/src/pet/utils/config.py new file mode 100644 index 0000000..d989601 --- /dev/null +++ b/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"}) diff --git a/src/pet/utils/other.py b/src/pet/utils/other.py new file mode 100644 index 0000000..ac8b329 --- /dev/null +++ b/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 diff --git a/src/pet/utils/save_and_load.py b/src/pet/utils/save_and_load.py new file mode 100644 index 0000000..e2eea05 --- /dev/null +++ b/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 diff --git a/utils/style_doc.py b/utils/style_doc.py new file mode 100644 index 0000000..0422ebe --- /dev/null +++ b/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 , and blocks. +_re_tip = re.compile("^\s*|\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)