Upload TFBilma

config.json

@@ -1,5 +1,4 @@
 {
-  "_name_or_path": "bilma_AR",
   "add_head": null,
   "architectures": [
     "Bilma"

configuration_bilma.py

@@ -6,7 +6,9 @@ class BilmaConfig(PretrainedConfig):
     def __init__(
         self,
         weights="AR",
-        include_top=True,
+        include_top = True,
+        add_head = None,
+        pooling = None,
         num_attention_heads: int = 4,
         num_hidden_layers: int = 2,
         seq_max_length: int = 280,
@@ -16,11 +18,20 @@ class BilmaConfig(PretrainedConfig):
         **kwargs,
     ):
         countries = ["AR"]
+        poolings = ["mean", "cls", "max"]
         if weights not in countries:
             raise ValueError(f"`weights` must be one of {countries}, got {weights}.")
+        if add_head is not None and include_top == True:
+            raise ValueError(f"To add a head, 'include_top' must be False")
+        if pooling is not None and include_top == True:
+            raise ValueError(f"To specify a pooling, 'include_top' must be False")
+        if pooling is not None and pooling not in poolings:
+            raise ValueError(f"`pooling` must be one of {poolings}, got {pooling}.")
         if weights is not None:
             self.weights = weights
             self.include_top = include_top
+            self.add_head = add_head
+            self.pooling = pooling
             self.num_attention_heads = 4
             self.num_hidden_layers = 2
             self.seq_max_length = 280
@@ -32,6 +43,8 @@ class BilmaConfig(PretrainedConfig):
 
         self.weights = weights
         self.include_top = include_top
+        self.add_head = add_head
+        self.pooling = pooling
         self.num_attention_heads = num_attention_heads
         self.num_hidden_layers = num_hidden_layers
         self.seq_max_length = seq_max_length

modeling_bilma.py

@@ -1,4 +1,5 @@
-from transformers import TFPreTrainedModel, PreTrainedTokenizer
+from transformers import TFPreTrainedModel, PreTrainedTokenizer, BatchEncoding
+
 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
@@ -38,6 +39,7 @@ class TFBilma(TFPreTrainedModel):
     def __init__(self, config):
         self.seq_max_length = config.seq_max_length
         self.include_top = config.include_top
+        self.add_head = config.add_head
         super().__init__(config)
 
         self.model = bilma(num_enc=config.num_hidden_layers,
@@ -47,7 +49,9 @@ class TFBilma(TFPreTrainedModel):
                            ff_dim=config.hidden_size,
                            vocab_size=config.vocab_size,
                            rate=config.hidden_dropout_prob,
-                           include_top = config.include_top)
+                           include_top = config.include_top,
+                           add_head = config.add_head,
+                           pooling = config.pooling)
             
     @property
     def dummy_inputs(self) -> Dict[str, tf.Tensor]:
@@ -70,13 +74,26 @@ class TFBilma(TFPreTrainedModel):
     
     
     def call(self, inputs):        
-        ins = tf.cast(inputs["input_ids"], tf.float32)
+        if isinstance(inputs, Dict) or isinstance(inputs, BatchEncoding):
+            ins = tf.cast(inputs["input_ids"], tf.float32)
+        else:
+            ins = inputs
         if self.include_top:
             output = {"logits":self.model(ins)}
         else:
-            output = {"last_hidden_state":self.model(ins)}
+            if self.add_head is None:
+                output = {"last_hidden_state":self.model(ins)}
+            else:
+                output = {"label":self.model(ins)}
         return output
     
+    def get_loss_function():
+        return loss_funtion()
+    
+    def get_acc_function():
+        return accuracy_function()
+    
+    
 # copied from bilma_model.py
 # --------------------------
 
@@ -105,7 +122,40 @@ def accuracy_function(ignore_id=0):
         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, include_top=True):
