fixing doc strings to follow hf doc-builder format

This commit is contained in:
Sourab Mangrulkar
2023-01-19 03:23:57 +05:30
parent 416a04008e
commit 246b2cffba
8 changed files with 91 additions and 86 deletions
+3 -3
View File
@@ -62,7 +62,7 @@ def get_peft_config(config_dict):
Returns a Peft config object from a dictionary.
Args:
config_dict (`Dict[str, Any]`):
config_dict (`Dict[str, Any]`): Dictionary containing the configuration parameters.
"""
return PEFT_TYPE_TO_CONFIG_MAPPING[config_dict["peft_type"]](**config_dict)
@@ -128,8 +128,8 @@ def get_peft_model(model, peft_config):
Returns a Peft model object from a model and a config.
Args:
model (`transformers.PreTrainedModel`):
peft_config (`transformers.PeftConfig`):
model ([`transformers.PreTrainedModel`]): Model to be wrapped.
peft_config ([`PeftConfig`]): Configuration object containing the parameters of the Peft model.
"""
model_config = model.config.to_dict()
+31 -28
View File
@@ -30,20 +30,23 @@ class PeftModel(torch.nn.Module):
Parameter-Efficient Fine-Tuning Model. Base model encompassing various Peft methods.
Args:
model (`PreTrainedModel`): The base transformer model used for Peft.
peft_config (`PeftConfig`): The configuration of the Peft model.
model ([`PreTrainedModel`]): The base transformer model used for Peft.
peft_config ([`PeftConfig`]): The configuration of the Peft model.
Attributes:
base_model (`PreTrainedModel`): The base transformer model used for Peft. peft_config (`PeftConfig`): The
configuration of the Peft model. modules_to_save (`list` of `str`): The list of sub-module names to save when
saving the model. prompt_encoder (`PromptEncoder`): The prompt encoder used for Peft if `peft_config.peft_type
!= PeftType.LORA`. prompt_tokens (`torch.Tensor`): The virtual prompt tokens used for Peft if
`peft_config.peft_type != PeftType.LORA`. transformer_backbone_name (`str`): The name of the transformer
backbone in the base model
if `peft_config.peft_type != PeftType.LORA`.
word_embeddings (`torch.nn.Embedding`): The word embeddings of the transformer backbone
in the base model if `peft_config.peft_type != PeftType.LORA`.
**Attributes**:
- **base_model** ([`PreTrainedModel`]) -- The base transformer model used for Peft.
- **peft_config** ([`PeftConfig`]) -- The configuration of the Peft model.
- **modules_to_save** (`list` of `str`) -- The list of sub-module names to save when
saving the model.
- **prompt_encoder** ([`PromptEncoder`]) -- The prompt encoder used for Peft if `peft_config.peft_type
!= PeftType.LORA`.
- **prompt_tokens** (`torch.Tensor`) -- The virtual prompt tokens used for Peft if
`peft_config.peft_type != PeftType.LORA`.
- **transformer_backbone_name** (`str`) -- The name of the transformer
backbone in the base model if `peft_config.peft_type != PeftType.LORA`.
- **word_embeddings** (`torch.nn.Embedding`) -- The word embeddings of the transformer backbone
in the base model if `peft_config.peft_type != PeftType.LORA`.
"""
def __init__(self, model, peft_config: PeftConfig):
@@ -92,7 +95,7 @@ class PeftModel(torch.nn.Module):
def get_prompt_embedding_to_save(self):
"""
Returns the prompt embedding to save when saving the model. Only applocable when `peft_config.peft_type !=
Returns the prompt embedding to save when saving the model. Only applicable when `peft_config.peft_type !=
PeftType.LORA`.
"""
prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(1, -1).to(self.device)
@@ -103,7 +106,7 @@ class PeftModel(torch.nn.Module):
def get_prompt(self, batch_size):
"""
Returns the virtual prompts to use for Peft. Only applocable when `peft_config.peft_type != PeftType.LORA`.
Returns the virtual prompts to use for Peft. Only applicable when `peft_config.peft_type != PeftType.LORA`.
"""
prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(batch_size, -1).to(self.device)
if self.peft_config.peft_type == PeftType.PREFIX_TUNING:
@@ -162,12 +165,12 @@ class PeftModelForSequenceClassification(PeftModel):
Peft model for sequence classification tasks.
Args:
model (`PreTrainedModel`): Base transformer model
peft_config (`PeftConfig`): Peft config.
model ([`PreTrainedModel`]): Base transformer model
peft_config ([`PeftConfig`]): Peft config.
Attributes:
config (`PretrainedConfig`): The configuration object of the base model. cls_layer_name (`str`): The name of
the classification layer.
**Attributes**:
- **config** ([`PretrainedConfig`]) -- The configuration object of the base model.
- **cls_layer_name** (`str`) -- The name of the classification layer.
Example::
@@ -332,8 +335,8 @@ class PeftModelForCausalLM(PeftModel):
Peft model for Causal LM
Args:
model (`PreTrainedModel`): Base transformer model
peft_config (`PeftConfig`): Peft config.
model ([`PreTrainedModel`]): Base transformer model
peft_config ([`PeftConfig`]): Peft config.
Example::
@@ -454,8 +457,8 @@ class PeftModelForSeq2SeqLM(PeftModel):
Peft model for Seq2Seq LM
Args:
model (`PreTrainedModel`): Base transformer model
peft_config (`PeftConfig`): Peft config.
model ([`PreTrainedModel`]): Base transformer model
peft_config ([`PeftConfig`]): Peft config.
Example::
@@ -606,12 +609,12 @@ class PeftModelForTokenClassification(PeftModel):
Peft model for sequence classification tasks.
Args:
model (`PreTrainedModel`): Base transformer model
peft_config (`PeftConfig`): Peft config.
model ([`PreTrainedModel`]): Base transformer model
peft_config ([`PeftConfig`]): Peft config.
Attributes:
config (`PretrainedConfig`): The configuration object of the base model. cls_layer_name (`str`): The name of
the classification layer.
**Attributes**:
- **config** ([`PretrainedConfig`]) -- The configuration object of the base model.
- **cls_layer_name** (`str`) -- The name of the classification layer.
Example::
+14 -14
View File
@@ -33,18 +33,18 @@ from ..utils import PeftConfig, PeftType, transpose
@dataclass
class LoraConfig(PeftConfig):
"""
This is the configuration class to store the configuration of a :class:`~peft.Lora`.
This is the configuration class to store the configuration of a [`~peft.Lora`].
Args:
r: (int): Lora attention dimension
target_modules ( list of str): The names of the modules to apply Lora to.
lora_alpha ( float): The alpha parameter for Lora scaling.
lora_dropout ( float): The dropout probability for Lora layers.
merge_weights ( bool):
r (`int`): Lora attention dimension
target_modules (`List[str]`): The names of the modules to apply Lora to.
lora_alpha (`float`): The alpha parameter for Lora scaling.
lora_dropout (`float`): The dropout probability for Lora layers.
merge_weights (`bool`):
Whether to merge the weights of the Lora layers with the base transformer model in `eval` mode.
fan_in_fan_out ( bool): Set this to True if the layer to replace stores weight like (fan_in, fan_out)
enable_lora ( list of bool): Used with `lora.MergedLinear`.
bias ( str): Bias type for Lora. Can be 'none', 'all' or 'lora_only'
fan_in_fan_out (`bool`): Set this to True if the layer to replace stores weight like (fan_in, fan_out)
enable_lora ( `List[bool]`): Used with `lora.MergedLinear`.
bias (`str`): Bias type for Lora. Can be 'none', 'all' or 'lora_only'
"""
r: int = field(default=8, metadata={"help": "Lora attention dimension"})
@@ -70,8 +70,8 @@ class LoraModel(torch.nn.Module):
Creates Low Rank Adapter (Lora) model from a pretrained transformers model.
Args:
model (`transformers.PreTrainedModel`): The model to be adapted.
config (`LoraConfig`): The configuration of the Lora model.
model ([`transformers.PreTrainedModel`]): The model to be adapted.
config ([`LoraConfig`]): The configuration of the Lora model.
Returns:
`torch.nn.Module`: The Lora model.
@@ -84,9 +84,9 @@ class LoraModel(torch.nn.Module):
lora_dropout=0.01, )
>>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") >>> lora_model = LoraModel(config, model)
Attributes:
model (`transformers.PreTrainedModel`): The model to be adapted. config (`LoraConfig`): The configuration of
the Lora model.
**Attributes**:
- **model** ([`transformers.PreTrainedModel`]) -- The model to be adapted.
- **peft_config** ([`LoraConfig`]): The configuration of the Lora model.
"""
def __init__(self, config, model):
+15 -12
View File
@@ -30,11 +30,11 @@ class PromptEncoderReparameterizationType(str, enum.Enum):
@dataclass
class PromptEncoderConfig(PromptLearningConfig):
"""
This is the configuration class to store the configuration of a :class:`~peft.PromptEncoder`.
This is the configuration class to store the configuration of a [`~peft.PromptEncoder`].
Args:
encoder_reparameterization_type
(Union[:class:`PromptEncoderReparameterizationType`, `str`]): The type of reparameterization to use.
(Union[[`PromptEncoderReparameterizationType`], `str`]): The type of reparameterization to use.
encoder_hidden_size (`int`): The hidden size of the prompt encoder.
encoder_num_layers (`int`): The number of layers of the prompt encoder.
encoder_dropout (`float`): The dropout probability of the prompt encoder.
@@ -68,7 +68,7 @@ class PromptEncoder(torch.nn.Module):
The prompt encoder network that is used to generate the virtual token embeddings for p-tuning.
Args:
config (:class:`PromptEncoderConfig`): The configuration of the prompt encoder.
config ([`PromptEncoderConfig`]): The configuration of the prompt encoder.
Example::
@@ -79,15 +79,18 @@ class PromptEncoder(torch.nn.Module):
)
>>> prompt_encoder = PromptEncoder(config)
Attributes:
embedding (:class:`~torch.nn.Embedding`): The embedding layer of the prompt encoder. mlp_head
(:class:`~torch.nn.Sequential`): The MLP head of the prompt encoder if `inference_mode=False`. lstm_head
(:class:`~torch.nn.LSTM`):
The LSTM head of the prompt encoder if `inference_mode=False` and `encoder_reparameterization_type="LSTM"`.
token_dim (`int`): The hidden embedding dimension of the base transformer model. input_size (`int`): The input
size of the prompt encoder. output_size (`int`): The output size of the prompt encoder. hidden_size (`int`):
The hidden size of the prompt encoder. total_virtual_tokens (`int`): The total number of virtual tokens of the
prompt encoder. encoder_type (Union[:class:`PromptEncoderReparameterizationType`, `str`]):
**Attributes**:
- **embedding** ([`~torch.nn.Embedding`]) -- The embedding layer of the prompt encoder.
- **mlp_head** ([`~torch.nn.Sequential`]) -- The MLP head of the prompt encoder if `inference_mode=False`.
- **lstm_head** ([`~torch.nn.LSTM`]) -- The LSTM head of the prompt encoder if `inference_mode=False` and
`encoder_reparameterization_type="LSTM"`.
- **token_dim** (`int`) -- The hidden embedding dimension of the base transformer model.
- **input_size** (`int`) -- The input size of the prompt encoder.
- **output_size** (`int`) -- The output size of the prompt encoder.
- **hidden_size** (`int`) -- The hidden size of the prompt encoder.
- **total_virtual_tokens** (`int`): The total number of virtual tokens of the
prompt encoder.
- **encoder_type** (Union[[`PromptEncoderReparameterizationType`], `str`]):
The encoder type of the prompt encoder.
+13 -14
View File
@@ -25,13 +25,12 @@ from ..utils import PeftType, PromptLearningConfig
@dataclass
class PrefixTuningConfig(PromptLearningConfig):
"""
This is the configuration class to store the configuration of a :class:`~peft.PrefixEncoder`.
This is the configuration class to store the configuration of a [`~peft.PrefixEncoder`].
Args:
encoder_hidden_size ( int): The hidden size of the prompt encoder.
prefix_projection ( bool): Whether to project the prefix embeddings.
postprocess_past_key_value_function (:
obj: Optional[Callable]): The function to postprocess the past key value.
encoder_hidden_size (`int`): The hidden size of the prompt encoder.
prefix_projection (`bool`): Whether to project the prefix embeddings.
postprocess_past_key_value_function (`Callable`, *optional*): The function to postprocess the past key value.
"""
encoder_hidden_size: int = field(
@@ -58,7 +57,7 @@ class PrefixEncoder(torch.nn.Module):
The torch.nn model to encode the prefix
Args:
config (:class:`PrefixTuningConfig`): The configuration of the prefix encoder.
config ([`PrefixTuningConfig`]): The configuration of the prefix encoder.
Example::
@@ -69,12 +68,12 @@ class PrefixEncoder(torch.nn.Module):
>>> prefix_encoder = PrefixEncoder(config)
Attributes:
embedding (`torch.nn.Embedding`):
The embedding layer of the prefix encoder. trans (`torch.nn.Sequential`): The
two-layer MLP to transform the prefix embeddings
if `prefix_projection` is `True`.
prefix_projection (`bool`): Whether to project the prefix embeddings.
**Attributes**:
- **embedding** (`torch.nn.Embedding`) --
The embedding layer of the prefix encoder.
- **transform** (`torch.nn.Sequential`) -- The
two-layer MLP to transform the prefix embeddings if `prefix_projection` is `True`.
- **prefix_projection** (`bool`) -- Whether to project the prefix embeddings.
Input shape: (batch_size, num_virtual_tokens)
@@ -91,7 +90,7 @@ class PrefixEncoder(torch.nn.Module):
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(
self.transform = 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),
@@ -102,7 +101,7 @@ class PrefixEncoder(torch.nn.Module):
def forward(self, prefix: torch.Tensor):
if self.prefix_projection:
prefix_tokens = self.embedding(prefix)
past_key_values = self.trans(prefix_tokens)
past_key_values = self.transform(prefix_tokens)
else:
past_key_values = self.embedding(prefix)
return past_key_values
+5 -5
View File
@@ -31,10 +31,10 @@ class PromptTuningInit(str, enum.Enum):
@dataclass
class PromptTuningConfig(PromptLearningConfig):
"""
This is the configuration class to store the configuration of a :class:`~peft.PromptEmbedding`.
This is the configuration class to store the configuration of a [`~peft.PromptEmbedding`].
Args:
prompt_tuning_init (Union[:class:`PromptTuningInit`, `str`]): The initialization of the prompt embedding.
prompt_tuning_init (Union[[`PromptTuningInit`], `str`]): The initialization of the prompt embedding.
prompt_tuning_init_text ( Optional[`str`]): The text to initialize the prompt embedding.
Only used if `prompt_tuning_init` is `TEXT`
tokenizer_name_or_path ( Optional[`str`]): The name or path of the tokenizer.
@@ -67,11 +67,11 @@ class PromptEmbedding(torch.nn.Module):
The model to encode virtual tokens into prompt embeddings.
Args:
config (:class:`PromptTuningConfig`): The configuration of the prompt embedding.
config ([`PromptTuningConfig`]): The configuration of the prompt embedding.
word_embeddings (`torch.nn.Module`): The word embeddings of the base transformer model.
Attributes:
embedding (`torch.nn.Embedding`): The embedding layer of the prompt embedding.
**Attributes**:
**embedding** (`torch.nn.Embedding`) -- The embedding layer of the prompt embedding.
Example::
+4 -4
View File
@@ -38,8 +38,8 @@ class PeftConfig:
This is the base configuration class to store the configuration of a :class:`~peft.PeftModel`.
Args:
peft_type (Union[:class:`~peft.utils.config.PeftType`, `str`]): The type of Peft method to use.
task_type (Union[:class:`~peft.utils.config.TaskType`, `str`]): The type of task to perform.
peft_type (Union[[`~peft.utils.config.PeftType`], `str`]): The type of Peft method to use.
task_type (Union[[`~peft.utils.config.TaskType`], `str`]): The type of task to perform.
inference_mode (`bool`, defaults to `False`): Whether to use the Peft model in inference mode.
"""
@@ -51,8 +51,8 @@ class PeftConfig:
@dataclass
class PromptLearningConfig(PeftConfig):
"""
This is the base configuration class to store the configuration of a Union[:class:`~peft.PrefixTuning`,
:class:`~peft.PromptEncoder`, :class:`~peft.PromptTuning`].
This is the base configuration class to store the configuration of a Union[[`~peft.PrefixTuning`],
[`~peft.PromptEncoder`], [`~peft.PromptTuning`]].
Args:
num_virtual_tokens (`int`): The number of virtual tokens to use.
+6 -6
View File
@@ -21,10 +21,10 @@ def get_peft_model_state_dict(model, state_dict=None):
Get the state dict of the Peft model.
Args:
model (`PeftModel`): The Peft model. When using torch.nn.DistributedDataParallel, DeepSpeed or FSDP,
the model should be teh underlying model/unwrapped model (i.e. model.module).
state_dict (:
obj:`dict`, `optional`): The state dict of the model. If not provided, the state dict of the model
model ([`PeftModel`]): The Peft model. When using torch.nn.DistributedDataParallel, DeepSpeed or FSDP,
the model should be the underlying model/unwrapped model (i.e. model.module).
state_dict (`dict`, *optional*, defaults to `None`):
The state dict of the model. If not provided, the state dict of the model
will be used.
"""
if state_dict is None:
@@ -64,7 +64,7 @@ def set_peft_model_state_dict(model, peft_model_state_dict):
Set the state dict of the Peft model.
Args:
model (`PeftModel`): The Peft model.
model ([`PeftModel`]): The Peft model.
peft_model_state_dict (`dict`): The state dict of the Peft model.
"""
@@ -81,7 +81,7 @@ def peft_model_load_and_dispatch(model, peft_model_state_dict, peft_config, max_
Load the Peft model state dict and dispatch the model to the correct device.
Args:
model (`PeftModel`): The Pre-trained base model which has already been sharded and dispatched
model ([`PeftModel`]): The Pre-trained base model which has already been sharded and dispatched
using `accelerate` functionalities.
peft_model_state_dict (`dict`): The state dict of the Peft model.
max_memory (`Dict`, *optional*):