四、文本摘要

发布于 2023-07-17 23:38:23 字数 27647 浏览 0 评论 0 收藏 0

这里我们将采用不同的方法并从头开始训练一个全新的因果语言模型。为了减小数据规模从而用于演示，我们将使用 Python 代码的子集专注于单行代码的补全（而不是补全完整的函数或类）。

3.1 数据集和数据处理

加载数据集：CodeParrot 数据集来自于 Google's BigQuery 数据集，使用了大约 180 GB 的 GitHub dump，包含大约 20M 个 Python 文件。创建过程：


xxxxxxxxxx
SELECT
  f.repo_name, f.path, c.copies, c.size, c.content, l.license
FROM
  `bigquery-public-data.github_repos.files` AS f
JOIN
  `bigquery-public-data.github_repos.contents` AS c
ON
  f.id = c.id
JOIN
  `bigquery-public-data.github_repos.licenses` AS l
ON
  f.repo_name = l.repo_name 
WHERE
  NOT c.binary AND ((f.path LIKE '%.py') AND (c.size BETWEEN 1024 AND 1048575))

为了演示的效果，我们进一步仅考虑与 Python 数据科学相关的子集。我们使用过滤函数：


xxxxxxxxxx
filters = ["pandas", "sklearn", "matplotlib", "seaborn"]


def any_keyword_in_string(string, keywords):
    for keyword in keywords:
        if keyword in string:
            return True
    return False

然后我们用这个过滤函数来流式地过滤数据集：


xxxxxxxxxx
def filter_streaming_dataset(dataset, filters):
    filtered_dict = defaultdict(list)
    total = 0
    for sample in tqdm(iter(dataset)):
        total += 1
        if any_keyword_in_string(sample["content"], filters):
            for k, v in sample.items():
                filtered_dict[k].append(v)
    print(f"{len(filtered_dict['content'])/total:.2%} of data after filtering.")
    return Dataset.from_dict(filtered_dict)
  
from datasets import load_dataset
split = "train"  # "valid"
data = load_dataset(f"transformersbook/codeparrot-{split}", split=split, streaming=True)
filtered_data = filter_streaming_dataset(data, filters) 
# 3.26% of data after filtering.

这个加载数据集并过滤的耗时非常长，可能需要数个小时。过滤之后保留了大约 3% 的原始数据，仍然高达 6GB ，包含大约 600k 个 python 文件。HuggingFace 提供好了过滤后的数据集：


xxxxxxxxxx
from datasets import load_dataset, DatasetDict


ds_train = load_dataset("huggingface-course/codeparrot-ds-train", split="train")
ds_valid = load_dataset("huggingface-course/codeparrot-ds-valid", split="validation")


raw_datasets = DatasetDict( { "train": ds_train,   "valid": ds_valid })
print(raw_datasets)
# DatasetDict({
#     train: Dataset({
#         features: ['repo_name', 'path', 'copies', 'size', 'content', 'license'],
#         num_rows: 606720
#     })
#     valid: Dataset({
#         features: ['repo_name', 'path', 'copies', 'size', 'content', 'license'],
#         num_rows: 3322
#     })
# })

Tokenization：第一步是对数据集进行 tokenization。我们的目标单行代码的补全，因此可以选择较短的上下文。这样做的好处是，我们可以更快地训练模型并且需要更少的内存。如果你需要更长的上下文（如，补全函数、类、或自动生成单元测试），那么需要设置较长的上下文。

这里我们选择上下文为 128 个 token （相比之下，GPT2 选择了 1024 个 token、GPT3 选择了 2048 个 token ）。大多数文档包含超过 128 个 token，如果简单地进行数据截断，那么将丢弃大多数的数据。相反，我们使用 return_overflowing_tokens 来执行 tokenizeation 从而返回几个块，并且还使用 return_length 来返回每个块的长度。通常最后一个块会小于上下文大小，我们会去掉这些块从而免于填充，因为这里的数据量已经足够用了。


xxxxxxxxxx
from transformers import AutoTokenizer


context_length = 128
tokenizer = AutoTokenizer.from_pretrained("huggingface-course/code-search-net-tokenizer")


