text:
"""
卷积理解
"""
import os
# os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
# os.environ['CUDA_VISIBLE_DEVICES'] = '-1'   # 设置在CPU上运行
import cv2
import numpy as np
import tensorflow as tf
from tensorflow import keras
import matplotlib.pyplot as plt
physical_devices = tf.config.experimental.list_physical_devices('GPU')
for gpu in physical_devices:
    tf.config.experimental.set_memory_growth(gpu, True)
image_root = '../data/Training/cats'

def show_conv():
    np.random.seed(1671)
    sobel_x = [[-1, -2, -1], [0, 0, 0], [1, 2, 1]]
    sobel_y = [[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]
    img = cv2.imread(os.path.join(image_root, 'cat.1.jpg'), 0)

    cv2.imshow('src', img)
    cv2.waitKey()
    input_array = img.reshape([1, img.shape[0], img.shape[1], 1])
    input_array = np.array(input_array, dtype=np.float64)
    model = keras.models.Sequential()
    conv1 = keras.layers.Conv2D(filters=1, kernel_size=3, strides=1, padding='same',
                                kernel_initializer=keras.initializers.Constant(value=sobel_x),
                                input_shape=input_array.shape[1:])
    conv2 = keras.layers.Conv2D(filters=1, kernel_size=3, strides=1, padding='same',
                                kernel_initializer=keras.initializers.Constant(value=sobel_y),
                                input_shape=input_array.shape[1:])
    model.add(conv1)
    model.summary()
    out_put = conv1(input_array)
    print(out_put.shape)
    cv2.imshow('x', np.array(out_put)[0, :, :, :])
    cv2.waitKey()

    deconv1 = keras.layers.Conv2DTranspose(filters=1, kernel_size=3, strides=1, padding='same',
                                           kernel_initializer=keras.initializers.Constant(
                                               value=np.linalg.pinv(sobel_x)),
                                           input_shape=np.array(out_put).shape[1:])
    deconv_x = deconv1(np.array(out_put))
    cv2.imshow('deconv_x', np.array(deconv_x)[0, :, :, :])
    cv2.waitKey()

    out_put = conv2(input_array)
    print(out_put.shape)
    cv2.imshow('y', np.array(out_put)[0, :, :, :])
    cv2.waitKey()

def spatial_attention():
    img = cv2.imread(os.path.join(image_root, 'cat.1.jpg'))
    cv2.imshow('1', img)
    input_array = img.reshape([1, img.shape[0], img.shape[1], img.shape[2]])
    input_array = np.array(input_array, dtype=np.float64)
    avg_out = tf.reduce_mean(input_array, 3)
    max_out = tf.reduce_max(input_array, 3)

    ave_img = np.array(avg_out, dtype=np.uint8).reshape(img.shape[0], img.shape[1], 1)
    cv2.imshow('ave', ave_img)

    max_img = np.array(max_out, dtype=np.uint8).reshape(img.shape[0], img.shape[1], 1)
    cv2.imshow('max', max_img)
    cv2.waitKey()
    # plt.imshow(show_img)
    # cv2.waitKey()

def conv_calculate():
    np.random.seed(1671)
    sobel_x = [[-1, -2, -1], [0, 0, 0], [1, 2, 1]]*3
    sobel_y = [[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]*3
    img = np.array(range(27))
    input_array = img.reshape([1, 3, 3, 3])
    input_array = np.array(input_array, dtype=np.float64)
    model = keras.models.Sequential()
    conv1 = keras.layers.Conv2D(filters=1, kernel_size=3, strides=1, padding='valid',
                                kernel_initializer=keras.initializers.Constant(value=sobel_x),
                                input_shape=input_array.shape[1:])
    conv2 = keras.layers.Conv2D(filters=1, kernel_size=3, strides=1, padding='same',
                                kernel_initializer=keras.initializers.Constant(value=sobel_y),
                                input_shape=input_array.shape[1:])
    model.add(conv1)
    model.summary()
    out_put = conv1(input_array)
    print(out_put.shape)
    print(np.array(out_put)[0, :, :, :])

    out_put = conv2(input_array)
    print(out_put.shape)
    # cv2.imshow('2', np.array(out_put)[0, :, :, :])
    # cv2.waitKey()

def map_a():
    img = cv2.imread('../data/Training/cats/cat.10.jpg', 0)
    cv2.imshow('1', img)
    # cv2.waitKey()

    img = img // 2
    cv2.imshow('2', img)
    cv2.waitKey()

if __name__ == r'__main__':
    # print(tf.__version__)
    show_conv()
    # conv_calculate()
    # spatial_attention()
    # map_a()


text:
import numpy as np

import os
import subprocess
import json

from timeit import default_timer as timer

repeats = 5

threads = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096]

directory = 'images/'

images = []

for dirpath,_,filenames in os.walk(directory):
    for f in filenames:
        if f.endswith('.jpg') and 'size' in f:
            images.append(os.path.abspath(os.path.join(dirpath, f)))

images.sort()

cuda_comand = "../Integral_Images_(final-term)/Code/CUDA/cuda -json"

cuda_results = np.zeros((repeats,len(images)))

start_time = timer()

j = 0
for image in images:
    print("Working on image: " + str(j))
    for i in range(repeats):
        com_cuda = cuda_comand + " -p " + image
        process = subprocess.Popen(com_cuda.split(), stdout=subprocess.PIPE)
        output, error = process.communicate()
        output = output.decode("utf-8")
        json_data = json.loads(output)
        cuda_results[i,j] = json_data["time"]
    j = j+1

end_time = timer()

print(cuda_results)

print("\n\nTime elapsed: " + "{0:.2f}".format((end_time - start_time)) + " s")

np.save("cuda_results_integral_max-threads.npy", cuda_results)

ext: py

max_stars_repo_name: LucaAngioloni/ParallelComputingExam

max_issues_repo_name: SqrtPapere/ParallelComputingExam

max_forks_repo_name: SqrtPapere/ParallelComputingExam

reason: import numpy


text:
import numpy as np
import pandas as pd
import mxnet as mx
import logging

# set to 0 to train on all available data
VALIDATION_SIZE = 1000

# create the training
dataset = pd.read_csv("./train.csv")
target = dataset[[0]].values.ravel()
train = dataset.iloc[:,1:].values

def get_lenet():
    data = mx.symbol.Variable('data')
    # first conv
    conv1 = mx.symbol.Convolution(data=data, kernel=(5,5), num_filter=20)
    tanh1 = mx.symbol.Activation(data=conv1, act_type="tanh")
    pool1 = mx.symbol.Pooling(data=tanh1, pool_type="max",
                              kernel=(3,3), stride=(2,2))
    # second conv
    conv2 = mx.symbol.Convolution(data=pool1, kernel=(3,3), num_filter=50)
    tanh2 = mx.symbol.Activation(data=conv2, act_type="tanh")
    pool2 = mx.symbol.Pooling(data=tanh2, pool_type="max",
                              kernel=(3,3), stride=(2,2))
    # first fullc
    flatten = mx.symbol.Flatten(data=pool2)
    fc1 = mx.symbol.FullyConnected(data=flatten, num_hidden=500)
    tanh3 = mx.symbol.Activation(data=fc1, act_type="tanh")
    # second fullc
    fc2 = mx.symbol.FullyConnected(data=tanh3, num_hidden=10)
    # loss
    lenet = mx.symbol.SoftmaxOutput(data=fc2, name='softmax')
    return lenet

# split dataset
val_data = train[:VALIDATION_SIZE].astype('float32')
val_label = target[:VALIDATION_SIZE]
train_data = train[VALIDATION_SIZE: , :].astype('float32')
train_label = target[VALIDATION_SIZE:]
train_data = np.array(train_data).reshape((-1, 1, 28, 28))
val_data = np.array(val_data).reshape((-1, 1, 28, 28))

# Normalize data
train_data[:] /= 256.0
val_data[:]/= 256.0

batch_size = 500
train_iter = mx.io.NDArrayIter(train_data, train_label, batch_size=batch_size, shuffle=True)
val_iter = mx.io.NDArrayIter(val_data, val_label, batch_size=batch_size)

# logging
head = '%(asctime)-15s Node[0] %(message)s'
logging.basicConfig(level=logging.DEBUG, format=head)

# create model
devs = mx.gpu(0)
network=get_lenet()
model = mx.model.FeedForward(
        ctx                = devs,
        symbol             = network,
        num_epoch          = 50,
        learning_rate      = 0.01,
        momentum           = 0.9,
        wd                 = 0.0001,
        initializer        = mx.init.Xavier(factor_type="in", magnitude=2.34)
        )

eval_metrics = ['accuracy']
model.fit(
	X=train_iter,
	eval_metric        = eval_metrics,
	eval_data	 = val_iter
	)

#predict
test = pd.read_csv("./test.csv").values
test_data = test.astype('float32')
test_data = np.array(test_data).reshape((-1, 1, 28, 28))
test_data[:]/= 256.0
test_iter = mx.io.NDArrayIter(test_data, batch_size=batch_size)

pred = model.predict(X = test_iter)
pred = np.argsort(pred)
np.savetxt('submission_lb.csv', np.c_[range(1,len(test)+1),pred[:,9]], delimiter=',', header = 'ImageId,Label', comments = '', fmt='%d')

max_stars_repo_name: lbin/mnist

max_issues_repo_name: lbin/mnist

max_forks_repo_name: lbin/mnist


text:
import os
import util
import deep_net_utils
from sklearn.linear_model import LogisticRegression
import random
import argparse
from sklearn.svm import SVC
import pickle
import numpy as np

parser = argparse.ArgumentParser()
# setting the default to be 'experiments/fallback' so we don't risk overwriting existing experiments
parser.add_argument('--model_dir', default='experiments/fallback',
                    help="Directory containing params.json")
parser.add_argument('--verbose', default=1,
                    help="Directory containing params.json")
parser.add_argument('--vis', default=0,
                    help="Directory containing params.json")
# You may evaluate on 'train', 'val', or 'test'.
# Default '' means that we are simply training, and metrics for train and val will be saved. Model will also be saved
parser.add_argument('--evaluate', default='',
                    help="Directory containing params.json")

def evaluate(args, params):
    split = args.evaluate
    data_dir = '../data/' + params.dataset
    X, Y = util.load_dataset_from_split_directory(data_dir + '/' + split + '/', params.img_dimensions, args.verbose)

    # Load model
    saved_model_path = os.path.join(
        args.model_dir, "finalized_model.sav")
    loaded_model = pickle.load(open(saved_model_path, 'rb'))

    # predit
    predicted = loaded_model.predict(X)

    # evaluate metrics
    confus_save_path = os.path.join(
        args.model_dir, "confus_f1_{}.json".format(split))
    util.compute_and_save_f1(predicted, Y, confus_save_path)

def fit_softmax_or_svm(X_train, Y_train, params):
    class_weight = None if params.class_weight == "none" else params.class_weight
    if params.class_weight == 'custom': class_weight = {0: 10, 1:1, 2: 3}
    if params.model_type == 'softmax':
        return LogisticRegression(
            random_state=0,
            max_iter=params.iter,
            multi_class='multinomial',
            solver='lbfgs',
            penalty=params.penalty,
            C=1/params.regularization_constant,
            class_weight=class_weight
        ).fit(X_train, Y_train)
    if params.model_type == 'svm':
        svm = SVC(
            C=1/params.regularization_contant,
            kernel=params.kernel,
            degree=params.degree,
            gamma=params.gamma,
            class_weight=class_weight,
            probability=True,
            random_state=42).fit(X_train, Y_train)
        return svm
    return None

def train(args, params):
    random.seed(229)

    # Check model_type is set (if 'model_type' is 'svm', 'kernel' is also set)
    assert(hasattr(params, 'model_type'))
    if params.model_type == 'svm': assert(hasattr(params, 'kernel'))

    data_container = util.load_dataset_for_split(args, params)
    X_train, Y_train = data_container['train']
    X_val, Y_val = data_container['val']
    X_test, Y_test = data_container['test']

    if args.verbose: print("We have {}, {}, and {} training examples for classes 0, 1, 2 respectively.".format(
        np.sum(np.where(Y_train == 0, 1, 0)),
        np.sum(np.where(Y_train == 1, 1, 0)),
        np.sum(np.where(Y_train == 2, 1, 0))
    ))

    model = fit_softmax_or_svm(X_train, Y_train, params)

    model_file_path = os.path.join(
        args.model_dir, "finalized_model.sav")
    pickle.dump(model, open(model_file_path, 'wb'))

    predicted_train = model.predict(X_train)
    predicted_val = model.predict(X_val)

    train_confus_save_path = os.path.join(
        args.model_dir, "confus_f1_train.json")
    val_confus_save_path = os.path.join(
        args.model_dir, "confus_f1_val.json")

    util.compute_and_save_f1(predicted_train, Y_train, train_confus_save_path)
    util.compute_and_save_f1(predicted_val, Y_val, val_confus_save_path)
    if params.model_type == 'softmax' and args.verbose: print("Model converged in {} iterations.".format(model.n_iter_))

    train_accur = model.score(X_train, Y_train)
    print("Acc on train: ", train_accur)

    val_accur = model.score(X_val, Y_val)
    print("Acc on validation: ", val_accur)

    if args.vis: util.visualize(X_train, Y_train)

if __name__ == "__main__":
    random.seed(229)
    args = parser.parse_args()

    # load hyperparameters
    json_path = os.path.join(args.model_dir, "params.json")
    assert os.path.isfile(
        json_path), "No json configuration file found at {}".format(json_path)
    params = deep_net_utils.Params(json_path)

    if args.evaluate == "":
        train(args, params)
    else:
        evaluate(args, params)


text:
""" Imitation Learning (IL) training. """

import time
import os
import random
import argparse
import pickle
from collections import OrderedDict

import numpy as np
import torch
import torch.optim as optim
import math

from src.environments import *
from models.feedforward import *
from utilities.hdf5_dataloader import *

# state dimensions
# var_dim is the dimension of each candidate variable's input, i.e., the fixed dimension of matrix C_t
# Tree_t is given by concatenation of two states, for a total dimension node_dim + mip_dim
state_dims = {
    'var_dim': 25,
    'node_dim': 8,
    'mip_dim': 53
}

if __name__ == '__main__':

    # parser definition
    parser = argparse.ArgumentParser(description='Parser for IL training experiments.')
    parser.add_argument(
        '--seed',
        type=int,
        default=0,
        help='Random seed for IL training experiment.'
    )
    parser.add_argument(
        '--policy_type',
        type=str,
        choices=['NoTreePolicy', 'TreeGatePolicy'],
        help='Type of policy to use.'
    )
    parser.add_argument(
        '--hidden_size',
        type=int,
        help='Hidden size of the branching policy network.'
    )
    parser.add_argument(
        '--depth',
        type=int,
        help='Depth of the branching policy network.'
    )
    parser.add_argument(
        '--dropout',
        type=float,
        default=0.0,
        help='Dropout parameter for the branching policy network.'
    )
    parser.add_argument(
        '--dim_reduce_factor',
        type=int,
        default=2,
        help='Dimension reduce factor of the branching policy network.'
    )
    parser.add_argument(
        '--infimum',
        type=int,
        default=8,
        help='Infimum parameter of the branching policy network.'
    )
    parser.add_argument(
        '--norm',
        type=str,
        default='none',
        help='Normalization type of the branching policy network.'
    )
    parser.add_argument(
        '--train_h5_path',
        type=str,
        help='Pathway to the train H5 file.'
    )
    parser.add_argument(
        '--val_h5_path',
        type=str,
        help='Pathway to the val H5 file.'
    )
    parser.add_argument(
        '--test_h5_path',
        type=str,
        help='Pathway to the test H5 file.'
    )
    parser.add_argument(
        '--out_dir',
        type=str,
        help='Directory to save the experimental results.'
    )
    parser.add_argument(
        '--use_gpu',
        default=False,
        action='store_true',
        help='Use gpu or not.'
    )
    parser.add_argument(
        '--opt',
        default='adam',
        type=str,
        help='Type of optimizer to use.'
    )
    parser.add_argument(
        '--lr',
        type=float,
        help='Learning rate.'
    )
    parser.add_argument(
        '--momentum',
        default=0.9,
        type=float,
        help='Momentum optimization parameter.'
    )
    parser.add_argument(
        '--weight_decay',
        default=1e-5,
        type=float,
        help='Weight decay optimization parameter.'
    )
    parser.add_argument(
        '--num_epochs',
        type=int,
        default=40,
        help='Number of training epochs.'
    )
    parser.add_argument(
        '--lr_decay_schedule',
        type=int,
        nargs='+',
        default=[20, 30],
        help='Learning rate decay schedule.'
    )
    parser.add_argument(
        '--lr_decay_factor',
        type=float,
        default=0.1,
        help='LR decay factor.'
    )
    parser.add_argument(
        '--train_batchsize',
        type=int,
        default=32,
        help='Training batchsize.')
    parser.add_argument(
        '--eval_batchsize',
        type=int,
        default=500,
        help='Evaluation batchsize.')
    parser.add_argument(
        '--top_k',
        type=int,
        nargs='+',
        default=[2, 3, 5, 10],
        help='In addition to top-1 generalization accuracy, we track top-k.'
    )
    args = parser.parse_args()
    print(args)

    # set all the random seeds
    seed = args.seed
    np.random.seed(seed)
    torch.manual_seed(seed)
    random.seed(seed)

    # setup output directory
    if not os.path.exists(args.out_dir):
        os.makedirs(args.out_dir, exist_ok=True)

    # use gpu or cpu
    if args.use_gpu:
        import torch.backends.cudnn as cudnn

        device = torch.device('cuda')
        cudnn.benchmark = True
    else:
        device = torch.device('cpu')

    # if final checkpoint exists exit the script
    chkpnt_path = os.path.join(args.out_dir, 'final_checkpoint.pth.tar')
    if os.path.isfile(chkpnt_path):
        print('Final checkpoint exists, experiment has already been run, exiting...')
        exit()
    elif os.path.isfile(os.path.join(args.out_dir, 'final_job_crashed_checkpoint.pth.tar')):
        print('Experiment previously crashed, exiting...')
        exit()

    # load a checkpoint path
    chkpnt_path = os.path.join(args.out_dir, 'checkpoint.pth.tar')
    if os.path.isfile(chkpnt_path):
        chkpnt = torch.load(chkpnt_path)
        epoch_start = chkpnt['epoch']
        print('Checkpoint loaded from path {}, starting at epoch {}...'.format(chkpnt_path, epoch_start))
    else:
        chkpnt = None
        epoch_start = 0

    # setup the policy
    if args.policy_type == 'NoTreePolicy':
        policy = NoTreePolicy(
            var_dim=state_dims['var_dim'],
            node_dim=state_dims['node_dim'],
            mip_dim=state_dims['mip_dim'],
            hidden_size=args.hidden_size,
            depth=args.depth,
            dropout=args.dropout,
            dim_reduce_factor=args.dim_reduce_factor,
            infimum=args.infimum,
            norm=args.norm,
        )
        policy_name = 'NoTreePolicy'
    elif args.policy_type == 'TreeGatePolicy':
        policy = TreeGatePolicy(
            var_dim=state_dims['var_dim'],
            node_dim=state_dims['node_dim'],
            mip_dim=state_dims['mip_dim'],
            hidden_size=args.hidden_size,
            depth=args.depth,
            dropout=args.dropout,
            dim_reduce_factor=args.dim_reduce_factor,
            infimum=args.infimum,
            norm=args.norm,
        )
        policy_name = 'TreeGatePolicy'
    else:
        raise ValueError('A valid policy should be set.')

    # set the policy into train mode
    policy.train()
    policy = policy.to(device)

    if args.opt == 'adam':
        optimizer = optim.Adam(
            policy.parameters(),
            lr=args.lr,
            betas=(args.momentum, 0.999),
            weight_decay=args.weight_decay
        )
        eps = np.finfo(np.float32).eps.item()
    else:
        raise ValueError('A valid optimizer should be set.')

    # specify a learning rate scheduler
    if args.lr_decay_schedule:
        lr_decay_schedule = args.lr_decay_schedule
        lr_decay_factor = args.lr_decay_factor
        scheduler = optim.lr_scheduler.MultiStepLR(optimizer, lr_decay_schedule, lr_decay_factor)
        use_scheduler = True
    else:
        use_scheduler = False

    # if checkpoint available, load the policy's and the optimizers parameters
    if chkpnt:
        policy.load_state_dict(chkpnt['state_dict'])
        optimizer.load_state_dict(chkpnt['optimizer'])
        scheduler.load_state_dict(chkpnt['scheduler'])
        exp_dict = chkpnt['exp_dict']
    else:
        # setup an experiment log
        exp_dict = OrderedDict()
        exp_dict['train_loss_epoch_arr'] = np.zeros(args.num_epochs, )
        exp_dict['val_loss_epoch_arr'] = np.zeros(args.num_epochs, )
        exp_dict['val_acc_epoch_arr'] = np.zeros(args.num_epochs, )
        for k in args.top_k:
            exp_dict['val_acc_top_{}_epoch_arr'.format(k)] = np.zeros(args.num_epochs, )
        exp_dict['val_nan_counter_epoch_arr'] = np.zeros(args.num_epochs, )
        exp_dict['test_loss'] = 0.0
        exp_dict['test_acc'] = 0.0
        for k in args.top_k:
            exp_dict['test_acc_top_{}'.format(k)] = 0.0
        exp_dict['test_nan_counter'] = 0

    # setup a train loader
    train_h5 = dataset_h5(
        h5_file=args.train_h5_path,
        node_dim=state_dims['node_dim'],
        mip_dim=state_dims['mip_dim'],
        var_dim=state_dims['var_dim']
    )
    num_train_batches = train_h5.__len__() // args.train_batchsize
    train_loader = DataLoader(
        dataset=train_h5,
        batch_size=args.train_batchsize,
        shuffle=True,
        collate_fn=collate_fn
    )

    # setup a val loader
    val_h5 = dataset_h5(
        h5_file=args.val_h5_path,
        node_dim=state_dims['node_dim'],
        mip_dim=state_dims['mip_dim'],
        var_dim=state_dims['var_dim']
    )
    num_val_batches = val_h5.__len__() // args.eval_batchsize
    val_loader = DataLoader(
        dataset=val_h5,
        batch_size=args.eval_batchsize,
        shuffle=False,
        collate_fn=collate_fn
    )

    # setup a test loader
    test_h5 = dataset_h5(
        h5_file=args.test_h5_path,
        node_dim=state_dims['node_dim'],
        mip_dim=state_dims['mip_dim'],
        var_dim=state_dims['var_dim']
    )
    num_test_batches = test_h5.__len__() // args.eval_batchsize
    test_loader = DataLoader(
        dataset=test_h5,
        batch_size=args.eval_batchsize,
        shuffle=False,
        collate_fn=collate_fn
    )

    # setup the loss
    criterion = nn.CrossEntropyLoss().to(device)

    # main training loop
    print('Starting training loop...\n')
    for i in range(epoch_start, args.num_epochs):
        # set the policy into train mode
        policy.train()
        start_time = time.time()
        start_time_process = time.process_time()

        running_loss = 0.0
        for batch in train_loader:
            optimizer.zero_grad()
            batch_loss = 0.0
            for idx, data_tuple in enumerate(batch):
                target, node, mip, grid = data_tuple
                target, node, mip, grid = target.to(device), node.to(device), mip.to(device), grid.to(device)
                logits = policy(grid, node, mip)
                logits = logits.transpose(1, 0)
                batch_loss += criterion(logits, target)
            batch_loss /= float(args.train_batchsize)
            batch_loss.backward()
            optimizer.step()
            running_loss += batch_loss.item()
        running_loss /= float(num_train_batches)

        if use_scheduler:
            scheduler.step()

        train_time = time.time() - start_time
        train_time_process = time.process_time() - start_time_process
        exp_dict['train_loss_epoch_arr'][i] = running_loss

        # set the policy into eval mode
        policy.eval()
        eval_start = time.time()
        eval_start_process = time.process_time()

        total_correct = 0
        top_k_correct = dict.fromkeys(args.top_k)
        for k in args.top_k:
            top_k_correct[k] = 0
        val_acc_top_k = dict.fromkeys(args.top_k)
        total_loss = 0.0
        nan_counter = 0

        with torch.no_grad():
            for batch in val_loader:
                for idx, data_tuple in enumerate(batch):
                    target, node, mip, grid = data_tuple
                    target, node, mip, grid = target.to(device), node.to(device), mip.to(device), grid.to(device)
                    logits = policy(grid, node, mip)
                    logits = logits.transpose(1, 0)
                    _loss = criterion(logits, target).item()
                    if math.isnan(_loss):
                        nan_counter += 1
                    else:
                        total_loss += _loss
                        _, predicted = torch.max(logits, 1)
                        total_correct += predicted.eq(target.item()).cpu().item()
                        grid_size = grid.size(0)
                        for k in args.top_k:
                            max_k = min(k, grid_size)  # Accounts for when grid_size is smaller than top_k
                            top_k_correct[k] += int(target.item() in logits.topk(max_k, dim=1).indices)
            if nan_counter < val_h5.n_data:
                val_loss = total_loss / float(val_h5.n_data - nan_counter)
                val_acc = total_correct / float(val_h5.n_data - nan_counter)
                for k in args.top_k:
                    val_acc_top_k[k] = top_k_correct[k] / float(val_h5.n_data - nan_counter)
            else:
                val_loss = np.nan
                val_acc = np.nan
                for k in args.top_k:
                    val_acc_top_k[k] = np.nan

                print('Model overflow on entire val set, hyperparameter configuration is ill-posed, killing the job.')

                # save the final checkpoint
                torch.save(
                    {'state_dict': policy.state_dict(),
                     'optimizer': optimizer.state_dict(),
                     'scheduler': scheduler.state_dict(),
                     'exp_dict': exp_dict,
                     'args': args,
                     'state_dims': state_dims
                     },
                    os.path.join(args.out_dir, 'final_job_crashed_checkpoint.pth.tar')
                )

                # exit
                exit()

            eval_time = time.time() - eval_start
            eval_time_process = time.process_time() - eval_start_process
            exp_dict['val_loss_epoch_arr'][i] = val_loss
            exp_dict['val_acc_epoch_arr'][i] = val_acc
            for k in args.top_k:
                exp_dict['val_acc_top_{}_epoch_arr'.format(k)][i] = val_acc_top_k[k]
            exp_dict['val_nan_counter_epoch_arr'][i] = nan_counter

        print(
            "[Epoch {:d}] Train loss: {:.4f}. Train time: {:.2f}sec. "
            "Val loss: {:.4f}. Val acc: {:.2f}%, Val acc top-{}: {:.2f}%, Val acc top-{}: {:.2f}%. Val time: {:.2f}sec.".format(
                i + 1, running_loss, train_time, val_loss, 100 * val_acc, args.top_k[0],
                100 * val_acc_top_k[args.top_k[0]],
                args.top_k[1], 100 * val_acc_top_k[args.top_k[1]], eval_time))

        # checkpoint
        torch.save(
            {'epoch': i + 1,
             'state_dict': policy.state_dict(),
             'optimizer': optimizer.state_dict(),
             'scheduler': scheduler.state_dict(),
             'exp_dict': exp_dict,
             'args': args,
             'state_dims': state_dims
             },
            os.path.join(args.out_dir, 'checkpoint.pth.tar')
        )

        # create per epoch save
        torch.save(
            {'epoch': i + 1,
             'state_dict': policy.state_dict(),
             'optimizer': optimizer.state_dict(),
             'scheduler': scheduler.state_dict(),
             'exp_dict': exp_dict,
             'args': args,
             'state_dims': state_dims
             },
            os.path.join(args.out_dir, 'epoch_{}_checkpoint.pth.tar'.format(i + 1))
        )

    # put the policy into eval/validation mode
    print('\nEvaluating on the test set...\n')
    policy.eval()

    total_correct = 0
    top_k_correct = dict.fromkeys(args.top_k)
    for k in args.top_k:
        top_k_correct[k] = 0
    test_acc_top_k = dict.fromkeys(args.top_k)
    total_loss = 0.0
    nan_counter = 0

    with torch.no_grad():
        for batch in test_loader:
            for idx, data_tuple in enumerate(batch):
                target, node, mip, grid = data_tuple
                target, node, mip, grid = target.to(device), node.to(device), mip.to(device), grid.to(device)
                logits = policy(grid, node, mip)
                logits = logits.transpose(1, 0)
                _loss = criterion(logits, target).item()
                if math.isnan(_loss):
                    nan_counter += 1
                else:
                    total_loss += _loss
                    _, predicted = torch.max(logits, 1)
                    total_correct += predicted.eq(target.item()).cpu().item()
                    grid_size = grid.size(0)
                    for k in args.top_k:
                        max_k = min(k, grid_size)  # Accounts for when grid_size is smaller than top_k
                        top_k_correct[k] += int(target.item() in logits.topk(max_k, dim=1).indices)
        if nan_counter < test_h5.n_data:
            test_loss = total_loss / float(test_h5.n_data - nan_counter)
            test_acc = total_correct / float(test_h5.n_data - nan_counter)
            for k in args.top_k:
                test_acc_top_k[k] = top_k_correct[k] / float(test_h5.n_data - nan_counter)
        else:
            test_loss = np.nan
            test_acc = np.nan
            for k in args.top_k:
                test_acc_top_k[k] = np.nan

        exp_dict['test_loss'] = test_loss
        exp_dict['test_acc'] = test_acc
        for k in args.top_k:
            exp_dict['test_acc_top_{}'.format(k)] = test_acc_top_k[k]
        exp_dict['test_nan_counter'] = nan_counter

    print('Test loss: {:.6f}, Test acc: {:.2f}%, Test acc top-{}: {:.2f}%, Test acc top-{}: {:.2f}%'.format(
        test_loss, 100 * test_acc, args.top_k[0], 100 * test_acc_top_k[args.top_k[0]], args.top_k[1],
        100 * test_acc_top_k[args.top_k[1]]))

    # save the final checkpoint
    torch.save(
        {'state_dict': policy.state_dict(),
         'optimizer': optimizer.state_dict(),
         'scheduler': scheduler.state_dict(),
         'exp_dict': exp_dict,
         'args': args,
         'state_dims': state_dims
         },
        os.path.join(args.out_dir, 'final_checkpoint.pth.tar')
    )


text:
# coding=utf-8
# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""PyTorch BERT model. """

from __future__ import annotations
import math
import os
import warnings
from dataclasses import dataclass
from typing import Optional, Tuple

import torch
import torch.utils.checkpoint
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from transformers.activations import ACT2FN
from transformers.file_utils import (
    ModelOutput,
    add_code_sample_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    replace_return_docstrings,
)
from transformers.modeling_outputs import (
    BaseModelOutputWithPastAndCrossAttentions,
    BaseModelOutputWithPoolingAndCrossAttentions,
    CausalLMOutputWithCrossAttentions,
    MaskedLMOutput,
    MultipleChoiceModelOutput,
    NextSentencePredictorOutput,
    QuestionAnsweringModelOutput,
    SequenceClassifierOutput,
    TokenClassifierOutput,
)
from transformers.modeling_utils import (
    PreTrainedModel,
    apply_chunking_to_forward,
    find_pruneable_heads_and_indices,
    prune_linear_layer,
)
from transformers.utils import logging

BERT_PRETRAINED_CONFIG_ARCHIVE_MAP = {
    "bert-base-uncased": "https://huggingface.co/bert-base-uncased/resolve/main/config.json",
    "bert-large-uncased": "https://huggingface.co/bert-large-uncased/resolve/main/config.json",
    "bert-base-cased": "https://huggingface.co/bert-base-cased/resolve/main/config.json",
    "bert-large-cased": "https://huggingface.co/bert-large-cased/resolve/main/config.json",
    "bert-base-multilingual-uncased": "https://huggingface.co/bert-base-multilingual-uncased/resolve/main/config.json",
    "bert-base-multilingual-cased": "https://huggingface.co/bert-base-multilingual-cased/resolve/main/config.json",
    "bert-base-chinese": "https://huggingface.co/bert-base-chinese/resolve/main/config.json",
    "bert-base-german-cased": "https://huggingface.co/bert-base-german-cased/resolve/main/config.json",
    "bert-large-uncased-whole-word-masking": "https://huggingface.co/bert-large-uncased-whole-word-masking/resolve/main/config.json",
    "bert-large-cased-whole-word-masking": "https://huggingface.co/bert-large-cased-whole-word-masking/resolve/main/config.json",
    "bert-large-uncased-whole-word-masking-finetuned-squad": "https://huggingface.co/bert-large-uncased-whole-word-masking-finetuned-squad/resolve/main/config.json",
    "bert-large-cased-whole-word-masking-finetuned-squad": "https://huggingface.co/bert-large-cased-whole-word-masking-finetuned-squad/resolve/main/config.json",
    "bert-base-cased-finetuned-mrpc": "https://huggingface.co/bert-base-cased-finetuned-mrpc/resolve/main/config.json",
    "bert-base-german-dbmdz-cased": "https://huggingface.co/bert-base-german-dbmdz-cased/resolve/main/config.json",
    "bert-base-german-dbmdz-uncased": "https://huggingface.co/bert-base-german-dbmdz-uncased/resolve/main/config.json",
    "cl-tohoku/bert-base-japanese": "https://huggingface.co/cl-tohoku/bert-base-japanese/resolve/main/config.json",
    "cl-tohoku/bert-base-japanese-whole-word-masking": "https://huggingface.co/cl-tohoku/bert-base-japanese-whole-word-masking/resolve/main/config.json",
    "cl-tohoku/bert-base-japanese-char": "https://huggingface.co/cl-tohoku/bert-base-japanese-char/resolve/main/config.json",
    "cl-tohoku/bert-base-japanese-char-whole-word-masking": "https://huggingface.co/cl-tohoku/bert-base-japanese-char-whole-word-masking/resolve/main/config.json",
    "TurkuNLP/bert-base-finnish-cased-v1": "https://huggingface.co/TurkuNLP/bert-base-finnish-cased-v1/resolve/main/config.json",
    "TurkuNLP/bert-base-finnish-uncased-v1": "https://huggingface.co/TurkuNLP/bert-base-finnish-uncased-v1/resolve/main/config.json",
    "wietsedv/bert-base-dutch-cased": "https://huggingface.co/wietsedv/bert-base-dutch-cased/resolve/main/config.json",
    # See all BERT models at https://huggingface.co/models?filter=bert
}

from transformers.configuration_utils import PretrainedConfig
class BertConfig(PretrainedConfig):
    model_type = "bert"
    def __init__(
        self,
        vocab_size=30522,
        hidden_size=768,
        num_hidden_layers=12,
        num_attention_heads=12,
        intermediate_size=3072,
        hidden_act="gelu",
        hidden_dropout_prob=0.1,
        attention_probs_dropout_prob=0.1,
        max_position_embeddings=512,
        type_vocab_size=2,
        initializer_range=0.02,
        layer_norm_eps=1e-12,
        pad_token_id=0,
        gradient_checkpointing=False,
        position_embedding_type="absolute",
        use_cache=True,
        **kwargs
    ):
        super().__init__(pad_token_id=pad_token_id, **kwargs)

        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.hidden_act = hidden_act
        self.intermediate_size = intermediate_size
        self.hidden_dropout_prob = hidden_dropout_prob
        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.max_position_embeddings = max_position_embeddings
        self.type_vocab_size = type_vocab_size
        self.initializer_range = initializer_range
        self.layer_norm_eps = layer_norm_eps
        self.gradient_checkpointing = gradient_checkpointing
        self.position_embedding_type = position_embedding_type
        self.use_cache = use_cache

        self.cross_index = 8

logger = logging.get_logger(__name__)

_CONFIG_FOR_DOC = "BertConfig"
_TOKENIZER_FOR_DOC = "BertTokenizer"

BERT_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "bert-base-uncased",
    "bert-large-uncased",
    "bert-base-cased",
    "bert-large-cased",
    "bert-base-multilingual-uncased",
    "bert-base-multilingual-cased",
    "bert-base-chinese",
    "bert-base-german-cased",
    "bert-large-uncased-whole-word-masking",
    "bert-large-cased-whole-word-masking",
    "bert-large-uncased-whole-word-masking-finetuned-squad",
    "bert-large-cased-whole-word-masking-finetuned-squad",
    "bert-base-cased-finetuned-mrpc",
    "bert-base-german-dbmdz-cased",
    "bert-base-german-dbmdz-uncased",
    "cl-tohoku/bert-base-japanese",
    "cl-tohoku/bert-base-japanese-whole-word-masking",
    "cl-tohoku/bert-base-japanese-char",
    "cl-tohoku/bert-base-japanese-char-whole-word-masking",
    "TurkuNLP/bert-base-finnish-cased-v1",
    "TurkuNLP/bert-base-finnish-uncased-v1",
    "wietsedv/bert-base-dutch-cased",
    # See all BERT models at https://huggingface.co/models?filter=bert
]

