<reponame>xiongma/bert2tf<filename>examples/p_tuning/main.py
import numpy as np
import tensorflow as tf
import tensorflow.keras.backend as K
from tensorflow.python.keras.metrics import MeanMetricWrapper

from bert2tf import BertTokenizer, RobertaPreTraining, Adam
from bert2tf.executors.models.bert import BertEmbedding

max_length = 128
batch_size = 32
masking = True

strategy = tf.distribute.MirroredStrategy()
batch_size = batch_size * strategy.num_replicas_in_sync

tokenizer = BertTokenizer(
    '../../tests/resources/pre_models/roberta_wwm_ext/vocab.txt',
    do_lower_case=True)

mask_idx = 5
prefix_tokens = [f'[unused{idx}]' for idx in range(1, 9)]
prefix_tokens.insert(mask_idx - 1, '[MASK]')
prefix_ids = tokenizer.convert_tokens_to_ids(prefix_tokens)
pos_id = tokenizer.convert_tokens_to_ids(['很'])[0]
neg_id = tokenizer.convert_tokens_to_ids(['不'])[0]


def random_masking(token_ids):
    rands = np.random.random(len(token_ids))
    source, target = [], []
    for r, t in zip(rands, token_ids):
        if r < 0.15 * 0.8:
            source.append(tokenizer.convert_tokens_to_ids(['[MASK]'])[0])
            target.append(t)
        elif r < 0.15 * 0.9:
            source.append(t)
            target.append(t)
        elif r < 0.15:
            source.append(np.random.choice(len(tokenizer.vocabs) - 1) + 1)
            target.append(t)
        else:
            source.append(t)
            target.append(0)
    return source, target


def padding(inputs):
    max_seq_length = max([len(x) for x in inputs])

    for x in inputs:
        while len(x) < max_seq_length:
            x.append(0)

    return inputs


def load_data(path):
    items = []
    with open(path, 'r', encoding='utf-8') as f:
        for line in f.readlines():
            line = line.strip().replace('\n', '')
            text, label = line.split('\t')
            label = int(label)

            items.append((text, label))

    return items


class PTuningBertEmbedding(BertEmbedding):
    def call(self, inputs, **kwargs):
        # stop gradients
        word_embeddings = self.word_embeddings
        sg_word_embeddings = K.stop_gradient(self.word_embeddings)

        mask = np.zeros((K.int_shape(self.word_embeddings)[0], 1))
        mask[1:9] += 1
        self.word_embeddings = word_embeddings * mask + sg_word_embeddings * (1 - mask)

        outputs = super(PTuningBertEmbedding, self).call(inputs, **kwargs)

        # reset gradients
        self.word_embeddings = word_embeddings

        return outputs


class PTuningBert(RobertaPreTraining):
    def create_embedding_layer(self, config):
        return PTuningBertEmbedding(config.hidden_size,
                                    config.type_vocab_size,
                                    vocab_size=config.vocab_size,
                                    initializer_range=config.initializer_range,
                                    hidden_dropout_prob=config.hidden_dropout_prob,
                                    max_position_embeddings=config.max_position_embeddings,
                                    name='embeddings')


class Loss(tf.keras.losses.Loss):
    def call(self, y_true, y_pred, **kwargs):
        y_mask = K.cast(K.not_equal(y_true, 0), K.floatx())

        y_pred = K.softmax(y_pred)
        loss = K.sparse_categorical_crossentropy(y_true, y_pred)
        loss = K.sum(loss * y_mask) / K.sum(y_mask)
        return loss


def accuracy(y_true, y_pred, **kwargs):
    y_mask = K.cast(K.not_equal(y_true, 0), K.floatx())
    accuracy = tf.metrics.sparse_categorical_accuracy(y_true, y_pred)
    accuracy = K.sum(accuracy * y_mask) / K.sum(y_mask)

    return accuracy


