Spaces:

AT-VN-Research-Group
/

PUDMED-Sentence-Classification

Sleeping

PUDMED-Sentence-Classification / src /models /penta_embeddings_model.py

src

e8e72fb almost 3 years ago

9.76 kB

import tensorflow as tf
from tensorflow.keras import layers, Sequential
import tensorflow_hub as hub
from tensorflow.keras.utils import plot_model
import sys
import os


current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.abspath(os.path.join(current_dir, os.pardir, os.pardir))
sys.path.append(project_root)
from config.configs import *
params = Params()


class PentaEmbeddingModel():
    def __init__(self, word_vectorizer = None, char_vectorizer = None, word_embed = None, char_embed = None, pretrained_embedding = None, glove_embed = None, bert_process = None, bert_layer = None, num_classes = 5):
        super(PentaEmbeddingModel, self).__init__()
        """
        word_vectorizer: Word-level vectorizer
        char_vectorizer: Char-level vectorizer
        word_embed: Word-level embedding layer
        char_embed: Char-level embedding layer
        pretrained_embedding: "bert", "glove" or None. Default: None
        glove_embed: glove embedding layer 
        bert_process: BERT input processing layer
        bert_layer: BERT embedding layer
        num_classes: Number of classes. Default: 5 ("Do not change")
        """

        # Params
        self.pretrained_embedding = pretrained_embedding
        self.num_classes = params.NUM_CLASSES
        self.word_output_dim = params.WORD_OUTPUT_DIM
        self.char_output_dim = params.CHAR_OUTPUT_DIM
        self.concate_dim = self.word_output_dim + 2 * self.char_output_dim

        self.line_ids_input_dim = params.LINE_IDS_DEPTH
        self.length_lines_input_dim = params.LENGTH_LINES_DEPTH
        self.total_lines_input_dim = params.TOTAL_LINES_DEPTH


        # Vectorizer
        self.word_vectorizer = word_vectorizer
        self.char_vectorizer = char_vectorizer
        
        # Embedding
        self.word_embed = word_embed
        self.char_embed = char_embed
        self.glove_embed = glove_embed
        self.bert_process = bert_process
        self.bert_layer = bert_layer

        # Layers
        self.word_biLSTM = layers.Bidirectional(layers.LSTM(int(self.word_output_dim / 2)))
        self.char_biLSTM = layers.Bidirectional(layers.LSTM(int(self.char_output_dim)))
        self.concat_biLSTM = layers.Bidirectional(layers.LSTM(int(self.concate_dim)))
        self.concatenate = layers.Concatenate()
        self.dense_classes = layers.Dense(self.num_classes, activation = "softmax")
        self.dropout = layers.Dropout(0.5)



#---------- First level branch----------------------
    def word_level_branch(self, word_input):
        """
        Word-token embedding branch
        Pretrained BERT don't need vectorization layer
        """
        if str(self.pretrained_embedding).lower() == "bert":
            # Pretrained Bert embeddings
            bert_input = self.bert_process(word_input)
            bert_output = self.bert_layer(bert_input, training = False)
            word_embeddings = bert_output['sequence_output']
        else:
            if (self.word_vectorizer):    
                if (str(self.pretrained_embedding).lower() == "glove"):
                    # Get glove embedding
                    word_vectors = self.word_vectorizer(word_input)
                    word_embeddings = self.glove_embed(word_vectors)
                else:
                    # Original word_embeddings
                    word_vectors = self.word_vectorizer(word_input)
                    word_embeddings = self.word_embed(word_vectors)
            else:
                raise Exception("Please provide word vectorizer.")
        x = layers.Conv1D(64, kernel_size=5, padding="same", activation="relu")(word_embeddings)
        x = layers.Dense(128, activation = "relu")(x)
        x = layers.BatchNormalization()(x)
        word_outputs = self.word_biLSTM(x)
        return word_outputs


    def char_level_branch(self, char_input):
        """
        arg: 
        - char_input: char-level tokens embedding
        """
        char_vectors = self.char_vectorizer(char_input)
        char_embeddings = self.char_embed(char_vectors)
        x = self.char_biLSTM(char_embeddings)
        return x


    def line_ids_branch(self, line_id_in):
        """
        arg: 
        - line_id_in: line_ids one-hot embedding
        """
        x = layers.Dense(64, activation = "relu")(line_id_in)
        x = layers.BatchNormalization()(x)
        return x



    def length_lines_branch(self, length_line_in):
        """
        arg:
        - length_line_in: length_lines one-hot embedding
        """
        x = layers.Dense(64, activation = "relu")(length_line_in)
        x = layers.BatchNormalization()(x)
        return x


    def total_lines_branch(self, total_line_in):
        """
        arg:
        -total_lines_in: total_lines one-hot embedding
        """
        x = layers.Dense(64, activation = "relu")(total_line_in)
        x = layers.BatchNormalization()(x)
        return x