##**********  如果 return_overflowing_tokens 未设置  **********
outputs = tokenizer(
    raw_datasets["train"][:2]["content"],
    truncation=True, 
    max_length=context_length,
)


print(f"Input IDs length: {len(outputs['input_ids'])}")
# Input IDs length: 2


##**********  如果 return_overflowing_tokens 被设置  **********
outputs = tokenizer(
    raw_datasets["train"][:2]["content"],
    truncation=True,
    max_length=context_length,
    return_overflowing_tokens=True,
    return_length=True,
)


print(f"Input IDs length: {len(outputs['input_ids'])}")
# Input IDs length: 34
print(f"Input chunk lengths: {(outputs['length'])}")
# Input chunk lengths: [128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 117, 
#                       128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 41]
print(f"Chunk mapping: {outputs['overflow_to_sample_mapping']}")
# Chunk mapping: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
#                 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]

可以看到：从两个样本中我们得到了 34 个块，其中：

outputs['input_ids'] 存放每个块的数据。
outputs['length'] 存放每个块的长度。可以看到，每个文档的末尾的块，它的长度都小于 128 （分别为 117 和 41）。由于这种不足 128 的块的占比很小，因此我们可以将它们丢弃。
outputs['overflow_to_sample_mapping'] 存放每个块属于哪个样本。

然后我们把上述代码封装到一个函数中，并在 Dataset.map() 中调用：


xxxxxxxxxx
def tokenize(element):
    outputs = tokenizer(
        element["content"],
        truncation=True,
        max_length=context_length,
        return_overflowing_tokens=True,
        return_length=True,
    )
    input_batch = []
    for length, input_ids in zip(outputs["length"], outputs["input_ids"]):
        if length == context_length: # 过滤掉长度小于 context_length 的块
            input_batch.append(input_ids)
    return {"input_ids": input_batch}


tokenized_datasets = raw_datasets.map(
    tokenize, batched=True, 
    remove_columns=raw_datasets["train"].column_names # 我们只需要 input_ids 列，因此移除所有其它的列
)
print(tokenized_datasets)
# DatasetDict({
#     train: Dataset({
#         features: ['input_ids'],
#         num_rows: 16702061
#     })
#     valid: Dataset({
#         features: ['input_ids'],
#         num_rows: 93164
#     })
# })

我们现在有 16.70M 样本，每个样本有 128 个 token ，总计相当于大约 2.1B tokens 。作为参考，OpenAI 的 GPT-3 和 Codex 模型分别在 30B 和 100B 个 token 上训练，其中 Codex 模型从 GPT-3 checkpoint 初始化。

3.2 使用 Trainer API 微调模型

我们初始化一个 GPT-2 模型。我们采用与 GPT-2 相同的配置，并确保词表规模与 tokenizer 规模相匹配，然后设置 bos_token_id 和 eos_token_id 。利用该配置，我们加载一个新模型。注意，这是我们首次不使用 from_pretrained() 函数，因为我们实际上是在自己初始化模型：


xxxxxxxxxx
from transformers import AutoTokenizer, GPT2LMHeadModel, AutoConfig


config = AutoConfig.from_pretrained(
    "gpt2",
    vocab_size=len(tokenizer),
    n_ctx=context_length,
    bos_token_id=tokenizer.bos_token_id, # 句子开始的 token
    eos_token_id=tokenizer.eos_token_id, # 句子结束的 token
)


model = GPT2LMHeadModel(config)
model_size = sum(t.numel() for t in model.parameters())
print(f"GPT-2 size: {model_size/1000**2:.1f}M parameters")
# GPT-2 size: 124.2M parameters

该模型有 124.2 M 参数。

