modeling_bilma.py

from transformers import TFPreTrainedModel
from tensorflow.keras.models import Model, load_model, Sequential
from tensorflow.keras.layers import Layer, Dense, concatenate, Input, add, Dropout, LayerNormalization, MultiHeadAttention, Embedding
import tensorflow as tf
import numpy as np

import re
import unicodedata

from configuration_bilma import BilmaConfig

# copied from preprocessing.py
BLANK = ' '

RE_OPS = re.I | re.M | re.S
RE_USR = re.compile(r"""@\S+""", RE_OPS)
RE_TAG = re.compile(r"""#\S+""", RE_OPS)
RE_URL = re.compile(r"""(http|ftp|https)://\S+""", RE_OPS)
RE_NUM = re.compile(r"""[-+]?\d+\.?\d*""", RE_OPS)

SYMBOLS_ = "()[]¿?¡!{}~<>|"
SYMBOLS = set(";:,.@\\-\"/" + SYMBOLS_)



# ------------------
# Class declaration
# ------------------


class Bilma(TFPreTrainedModel):
    config_class = BilmaConfig
    main_input_name = "input_ids"

    def __init__(self, config):
        super().__init__(config)
        #if config.weights == "spanish":
        #    my_resources = importlib_resources.files("hf_bilma") 
        #    model_file = str((my_resources / "bilma_dataset_small_epoch-1_part-60.h5").joinpath())
        #    self.model = bm.load(model_file)
        #else:
        self.model = bilma(num_enc=config.num_encoders,
                           embed_dim=config.embedding_dim, 
                           max_length=config.max_length,
                           num_heads=config.num_attention_heads,
                           ff_dim=config.embedding_dim,
                           vocab_size=config.vocab_size,
                           rate=config.drop_rate)
        

    def call(self, tensor):
        return self.model(tensor)
    

# 
# Copied from transformer_text.py   
# -------------------------------

class EncoderBlock(Layer):
    def __init__(self, patch_dim, num_heads, ff_dim, rate=0.1, **kwargs):
        super(EncoderBlock, self).__init__(**kwargs)
        self.p_d = patch_dim
        self.n_h = num_heads
        self.f_d = ff_dim
        self.rate = rate
        
        self.att = MultiHeadAttention(num_heads=num_heads, key_dim=patch_dim)
        self.ffn = Sequential(
            #[Conv1D(ff_dim, kernel_size=1, activation=tf.nn.gelu), 
            # Conv1D(patch_dim, kernel_size=1),]
            [Dense(ff_dim, activation=tf.nn.gelu), 
             Dense(patch_dim),]
        )
        #self.layernorm0 = LayerNormalization(epsilon=1e-6)
        self.layernorm1 = LayerNormalization(epsilon=1e-6)
        self.layernorm2 = LayerNormalization(epsilon=1e-6)
        self.dropout1 = Dropout(rate)
        self.dropout2 = Dropout(rate)
        
    def get_config(self):
        config = super(EncoderBlock, self).get_config()
        config.update({"patch_dim":self.p_d, "num_heads":self.n_h, "ff_dim":self.f_d, "rate":self.rate})
        return config

    def call(self, inputs, training=False):
        #inputs = self.layernorm0(inputs)
        attn_output = self.att(inputs, inputs)
        attn_output = self.dropout1(attn_output, training=training)
        out1 = self.layernorm1(add([inputs, attn_output]))
        ffn_output = self.ffn(out1)
        ffn_output = self.dropout2(ffn_output, training=training)
        return self.layernorm2(add([out1, ffn_output]))
    

class DecoderBlock(Layer):
    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1, **kwargs):
        super(DecoderBlock, self).__init__(**kwargs)
        self.e_d = embed_dim
        self.n_h = num_heads
        self.f_d = ff_dim
        self.rate = rate
        
        self.att1 = MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
        self.att2 = MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
        self.ffn = Sequential(
            #[Conv1D(ff_dim, kernel_size=1, activation=tf.nn.gelu), 
            # Conv1D(embed_dim, kernel_size=1),]
            [Dense(ff_dim, activation=tf.nn.gelu), 
             Dense(embed_dim),]
        )
        self.layernorm1 = LayerNormalization(epsilon=1e-6)
        self.layernorm2 = LayerNormalization(epsilon=1e-6)
        self.dropout1 = Dropout(rate)
        self.dropout2 = Dropout(rate)
        self.dropout3 = Dropout(rate)
        
    def get_config(self):
        config = super(DecoderBlock, self).get_config()
        config.update({"embed_dim":self.e_d, "num_heads":self.n_h, "ff_dim":self.f_d, "rate":self.rate})
        return config

    def call(self, inputs, encoder_output, look_ahead_mask, padding_mask, training=None):
        y, attn_output1 = self.att1(inputs, inputs, attention_mask=look_ahead_mask, return_attention_scores=True)
        y = self.dropout1(y, training=training)
        y = add([inputs, y])                
        out1 = self.layernorm1(y)
        
        y, attn_encoder = self.att2(out1, encoder_output, attention_mask=padding_mask, return_attention_scores=True)
        y = self.dropout2(y, training=training)
        y = add([out1, y])                
        out2 = self.layernorm1(y)
        
        ffn_output = self.ffn(out2)
        ffn_output = self.dropout3(ffn_output, training=training)
        final_output =  self.layernorm2(out2 + ffn_output)
        
        return final_output, attn_output1, attn_encoder