def data_gen(items, training=False):
    batch_input_ids, batch_segment_ids, batch_input_masks, batch_target_ids = [], [], [], []

    while True:
        from sklearn.utils import shuffle
        items = shuffle(items)
        for text, label in items:
            input_ids, input_masks, segment_ids = tokenizer.encode(text, max_length=max_length)

            if label != 2:
                input_ids = input_ids[:1] + prefix_ids + input_ids[1:]
                input_masks = input_masks[:1] + [1] * len(prefix_ids) + input_masks[1:]
                segment_ids = segment_ids[:1] + [0] * len(prefix_ids) + segment_ids[1:]

            input_ids, target_ids = random_masking(input_ids)

            if label < 2:
                input_ids[mask_idx] = tokenizer.convert_tokens_to_ids(['[MASK]'])[0]
                target_ids[mask_idx] = neg_id if label == 0 else pos_id

            batch_input_ids.append(input_ids)
            batch_segment_ids.append(segment_ids)
            batch_input_masks.append(input_masks)
            batch_target_ids.append(target_ids)

            if len(batch_input_ids) == batch_size:
                batch_input_ids = padding(batch_input_ids)
                batch_segment_ids = padding(batch_segment_ids)
                batch_input_masks = padding(batch_input_masks)
                batch_target_ids = padding(batch_target_ids)

                yield (np.array(batch_input_ids), np.array(batch_input_masks), np.array(batch_segment_ids)), \
                      (np.array(batch_target_ids),)
                batch_input_ids, batch_segment_ids, batch_input_masks, batch_target_ids = [], [], [], []

        if not training:
            # for file end
            batch_input_ids = padding(batch_input_ids)
            batch_segment_ids = padding(batch_segment_ids)
            batch_input_masks = padding(batch_input_masks)
            batch_target_ids = padding(batch_target_ids)

            yield (np.array(batch_input_ids), np.array(batch_input_masks), np.array(batch_segment_ids)), \
                  (np.array(batch_target_ids),)

            yield None, None


train_data = load_data('sentiment/sentiment.train.data')
test_data = load_data('sentiment/sentiment.test.data')
valid_data = load_data('sentiment/sentiment.valid.data')

train_data_generator = data_gen(train_data, training=True)
test_data_generator = data_gen(test_data)
valid_data_generator = data_gen(valid_data)

with strategy.scope():
    model = PTuningBert(
        config='../../tests/resources/pre_models/roberta_wwm_ext/bert_config.json',
        pretrained_weights_path='../../tests/resources/pre_models/roberta_wwm_ext/bert_model.ckpt')

    # set variables to non-trainable exclude word embeddings
    for layer in model.layers:
        layer.trainable = False
    embedding_layer = model.layers[0]
    embedding_layer.trainable = True
    embedding_layer.position_embeddings.trainable = False
    embedding_layer.token_type_embeddings.trainable = False
    embedding_layer.ln.trainable = False

    model.compile(optimizer=Adam(6e-4),
                  loss=Loss(),
                  metrics=[MeanMetricWrapper(fn=accuracy, name='accuracy')])


def evaluate(generator):
    total, right = 0., 0.
    for inputs, y_true in generator:
        if not inputs and not y_true:
            break

        y_pred = model(inputs)
        y_true = y_true[0]

        y_pred1 = tf.expand_dims(y_pred[:, mask_idx, neg_id], axis=-1)
        y_pred2 = tf.expand_dims(y_pred[:, mask_idx, pos_id], axis=-1)
        y_pred = tf.concat([y_pred1, y_pred2], axis=-1)
        y_pred = tf.argmax(y_pred, axis=1).numpy()

        y_true = (y_true[:, mask_idx] == pos_id).astype(int)

        total += len(y_true)
        right += (y_true == y_pred).sum()

    return right / total


class Evaluator(tf.keras.callbacks.Callback):
    def __init__(self):
        self.best_val_acc = 0.

    def on_epoch_end(self, epoch, logs=None):
        val_acc = evaluate(valid_data_generator)
        if val_acc > self.best_val_acc:
            self.best_val_acc = val_acc
            model.save_weights('best_model_bert.weights')
        test_acc = evaluate(test_data_generator)
        print(
            u'val_acc: %.5f, best_val_acc: %.5f, test_acc: %.5f\n' %
            (val_acc, self.best_val_acc, test_acc)
        )


model.summary()

model.fit(x=train_data_generator,
          steps_per_epoch=len(train_data) / batch_size + 1,
          epochs=1000,
          callbacks=[Evaluator()])