在开始训练之前，我们需要设置一个负责创建 batch 的 data collator 。我们可以使用 DataCollatorForLanguageModeling ，它是专为语言建模而设计。除了 batch 和 padding ，它还负责创建语言模型的标签：在因果语言建模中，input 也用作 label （只是右移一个位置），并且这个 data collator 在训练期间动态地创建 label ，所以我们不需要复制 input_ids。
注意 DataCollatorForLanguageModeling 支持掩码语言建模 (Masked Language Model: MLM ) 和因果语言建模 (Causal Language Model: CLM )。默认情况下它为 MLM 准备数据，但我们可以通过设置 mlm=False 参数切换到 CLM ：
```
xxxxxxxxxx
from transformers import DataCollatorForLanguageModeling
print(tokenizer.pad_token, tokenizer.eos_token)
# None <|endoftext|>
tokenizer.pad_token = tokenizer.eos_token # 利用 eos_token 来填充
data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)
```
用法示例：
```
xxxxxxxxxx
out = data_collator([tokenized_datasets["train"][i] for i in range(5)])
for key in out:
    print(f"{key} shape: {out[key].shape}")
# input_ids shape: torch.Size([5, 128])
# attention_mask shape: torch.Size([5, 128])
# labels shape: torch.Size([5, 128])
```

接下来就是配置 TrainingArguments 并启动 Trainer 。我们将使用余弦学习率，进行一些 warmup ，设置有效 batch size = 256 （ per_device_train_batch_size * gradient_accumulation_steps ）。gradient_accumulation_steps 用于梯度累积，它通过多次前向传播和反向传播但是只有一次梯度更新，从而实现 large batch size 的效果。当我们使用Accelerate 手动创建训练循环时，我们将看到这一点。


xxxxxxxxxx
from transformers import Trainer, TrainingArguments


args = TrainingArguments(
    output_dir="codeparrot-ds",
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    evaluation_strategy="steps", # 每隔若干个 step 执行一次评估
    save_strategy="epoch",       # 每个 epoch 保存一次
    eval_steps=5_000,
    logging_steps=5_000,
    gradient_accumulation_steps=8, # 梯度累积
    num_train_epochs=1,
    weight_decay=0.1,
    warmup_steps=1_000,
    lr_scheduler_type="cosine",
    learning_rate=5e-4,
    save_steps=5_000,
    fp16=True,
    push_to_hub=False, # 暂时不推送到 HuggingFace Hub
)


trainer = Trainer(
    model=model,
    tokenizer=tokenizer,
    args=args,
    data_collator=data_collator,
    train_dataset=tokenized_datasets["train"].select(range(100000)), # 为演示目的，仅用 10 万个 batch
    eval_dataset=tokenized_datasets["valid"].select(range(10000)),   
)


trainer.train()
# TrainOutput(
#    global_step=390, 
#    training_loss=4.230728540665064, 
#    metrics={
#        'train_runtime': 235.341, 
#        'train_samples_per_second': 424.915, 
#        'train_steps_per_second': 1.657, 
#        'total_flos': 6521849118720000.0, 
#        'train_loss': 4.230728540665064, 
#        'epoch': 1.0}
#  )


tokenizer.save_pretrained("codeparrot-ds") # 保存 tokenizer
# trainer.push_to_hub() # 训练完成后，将模型和 tokenizer 推送到 HuggingFace Hub

使用模型：现在可以通过 Transformers 的 pipeline 来调用微调后的模型：


xxxxxxxxxx
import torch
from transformers import pipeline


device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
pipe = pipeline(
    "text-generation", model="./codeparrot-ds/", device=device # 使用 GPU 加速生成过程
)


txt = """\
# create some data
x = np.random.randn(100)
y = np.random.randn(100)


# create scatter plot with x, y
"""
result = pipe(txt, num_return_sequences=1)
print(result) 
# [{'generated_text': '# create some data\nx = np.random.randn(100)\ny = np.random.randn(100)\n\n# create scatter plot with x, y\ny = rng.randn(n_samples, (1,'}]
print([0]["generated_text"]) # 因为训练不充分，这里的结果仅供参考

3.3 自定义训练过程

有时我们想要完全控制训练循环，或者我们想要进行一些特殊的更改，这时我们可以利用 Accelerate 来进行自定义的训练过程。

众所周知，数学科学的 package 中有一些关键字，如 plt, pd, sk, fit, predict 等等。我们仅关注那些表示为单个 token 的关键字，并且我们还关注那些带有一个空格作为前缀的关键字版本。


xxxxxxxxxx
from transformers import AutoTokenizer


tokenizer = AutoTokenizer.from_pretrained("huggingface-course/code-search-net-tokenizer")