def load_tf_weights_in_bert(model, config, tf_checkpoint_path):
    """Load tf checkpoints in a pytorch model."""
    try:
        import re

        import numpy as np
        import tensorflow as tf
    except ImportError:
        logger.error(
            "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "
            "https://www.tensorflow.org/install/ for installation instructions."
        )
        raise
    tf_path = os.path.abspath(tf_checkpoint_path)
    logger.info("Converting TensorFlow checkpoint from {}".format(tf_path))
    # Load weights from TF model
    init_vars = tf.train.list_variables(tf_path)
    names = []
    arrays = []
    for name, shape in init_vars:
        logger.info("Loading TF weight {} with shape {}".format(name, shape))
        array = tf.train.load_variable(tf_path, name)
        names.append(name)
        arrays.append(array)

    for name, array in zip(names, arrays):
        name = name.split("/")
        # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v
        # which are not required for using pretrained model
        if any(
            n in ["adam_v", "adam_m", "AdamWeightDecayOptimizer", "AdamWeightDecayOptimizer_1", "global_step"]
            for n in name
        ):
            logger.info("Skipping {}".format("/".join(name)))
            continue
        pointer = model
        for m_name in name:
            if re.fullmatch(r"[A-Za-z]+_\d+", m_name):
                scope_names = re.split(r"_(\d+)", m_name)
            else:
                scope_names = [m_name]
            if scope_names[0] == "kernel" or scope_names[0] == "gamma":
                pointer = getattr(pointer, "weight")
            elif scope_names[0] == "output_bias" or scope_names[0] == "beta":
                pointer = getattr(pointer, "bias")
            elif scope_names[0] == "output_weights":
                pointer = getattr(pointer, "weight")
            elif scope_names[0] == "squad":
                pointer = getattr(pointer, "classifier")
            else:
                try:
                    pointer = getattr(pointer, scope_names[0])
                except AttributeError:
                    logger.info("Skipping {}".format("/".join(name)))
                    continue
            if len(scope_names) >= 2:
                num = int(scope_names[1])
                pointer = pointer[num]
        if m_name[-11:] == "_embeddings":
            pointer = getattr(pointer, "weight")
        elif m_name == "kernel":
            array = np.transpose(array)
        try:
            assert (
                pointer.shape == array.shape
            ), f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched"
        except AssertionError as e:
            e.args += (pointer.shape, array.shape)
            raise
        logger.info("Initialize PyTorch weight {}".format(name))
        pointer.data = torch.from_numpy(array)
    return model

