TRL 文档

使用训练后的模型

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

使用训练后的模型

一旦你使用 SFTTrainer、PPOTrainer 或 DPOTrainer 训练了一个模型,你将拥有一个微调模型,可用于文本生成。在本节中,我们将介绍加载微调模型和生成文本的过程。如果你需要使用训练好的模型运行推理服务器,你可以探索诸如 text-generation-inference 等库。

加载和生成

如果你已经完全微调了一个模型,即没有使用 PEFT,你可以像加载 transformers 中的任何其他语言模型一样简单地加载它。例如,PPO 训练期间训练的 value head 不再需要,如果你使用原始 transformer 类加载模型,它将被忽略

from transformers import AutoTokenizer, AutoModelForCausalLM

model_name_or_path = "kashif/stack-llama-2" #path/to/your/model/or/name/on/hub
device = "cpu" # or "cuda" if you have a GPU

model = AutoModelForCausalLM.from_pretrained(model_name_or_path).to(device)
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)

inputs = tokenizer.encode("This movie was really", return_tensors="pt").to(device)
outputs = model.generate(inputs)
print(tokenizer.decode(outputs[0]))

或者你也可以使用 pipeline

from transformers import pipeline

model_name_or_path = "kashif/stack-llama-2" #path/to/your/model/or/name/on/hub
pipe = pipeline("text-generation", model=model_name_or_path)
print(pipe("This movie was really")[0]["generated_text"])

使用适配器 PEFT

from peft import PeftConfig, PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

base_model_name = "kashif/stack-llama-2" #path/to/your/model/or/name/on/hub"
adapter_model_name = "path/to/my/adapter"

model = AutoModelForCausalLM.from_pretrained(base_model_name)
model = PeftModel.from_pretrained(model, adapter_model_name)

tokenizer = AutoTokenizer.from_pretrained(base_model_name)

你也可以将适配器合并到基础模型中,这样你就可以像使用普通的 transformers 模型一样使用该模型,但是 checkpoint 会大得多

model = AutoModelForCausalLM.from_pretrained(base_model_name)
model = PeftModel.from_pretrained(model, adapter_model_name)

model = model.merge_and_unload()
model.save_pretrained("merged_adapters")

一旦你加载了模型,并且合并了适配器或者将它们分开放在顶部,你就可以像上面概述的普通模型一样运行生成。

< > GitHub 上更新