import numpy as np
import tensorflow as tf
import tensorflow.keras.backend as K
from tensorflow.python.keras.metrics import MeanMetricWrapper

from bert2tf import BertTokenizer, RobertaPreTraining, Adam
from bert2tf.executors.models.bert import BertEmbedding

max_length = 128
batch_size = 32
masking = True

strategy = tf.distribute.MirroredStrategy()
batch_size = batch_size * strategy.num_replicas_in_sync

tokenizer = BertTokenizer(
    '../../tests/resources/pre_models/roberta_wwm_ext/vocab.txt',
    do_lower_case=True)

mask_idx = 5
prefix_tokens = [f'[unused{idx}]' for idx in range(1, 9)]
prefix_tokens.insert(mask_idx - 1, '[MASK]')
prefix_ids = tokenizer.convert_tokens_to_ids(prefix_tokens)
pos_id = tokenizer.convert_tokens_to_ids(['很'])[0]
neg_id = tokenizer.convert_tokens_to_ids(['不'])[0]


def random_masking(token_ids):
    rands = np.random.random(len(token_ids))
    source, target = [], []
    for r, t in zip(rands, token_ids):
        if r < 0.15 * 0.8:
            source.append(tokenizer.convert_tokens_to_ids(['[MASK]'])[0])
            target.append(t)
        elif r < 0.15 * 0.9:
            source.append(t)
            target.append(t)
        elif r < 0.15:
            source.append(np.random.choice(len(tokenizer.vocabs) - 1) + 1)
            target.append(t)
        else:
            source.append(t)
            target.append(0)
    return source, target


def padding(inputs):
    max_seq_length = max([len(x) for x in inputs])

    for x in inputs:
        while len(x) < max_seq_length:
            x.append(0)

    return inputs


def load_data(path):
    items = []
    with open(path, 'r', encoding='utf-8') as f:
        for line in f.readlines():
            line = line.strip().replace('\n', '')
            text, label = line.split('\t')
            label = int(label)

            items.append((text, label))

    return items


class PTuningBertEmbedding(BertEmbedding):
    def call(self, inputs, **kwargs):
        # stop gradients
        word_embeddings = self.word_embeddings
        sg_word_embeddings = K.stop_gradient(self.word_embeddings)

        mask = np.zeros((K.int_shape(self.word_embeddings)[0], 1))
        mask[1:9] += 1
        self.word_embeddings = word_embeddings * mask + sg_word_embeddings * (1 - mask)

        outputs = super(PTuningBertEmbedding, self).call(inputs, **kwargs)

        # reset gradients
        self.word_embeddings = word_embeddings

        return outputs


class PTuningBert(RobertaPreTraining):
    def create_embedding_layer(self, config):
        return PTuningBertEmbedding(config.hidden_size,
                                    config.type_vocab_size,
                                    vocab_size=config.vocab_size,
                                    initializer_range=config.initializer_range,
                                    hidden_dropout_prob=config.hidden_dropout_prob,
                                    max_position_embeddings=config.max_position_embeddings,
                                    name='embeddings')


class Loss(tf.keras.losses.Loss):
    def call(self, y_true, y_pred, **kwargs):
        y_mask = K.cast(K.not_equal(y_true, 0), K.floatx())

        y_pred = K.softmax(y_pred)
        loss = K.sparse_categorical_crossentropy(y_true, y_pred)
        loss = K.sum(loss * y_mask) / K.sum(y_mask)
        return loss


def accuracy(y_true, y_pred, **kwargs):
    y_mask = K.cast(K.not_equal(y_true, 0), K.floatx())
    accuracy = tf.metrics.sparse_categorical_accuracy(y_true, y_pred)
    accuracy = K.sum(accuracy * y_mask) / K.sum(y_mask)

    return accuracy


