TRL 文档

模型

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

模型

TRL 通过 `AutoModelForCausalLMWithValueHead` 类支持 transformers 中的所有解码器模型架构,如 GPT-2、OPT 和 GPT-Neo。此外,通过 `AutoModelForSeq2SeqLMWithValueHead`,您可以使用编码器-解码器架构,如 T5。TRL 还需要参考模型,它们是所训练模型的冻结副本。使用 `create_reference_model`,您可以轻松创建一个冻结副本,并在两个模型之间共享层以节省内存。

PreTrainedModelWrapper

class trl.PreTrainedModelWrapper

< >

( pretrained_model = None score_module = None supports_rm_adapter = False rm_adapter_name = None **kwargs )

参数

  • pretrained_model (transformers.PreTrainedModel) — 要包装的模型。
  • parent_class (transformers.PreTrainedModel) — 要包装的模型的父类。
  • supported_args (list) — 包装器类支持的参数列表。

一个围绕 (`transformers.PreTrainedModel`) 的包装类,以兼容 (`~transformers.PreTrained`) 类,从而保留 (`~transformers.PreTrainedModel`) 类的某些属性和方法。

add_and_load_reward_modeling_adapter

< >

( pretrained_model adapter_model_id adapter_name = 'reward_model_adapter' token = None )

添加并加载一个奖励建模适配器。仅当模型是 `PeftModel` 并且您已使用 `reward_modeling_adapter_id` 参数(指向奖励建模适配器的ID)初始化模型时,才能使用此方法。后者还需要包含分数头才能产生奖励。

compute_reward_score

< >

( input_ids attention_mask = None **kwargs )

计算给定输入的奖励分数。该方法首先必须启用适配器,然后计算奖励分数。之后,模型将禁用奖励建模适配器并重新启用默认的 ppo 适配器。

from_pretrained

< >

( pretrained_model_name_or_path *model_args **kwargs )

参数

  • pretrained_model_name_or_path (str or transformers.PreTrainedModel) — 预训练模型的路径或其名称。
  • *model_args (list, optional)) — 传递给底层模型的 `from_pretrained` 方法的额外位置参数。
  • **kwargs (dict, optional) — 传递给底层模型的 `from_pretrained` 方法的额外关键字参数。我们还会预处理 kwargs 以提取特定于 `transformers.PreTrainedModel` 类和 trl 模型的参数。kwargs 还支持 `peft` 库中的 `prepare_model_for_kbit_training` 参数。

从 `transformers` 中的预训练模型实例化一个新模型。预训练模型使用 `transformers.PreTrainedModel` 类的 `from_pretrained` 方法加载。特定于 `transformers.PreTrainedModel` 类的参数会传递给此方法并从 `kwargs` 参数中过滤掉。

post_init

< >

( *args **kwargs )

后初始化方法。此方法在模型实例化并从检查点加载后调用。它可用于执行其他操作,如加载 state_dict。

push_to_hub

< >

( *args **kwargs )

参数

  • *args (list, optional) — 传递给底层模型的 `push_to_hub` 方法的位置参数。
  • **kwargs (dict, optional) — 传递给底层模型的 `push_to_hub` 方法的关键字参数。

将预训练模型推送到 Hub。此方法是 `transformers.PreTrainedModel.push_to_hub` 的包装器。有关更多信息,请参阅 `transformers.PreTrainedModel.push_to_hub` 的文档。

save_pretrained

< >

( *args **kwargs )

参数

  • *args (list, optional) — 传递给底层模型的 `save_pretrained` 方法的位置参数。
  • **kwargs (dict, optional) — 传递给底层模型的 `save_pretrained` 方法的关键字参数。

将预训练模型保存到目录。此方法是 `transformers.PreTrainedModel.save_pretrained` 的包装器。有关更多信息,请参阅 `transformers.PreTrainedModel.save_pretrained` 的文档。

state_dict

< >

( *args **kwargs )

返回预训练模型的 state_dict。

AutoModelForCausalLMWithValueHead

class trl.AutoModelForCausalLMWithValueHead

< >

( pretrained_model **kwargs )

一个自回归模型,除了语言模型头外,还有一个价值头。这个类继承自 `~trl.PreTrainedModelWrapper` 并包装了一个 `transformers.PreTrainedModel` 类。这个包装类支持诸如 `from_pretrained`、`push_to_hub` 和 `generate` 等经典函数。要调用被包装模型的方法,只需操作该类的 `pretrained_model` 属性即可。

