from transformers import AutoModelForMaskedLM

model_checkpoint = "distilbert-base-uncased"

model = AutoModelForMaskedLM.from_pretrained(model_checkpoint)

distilbert_num_parameters = model.num_parameters() / 1_000_000

print(f"'>>> DistilBERT number of parameters: {round(distilbert_num_parameters)}M'")

print(f"'>>> BERT number of parameters: 110M'")

'>>> DistilBERT number of parameters: 67M'

'>>> BERT number of parameters: 110M'

text = "This is a great [MASK]."

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)

import torch

inputs = tokenizer(text, return_tensors="pt")

token_logits = model(**inputs).logits

# Find the location of [MASK] and extract its logits

mask_token_index = torch.where(inputs["input_ids"] == tokenizer.mask_token_id)[1]

mask_token_logits = token_logits[0, mask_token_index, :]

# Pick the [MASK] candidates with the highest logits

top_5_tokens = torch.topk(mask_token_logits, 5, dim=1).indices[0].tolist()

for token in top_5_tokens:

    print(f"'>>> {text.replace(tokenizer.mask_token, tokenizer.decode([token]))}'")

inputs

{'input_ids': tensor([[ 101, 2023, 2003, 1037, 2307,  103, 1012,  102]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1]])}

tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])

['[CLS]', 'this', 'is', 'a', 'great', '[MASK]', '.', '[SEP]']

model(**inputs)

MaskedLMOutput(loss=None, logits=tensor([[[ -5.5882,  -5.5868,  -5.5958,  ...,  -4.9448,  -4.8174,  -2.9905],

         [-11.9031, -11.8872, -12.0623,  ..., -10.9570, -10.6464,  -8.6324],

         [-11.9604, -12.1520, -12.1279,  ..., -10.0218,  -8.6074,  -8.0971],

         ...,

         [ -4.8228,  -4.6268,  -5.1041,  ...,  -4.2771,  -5.0184,  -3.9428],

         [-11.2945, -11.2388, -11.3857,  ...,  -9.2063,  -9.3411,  -6.1505],

         [ -9.5213,  -9.4632,  -9.5022,  ...,  -8.6561,  -8.4908,  -4.6903]]],

       grad_fn=<ViewBackward0>), hidden_states=None, attentions=None)

model(**inputs)['logits']和model(**inputs).logits相同

model(**inputs)['logits'].shape

torch.Size([1, 8, 30522])

torch.where(inputs["input_ids"] == tokenizer.mask_token_id)

(tensor([0]), tensor([5])) # 列出横纵坐标

mask_token_logits.shape

torch.Size([1, 30522])

torch.topk(mask_token_logits, 5, dim=1)

torch.return_types.topk(

values=tensor([[7.0727, 6.6514, 6.6425, 6.2530, 5.8618]], grad_fn=<TopkBackward0>),

indices=tensor([[3066, 3112, 6172, 2801, 8658]]))

# 上面应该是对 mask_token_logits 的 topk 进行了优化，所以返回值中有 indices，这表示词表中 token 对应的数字

mask_token_logits.topk(3, dim=1)

'>>> This is a great deal.'

'>>> This is a great success.'

'>>> This is a great adventure.'

'>>> This is a great idea.'

'>>> This is a great feat.'

from datasets import load_dataset

imdb_dataset = load_dataset("imdb")

imdb_dataset

DatasetDict({

    train: Dataset({

        features: ['text', 'label'],

        num_rows: 25000

    })

    test: Dataset({

        features: ['text', 'label'],

        num_rows: 25000

    })

    unsupervised: Dataset({

        features: ['text', 'label'],

        num_rows: 50000

    })

})

sample = imdb_dataset["train"].shuffle(seed=42).select(range(3))

for row in sample:

    print(f"\n'>>> Review: {row['text']}'")

    print(f"'>>> Label: {row['label']}'")

'>>> Review: This is your typical Priyadarshan movie--a bunch of loony characters out on some silly mission. His signature climax has the entire cast of the film coming together and fighting each other in some crazy moshpit over hidden money. Whether it is a winning lottery ticket in Malamaal Weekly, black money in Hera Pheri, "kodokoo" in Phir Hera Pheri, etc., etc., the director is becoming ridiculously predictable. Don\'t get me wrong; as clichéd and preposterous his movies may be, I usually end up enjoying the comedy. However, in most his previous movies there has actually been some good humor, (Hungama and Hera Pheri being noteworthy ones). Now, the hilarity of his films is fading as he is using the same formula over and over again.<br /><br />Songs are good. Tanushree Datta looks awesome. Rajpal Yadav is irritating, and Tusshar is not a whole lot better. Kunal Khemu is OK, and Sharman Joshi is the best.'

'>>> Label: 0'

