BGE-M3 모델 Dense, LayerNorm만으로 처음부터 구현하기

class BGEM3TensorFlow(tf.keras.Model):

def __init__(self,...):

   # Word embeddings
  self.word_embedding = tf.keras.layers.Embedding( 
		  input_dim=250002,
      output_dim=1024,
	)
  
  # Position embeddings
  self.position_embedding = tf.keras.layers.Embedding(
      input_dim=8194,
      output_dim=1024,
  )
  
  # Token type embeddings
  self.token_type_embedding = tf.keras.layers.Embedding(
      input_dim=1,
      output_dim=1024,
  )
  
  # Layer normalization and dropout
  self.layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5)
  #self.dropout = layers.Dropout(rate=0.1)

def call(self, ..)

	self.inputs_embeds = tf.gather(params=self.weight, indices=input_ids)
	self.position_embeds = tf.gather(params=self.position_embeddings, indices=position_ids)
	self.token_type_embeds = tf.gather(params=self.token_type_embeddings, indices=token_type_ids)

	embedding_output = inputs_embeds + position_embeds + token_type_embeds
	embedding_output = self.layerNorm(embedding_output)

class TransformerBlock(tf.keras.layers.Layer):

# Query, Key, Value 가중치
encoder.layer.0.attention.self.query | shape: torch.Size([1024, 1024])
encoder.layer.0.attention.self.key | shape: torch.Size([1024, 1024])
encoder.layer.0.attention.self.value| shape: torch.Size([1024, 1024])
# Attention 출력 처리
encoder.layer.0.attention.output.dense | shape: torch.Size([1024, 1024])
encoder.layer.0.attention.output.LayerNorm | shape: torch.Size([1024])

# 중간 레이어
encoder.layer.0.intermediate.dense | shape: torch.Size([4096, 1024]) # 확장
encoder.layer.0.output.dense | shape: torch.Size([1024, 4096]) # 축소

# 최종 정규화
encoder.layer.0.output.LayerNorm | shape: torch.Size([1024])

def __init__(self, ...):
    
    self.wq = tf.keras.layers.Dense(1024)
    self.wk = tf.keras.layers.Dense(1024)
    self.wv = tf.keras.layers.Dense(1024)
    self.dense = tf.keras.layers.Dense(1024)
    self.attlayerNorm = tf.keras.layers.LayerNormalization(epsilon=1e-5)
		
    self.intermediate = tf.keras.layers.Dense(4096)
    self.output_dense = tf.keras.layers.Dense(1024)
    self.output_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5)

def call(self, ..)

	inputs = embedding_output

	# Query, Key, Value를 계산
	q = self.wq(inputs)  # (batch_size, seq_len, d_model)
	k = self.wk(inputs)  # (batch_size, seq_len, d_model)
	v = self.wv(inputs)  # (batch_size, seq_len, d_model)
	
	# 다중 헤드로 분리
	q = self.split_heads(q, batch_size, 16, 64)  # (batch_size, num_heads, seq_len_q, depth)
	k = self.split_heads(k, batch_size, 16, 64)  # (batch_size, num_heads, seq_len_k, depth)
	v = self.split_heads(v, batch_size, 16, 64)  # (batch_size, num_heads, seq_len_v, depth)
	
	def split_heads(self, x, batch_size, num_heads, depth):
	  x = tf.reshape(x, (batch_size, -1, num_heads, depth))
	  return tf.transpose(x, perm=[0, 2, 1, 3])  # (batch_size, num_heads, seq_len, depth)