class BertEmbeddings(nn.Module):
    """Construct the embeddings from word, position and token_type embeddings."""

    def __init__(self, config):
        super().__init__()
        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)

        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
        # any TensorFlow checkpoint file
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
        self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))
        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")

    def forward(
        self, input_ids=None, token_type_ids=None, position_ids=None, inputs_embeds=None, past_key_values_length=0
    ):
        if input_ids is not None:
            input_shape = input_ids.size()
        else:
            input_shape = inputs_embeds.size()[:-1]

        seq_length = input_shape[1]

        if position_ids is None:
            position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]

        if token_type_ids is None:
            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)

        if inputs_embeds is None:
            inputs_embeds = self.word_embeddings(input_ids)
        token_type_embeddings = self.token_type_embeddings(token_type_ids)

        embeddings = inputs_embeds + token_type_embeddings
        if self.position_embedding_type == "absolute":
            position_embeddings = self.position_embeddings(position_ids)
            embeddings += position_embeddings
        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)
        return embeddings

class BertSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
            raise ValueError(
                "The hidden size (%d) is not a multiple of the number of attention "
                "heads (%d)" % (config.hidden_size, config.num_attention_heads)
            )

        self.num_attention_heads = config.num_attention_heads
        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
        self.all_head_size = self.num_attention_heads * self.attention_head_size

        self.query = nn.Linear(config.hidden_size, self.all_head_size)
        self.key = nn.Linear(config.hidden_size, self.all_head_size)
        self.value = nn.Linear(config.hidden_size, self.all_head_size)

        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
            self.max_position_embeddings = config.max_position_embeddings
            self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)

        self.is_decoder = config.is_decoder

    def transpose_for_scores(self, x):
        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
        x = x.view(*new_x_shape)
        return x.permute(0, 2, 1, 3)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_value=None,
        output_attentions=False,
    ):
        mixed_query_layer = self.query(hidden_states)

        # If this is instantiated as a cross-attention module, the keys
        # and values come from an encoder; the attention mask needs to be
        # such that the encoder's padding tokens are not attended to.
        is_cross_attention = encoder_hidden_states is not None

        if is_cross_attention and past_key_value is not None:
            # reuse k,v, cross_attentions
            key_layer = past_key_value[0]
            value_layer = past_key_value[1]
            attention_mask = encoder_attention_mask
        elif is_cross_attention:
            key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
            value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
            attention_mask = encoder_attention_mask
        elif past_key_value is not None:
            key_layer = self.transpose_for_scores(self.key(hidden_states))
            value_layer = self.transpose_for_scores(self.value(hidden_states))
            key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
            value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
        else:
            key_layer = self.transpose_for_scores(self.key(hidden_states))
            value_layer = self.transpose_for_scores(self.value(hidden_states))

        query_layer = self.transpose_for_scores(mixed_query_layer)

        if self.is_decoder:
            # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
            # Further calls to cross_attention layer can then reuse all cross-attention
            # key/value_states (first "if" case)
            # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
            # all previous decoder key/value_states. Further calls to uni-directional self-attention
            # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
            # if encoder bi-directional self-attention `past_key_value` is always `None`
            past_key_value = (key_layer, value_layer)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))

        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
            seq_length = hidden_states.size()[1]
            position_ids_l = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)
            position_ids_r = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(1, -1)
            distance = position_ids_l - position_ids_r
            positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)
            positional_embedding = positional_embedding.to(dtype=query_layer.dtype)  # fp16 compatibility

            if self.position_embedding_type == "relative_key":
                relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
                attention_scores = attention_scores + relative_position_scores
            elif self.position_embedding_type == "relative_key_query":
                relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
                relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)
                attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key

        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
        if attention_mask is not None:
            # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
            attention_scores = attention_scores + attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = nn.Softmax(dim=-1)(attention_scores)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        context_layer = torch.matmul(attention_probs, value_layer)

        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(*new_context_layer_shape)

        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)

        if self.is_decoder:
            outputs = outputs + (past_key_value,)
        return outputs