# ----------------Second-level layer----------------------

    def word_char_block(self, word_char_concat):
        """
        Blocks for word-level tokens, char-level tokens
        arg: - word_char_concat: word-level, char-level tokens
        """
        word_char_concat = tf.expand_dims(word_char_concat, axis = 1)
        # LSTM layer for first two concate embeddings
        lstm_concat = self.concat_biLSTM(word_char_concat)
        lstm_concat = layers.Dense(256, activation = "relu")(lstm_concat)
        lstm_concat = layers.Dropout(0.5) (lstm_concat)

        return lstm_concat, lstm_concat.shape

# ---------------Third-level layer-----------------------------

    def sequence_opt_layer(self, total_embed):
        """
        Context-position enrichment layer
        arg:
        - total_embed: 5 concatenated input embedding
        """
        total_embed = tf.expand_dims(total_embed, axis = 1)
        bilstm_out = layers.Bidirectional(layers.LSTM(int(total_embed.shape[-1] / 2)))(total_embed)
        return bilstm_out


    def fcn(self, total_embed):
        """
        Feed forward FCN
        """
        x = layers.Dense(64, activation = "relu", input_shape = (total_embed.shape[1], ))(total_embed)
        x = layers.BatchNormalization()(x)
        x = layers.Dropout(0.5)(x)
        x = layers.Dense(self.num_classes, activation = "softmax")(x)
        return x


    def _get_model(self):
        
        # Penta embedding inputs
        word_inputs = layers.Input(shape = [], dtype = tf.string, name = "token_input")
        char_inputs = layers.Input(shape = (1, ), dtype = tf.string, name = "char_input")
        line_ids_inputs = layers.Input(shape = (self.line_ids_input_dim, ), name = "line_ids_input")
        length_lines_inputs = layers.Input(shape = (self.length_lines_input_dim, ), name = "length_lines_input")
        total_lines_inputs = layers.Input(shape = (self.total_lines_input_dim, ), name = "total_lines_input")

        #-----------------------------------------------
        # Branch outputs

        # Word-level
        word_level_output = self.word_level_branch(word_inputs)

        # Char-level branch
        char_level_output = self.char_level_branch(char_inputs)

        #line_ids, length_lines, total_lines branch
        line_ids_output = self.line_ids_branch(line_ids_inputs)
        length_lines_output  = self.length_lines_branch(length_lines_inputs)
        total_lines_output = self.total_lines_branch(total_lines_inputs)
        #---------------------------------------------------------

        #Concate two embeddings
        word_char_concat = self.concatenate([word_level_output,char_level_output])

        # Pass to word_char_block
        word_char_output, word_char_output_shape = self.word_char_block(word_char_concat)
        #--------------------------------------------------------

        # Concanate last three input
        position_embed = self.concatenate([length_lines_output, line_ids_output, total_lines_output])

        # Concatnate 5 input
        total_embed = self.concatenate([word_char_output, position_embed])

        #Sequence label opt layers
        total_embed = self.sequence_opt_layer(total_embed)
        # FCN

        output_layer = self.fcn(total_embed)

        model= tf.keras.Model(inputs=[word_inputs, char_inputs,line_ids_inputs, length_lines_inputs, total_lines_inputs],
                         outputs= output_layer,
                         name="penta_embeddings_model")

        model.compile(optimizer = params.OPTIMIZER, loss = params.LOSS, metrics = params.METRICS)
        return model



    def _define_checkpoint(self):
        """
        Define checkpoint for model
        """
        if str(self.pretrained_embedding).lower() == "glove":
            if not os.path.exists(params.PENTA_GLOVE_MODEL_DIR):
                os.makedirs(params.PENTA_GLOVE_MODEL_DIR)
            checkpoint_dir = params.PENTA_GLOVE_MODEL_DIR
        elif str(self.pretrained_embedding).lower() == "bert":
            if not os.path.exists(params.PENTA_BERT_MODEL_DIR):
                os.makedirs(params.PENTA_BERT_MODEL_DIR)
            checkpoint_dir = params.PENTA_BERT_MODEL_DIR
        else:
            if not os.path.exists(params.PENTA_NOR_MODEL_DIR):
                os.makedirs(params.PENTA_NOR_MODEL_DIR)
            checkpoint_dir = params.PENTA_NOR_MODEL_DIR
    
        checkpoint= tf.keras.callbacks.ModelCheckpoint(
        filepath = checkpoint_dir + '/best_model.ckpt',
        monitor = "val_categorical_accuracy",
        save_best_only = True,
        save_weights_only = True,
        verbose = 1
        )
        print("Create checkpoint for penta embeddings model at: ", checkpoint_dir)
        return checkpoint
    

    def _plot_model(self, model):
        plot_model(model)
        return 

	import tensorflow as tf
	from tensorflow.keras import layers, Sequential
	import tensorflow_hub as hub
	from tensorflow.keras.utils import plot_model
	import sys
	import os


	current_dir = os.path.dirname(os.path.abspath(__file__))
	project_root = os.path.abspath(os.path.join(current_dir, os.pardir, os.pardir))
	sys.path.append(project_root)
	from config.configs import *
	params = Params()


	class PentaEmbeddingModel():
	def __init__(self, word_vectorizer = None, char_vectorizer = None, word_embed = None, char_embed = None, pretrained_embedding = None, glove_embed = None, bert_process = None, bert_layer = None, num_classes = 5):
	super(PentaEmbeddingModel, self).__init__()
	"""
	word_vectorizer: Word-level vectorizer
	char_vectorizer: Char-level vectorizer
	word_embed: Word-level embedding layer
	char_embed: Char-level embedding layer
	pretrained_embedding: "bert", "glove" or None. Default: None
	glove_embed: glove embedding layer
	bert_process: BERT input processing layer
	bert_layer: BERT embedding layer
	num_classes: Number of classes. Default: 5 ("Do not change")
	"""

	# Params
	self.pretrained_embedding = pretrained_embedding
	self.num_classes = params.NUM_CLASSES
	self.word_output_dim = params.WORD_OUTPUT_DIM
	self.char_output_dim = params.CHAR_OUTPUT_DIM
	self.concate_dim = self.word_output_dim + 2 * self.char_output_dim

	self.line_ids_input_dim = params.LINE_IDS_DEPTH
	self.length_lines_input_dim = params.LENGTH_LINES_DEPTH
	self.total_lines_input_dim = params.TOTAL_LINES_DEPTH


	# Vectorizer
	self.word_vectorizer = word_vectorizer
	self.char_vectorizer = char_vectorizer

	# Embedding
	self.word_embed = word_embed
	self.char_embed = char_embed
	self.glove_embed = glove_embed
	self.bert_process = bert_process
	self.bert_layer = bert_layer

	# Layers
	self.word_biLSTM = layers.Bidirectional(layers.LSTM(int(self.word_output_dim / 2)))
	self.char_biLSTM = layers.Bidirectional(layers.LSTM(int(self.char_output_dim)))
	self.concat_biLSTM = layers.Bidirectional(layers.LSTM(int(self.concate_dim)))
	self.concatenate = layers.Concatenate()
	self.dense_classes = layers.Dense(self.num_classes, activation = "softmax")
	self.dropout = layers.Dropout(0.5)



	#---------- First level branch----------------------
	def word_level_branch(self, word_input):
	"""
	Word-token embedding branch
	Pretrained BERT don't need vectorization layer
	"""
	if str(self.pretrained_embedding).lower() == "bert":
	# Pretrained Bert embeddings
	bert_input = self.bert_process(word_input)
	bert_output = self.bert_layer(bert_input, training = False)
	word_embeddings = bert_output['sequence_output']
	else:
	if (self.word_vectorizer):
	if (str(self.pretrained_embedding).lower() == "glove"):
	# Get glove embedding
	word_vectors = self.word_vectorizer(word_input)
	word_embeddings = self.glove_embed(word_vectors)
	else:
	# Original word_embeddings
	word_vectors = self.word_vectorizer(word_input)
	word_embeddings = self.word_embed(word_vectors)
	else:
	raise Exception("Please provide word vectorizer.")
	x = layers.Conv1D(64, kernel_size=5, padding="same", activation="relu")(word_embeddings)
	x = layers.Dense(128, activation = "relu")(x)
	x = layers.BatchNormalization()(x)
	word_outputs = self.word_biLSTM(x)
	return word_outputs


	def char_level_branch(self, char_input):
	"""
	arg:
	- char_input: char-level tokens embedding
	"""
	char_vectors = self.char_vectorizer(char_input)
	char_embeddings = self.char_embed(char_vectors)
	x = self.char_biLSTM(char_embeddings)
	return x


	def line_ids_branch(self, line_id_in):
	"""
	arg:
	- line_id_in: line_ids one-hot embedding
	"""
	x = layers.Dense(64, activation = "relu")(line_id_in)
	x = layers.BatchNormalization()(x)
	return x



	def length_lines_branch(self, length_line_in):
	"""
	arg:
	- length_line_in: length_lines one-hot embedding
	"""
	x = layers.Dense(64, activation = "relu")(length_line_in)
	x = layers.BatchNormalization()(x)
	return x


	def total_lines_branch(self, total_line_in):
	"""
	arg:
	-total_lines_in: total_lines one-hot embedding
	"""
	x = layers.Dense(64, activation = "relu")(total_line_in)
	x = layers.BatchNormalization()(x)
	return x

	# ----------------Second-level layer----------------------

	def word_char_block(self, word_char_concat):
	"""
	Blocks for word-level tokens, char-level tokens
	arg: - word_char_concat: word-level, char-level tokens
	"""
	word_char_concat = tf.expand_dims(word_char_concat, axis = 1)
	# LSTM layer for first two concate embeddings
	lstm_concat = self.concat_biLSTM(word_char_concat)
	lstm_concat = layers.Dense(256, activation = "relu")(lstm_concat)
	lstm_concat = layers.Dropout(0.5) (lstm_concat)

	return lstm_concat, lstm_concat.shape

	# ---------------Third-level layer-----------------------------

	def sequence_opt_layer(self, total_embed):
	"""
	Context-position enrichment layer
	arg:
	- total_embed: 5 concatenated input embedding
	"""
	total_embed = tf.expand_dims(total_embed, axis = 1)
	bilstm_out = layers.Bidirectional(layers.LSTM(int(total_embed.shape[-1] / 2)))(total_embed)
	return bilstm_out


	def fcn(self, total_embed):
	"""
	Feed forward FCN
	"""
	x = layers.Dense(64, activation = "relu", input_shape = (total_embed.shape[1], ))(total_embed)
	x = layers.BatchNormalization()(x)
	x = layers.Dropout(0.5)(x)
	x = layers.Dense(self.num_classes, activation = "softmax")(x)
	return x


	def _get_model(self):

	# Penta embedding inputs
	word_inputs = layers.Input(shape = [], dtype = tf.string, name = "token_input")
	char_inputs = layers.Input(shape = (1, ), dtype = tf.string, name = "char_input")
	line_ids_inputs = layers.Input(shape = (self.line_ids_input_dim, ), name = "line_ids_input")
	length_lines_inputs = layers.Input(shape = (self.length_lines_input_dim, ), name = "length_lines_input")
	total_lines_inputs = layers.Input(shape = (self.total_lines_input_dim, ), name = "total_lines_input")

	#-----------------------------------------------
	# Branch outputs

	# Word-level
	word_level_output = self.word_level_branch(word_inputs)

	# Char-level branch
	char_level_output = self.char_level_branch(char_inputs)

	#line_ids, length_lines, total_lines branch
	line_ids_output = self.line_ids_branch(line_ids_inputs)
	length_lines_output = self.length_lines_branch(length_lines_inputs)
	total_lines_output = self.total_lines_branch(total_lines_inputs)
	#---------------------------------------------------------

	#Concate two embeddings
	word_char_concat = self.concatenate([word_level_output,char_level_output])

	# Pass to word_char_block
	word_char_output, word_char_output_shape = self.word_char_block(word_char_concat)
	#--------------------------------------------------------

	# Concanate last three input
	position_embed = self.concatenate([length_lines_output, line_ids_output, total_lines_output])

	# Concatnate 5 input
	total_embed = self.concatenate([word_char_output, position_embed])

	#Sequence label opt layers
	total_embed = self.sequence_opt_layer(total_embed)
	# FCN

	output_layer = self.fcn(total_embed)

	model= tf.keras.Model(inputs=[word_inputs, char_inputs,line_ids_inputs, length_lines_inputs, total_lines_inputs],
	outputs= output_layer,
	name="penta_embeddings_model")

	model.compile(optimizer = params.OPTIMIZER, loss = params.LOSS, metrics = params.METRICS)
	return model



	def _define_checkpoint(self):
	"""
	Define checkpoint for model
	"""
	if str(self.pretrained_embedding).lower() == "glove":
	if not os.path.exists(params.PENTA_GLOVE_MODEL_DIR):
	os.makedirs(params.PENTA_GLOVE_MODEL_DIR)
	checkpoint_dir = params.PENTA_GLOVE_MODEL_DIR
	elif str(self.pretrained_embedding).lower() == "bert":
	if not os.path.exists(params.PENTA_BERT_MODEL_DIR):
	os.makedirs(params.PENTA_BERT_MODEL_DIR)
	checkpoint_dir = params.PENTA_BERT_MODEL_DIR
	else:
	if not os.path.exists(params.PENTA_NOR_MODEL_DIR):
	os.makedirs(params.PENTA_NOR_MODEL_DIR)
	checkpoint_dir = params.PENTA_NOR_MODEL_DIR

	checkpoint= tf.keras.callbacks.ModelCheckpoint(
	filepath = checkpoint_dir + '/best_model.ckpt',
	monitor = "val_categorical_accuracy",
	save_best_only = True,
	save_weights_only = True,
	verbose = 1
	)
	print("Create checkpoint for penta embeddings model at: ", checkpoint_dir)
	return checkpoint


	def _plot_model(self, model):
	plot_model(model)
	return