keytoken_ids = []
for keyword in [ "plt", "pd", "sk", "fit", "predict", 
                " plt", " pd", " sk", " fit", " predict", 
]:
    ids = tokenizer([keyword]).input_ids[0]
    if len(ids) == 1: # 仅考虑单个 token 的关键字
        keytoken_ids.append(ids[0])
    else:
        print(f"Keyword has not single token: {keyword}")
print(keytoken_ids)
# [8436, 4289, 1201, 2770, 5431, 2564, 2604, 2110, 2872, 4969]

我们可以计算每个样本的损失，并计算每个样本中所有关键字的出现次数。然后我们以这个出现次数为权重来对样本的损失函数进行加权，使得模型更加注重那些具有多个关键字的样本。这是一个自定义的损失函数，它将输入序列、logits、以及我们刚刚选择的关键字 token 作为输入，然后输出关键字频率加权的损失函数：


xxxxxxxxxx
from torch.nn import CrossEntropyLoss
import torch


def keytoken_weighted_loss(inputs, logits, keytoken_ids, alpha=1.0):
    shift_labels = inputs[..., 1:].contiguous()     # input 序列第 i+1 位就是第 i 个标签
    shift_logits = logits[..., :-1, :].contiguous() # 最后一个位置不需要预测，因为没有 label
    
    loss_fct = CrossEntropyLoss(reduce=False) # 用于计算 per-token 的损失
    loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
    loss_per_sample = loss.view(shift_logits.size(0), shift_logits.size(1)).mean(axis=1) # 该样本的平均损失
 
    weights = torch.stack([(inputs == kt).float() for kt in keytoken_ids]).sum( # 每个样本出现的所有关键字的数量
        axis=[0, 2]
    )
    weights = alpha * (1.0 + weights)
    weighted_loss = (loss_per_sample * weights).mean() # 计算 batch 的加权平均损失
    return weighted_loss

加载数据集：


xxxxxxxxxx
from datasets import load_dataset, DatasetDict
context_length = 128
ds_train = load_dataset("huggingface-course/codeparrot-ds-train", split="train")
ds_valid = load_dataset("huggingface-course/codeparrot-ds-valid", split="validation")


raw_datasets = DatasetDict( { "train": ds_train,   "valid": ds_valid })
def tokenize(element):
    outputs = tokenizer(
        element["content"],
        truncation=True,
        max_length=context_length,
        return_overflowing_tokens=True,
        return_length=True,
    )
    input_batch = []
    for length, input_ids in zip(outputs["length"], outputs["input_ids"]):
        if length == context_length: # 过滤掉长度小于 context_length 的块
            input_batch.append(input_ids)
    return {"input_ids": input_batch}


tokenized_datasets = raw_datasets.map(
    tokenize, batched=True, 
    remove_columns=raw_datasets["train"].column_names # 我们只需要 input_ids 列，因此移除所有其它的列
)


tokenized_dataset.set_format("torch") # 设置为 torch 格式
train_dataloader = DataLoader(tokenized_dataset["train"], batch_size=32, shuffle=True)
eval_dataloader = DataLoader(tokenized_dataset["valid"], batch_size=32)

设置 weight-decay：我们对参数进行分组，以便优化器知道哪些将获得额外的 weight-decay。通常，所有的 bias 项和 LayerNorm weight 都不需要 weight-decay：


xxxxxxxxxx
weight_decay = 0.1


def get_grouped_params(model, no_decay=["bias", "LayerNorm.weight"]):
    params_with_wd, params_without_wd = [], []
    for name, param in model.named_parameters():
        if any(nd in name for nd in no_decay): # 判断 bias 字符串是否出现在 parameter.name 中，因为 parameter.name 可能为 attention1.bias1
            params_without_wd.append(param)
        else:
            params_with_wd.append(param)
    return [
        {"params": params_with_wd, "weight_decay": weight_decay},
        {"params": params_without_wd, "weight_decay": 0.0},
    ]

评估函数：由于我们希望在训练期间定期地在验证集上评估模型，因此我们也为此编写一个函数。它只是运行 eval_dataloader 并收集跨进程的所有损失函数值：