class BertSelfOutput(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states, input_tensor):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        return hidden_states

class BertAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.self = BertSelfAttention(config)
        self.output = BertSelfOutput(config)
        self.pruned_heads = set()

    def prune_heads(self, heads):
        if len(heads) == 0:
            return
        heads, index = find_pruneable_heads_and_indices(
            heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
        )

        # Prune linear layers
        self.self.query = prune_linear_layer(self.self.query, index)
        self.self.key = prune_linear_layer(self.self.key, index)
        self.self.value = prune_linear_layer(self.self.value, index)
        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)

        # Update hyper params and store pruned heads
        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
        self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
        self.pruned_heads = self.pruned_heads.union(heads)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_value=None,
        output_attentions=False,
    ):
        self_outputs = self.self(
            hidden_states,
            attention_mask,
            head_mask,
            encoder_hidden_states,
            encoder_attention_mask,
            past_key_value,
            output_attentions,
        )
        attention_output = self.output(self_outputs[0], hidden_states)
        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
        return outputs

class BertIntermediate(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
        if isinstance(config.hidden_act, str):
            self.intermediate_act_fn = ACT2FN[config.hidden_act]
        else:
            self.intermediate_act_fn = config.hidden_act

    def forward(self, hidden_states):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.intermediate_act_fn(hidden_states)
        return hidden_states

class BertOutput(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states, input_tensor):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        return hidden_states

class BertLayer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.chunk_size_feed_forward = config.chunk_size_feed_forward
        self.seq_len_dim = 1
        self.attention = BertAttention(config)
        self.is_decoder = config.is_decoder
        self.add_cross_attention = config.add_cross_attention
        if self.add_cross_attention:
            assert self.is_decoder, f"{self} should be used as a decoder model if cross attention is added"
            self.crossattention = BertAttention(config)
        self.intermediate = BertIntermediate(config)
        self.output = BertOutput(config)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_value=None,
        output_attentions=False,
    ):
        # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
        self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
        self_attention_outputs = self.attention(
            hidden_states,
            attention_mask,
            head_mask,
            output_attentions=output_attentions,
            past_key_value=self_attn_past_key_value,
        )
        attention_output = self_attention_outputs[0]

        # if decoder, the last output is tuple of self-attn cache
        if self.is_decoder:
            outputs = self_attention_outputs[1:-1]
            present_key_value = self_attention_outputs[-1]
        else:
            outputs = self_attention_outputs[1:]  # add self attentions if we output attention weights

        cross_attn_present_key_value = None
        if self.is_decoder and encoder_hidden_states is not None:
            assert hasattr(
                self, "crossattention"
            ), f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers by setting `config.add_cross_attention=True`"

            # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple
            cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
            cross_attention_outputs = self.crossattention(
                attention_output,
                attention_mask,
                head_mask,
                encoder_hidden_states,
                encoder_attention_mask,
                cross_attn_past_key_value,
                output_attentions,
            )
            attention_output = cross_attention_outputs[0]
            outputs = outputs + cross_attention_outputs[1:-1]  # add cross attentions if we output attention weights

            # add cross-attn cache to positions 3,4 of present_key_value tuple
            cross_attn_present_key_value = cross_attention_outputs[-1]
            present_key_value = present_key_value + cross_attn_present_key_value

        layer_output = apply_chunking_to_forward(
            self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
        )
        outputs = (layer_output,) + outputs

        # if decoder, return the attn key/values as the last output
        if self.is_decoder:
            outputs = outputs + (present_key_value,)

        return outputs

    def feed_forward_chunk(self, attention_output):
        intermediate_output = self.intermediate(attention_output)
        layer_output = self.output(intermediate_output, attention_output)
        return layer_output

class BertEncoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.num_hidden_layers)])

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_values=None,
        use_cache=None,
        output_attentions=False,
        output_hidden_states=False,
        return_dict=True,
    ):
        all_hidden_states = () if output_hidden_states else None
        all_self_attentions = () if output_attentions else None
        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None

        next_decoder_cache = () if use_cache else None
        for i, layer_module in enumerate(self.layer):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            layer_head_mask = head_mask[i] if head_mask is not None else None
            past_key_value = past_key_values[i] if past_key_values is not None else None

            if getattr(self.config, "gradient_checkpointing", False) and self.training:

                if use_cache:
                    logger.warn(
                        "`use_cache=True` is incompatible with `config.gradient_checkpointing=True`. Setting "
                        "`use_cache=False`..."
                    )
                    use_cache = False

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        return module(*inputs, past_key_value, output_attentions)

                    return custom_forward

                layer_outputs = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(layer_module),
                    hidden_states,
                    attention_mask,
                    layer_head_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                )
            else:
                layer_outputs = layer_module(
                    hidden_states,
                    attention_mask,
                    layer_head_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    past_key_value,
                    output_attentions,
                )

            hidden_states = layer_outputs[0]
            if use_cache:
                next_decoder_cache += (layer_outputs[-1],)
            if output_attentions:
                all_self_attentions = all_self_attentions + (layer_outputs[1],)
                if self.config.add_cross_attention:
                    all_cross_attentions = all_cross_attentions + (layer_outputs[2],)

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return tuple(
                v
                for v in [
                    hidden_states,
                    next_decoder_cache,
                    all_hidden_states,
                    all_self_attentions,
                    all_cross_attentions,
                ]
                if v is not None
            )
        return BaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,
            past_key_values=next_decoder_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
            cross_attentions=all_cross_attentions,
        )