class Encoder(Layer):
    def __init__(self, n, embed_dim, max_length, num_heads, ff_dim, rate=0.1, **kwargs):
        super(Encoder, self).__init__(**kwargs)
        self.n = n        
        self.embed_dim = embed_dim
        self.max_length = max_length
        self.n_h = num_heads
        self.f_d = ff_dim
        self.rate = rate
        self._layers = [EncoderBlock(embed_dim, num_heads, ff_dim, rate=0.1) for _ in range(n)]
        self.pe = positional_encoding(self.max_length, self.embed_dim)
        
    def get_config(self):
        config = super(Encoder, self).get_config()
        config.update({"n": self.n, "embed_dim":self.embed_dim, "max_length": self.max_length, "num_heads":self.n_h, "ff_dim":self.f_d, "rate":self.rate})
        return config
    
    def call(self, x, training=False):
        x *= tf.math.sqrt(tf.cast(self.embed_dim, tf.float32))
        x = x + self.pe[:, :tf.shape(x)[1], :]
        for layer in self._layers:
            x = layer(x, training)
        return x

    
class Decoder(Layer):
    def __init__(self, n, embed_dim, max_length, num_heads, ff_dim, rate=0.1, **kwargs):
        super(Decoder, self).__init__(**kwargs)
        self.n = n
        self.embed_dim = embed_dim
        self.max_length = max_length
        self.n_h = num_heads
        self.f_d = ff_dim
        self.rate = rate
        self._layers = [DecoderBlock(embed_dim, num_heads, ff_dim, rate=0.1) for _ in range(n)]
        self.pe = positional_encoding(self.max_length, self.embed_dim)
    
    def get_config(self):
        config = super(Decoder, self).get_config()
        config.update({"n": self.n, "embed_dim":self.embed_dim, "max_length": self.max_length, "num_heads":self.n_h, "ff_dim":self.f_d, "rate":self.rate})
        return config
    
    def call(self, x, encoder_output, look_ahead_mask, padding_mask, training):      
        x *= tf.math.sqrt(tf.cast(self.embed_dim, tf.float32))
        x = x + self.pe[:, :tf.shape(x)[1], :]
        
        for layer in self._layers:
            x, self_att, enc_att = layer(x, encoder_output, look_ahead_mask, padding_mask, training)

        return x




# =========================================
#   M A S K S 
# =========================================
def create_padding_mask(seq):
    """
    For self-attention
    seq shape(bs, max_length, emb_dim)
    output shape (bs, max_length, max_length)
    """
    mask = tf.cast(tf.not_equal(seq, 0), tf.bool)
    mask = tf.reduce_any(mask, 2)
    mask = tf.repeat(mask, seq.shape[1], 0)
    mask = tf.reshape(mask, (-1,seq.shape[1], seq.shape[1]))
    return tf.cast(mask, tf.float32)


def create_cross_padding_mask(seq, target_seq):
    """
    For cross-attention
    seq shape(bs, k, image_features)
    target_seq(bs, max_length, emb_dim)
    output shape (bs, max_length, k)
    """
    mask = tf.cast(tf.not_equal(target_seq, 0), tf.bool)
    mask = tf.reduce_any(mask, 2)
    mask = tf.repeat(mask, seq.shape[1], 0)
    mask = tf.reshape(mask, (-1, tf.shape(seq)[1], tf.shape(target_seq)[1]))
    mask = tf.transpose(mask, [0, 2, 1])
    return mask


def create_look_ahead_mask(seq):
    """
    seq shape(bs, max_length, emb_dim)
    output 2D matrix of shape (bs, max_length, max_length) with ones on the diagonal and below.
    """
    size = seq.shape[1]
    mask = tf.linalg.band_part(tf.ones((size, size)), -1, 0)
    mask = tf.expand_dims(mask, 0)
    mask = tf.repeat(mask, tf.shape(seq)[0], 0)
    return mask


def create_masks(seq, target_seq):
    decoder_mask = create_padding_mask(target_seq)
    decoder_mask *= create_look_ahead_mask(target_seq)
    cross_att_mask = create_cross_padding_mask(seq, target_seq)
    return decoder_mask, cross_att_mask
        
    