'>>> Review: Okay, the story makes no sense, the characters lack any dimensionally, the best dialogue is ad-libs about the low quality of movie, the cinematography is dismal, and only editing saves a bit of the muddle, but Sam" Peckinpah directed the film. Somehow, his direction is not enough. For those who appreciate Peckinpah and his great work, this movie is a disappointment. Even a great cast cannot redeem the time the viewer wastes with this minimal effort.<br /><br />The proper response to the movie is the contempt that the director San Peckinpah, James Caan, Robert Duvall, Burt Young, Bo Hopkins, Arthur Hill, and even Gig Young bring to their work. Watch the great Peckinpah films. Skip this mess.'

'>>> Label: 0'

'>>> Review: I saw this movie at the theaters when I was about 6 or 7 years old. I loved it then, and have recently come to own a VHS version. <br /><br />My 4 and 6 year old children love this movie and have been asking again and again to watch it. <br /><br />I have enjoyed watching it again too. Though I have to admit it is not as good on a little TV.<br /><br />I do not have older children so I do not know what they would think of it. <br /><br />The songs are very cute. My daughter keeps singing them over and over.<br /><br />Hope this helps.'

'>>> Label: 1'

def tokenize_function(examples):

    result = tokenizer(examples["text"])

    if tokenizer.is_fast:

        result["word_ids"] = [result.word_ids(i) for i in range(len(result["input_ids"]))]

    return result

# Use batched=True to activate fast multithreading!

tokenized_datasets = imdb_dataset.map(

    tokenize_function, batched=True, remove_columns=["text", "label"]

)

tokenized_datasets

DatasetDict({

    train: Dataset({

        features: ['attention_mask', 'input_ids', 'word_ids'],

        num_rows: 25000

    })

    test: Dataset({

        features: ['attention_mask', 'input_ids', 'word_ids'],

        num_rows: 25000

    })

    unsupervised: Dataset({

        features: ['attention_mask', 'input_ids', 'word_ids'],

        num_rows: 50000

    })

})

tokenizer.model_max_length

512

chunk_size = 128

# Slicing produces a list of lists for each feature

tokenized_samples = tokenized_datasets["train"][:3]

for idx, sample in enumerate(tokenized_samples["input_ids"]):

    print(f"'>>> Review {idx} length: {len(sample)}'")

'>>> Review 0 length: 200'

'>>> Review 1 length: 559'

'>>> Review 2 length: 192'

concatenated_examples = {

    k: sum(tokenized_samples[k], []) for k in tokenized_samples.keys()

}

total_length = len(concatenated_examples["input_ids"])

print(f"'>>> Concatenated reviews length: {total_length}'")

'>>> Concatenated reviews length: 951'

chunks = {

    k: [t[i : i + chunk_size] for i in range(0, total_length, chunk_size)]

    for k, t in concatenated_examples.items()

}

for chunk in chunks["input_ids"]:

    print(f"'>>> Chunk length: {len(chunk)}'")

'>>> Chunk length: 128'

'>>> Chunk length: 128'

'>>> Chunk length: 128'

'>>> Chunk length: 128'

'>>> Chunk length: 128'

'>>> Chunk length: 128'

'>>> Chunk length: 128'

'>>> Chunk length: 55'

def group_texts(examples):

    # Concatenate all texts

    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}

    # Compute length of concatenated texts

    total_length = len(concatenated_examples[list(examples.keys())[0]])

    # We drop the last chunk if it's smaller than chunk_size

    total_length = (total_length // chunk_size) * chunk_size

    # Split by chunks of max_len

    result = {

        k: [t[i : i + chunk_size] for i in range(0, total_length, chunk_size)]

        for k, t in concatenated_examples.items()

    }

    # Create a new labels column

    result["labels"] = result["input_ids"].copy()

    return result

lm_datasets = tokenized_datasets.map(group_texts, batched=True)

lm_datasets

DatasetDict({

    train: Dataset({

        features: ['attention_mask', 'input_ids', 'labels', 'word_ids'],

        num_rows: 61289

    })

    test: Dataset({

        features: ['attention_mask', 'input_ids', 'labels', 'word_ids'],

        num_rows: 59905

    })

    unsupervised: Dataset({

        features: ['attention_mask', 'input_ids', 'labels', 'word_ids'],

        num_rows: 122963

    })

})

tokenizer.decode(lm_datasets["train"][1]["input_ids"])

".... at.......... high. a classic line : inspector : i'm here to sack one of your teachers. student : welcome to bromwell high. i expect that many adults of my age think that bromwell high is far fetched. what a pity that it isn't! [SEP] [CLS] homelessness ( or houselessness as george carlin stated ) has been an issue for years but never a plan to help those on the street that were once considered human who did everything from going to school, work, or vote for the matter. most people think of the homeless"

tokenizer.decode(lm_datasets["train"][1]["labels"])

".... at.......... high. a classic line : inspector : i'm here to sack one of your teachers. student : welcome to bromwell high. i expect that many adults of my age think that bromwell high is far fetched. what a pity that it isn't! [SEP] [CLS] homelessness ( or houselessness as george carlin stated ) has been an issue for years but never a plan to help those on the street that were once considered human who did everything from going to school, work, or vote for the matter. most people think of the homeless"

from transformers import DataCollatorForLanguageModeling

data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=0.15)