xxxxxxxxxx
def evaluate():
    model.eval()
    losses = []
    for step, batch in enumerate(eval_dataloader):
        with torch.no_grad():
            outputs = model(batch["input_ids"], labels=batch["input_ids"])
            
        losses.append(accelerator.gather(outputs.loss)) # 跨进程收集每个样本的 loss
    loss = torch.mean(torch.cat(losses))
    try:
        perplexity = torch.exp(loss) # 计算困惑度
    except OverflowError:
        perplexity = float("inf")
    return loss.item(), perplexity.item()

这个评估函数用于获取损失函数值、以及困惑度。

完整的训练过程：


xxxxxxxxxx
from transformers import AutoTokenizer
from datasets import load_dataset, DatasetDict
from transformers import AutoTokenizer, GPT2LMHeadModel, AutoConfig
from torch.optim import AdamW
from accelerate import Accelerator
from transformers import get_scheduler
from huggingface_hub import Repository, get_full_repo_name
from tqdm.notebook import tqdm
from torch.utils.data import DataLoader


##************** 加载数据集 ********************
tokenizer = AutoTokenizer.from_pretrained("huggingface-course/code-search-net-tokenizer")
context_length = 128
ds_train = load_dataset("huggingface-course/codeparrot-ds-train", split="train")
ds_valid = load_dataset("huggingface-course/codeparrot-ds-valid", split="validation")


raw_datasets = DatasetDict( { "train": ds_train,   "valid": ds_valid })
def tokenize(element):
    outputs = tokenizer(
        element["content"],
        truncation=True,
        max_length=context_length,
        return_overflowing_tokens=True,
        return_length=True,
    )
    input_batch = []
    for length, input_ids in zip(outputs["length"], outputs["input_ids"]):
        if length == context_length: # 过滤掉长度小于 context_length 的块
            input_batch.append(input_ids)
    return {"input_ids": input_batch}


tokenized_datasets = raw_datasets.map(
    tokenize, batched=True, 
    remove_columns=raw_datasets["train"].column_names # 我们只需要 input_ids 列，因此移除所有其它的列
)


