Exemplo n.º 1
0
def model_save():
    model = TF2BertModel(9, 128, 'models/albert_crf')
    model.load()
    vocab = TF2Tokenizer()
    export_path = 'models/albert_crf/saved_model/1'
    encoded = model.predict(vocab.tokenize_to_ids(['你好!'], model.max_seq_len))
    tf.saved_model.save(model,
                        export_path,
                        signatures=model.call.get_concrete_function(
                            tf.TensorSpec(shape=[None, 2, 128],
                                          dtype=tf.int32,
                                          name="input")))
    model.save(os.path.join(model.model_dir, 'saved_model'), save_format='tf')
    np.save('models/albert_crf/transition.npy',
            model.transition_params.numpy())
Exemplo n.º 2
0
        outs = tf.nn.softmax(x)
        # outs = tf.nn.softmax(x) * mask
        return outs

    def load_pretrained(self, pretrained_ckpt):
        self.build(input_shape=(None, 2, self._max_seq_len))
        bert.load_albert_weights(self.l_bert, pretrained_ckpt)

    def load(self, weight_file):
        self.build(input_shape=(None, 2, self._max_seq_len))
        self.load_weights(weight_file)


if __name__ == '__main__':
    model = TF2BertModel(9, 128, 'models/alBERT_only_NER')
    model.load(os.path.join(model.model_dir, 'ckpt_best'))
    # model.load_pretrained('models/albert_base_zh/albert_model.ckpt')

    vocab = TF2Tokenizer()
    encoded = model.predict(
        vocab.tokenize_to_ids(['马云,祖籍浙江省绍兴嵊州市谷来镇,后父母移居杭州,出生于浙江省杭州市,中国企业家,中国共产党党员。曾为亚洲首富、阿里巴巴集团董事局主席。'], model.max_seq_len))
    print(encoded.shape)
    print(np.argmax(encoded, axis=-1))

    # from losses import masked_sparse_categorical_crossentropy
    #
    # label = np.argmax(encoded, axis=-1)
    # label[label == 8] = 9
    # print(masked_sparse_categorical_crossentropy(tf.convert_to_tensor(label),
    #                                              tf.convert_to_tensor(encoded)))