class BertFrontEncoder(nn.Module):
    def __init__(self, config: BertConfig):
        super().__init__()
        self.config: BertConfig = config
        self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.cross_index)])

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_values=None,
        use_cache=None,
        output_attentions=False,
        output_hidden_states=False,
    ):
        all_hidden_states = () if output_hidden_states else None
        all_self_attentions = () if output_attentions else None
        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None

        next_decoder_cache = () if use_cache else None
        for i, layer_module in enumerate(self.layer):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            layer_head_mask = head_mask[i] if head_mask is not None else None
            past_key_value = past_key_values[i] if past_key_values is not None else None

            layer_outputs = layer_module(
                hidden_states,
                attention_mask,
                layer_head_mask,
                encoder_hidden_states,
                encoder_attention_mask,
                past_key_value,
                output_attentions,
            )

            hidden_states = layer_outputs[0]
            if use_cache:
                next_decoder_cache += (layer_outputs[-1],)
            if output_attentions:
                all_self_attentions = all_self_attentions + (layer_outputs[1],)
                if self.config.add_cross_attention:
                    all_cross_attentions = all_cross_attentions + (layer_outputs[2],)

        """
        hidden_states,
        attention_mask,
        layer_head_mask,
        encoder_hidden_states,
        encoder_attention_mask,
        past_key_value,
        output_attentions,

        all_hidden_states,
        all_self_attentions,
        all_cross_attentions,

        output_hidden_states,
        use_cache,
        head_mask,
        past_key_values
        """
        return dict(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            layer_head_mask=layer_head_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            past_key_value=past_key_value,
            output_attentions=output_attentions,

            all_hidden_states=all_hidden_states,
            all_self_attentions=all_self_attentions,
            all_cross_attentions=all_cross_attentions,

            output_hidden_states=output_hidden_states,
            use_cache=use_cache,
            head_mask=head_mask,
            past_key_values=past_key_values
        )