类属性

  • transformers_parent_class (transformers.PreTrainedModel) — 被包装模型的父类。对于这个类,应将其设置为 `transformers.AutoModelForCausalLM`。
  • supported_args (tuple) — 一个字符串元组,用于标识 `ValueHead` 类支持的参数。目前支持的参数有
    • summary_dropout_prob (float, optional, 默认为 None) — `ValueHead` 类的 dropout 概率。
    • v_head_initializer_range (float, optional, 默认为 0.2) — 如果选择了特定的初始化策略,则为 `ValueHead` 的初始化范围。
    • v_head_init_strategy (str, optional, 默认为 None) — `ValueHead` 的初始化策略。目前支持的策略有
      • `None` — 使用随机分布初始化 `ValueHead` 的权重。这是默认策略。
      • “normal” — 使用正态分布初始化 `ValueHead` 的权重。

__init__

< >

( pretrained_model **kwargs )

参数

  • pretrained_model (transformers.PreTrainedModel) — 要包装的模型。它应该是一个因果语言模型,例如 GPT2,或任何映射在 `AutoModelForCausalLM` 类中的模型。
  • kwargs (dict, optional) — 传递给 `ValueHead` 类的附加关键字参数。

初始化模型。

forward

< >

( input_ids = None past_key_values = None attention_mask = None return_past_key_values = False **kwargs )

参数

  • input_ids (torch.LongTensor of shape (batch_size, sequence_length)) — 词汇表中输入序列标记的索引。
  • past_key_values (tuple(tuple(torch.FloatTensor)), optional) — 包含由模型计算的预计算隐藏状态(注意力块中的键和值),以加速顺序解码(参见 *past_key_values* 输入)。
  • attention_mask (torch.FloatTensor of shape (batch_size, sequence_length), optional) — 掩码,用于避免在填充标记索引上执行注意力。掩码值选择在 [0, 1] 中:
    • 1 表示**未被屏蔽**的标记,
    • 0 表示**被屏蔽**的标记。
  • return_past_key_values (bool) — 指示是否应返回计算的隐藏状态的标志。
  • kwargs (dict, optional) — 传递给被包装模型的附加关键字参数。

对被包装的模型应用前向传播,并返回价值头的 logits。

生成

< >

( *args **kwargs )

参数

  • *args (list, optional) — 传递给被包装模型的 `generate` 方法的位置参数。
  • **kwargs (dict, optional) — 传递给被包装模型的 `generate` 方法的关键字参数。

被包装模型 `generate` 方法的一个简单包装器。有关支持的参数的更多信息,请参阅被包装模型的 `generate` 方法。

_init_weights

< >

( **kwargs )

参数

  • **kwargs (dict, optional) — 传递给 `ValueHead` 类的附加关键字参数。这些参数可以包含 `v_head_init_strategy` 参数以及 `v_head_initializer_range` 参数。

初始化价值头的权重。默认的初始化策略是随机的。用户可以通过在调用 `.from_pretrained` 时传递 `v_head_init_strategy` 参数来传递不同的初始化策略。支持的策略有

  • normal:使用正态分布初始化权重。

AutoModelForSeq2SeqLMWithValueHead

class trl.AutoModelForSeq2SeqLMWithValueHead

< >

( pretrained_model **kwargs )

参数

  • pretrained_model (transformers.PreTrainedModel) — 要包装的模型。它应该是一个因果语言模型,例如 GPT2,或任何映射在 `AutoModelForSeq2SeqLM` 类中的模型。
  • kwargs — 传递给 `ValueHead` 类的附加关键字参数。

一个带有价值头的 seq2seq 模型,除了语言模型头之外。这个类继承自 `~trl.PreTrainedModelWrapper` 并包装了一个 `transformers.PreTrainedModel` 类。这个包装类支持诸如 `from_pretrained` 和 `push_to_hub` 等经典函数,并提供一些额外的功能,如 `generate`。

__init__

< >

( pretrained_model **kwargs )

forward

< >

( input_ids = None past_key_values = None attention_mask = None return_past_key_values = False **kwargs )

生成

< >

( *args **kwargs )

我们在被包装的模型上调用 `generate`。

_init_weights

< >

( **kwargs )

我们初始化价值头的权重。

create_reference_model

trl.create_reference_model

< >

( model: PreTrainedModelWrapper num_shared_layers: typing.Optional[int] = None pattern: typing.Optional[str] = None ) PreTrainedModelWrapper

参数

  • model (PreTrainedModelWrapper) — 要复制的模型。
  • num_shared_layers (int, optional) — 两个模型之间共享并保持冻结的初始层数。
  • pattern (str, optional) — 共享层通过字符串模式选择(例如,对于 GPT2 是 “transformer.h.{layer}”),如果需要自定义模式,可以在此处传递。

返回

PreTrainedModelWrapper

创建一个模型的静态参考副本。请注意,模型将处于 .eval() 模式。

< > 在 GitHub 上更新