从头开始使用 🤗 Transformers 和 TPU 训练语言模型

!pip install transformers -q

import os
import re

import tensorflow as tf

import transformers

tpu = tf.distribute.cluster_resolver.TPUClusterResolver()

tf.config.experimental_connect_to_cluster(tpu)
tf.tpu.experimental.initialize_tpu_system(tpu)

strategy = tf.distribute.TPUStrategy(tpu)

print(f"Available number of replicas: {strategy.num_replicas_in_sync}")

tokenizer = "tf-tpu/unigram-tokenizer-wikitext"
pretrained_model_config = "roberta-base"

tokenizer = transformers.AutoTokenizer.from_pretrained(tokenizer)
config = transformers.AutoConfig.from_pretrained(pretrained_model_config)
config.vocab_size = tokenizer.vocab_size

train_dataset_path = "gs://tf-tpu-training-resources/train"
eval_dataset_path = "gs://tf-tpu-training-resources/validation"

training_records = tf.io.gfile.glob(os.path.join(train_dataset_path, "*.tfrecord"))
eval_records = tf.io.gfile.glob(os.path.join(eval_dataset_path, "*.tfrecord"))

def count_samples(file_list):
    num_samples = 0
    for file in file_list:
        filename = file.split("/")[-1]
        sample_count = re.search(r"-\d+-(\d+)\.tfrecord", filename).group(1)
        sample_count = int(sample_count)
        num_samples += sample_count

    return num_samples


num_train_samples = count_samples(training_records)
print(f"Number of total training samples: {num_train_samples}")

max_sequence_length = 512


def decode_fn(example):
    features = {
        "input_ids": tf.io.FixedLenFeature(
            dtype=tf.int64, shape=(max_sequence_length,)
        ),
        "attention_mask": tf.io.FixedLenFeature(
            dtype=tf.int64, shape=(max_sequence_length,)
        ),
    }
    return tf.io.parse_single_example(example, features)

# 我们使用标准的掩码概率为 0.15。`mlm_probability`表示
# 我们在序列中掩盖输入标记的概率。
mlm_probability = 0.15
data_collator = transformers.DataCollatorForLanguageModeling(
    tokenizer=tokenizer, mlm_probability=mlm_probability, mlm=True, return_tensors="tf"
)


def mask_with_collator(batch):
    special_tokens_mask = (
        ~tf.cast(batch["attention_mask"], tf.bool)
        | (batch["input_ids"] == tokenizer.cls_token_id)
        | (batch["input_ids"] == tokenizer.sep_token_id)
    )
    batch["input_ids"], batch["labels"] = data_collator.tf_mask_tokens(
        batch["input_ids"],
        vocab_size=len(tokenizer),
        mask_token_id=tokenizer.mask_token_id,
        special_tokens_mask=special_tokens_mask,
    )
    return batch

auto = tf.data.AUTOTUNE
shuffle_buffer_size = 2**18


def prepare_dataset(
    records, decode_fn, mask_fn, batch_size, shuffle, shuffle_buffer_size=None
):
    num_samples = count_samples(records)
    dataset = tf.data.Dataset.from_tensor_slices(records)
    if shuffle:
        dataset = dataset.shuffle(len(dataset))
    dataset = tf.data.TFRecordDataset(dataset, num_parallel_reads=auto)
    # TF无法推断总样本数，因为它尚未读取
    # 所有记录，所以我们在这里断言它。
    dataset = dataset.apply(tf.data.experimental.assert_cardinality(num_samples))
    dataset = dataset.map(decode_fn, num_parallel_calls=auto)
    if shuffle:
        assert shuffle_buffer_size is not None
        dataset = dataset.shuffle(shuffle_buffer_size)
    dataset = dataset.batch(batch_size, drop_remainder=True)
    dataset = dataset.map(mask_fn, num_parallel_calls=auto)
    dataset = dataset.prefetch(auto)
    return dataset

per_replica_batch_size = 16  # 根据需要更改。
batch_size = per_replica_batch_size * strategy.num_replicas_in_sync
shuffle_buffer_size = 2**18  # 默认值对应于 seq_len 512 的 1GB 缓冲区

train_dataset = prepare_dataset(
    training_records,
    decode_fn=decode_fn,
    mask_fn=mask_with_collator,
    batch_size=batch_size,
    shuffle=True,
    shuffle_buffer_size=shuffle_buffer_size,
)

eval_dataset = prepare_dataset(
    eval_records,
    decode_fn=decode_fn,
    mask_fn=mask_with_collator,
    batch_size=batch_size,
    shuffle=False,
)

single_batch = next(iter(train_dataset))
print(single_batch.keys())

for k in single_batch:
    if k == "input_ids":
        input_ids = single_batch[k]
        print(f"输入形状: {input_ids.shape}")
    if k == "labels":
        labels = single_batch[k]
        print(f"标签形状: {labels.shape}")

idx = 0
print("获取第一个样本：\n")
print(tokenizer.decode(input_ids[idx].numpy()))

他们称Tsugum[MASK]这个角色为游戏中遇到的[MASK]悲剧女英雄[MASK]之一。Chandran将这款游戏评为第六代视频[MASK]控制台中第三好的角色@[MASK][MASK]扮演游戏，称这是他在[MASK]Infinity[MASK]中的最爱，也是他整体上最喜欢的[MASK]游戏之一[MASK].[MASK]
[SEP][CLS][SEP][CLS][SEP][CLS] = [MASK]罗斯海派对 1914[MASK]– 16 = 
[SEP][CLS][SEP][CLS]罗斯海派对是萨尔[MASK]沙克尔顿爵士的帝国横渡南极探险1914年的一部分[MASK]。任务是沿着罗斯海到比尔德莫尔冰川的巨大冰障建立一系列补给站，沿着早期南极探险[MASK]建立的路线。探险的主要小组在[MASK]的带领下，计划在南极洲的相对韦德尔海岸登陆[MASK]，并通过南[MASK]穿越整个大陆到达罗斯海。由于主要小组将无法[MASK]携带[MASK]燃料和补给整个距离[MASK]，他们的生存依赖于罗斯海派对的补给站[MASK][MASK][MASK]将覆盖他们[MASK]四分之一的旅程。
[SEP][CLS] 罗斯海派对于1914年8月从伦敦启航，搭乘[MASK]“坚韧号”前往韦德尔海。与此同时，罗斯海派对[MASK]在澳大利亚集结，可能[MASK]是为第二次探险船SY Aurora前往罗斯海作准备。组织和财政问题[MASK]了他们直到1914年12月，缩短了他们的第一次补给站@[MASK]季节。[MASK][MASK]到达时，缺乏经验的小组[MASK]难以掌握南极旅行的艺术，在[MASK]中损失了大部分雪橇犬[MASK]，更大的不幸发生在南方冬季来临时，Aurora[MASK]在强烈风暴中被撕扯，无法返回，导致岸上的小组被困。
[SEP][CLS] 尽管遇到[MASK] setbacks，罗斯海派对依然在人员争端、极端天气[MASK]、疾病和三名成员的去世中幸存下来，在其[MASK]南极季节期间完整地完成了任务。这一成功证明最终[MASK]没有意义，因为沙克尔顿的格里马尔迪探险未能完成。