def create_masks_looking_ahead(seq, target_seq):
    decoder_mask = create_padding_mask(target_seq)
    cross_att_mask = create_cross_padding_mask(seq, target_seq)
    return decoder_mask, cross_att_mask
    
# =========================================
#   P O S I T I O N A L   E N C O D I N G
# =========================================
def get_angles(pos, i, d_model):
    angle_rates = 1 / np.power(10000, (2 * (i//2)) / np.float32(d_model))
    return pos * angle_rates

@tf.autograph.experimental.do_not_convert
def positional_encoding(position, d_model):
    angle_rads = get_angles(np.arange(position)[:, np.newaxis],
                          np.arange(d_model)[np.newaxis, :],
                          d_model)

    # apply sin to even indices in the array; 2i
    angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])

    # apply cos to odd indices in the array; 2i+1
    angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])

    pos_encoding = angle_rads[np.newaxis, ...]

    return tf.cast(pos_encoding, dtype=tf.float32)

class PatchEncoder(Layer):
    def __init__(self, num_patches, projection_dim, **kwargs):
        super(PatchEncoder, self).__init__(**kwargs)
        self.num_patches = num_patches
        self.projection_dim = projection_dim
        self.projection = Dense(units=projection_dim)
        self.position_embedding = Embedding(
            input_dim=num_patches, output_dim=projection_dim
        )
    
    def get_config(self):
        config = super(PatchEncoder, self).get_config()
        config.update({"num_patches": self.num_patches, "projection_dim":self.projection_dim})
        return config

    def call(self, patch):
        positions = tf.range(start=0, limit=self.num_patches, delta=1)
        encoded = self.projection(patch) + self.position_embedding(positions)
        return encoded
    
    

# Copied from preprocessing.py
# ----------------------------
def norm_chars(text):
    L = []
  
    for u in unicodedata.normalize('NFD', text):
        o = ord(u)
        if 0x300 <= o and o <= 0x036F:
            continue
           
        if u in ('\n', '\r', BLANK, '\t'):
            if len(L) == 0:
                continue

            u = BLANK
        
        if u in SYMBOLS:
            if len(L) > 0 and L[-1] != BLANK:
                L.append(BLANK)
            
            L.append(u)
            L.append(BLANK)
            continue
        
        L.append(u)

    return "".join(L)


def preprocess(text):
    text = RE_URL.sub("_url ", text)
    text = RE_USR.sub("_usr ", text)
    #text = RE_TAG.sub("_htag ", text)
    #text = RE_NUM.sub("0 ", text)
    text = re.sub(r"&amp;", "&", text)
    text = re.sub(r"&gt;", ">", text)
    text = re.sub(r"&lt;", "<", text)
    #text = norm_chars(text.lower())
    text = re.sub(r"j(a|e|i)[jaei]+", r"j\1j\1", text)
    text = re.sub(r"h(a|e|i)[haei]+", r"j\1j\1", text)
    return re.sub(r"\s+", BLANK, text)



# Copied from wordpiece_tokenizer_ex.py
# -------------------------------------

class Tokenizer():
    def __init__(self, vocab_file, unk_token="[UNK]", end_token="[END]", mask_token="[MASK]"):
        self.word2idx = {}
        self.idx2word = []
        c = 0
        with open(vocab_file, "r", encoding="utf8") as f:
            while True:
                line = f.readline()
                if not line:
                    break
                self.word2idx[line[0:-1]] = c
                self.idx2word.append(line[0:-1])
                c += 1
        self.n_jobs = 2
        self.UNK = unk_token
        self.END = end_token
        self.MASK = mask_token
        
    def split(self, s):
        split = []
        i = 0
        while i < len(s):
            for j in range(i, len(s)):
                if (i==j and s[j:j+6] == self.MASK):
                    split.append(self.MASK)
                    i = j + 6
                    break
                if (s[j].isalnum()):
                    continue
                if (j==i):
                    if (s[j] != " "):
                        split.append(s[i:j+1])
                    i = j + 1
                    break
                split.append(s[i:j])
                i = j
                break
            else:
                split.append(s[i:j+1])
                i=j+1
        return split
    
    def tokenize(self, S):        
        #return Parallel(n_jobs=self.n_jobs)(delayed(self._tokenize)(s) for s in S)
        return [self._tokenize(s) for s in S]
    
    def detokenize(self, S, human_readable=True):        
        #return Parallel(n_jobs=self.n_jobs)(delayed(self._detokenize)(s) for s in S)
        return [self._detokenize(s, human_readable=human_readable) for s in S]
    
    def _tokenize(self, s):
        tokens = []
        s = s.rstrip('\n')
        for w in self.split(s):
            if w in self.word2idx:
                tokens.append(self.word2idx[w])
            else:
                if (len(w)==1):
                    tokens.append(self.word2idx["[UNK]"])
                    continue
                                                
                subtoken = []
                l = 0
                while len(w)>l:
                    
                    for i in range(len(w),l-1,-1):
                        if (w[0: i] in self.word2idx):
                            subtoken.append(self.word2idx[w[0: i]])
                            break
                    if (i == l):
                        subtoken = [self.word2idx["[UNK]"]]                
                        break
                    w = "##" + w[i: ]
                    l = 2
                tokens += subtoken
        return tokens
    
    
    def _detokenize(self, tokens, human_readable=True):
        sentence = []
        start = 0 if human_readable == False else 1
        
        for t in tokens[start:]:
            c = self.idx2word[t]
            if (human_readable and c == self.END):
                break
            sentence.append(c)
        return sentence
    
    
    