batch_size = 32
tokenized_dataset.set_format("torch") # 设置为 torch 格式
train_dataloader = DataLoader(tokenized_dataset["train"].select(range(100000), 
                     batch_size=batch_size, shuffle=True) # 为演示方便，用了更少的数据
eval_dataloader = DataLoader(tokenized_dataset["valid"].select(range(10000), 
                     batch_size=batch_size) # 为演示方便，用了更少的数据


##************** 定义加权损失函数 ***************
keytoken_ids = []
for keyword in [ "plt", "pd", "sk", "fit", "predict", 
                " plt", " pd", " sk", " fit", " predict", 
]:
    ids = tokenizer([keyword]).input_ids[0]
    if len(ids) == 1: # 仅考虑单个 token 的关键字
        keytoken_ids.append(ids[0])
    else:
        print(f"Keyword has not single token: {keyword}")
def keytoken_weighted_loss(inputs, logits, keytoken_ids, alpha=1.0):
    shift_labels = inputs[..., 1:].contiguous()     # input 序列第 i+1 位就是第 i 个标签
    shift_logits = logits[..., :-1, :].contiguous() # 最后一个位置不需要预测，因为没有 label
    
    loss_fct = CrossEntropyLoss(reduce=False) # 用于计算 per-token 的损失
    loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
    loss_per_sample = loss.view(shift_logits.size(0), shift_logits.size(1)).mean(axis=1) # 该样本的平均损失
 
    weights = torch.stack([(inputs == kt).float() for kt in keytoken_ids]).sum( # 每个样本出现的所有关键字的数量
        axis=[0, 2]
    )
    weights = alpha * (1.0 + weights)
    weighted_loss = (loss_per_sample * weights).mean() # 计算 batch 的加权平均损失
    return weighted_loss
  
##************** 定义评估函数 ***************                         
def evaluate():
    model.eval()
    losses = []
    for step, batch in enumerate(eval_dataloader):
        with torch.no_grad():
            outputs = model(batch["input_ids"], labels=batch["input_ids"])
            
        losses.append(accelerator.gather(outputs.loss)) # 跨进程收集每个样本的 loss
    loss = torch.mean(torch.cat(losses))
    try:
        perplexity = torch.exp(loss) # 计算困惑度
    except OverflowError:
        perplexity = float("inf")
    return loss.item(), perplexity.item()
  
##****************** 定义 weight-decay ****************
weight_decay = 0.1


def get_grouped_params(model, no_decay=["bias", "LayerNorm.weight"]):
    params_with_wd, params_without_wd = [], []
    for name, param in model.named_parameters():
        if any(nd in name for nd in no_decay): # 判断 bias 字符串是否出现在 parameter.name 中，因为 parameter.name 可能为 attention1.bias1
            params_without_wd.append(param)
        else:
            params_with_wd.append(param)
    return [
        {"params": params_with_wd, "weight_decay": weight_decay},
        {"params": params_without_wd, "weight_decay": 0.0},
    ]
##***************** 配置模型及其训练组件 *********************
config = AutoConfig.from_pretrained(
    "gpt2",
    vocab_size=len(tokenizer),
    n_ctx=context_length,
    bos_token_id=tokenizer.bos_token_id, # 句子开始的 token
    eos_token_id=tokenizer.eos_token_id, # 句子结束的 token
)


model = GPT2LMHeadModel(config)
optimizer = AdamW(get_grouped_params(model), lr=5e-4) # 使用 weight-decay 
accelerator = Accelerator(mixed_precision='fp16')
model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
    model, optimizer, train_dataloader, eval_dataloader
)
num_train_epochs = 1
num_update_steps_per_epoch = len(train_dataloader) # 必须在 accelerator.prepare() 之后执行，因为  accelerator.prepare 可能会改变 dataloader
num_training_steps = num_train_epochs * num_update_steps_per_epoch
lr_scheduler = get_scheduler(
    name="linear",
    optimizer=optimizer,
    num_warmup_steps=1_000,
    num_training_steps=num_training_steps,
)
##***************** 创建 Repository *********************
model_name = "codeparrot-ds-accelerate"
# repo_name = get_full_repo_name(model_name) #用于推送到 HuggingFace Hub
output_dir = "codeparrot-ds-accelerate"
# repo = Repository(output_dir, clone_from=repo_name) #用于推送到 HuggingFace Hub


##***************** Training Loop *********************
evaluate() # 先评估下结果，看看未训练的模型的效果
gradient_accumulation_steps = 8 # 每隔 8 个 step 来累积一次梯度 (可以通过 accumulator 的 gradient_accumulation_steps 选项来优化这里的梯度累积代码)
eval_steps = 5_000 # 每隔 eval_steps * gradient_accumulation_steps 步时评估一次


model.train()
completed_steps = 0 # 存放梯度更新的次数


for epoch in range(num_train_epochs):
    for step, batch in tqdm(
        enumerate(train_dataloader, start=1), total=len(train_dataloader)
    ):
        logits = model(batch["input_ids"]).logits
        loss = keytoken_weighted_loss(batch["input_ids"], logits, keytoken_ids) # 使用自定义的损失函数
        if step % 100 == 0:
            accelerator.print(
                {
                    "lr": lr_scheduler.get_last_lr(),
                    "samples": step * batch_size,
                    "steps": completed_steps, # 梯度更新的次数
                    "loss/train": loss.item(),
                }
            )
        loss = loss / gradient_accumulation_steps # 缩放损失从而对梯度取平均
        accelerator.backward(loss)
        if step % gradient_accumulation_steps == 0: # 执行梯度更新
            accelerator.clip_grad_norm_(model.parameters(), 1.0) # 梯度范数裁剪
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()
            completed_steps += 1
        if (step % (eval_steps * gradient_accumulation_steps)) == 0:
            eval_loss, perplexity = evaluate()
            accelerator.print({"loss/eval": eval_loss, "perplexity": perplexity})
            model.train() # 评估完成之后，需要设置为 training 模式
                        
accelerator.wait_for_everyone()
evaluate() # 再次评估下结果，看看训练好的模型的效果 
# (4.45915412902832, 86.41439056396484)
unwrapped_model = accelerator.unwrap_model(model)
unwrapped_model.save_pretrained(output_dir, save_function=accelerator.save)
if accelerator.is_main_process:
    tokenizer.save_pretrained(output_dir)
    # repo.push_to_hub(
    #  commit_message=f"Training in progress step {step}", blocking=False
    # )

分享到QQ

分享到微博