# Scaled Dot-Product Attention Score
dk = tf.cast(math.sqrt(1024 // 16), tf.float32)
attention_scores = tf.matmul(q, k, transpose_b=True)  # (batch_size, num_heads, seq_len_q, seq_len_k)
attention_scores = tf.divide(attention_scores, dk)

# Attention mask 
attention_scores = tf.add(attention_scores, mask)
attention_probs = tf.nn.softmax(attention_scores + 1e-9, axis=-1)

# Attention 결과 계산
attention_output = tf.matmul(attention_probs, v)  # (batch_size, num_heads, seq_len_q, depth)
attention_output = tf.transpose(attention_output,
                                perm=[0, 2, 1, 3])  # (batch_size, seq_len_q, num_heads, depth)
attention_output = tf.reshape(attention_output,
                              (batch_size, -1, self.d_model))  # (batch_size, seq_len_q, d_model)
                              
# Attention 최종 Dense Layer를 통과
attention_output = self.dense(attention_output)  # (batch_size, seq_len_q, d_model)
# 첫번째 Residual 연산
attention_output = self.attlayerNorm(inputs=attention_output + inputs)

# intermediate layer
intermediate_output = self.intermediate(attention_output)
intermediate_output = self.gelu_approx(intermediate_output)  # Use exact GELU approximation
layer_output = self.output_dense(intermediate_output)
layer_output = self.output_dropout(layer_output, training=training)

# 두번째 Residual 연산
output = layer_output + attention_output
output = self.output_norm(output)

# GELU Approximation
def gelu_approx(self, x):
    x = tf.convert_to_tensor(x)
    cdf = 0.5 * (1.0 + tf.math.erf(x / tf.cast(tf.sqrt(2.0), x.dtype)))

    return x * cdf

encoder_layers = []
for i in range(self.num_layers):
    layer = TransformerBlock(
        d_model=1024,
        num_heads=16,
        intermediate_size=4096,
        name=f"encoder.layer.{i}"
    )
    encoder_layers.append(layer)

pooler.dense.weight            | shape: torch.Size([1024, 1024])
pooler.dense.bias              | shape: torch.Size([1024])

def call(self, input_ids, ..)
	
	inputs_embeds = tf.gather(params=self.weight, indices=input_ids)
	input_shape = shape_list(inputs_embeds)[:-1]
	
	position_ids = self.create_position_ids_from_input_ids(input_ids=input_ids))
	position_embeds = tf.gather(params=self.position_embeddings, indices=position_ids)

	token_type_ids = tf.fill(dims=input_shape, value=0)
	token_type_embeds = tf.gather(params=self.token_type_embeddings, indices=token_type_ids)

	embedding_output = inputs_embeds + position_embeds + token_type_embeds
	embedding_output = self.layerNorm(embedding_output)
	
  def create_position_ids_from_input_ids(self, input_ids, past_key_values_length=0, padding_idx=1):
    mask = tf.cast(tf.math.not_equal(input_ids, padding_idx), dtype=input_ids.dtype)
    incremental_indices = (tf.math.cumsum(mask, axis=1) + past_key_values_length) * mask
    return incremental_indices + padding_idx

attention_mask_shape = shape_list(attention_mask)
extended_attention_mask = tf.reshape(
    attention_mask, (attention_mask_shape[0], 1, 1, attention_mask_shape[1])
)

extended_attention_mask = tf.cast(extended_attention_mask, dtype=embedding_output.dtype)
one_cst = tf.constant(1.0, dtype=embedding_output.dtype)
ten_thousand_cst = tf.constant(-10000.0, dtype=embedding_output.dtype)

# embedding_output = inputs_embeds + position_embeds + token_type_embeds

hidden_states = embedding_output
# Pass through encoder layers
for layer in self.encoder_layers:
    hidden_states = layer(
        hidden_states,
        attention_mask=attention_mask,
    )

pooled_output = hidden_states[:, 0, :] # (batch, seq_len, hidden_dim)

self.colbert_linear = tf.keras.layers.Dense(
	  units=self.d_model,
)


colbert_vecs = self.colbert_linear(hidden_states[:, 1:])
colbert_vecs = colbert_vecs * tf.cast(attention_mask_origin[:, 1:][:, :, None], dtype=tf.float32)

model_path = "./bge-m3/colbert_linear.pt"

colbert_model = torch.load(model_path, map_location=device, weights_only=True)

colbert_weights = colbert_model['weight']
colbert_bias = colbert_model['bias']


tf_model.colbert_linear.set_weights([
    colbert_weights.numpy().T,
    colbert_bias.numpy()
])

  outputs = {
      "dense_vecs": pooled_output, #[batch, hidden_size]
      "colbert_vecs" : colbert_vecs, #[hidden, seq_len, hidden_size]
  }
return outputs

def save_model_with_tokenizer(model, tokenizer, save_path):
    """Save both model and tokenizer"""
    os.makedirs(save_path, exist_ok=True)
    model_save_path = os.path.join(save_path, 'model')

    # Ensure model is built by calling it with dummy inputs
    dummy_inputs = {
        'input_ids': tf.zeros((2, 11), dtype=tf.int32),
        'attention_mask': tf.ones((2, 11), dtype=tf.int32)
    }
    _ = model(dummy_inputs, training=False, output_hidden_states=True)

    # Define serving signature
    @tf.function(input_signature=[
        tf.TensorSpec(shape=[None, None], dtype=tf.int32, name='input_ids'),
        tf.TensorSpec(shape=[None, None], dtype=tf.int32, name='attention_mask')
    ])

    def serving_fn(input_ids, attention_mask):
    
        inputs = {
            'input_ids': input_ids,
            'attention_mask': attention_mask
        }

      outputs = model(inputs=inputs, training=False)

        return {
            'dense_vecs': outputs['dense_vecs']
            'colbert_vecs': outputs['colbert_vecs']
        }


    # Save model
    tf.saved_model.save(
        model,
        model_save_path,
        signatures={'serving_default': serving_fn}
    )

    # Save tokenizer
    tokenizer.save_pretrained(save_path)

    return model_save_path