class BertBackEncoder(nn.Module):
    def __init__(self, config: BertConfig):
        super().__init__()
        self.config = config
        self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.num_hidden_layers - config.cross_index)])

    def forward(
        self,
        hidden_states,
        attention_mask,
        layer_head_mask,
        encoder_hidden_states,
        encoder_attention_mask,
        past_key_value,
        output_attentions,

        all_hidden_states,
        all_self_attentions,
        all_cross_attentions,

        output_hidden_states,
        use_cache,
        head_mask,
        past_key_values
    ):

        all_hidden_states = () if output_hidden_states else None
        all_self_attentions = () if output_attentions else None
        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None

        next_decoder_cache = () if use_cache else None
        for i, layer_module in enumerate(self.layer):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            layer_head_mask = head_mask[i] if head_mask is not None else None
            past_key_value = past_key_values[i] if past_key_values is not None else None

            layer_outputs = layer_module(
                hidden_states,
                attention_mask,
                layer_head_mask,
                encoder_hidden_states,
                encoder_attention_mask,
                past_key_value,
                output_attentions,
            )

            hidden_states = layer_outputs[0]
            if use_cache:
                next_decoder_cache += (layer_outputs[-1],)
            if output_attentions:
                all_self_attentions = all_self_attentions + (layer_outputs[1],)
                if self.config.add_cross_attention:
                    all_cross_attentions = all_cross_attentions + (layer_outputs[2],)

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        return BaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,
            past_key_values=next_decoder_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
            cross_attentions=all_cross_attentions,
        )