samples = [lm_datasets["train"][i] for i in range(2)]

for sample in samples:

    _ = sample.pop("word_ids")

for chunk in data_collator(samples)["input_ids"]:

    print(f"\n'>>> {tokenizer.decode(chunk)}'")

'>>> [CLS] bromwell [MASK] is a cartoon comedy. it ran at the same [MASK] as some other [MASK] about school life, [MASK] as " teachers ". [MASK] [MASK] [MASK] in the teaching [MASK] lead [MASK] to believe that bromwell high\'[MASK] satire is much closer to reality than is " teachers ". the scramble [MASK] [MASK] financially, the [MASK]ful students whogn [MASK] right through [MASK] pathetic teachers\'pomp, the pettiness of the whole situation, distinction remind me of the schools i knew and their students. when i saw [MASK] episode in [MASK] a student repeatedly tried to burn down the school, [MASK] immediately recalled. [MASK]...'

'>>> .... at.. [MASK]... [MASK]... high. a classic line plucked inspector : i\'[MASK] here to [MASK] one of your [MASK]. student : welcome to bromwell [MASK]. i expect that many adults of my age think that [MASK]mwell [MASK] is [MASK] fetched. what a pity that it isn\'t! [SEP] [CLS] [MASK]ness ( or [MASK]lessness as george 宇in stated )公 been an issue for years but never [MASK] plan to help those on the street that were once considered human [MASK] did everything from going to school, [MASK], [MASK] vote for the matter. most people think [MASK] the homeless'

import collections

import numpy as np

from transformers import default_data_collator

wwm_probability = 0.2

def whole_word_masking_data_collator(features):

    for feature in features:

        word_ids = feature.pop("word_ids")

        # Create a map between words and corresponding token indices

        mapping = collections.defaultdict(list)

        current_word_index = -1

        current_word = None

        for idx, word_id in enumerate(word_ids):

            if word_id is not None:

                if word_id != current_word:

                    current_word = word_id

                    current_word_index += 1

                mapping[current_word_index].append(idx)

        # Randomly mask words

        mask = np.random.binomial(1, wwm_probability, (len(mapping),))

        input_ids = feature["input_ids"]

        labels = feature["labels"]

        new_labels = [-100] * len(labels)

        for word_id in np.where(mask)[0]:

            word_id = word_id.item()

            for idx in mapping[word_id]:

                new_labels[idx] = labels[idx]

                input_ids[idx] = tokenizer.mask_token_id

        feature["labels"] = new_labels

    return default_data_collator(features)

samples = [lm_datasets["train"][i] for i in range(2)]

batch = whole_word_masking_data_collator(samples)

for chunk in batch["input_ids"]:

    print(f"\n'>>> {tokenizer.decode(chunk)}'")

'>>> [CLS] bromwell high is a cartoon comedy [MASK] it ran at the same time as some other programs about school life, such as " teachers ". my 35 years in the teaching profession lead me to believe that bromwell high\'s satire is much closer to reality than is " teachers ". the scramble to survive financially, the insightful students who can see right through their pathetic teachers\'pomp, the pettiness of the whole situation, all remind me of the schools i knew and their students. when i saw the episode in which a student repeatedly tried to burn down the school, i immediately recalled.....'

'>>> .... [MASK] [MASK] [MASK] [MASK]....... high. a classic line : inspector : i\'m here to sack one of your teachers. student : welcome to bromwell high. i expect that many adults of my age think that bromwell high is far fetched. what a pity that it isn\'t! [SEP] [CLS] homelessness ( or houselessness as george carlin stated ) has been an issue for years but never a plan to help those on the street that were once considered human who did everything from going to school, work, or vote for the matter. most people think of the homeless'

train_size = 10_000

test_size = int(0.1 * train_size)

downsampled_dataset = lm_datasets["train"].train_test_split(

    train_size=train_size, test_size=test_size, seed=42

)

downsampled_dataset

DatasetDict({

    train: Dataset({

        features: ['attention_mask', 'input_ids', 'labels', 'word_ids'],

        num_rows: 10000

    })

    test: Dataset({

        features: ['attention_mask', 'input_ids', 'labels', 'word_ids'],

        num_rows: 1000

    })

})

from transformers import TrainingArguments

batch_size = 64

# Show the training loss with every epoch

logging_steps = len(downsampled_dataset["train"]) // batch_size

model_name = model_checkpoint.split("/")[-1]

training_args = TrainingArguments(

    output_dir=f"{model_name}-finetuned-imdb",

    overwrite_output_dir=True,

    evaluation_strategy="epoch",

    learning_rate=2e-5,

    weight_decay=0.01,

    per_device_train_batch_size=batch_size,

    per_device_eval_batch_size=batch_size,

    push_to_hub=True,

    fp16=True,

    logging_steps=logging_steps,

)

import math

eval_results = trainer.evaluate()

print(f">>> Perplexity: {math.exp(eval_results['eval_loss']):.2f}")

Perplexity: 21.75

trainer.train()