def data_gen(items, training=False):
    batch_input_ids, batch_segment_ids, batch_input_masks, batch_target_ids = [], [], [], []

    while True:
        from sklearn.utils import shuffle
        items = shuffle(items)
        for text, label in items:
            input_ids, input_masks, segment_ids = tokenizer.encode(text, max_length=max_length)

            if label != 2:
                input_ids = input_ids[:1] + prefix_ids + input_ids[1:]
                input_masks = input_masks[:1] + [1] * len(prefix_ids) + input_masks[1:]
                segment_ids = segment_ids[:1] + [0] * len(prefix_ids) + segment_ids[1:]

            input_ids, target_ids = random_masking(input_ids)

            if label < 2:
                input_ids[mask_idx] = tokenizer.convert_tokens_to_ids(['[MASK]'])[0]
                target_ids[mask_idx] = neg_id if label == 0 else pos_id

            batch_input_ids.append(input_ids)
            batch_segment_ids.append(segment_ids)
            batch_input_masks.append(input_masks)
            batch_target_ids.append(target_ids)

            if len(batch_input_ids) == batch_size:
                batch_input_ids = padding(batch_input_ids)
                batch_segment_ids = padding(batch_segment_ids)
                batch_input_masks = padding(batch_input_masks)
                batch_target_ids = padding(batch_target_ids)

                yield (np.array(batch_input_ids), np.array(batch_input_masks), np.array(batch_segment_ids)), \
                      (np.array(batch_target_ids),)
                batch_input_ids, batch_segment_ids, batch_input_masks, batch_target_ids = [], [], [], []

        if not training:
            # for file end
            batch_input_ids = padding(batch_input_ids)
            batch_segment_ids = padding(batch_segment_ids)
            batch_input_masks = padding(batch_input_masks)
            batch_target_ids = padding(batch_target_ids)

            yield (np.array(batch_input_ids), np.array(batch_input_masks), np.array(batch_segment_ids)), \
                  (np.array(batch_target_ids),)

            yield None, None


train_data = load_data('sentiment/sentiment.train.data')
test_data = load_data('sentiment/sentiment.test.data')
valid_data = load_data('sentiment/sentiment.valid.data')

train_data_generator = data_gen(train_data, training=True)
test_data_generator = data_gen(test_data)
valid_data_generator = data_gen(valid_data)

with strategy.scope():
    model = PTuningBert(
        config='../../tests/resources/pre_models/roberta_wwm_ext/bert_config.json',
        pretrained_weights_path='../../tests/resources/pre_models/roberta_wwm_ext/bert_model.ckpt')

    # set variables to non-trainable exclude word embeddings
    for layer in model.layers:
        layer.trainable = False
    embedding_layer = model.layers[0]
    embedding_layer.trainable = True
    embedding_layer.position_embeddings.trainable = False
    embedding_layer.token_type_embeddings.trainable = False
    embedding_layer.ln.trainable = False

    model.compile(optimizer=Adam(6e-4),
                  loss=Loss(),
                  metrics=[MeanMetricWrapper(fn=accuracy, name='accuracy')])


def evaluate(generator):
    total, right = 0., 0.
    for inputs, y_true in generator:
        if not inputs and not y_true:
            break

        y_pred = model(inputs)
        y_true = y_true[0]

        y_pred1 = tf.expand_dims(y_pred[:, mask_idx, neg_id], axis=-1)
        y_pred2 = tf.expand_dims(y_pred[:, mask_idx, pos_id], axis=-1)
        y_pred = tf.concat([y_pred1, y_pred2], axis=-1)
        y_pred = tf.argmax(y_pred, axis=1).numpy()

        y_true = (y_true[:, mask_idx] == pos_id).astype(int)

        total += len(y_true)
        right += (y_true == y_pred).sum()

    return right / total


class Evaluator(tf.keras.callbacks.Callback):
    def __init__(self):
        self.best_val_acc = 0.

    def on_epoch_end(self, epoch, logs=None):
        val_acc = evaluate(valid_data_generator)
        if val_acc > self.best_val_acc:
            self.best_val_acc = val_acc
            model.save_weights('best_model_bert.weights')
        test_acc = evaluate(test_data_generator)
        print(
            u'val_acc: %.5f, best_val_acc: %.5f, test_acc: %.5f\n' %
            (val_acc, self.best_val_acc, test_acc)
        )


model.summary()

model.fit(x=train_data_generator,
          steps_per_epoch=len(train_data) / batch_size + 1,
          epochs=1000,
          callbacks=[Evaluator()])
# stop gradients # reset gradients # for file end # set variables to non-trainable exclude word embeddings