class BertPooler(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.activation = nn.Tanh()

    def forward(self, hidden_states):
        # We "pool" the model by simply taking the hidden state corresponding
        # to the first token.
        first_token_tensor = hidden_states[:, 0]
        pooled_output = self.dense(first_token_tensor)
        pooled_output = self.activation(pooled_output)
        return pooled_output

class BertPredictionHeadTransform(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        if isinstance(config.hidden_act, str):
            self.transform_act_fn = ACT2FN[config.hidden_act]
        else:
            self.transform_act_fn = config.hidden_act
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)

    def forward(self, hidden_states):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.transform_act_fn(hidden_states)
        hidden_states = self.LayerNorm(hidden_states)
        return hidden_states

class BertLMPredictionHead(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.transform = BertPredictionHeadTransform(config)

        # The output weights are the same as the input embeddings, but there is
        # an output-only bias for each token.
        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        self.bias = nn.Parameter(torch.zeros(config.vocab_size))

        # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
        self.decoder.bias = self.bias

    def forward(self, hidden_states):
        hidden_states = self.transform(hidden_states)
        hidden_states = self.decoder(hidden_states)
        return hidden_states

class BertOnlyMLMHead(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.predictions = BertLMPredictionHead(config)

    def forward(self, sequence_output):
        prediction_scores = self.predictions(sequence_output)
        return prediction_scores

class BertOnlyNSPHead(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.seq_relationship = nn.Linear(config.hidden_size, 2)

    def forward(self, pooled_output):
        seq_relationship_score = self.seq_relationship(pooled_output)
        return seq_relationship_score

class BertPreTrainingHeads(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.predictions = BertLMPredictionHead(config)
        self.seq_relationship = nn.Linear(config.hidden_size, 2)

    def forward(self, sequence_output, pooled_output):
        prediction_scores = self.predictions(sequence_output)
        seq_relationship_score = self.seq_relationship(pooled_output)
        return prediction_scores, seq_relationship_score

class BertPreTrainedModel(PreTrainedModel):
    """
    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
    models.
    """

    config_class = BertConfig
    load_tf_weights = load_tf_weights_in_bert
    base_model_prefix = "bert"
    _keys_to_ignore_on_load_missing = [r"position_ids"]

    def _init_weights(self, module):
        """ Initialize the weights """
        if isinstance(module, (nn.Linear, nn.Embedding)):
            # Slightly different from the TF version which uses truncated_normal for initialization
            # cf https://github.com/pytorch/pytorch/pull/5617
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)
        if isinstance(module, nn.Linear) and module.bias is not None:
            module.bias.data.zero_()

@dataclass
class BertForPreTrainingOutput(ModelOutput):
    """
    Output type of :class:`~transformers.BertForPreTraining`.

    Args:
        loss (`optional`, returned when ``labels`` is provided, ``torch.FloatTensor`` of shape :obj:`(1,)`):
            Total loss as the sum of the masked language modeling loss and the next sequence prediction
            (classification) loss.
        prediction_logits (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`):
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
        seq_relationship_logits (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, 2)`):
            Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation
            before SoftMax).
        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``):
            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True`` is passed or when ``config.output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape :obj:`(batch_size, num_heads,
            sequence_length, sequence_length)`.

            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
            heads.
    """

    loss: Optional[torch.FloatTensor] = None
    prediction_logits: torch.FloatTensor = None
    seq_relationship_logits: torch.FloatTensor = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None

class BertModel(BertPreTrainedModel):
    def __init__(self, config, add_pooling_layer=True):
        super().__init__(config)
        self.config = config

        self.embeddings = BertEmbeddings(config)
        self.encoder = BertEncoder(config)

        self.pooler = BertPooler(config) if add_pooling_layer else None

        self.init_weights()

    def get_input_embeddings(self):
        return self.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.embeddings.word_embeddings = value

    def _prune_heads(self, heads_to_prune):
        """
        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
        class PreTrainedModel
        """
        for layer, heads in heads_to_prune.items():
            self.encoder.layer[layer].attention.prune_heads(heads)

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_values=None,
        use_cache=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        r"""
        encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
            the model is configured as a decoder.
        encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
            the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:

            - 1 for tokens that are **not masked**,
            - 0 for tokens that are **masked**.
        past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.

            If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`
            (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`
            instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.
        use_cache (:obj:`bool`, `optional`):
            If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up
            decoding (see :obj:`past_key_values`).
        """
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if self.config.is_decoder:
            use_cache = use_cache if use_cache is not None else self.config.use_cache
        else:
            use_cache = False

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            input_shape = input_ids.size()
            batch_size, seq_length = input_shape
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
            batch_size, seq_length = input_shape
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        device = input_ids.device if input_ids is not None else inputs_embeds.device

        # past_key_values_length
        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0

        if attention_mask is None:
            attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)
        if token_type_ids is None:
            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)

        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
        # ourselves in which case we just need to make it broadcastable to all heads.
        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape, device)

        # If a 2D or 3D attention mask is provided for the cross-attention
        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
        if self.config.is_decoder and encoder_hidden_states is not None:
            encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
            if encoder_attention_mask is None:
                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
            encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
        else:
            encoder_extended_attention_mask = None

        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x n_heads x N x N
        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)

        embedding_output = self.embeddings(
            input_ids=input_ids,
            position_ids=position_ids,
            token_type_ids=token_type_ids,
            inputs_embeds=inputs_embeds,
            past_key_values_length=past_key_values_length,
        )
        named_parameters = list(self.embeddings.named_parameters())
        encoder_outputs = self.encoder(
            embedding_output,
            attention_mask=extended_attention_mask,
            head_mask=head_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_extended_attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        sequence_output = encoder_outputs[0]
        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None

        if not return_dict:
            return (sequence_output, pooled_output) + encoder_outputs[1:]

        return BaseModelOutputWithPoolingAndCrossAttentions(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
            past_key_values=encoder_outputs.past_key_values,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
            cross_attentions=encoder_outputs.cross_attentions,
        )

class CrossBertModel(BertPreTrainedModel):
    """
    The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
    cross-attention is added between the self-attention layers, following the architecture described in `Attention is
    all you need <https://arxiv.org/abs/1706.03762>`__ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
    Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.

    To behave as an decoder the model needs to be initialized with the :obj:`is_decoder` argument of the configuration
    set to :obj:`True`. To be used in a Seq2Seq model, the model needs to initialized with both :obj:`is_decoder`
    argument and :obj:`add_cross_attention` set to :obj:`True`; an :obj:`encoder_hidden_states` is then expected as an
    input to the forward pass.
    """

    def __init__(self, config, add_pooling_layer=True):
        super().__init__(config)
        self.config: BertConfig = config

        self.embeddings = BertEmbeddings(config)
        self.front_encoder = BertFrontEncoder(config)
        self.back_encoder = BertBackEncoder(config)

        self.pooler = BertPooler(config) if add_pooling_layer else None

        self.init_weights()

    @staticmethod
    def from_state_dict(file: str, config: BertConfig) -> CrossBertModel:
        def load_state_dict():
            state_dict = torch.load(file, map_location="cpu")

            # 2. transfer old key to new key
            old_keys = []
            new_keys = []
            for key in state_dict.keys():
                new_key = None
                if "gamma" in key:
                    new_key = key.replace("gamma", "weight")
                if "beta" in key:
                    new_key = key.replace("beta", "bias")
                if new_key:
                    new_key = new_key.replace("bert.", "")

                    old_keys.append(key)
                    new_keys.append(new_key)
                else:
                    new_key = key.replace("bert.", "")
                    new_keys.append(new_key)
                    old_keys.append(key)

            for old_key, new_key in zip(old_keys, new_keys):

                new_key = new_key.replace("bert.", "")

                state_dict[new_key] = state_dict.pop(old_key)

            # copy state_dict so _load_from_state_dict can modify it
            metadata = getattr(state_dict, "_metadata", None)
            state_dict = state_dict.copy()

            if metadata is not None:
                state_dict._metadata = metadata
            return state_dict, metadata
        state_dict, metadata = load_state_dict()

        model = CrossBertModel(config)
        model.load_state_dict(state_dict, metadata)
        return model

    def get_input_embeddings(self):
        return self.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.embeddings.word_embeddings = value

    def _prune_heads(self, heads_to_prune):
        """
        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
        class PreTrainedModel
        """
        for layer, heads in heads_to_prune.items():
            self.front_encoder.layer[layer].attention.prune_heads(heads)

        for layer, heads in heads_to_prune.items():
            self.back_encoder.layer[layer].attention.prune_heads(heads)

    def load_state_dict(self, state_dict, metadata):
        missing_keys = []
        unexpected_keys = []
        error_msgs = []

        def load(module: nn.Module, prefix="", sub_state_dict=None):
            local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
            module._load_from_state_dict(
                sub_state_dict,
                prefix,
                local_metadata,
                True,
                missing_keys,
                unexpected_keys,
                error_msgs,
            )
            for name, child in module._modules.items():
                if child is not None:
                    load(child, prefix + name + ".", sub_state_dict)

        from collections import OrderedDict
        front_encoder_state_dict = OrderedDict()
        for index in range(self.config.cross_index):
            for name, parameter in state_dict.items():
                if name.startswith(f'encoder.layer.{index}.'):
                    front_encoder_parameter_name = name.replace('encoder.', '')
                    front_encoder_state_dict[front_encoder_parameter_name] = parameter

        back_encoder_state_dict = OrderedDict()
        for index in range(self.config.num_hidden_layers - self.config.cross_index):
            for name, parameter in state_dict.items():
                if name.startswith(f'encoder.layer.{index + 8}.'):
                    back_encoder_parameter_name = name.replace(f'encoder.layer.{index + 8}.', f'layer.{index}.')
                    back_encoder_state_dict[back_encoder_parameter_name] = parameter

        # load front & back encoder weight
        load(self.front_encoder, "", front_encoder_state_dict)
        load(self.back_encoder, "", back_encoder_state_dict)

        # load embedding weight
        load(self.embeddings, 'embeddings.', state_dict)

        # load pooler weight
        load(self.pooler, 'pooler.', state_dict)

    def forward_on_front_transformers(self,input_ids=None,
            attention_mask=None,
            token_type_ids=None,
            position_ids=None,
            head_mask=None,
            inputs_embeds=None,
            encoder_hidden_states=None,
            encoder_attention_mask=None,
            past_key_values=None,
            use_cache=None,
            output_attentions=None,
            output_hidden_states=None,
            return_dict=None,
        ):

        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if self.config.is_decoder:
            use_cache = use_cache if use_cache is not None else self.config.use_cache
        else:
            use_cache = False

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            input_shape = input_ids.size()
            batch_size, seq_length = input_shape
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
            batch_size, seq_length = input_shape
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        device = input_ids.device if input_ids is not None else inputs_embeds.device

        # past_key_values_length
        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0

        if attention_mask is None:
            attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)
        if token_type_ids is None:
            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)

        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
        # ourselves in which case we just need to make it broadcastable to all heads.
        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape, device)

        # If a 2D or 3D attention mask is provided for the cross-attention
        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
        if self.config.is_decoder and encoder_hidden_states is not None:
            encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
            if encoder_attention_mask is None:
                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
            encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
        else:
            encoder_extended_attention_mask = None

        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x n_heads x N x N
        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)

        embedding_output = self.embeddings(
            input_ids=input_ids,
            position_ids=position_ids,
            token_type_ids=token_type_ids,
            inputs_embeds=inputs_embeds,
            past_key_values_length=past_key_values_length,
        )
        named_parameters = list(self.embeddings.named_parameters())

        front_transformer_output = self.front_encoder(
            embedding_output,
            attention_mask=extended_attention_mask,
            head_mask=head_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_extended_attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
        )
        return front_transformer_output

    def forward_on_back_transfomers(self, **kwargs):
        encoder_outputs = self.back_encoder(**kwargs)

        sequence_output = encoder_outputs[0]
        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None

        return BaseModelOutputWithPoolingAndCrossAttentions(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
            past_key_values=encoder_outputs.past_key_values,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
            cross_attentions=encoder_outputs.cross_attentions,
        )

    def forward(self, **kwargs):
        front_encoder_output = self.forward_on_front_transformers(**kwargs)
        back_encoder_output = self.forward_on_back_transfomers(**front_encoder_output)
        return back_encoder_output

class CrossBertModelTemp(nn.Module):
    def __init__(self, config: BertConfig, state_dict, metadata):
        super().__init__()
        self.left_bert = BertModel(config)

        self.load_state_dict(state_dict, self.left_bert, metadata)

        self.right_bert = BertModel(config)
        self.load_state_dict(state_dict, self.right_bert, metadata)

    def load_state_dict(self, state_dict, model, metadata):
        missing_keys = []
        unexpected_keys = []
        error_msgs = []

        def load(module: nn.Module, prefix=""):
            local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
            module._load_from_state_dict(
                state_dict,
                prefix,
                local_metadata,
                True,
                missing_keys,
                unexpected_keys,
                error_msgs,
            )
            for name, child in module._modules.items():
                if child is not None:
                    load(child, prefix + name + ".")

        load(model)

    def forward(self, **inputs):
        left_output = self.left_bert(**inputs)
        right_putput = self.right_bert(**inputs)

        output = left_output + right_putput
        return output

max_stars_repo_name: bupt-nlp/sequence-labeling-models

max_issues_repo_name: bupt-nlp/sequence-labeling-models

max_forks_repo_name: bupt-nlp/sequence-labeling-models


text:
import random
from typing import Union

import numpy as np
import pandas as pd
from streamad.base import BaseDetector

class RandomDetector(BaseDetector):
    """Return random anomaly score. A minimum score for benchmark."""

    def __init__(self):
        super().__init__()

    def fit(self, X: np.ndarray) -> None:
        """Detector fit current observation from StreamGenerator.

        Args:
            X (np.ndarray): The data value of current observation from StreamGenerator.
        """
        return self

    def score(self, X: np.ndarray) -> float:
        """Abstract method: Detector predict the probability of anomaly for current observation form StreamGenerator.

        Args:
            X (np.ndarray): The data value of current observation from StreamGenerator.

        Returns:
            float: Anomaly probability.
        """

        return random.random()