# copied from bilma_model.py
# --------------------------

def loss_function(ignore_id=0):
    loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction='none')
    def loss(real, pred):
        mask = tf.math.logical_not(tf.math.equal(real, ignore_id))
        loss_ = loss_object(real, pred)
        mask = tf.cast(mask, dtype=loss_.dtype)
        loss_ *= mask
        sum_ = tf.reduce_sum(mask,axis=1)
        
        loss_ = tf.math.divide_no_nan(tf.reduce_sum(loss_, axis=1), sum_)
        return loss_
    return loss

def accuracy_function(ignore_id=0):
    def acc_mlm(real, pred):
        accuracies = tf.equal(tf.cast(real, tf.int64), tf.argmax(pred, axis=2))

        mask = tf.math.logical_not(tf.math.equal(real, ignore_id))
        accuracies = tf.math.logical_and(mask, accuracies)

        accuracies = tf.cast(accuracies, dtype=tf.float32)
        mask = tf.cast(mask, dtype=tf.float32)
        return tf.math.divide_no_nan(tf.reduce_sum(accuracies), tf.reduce_sum(mask))
    return acc_mlm

def bilma(num_enc=6, embed_dim=300, max_length=50, num_heads=6, ff_dim=512, vocab_size=9739, rate=0.1):
    capt_inputs_ids = Input(shape=(max_length, ), name='capt_input')
    capt_embedding = Embedding(vocab_size, embed_dim, mask_zero=False, name="embedding")
    capt_inputs = capt_embedding(capt_inputs_ids)
    
    enc = Encoder(num_enc, embed_dim, max_length, num_heads, ff_dim, rate=rate)
    enc_output = enc(capt_inputs)
    fin_output = Dense(vocab_size, use_bias=True)(enc_output)
    
    caption_model = Model(inputs=capt_inputs_ids, outputs=[fin_output])
    return caption_model

def load(model_file):
    custom_objects={"EncoderBlock": EncoderBlock, 
                    "Encoder": Encoder,
                    "loss": loss_function(),
                    "acc_mlm":accuracy_function(),
                   }
    return load_model(model_file, custom_objects=custom_objects)

class tokenizer():
    def __init__(self, vocab_file, max_length):
        self.tokenizer = Tokenizer(vocab_file)
        self.emo_labels = "❤👌👏💔😄😊😌😍😒😘😡😢😭🤔🥺"
        self.max_length = max_length
        self.START = 2
        self.END = 3
        self.PAD = 0
        self.MASK = 4  
        
    def tokenize(self, text):
        text = [preprocess(t) for t in text]
        tokens = tf.ragged.constant(self.tokenizer.tokenize(text),  tf.int32)        
        count, _ = tokens.bounding_shape()
        starts = tf.fill([count,1], self.START)
        ends = tf.fill([count,1], self.END)
        tokens = tf.concat([starts, tokens[:, 0: self.max_length - 2], ends], axis=1)
        tokens = tokens.to_tensor(self.PAD, shape=(len(text), self.max_length))
        return tokens.numpy()
    
    def detokenize(self, tokens, human_readable=True):                
        words = self.tokenizer.detokenize(tokens, human_readable=human_readable)
        if (human_readable==True):
            return [" ".join(w) for w in words]
        text = tf.strings.reduce_join(words, separator=' ', axis=-1)
        return text
    
    def top_k(self, predictions, positions, k=10):
        top = []
        for p, m in zip(predictions, positions):
            top_k = self.detokenize([tf.argsort(p[m])[-k:][::-1]], False).numpy()[0].decode('utf8').split()
            top.append(top_k)
        return top

    def decode_emo(self, predictions):
        emo = tf.argmax(predictions, axis=-1)
        return [self.emo_labels[i] for i in emo]