+def mean_vectors(inputs, enc_vectors, max_length):
+    p = tf.where(inputs == 3)
+    pos = tf.transpose(p)[1]
+    C = tf.sequence_mask(pos, maxlen=max_length, dtype=tf.float32)
+    C = tf.reshape(C, (-1, max_length, 1))
+    S = tf.reduce_sum(enc_vectors * C, 1)
+    x = S / tf.expand_dims(tf.cast(pos, tf.float32), (1))
+    return x
+
+def mean_diff_vectors(inputs, enc_vectors, max_length):
+    p = tf.where(inputs == 3)
+    pos = tf.transpose(p)[1]
+    C = tf.sequence_mask(pos, maxlen=max_length, dtype=tf.float32)
+    C = tf.reshape(C, (-1, max_length, 1))
+    vecs = enc_vectors * C
+    S = tf.reduce_sum(vecs, 1)
+    mu = S / tf.expand_dims(tf.cast(pos, tf.float32), (1))
+    x = tf.reduce_sum(mu - vecs, 1) / tf.expand_dims(tf.cast(pos, tf.float32), (1))
+    return x
+
+def max_vectors(inputs, enc_vectors, max_length):
+    p = tf.where(inputs == 3)
+    pos = tf.transpose(p)[1]
+    C = tf.sequence_mask(pos, maxlen=max_length, dtype=tf.float32)
+    C = tf.reshape(C, (-1, max_length, 1))
+    x = tf.reduce_max(enc_vectors * C, 1)
+    return x
+
+def cls_vectors(inputs, enc_vectors, max_length):
+    x = tf.squeeze(enc_vectors[:, 0:1, :], axis=1)
+    return x
+
+
+def bilma(num_enc=6, embed_dim=300, max_length=50, num_heads=6, ff_dim=512, vocab_size=9739, rate=0.1, include_top=True, add_head=None, pooling=None):
     capt_inputs_ids = Input(shape=(max_length, ), name='input_ids')
     capt_embedding = Embedding(vocab_size, embed_dim, mask_zero=False, name="bilma/embedding")
     capt_inputs = capt_embedding(capt_inputs_ids)
@@ -115,9 +165,22 @@ def bilma(num_enc=6, embed_dim=300, max_length=50, num_heads=6, ff_dim=512, voca
     if include_top:
         fin_output = Dense(vocab_size, use_bias=True, name="bilma/dense_final")(enc_output)
     else:
-        fin_output = enc_output
+        x = enc_output
+        if pooling == "mean":
+            x = mean_vectors(capt_inputs_ids, x, max_length)
+        elif pooling == "cls":                
+            x = cls_vectors(capt_inputs_ids, x, max_length)
+        elif pooling == "max":
+            x = max_vectors(capt_inputs_ids, x, max_length)
+        
+        if add_head is None:
+            fin_output = x
+        else:            
+            for i, m in enumerate(add_head[:-1]):
+                x = Dense(m, use_bias=True, activation="relu", name=f"bilma/dense_ex_{i}")(x)
+            fin_output = Dense(add_head[-1], use_bias=True, activation="softmax", name=f"bilma/dense_ex_final")(x)
     
-    caption_model = Model(inputs=capt_inputs_ids, outputs=[fin_output], name="bilma_model")
+    caption_model = Model(inputs=capt_inputs_ids, outputs=fin_output, name="bilma_model")
     return caption_model
 
 def load(model_file):
@@ -132,7 +195,6 @@ def load(model_file):
 # 
 # Copied from transformer_text.py   
 # -------------------------------
-
 class EncoderBlock(Layer):
     def __init__(self, layer_num, patch_dim, num_heads, ff_dim, rate=0.1, **kwargs):
         super(EncoderBlock, self).__init__(**kwargs)
@@ -214,7 +276,6 @@ class DecoderBlock(Layer):
         
         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)

tf_model.h5

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6f092e461b986ede156fd67fcfe8b28a7f360bbadb52b6dcefd19b913575865c
+oid sha256:e8ca072bb7a05a66b756aea19821a7003fe14556d2331b00418c187ad4e14701
 size 156875820