# Source: content/notes/nlp/pretrained-model-families.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import os
os.environ["USE_TF"] = "0"  # select the PyTorch backend before importing Transformers
import torch
from transformers import BertConfig, BertForMaskedLM, GPT2Config, GPT2LMHeadModel

torch.set_num_threads(1)
torch.manual_seed(7)
ids = torch.tensor([[1, 5, 9, 2]])
bert = BertForMaskedLM(BertConfig(vocab_size=16, hidden_size=8,
    num_hidden_layers=1, num_attention_heads=2, intermediate_size=16,
    hidden_dropout_prob=0, attention_probs_dropout_prob=0))
masked = ids.clone()
masked[0, 2] = 3
labels = torch.full_like(ids, -100)
labels[0, 2] = ids[0, 2]
out = bert(masked, labels=labels)
manual = torch.nn.functional.cross_entropy(out.logits[:, 2], ids[:, 2])
torch.testing.assert_close(out.loss, manual)
gpt = GPT2LMHeadModel(GPT2Config(vocab_size=16, n_embd=8, n_layer=1,
    n_head=2, n_positions=8, resid_pdrop=0, embd_pdrop=0, attn_pdrop=0))
causal = gpt(ids, labels=ids)
manual = torch.nn.functional.cross_entropy(
    causal.logits[:, :-1].reshape(-1, 16), ids[:, 1:].reshape(-1))
torch.testing.assert_close(causal.loss, manual)
assert out.logits.shape == causal.logits.shape == (1, 4, 16)
print("MLM position selection and causal label shift verified")