Add files using upload-large-folder tool

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

.ipynb_checkpoints/README-checkpoint.md

@@ -0,0 +1,77 @@
+---
+base_model: []
+library_name: transformers
+tags:
+- mergekit
+- merge
+
+---
+# output_merge_multiSFT_v6_1
+
+This is a merge of pre-trained language models created using [mergekit](https://github.com/cg123/mergekit).
+
+## Merge Details
+### Merge Method
+
+This model was merged using the [SLERP](https://en.wikipedia.org/wiki/Slerp) merge method.
+
+### Models Merged
+
+The following models were included in the merge:
+* /workspace/storage-shared/cnm/embedding/code/llama_1B_legal_retrieval_v2_2.3M_clean_150726_finetune2_07150726/checkpoint-8000
+* /workspace/storage-shared/cnm/huyhq21_v2/mergekit/output_slerp_1B_data_v2_clean_v6
+
+### Configuration
+
+The following YAML configuration was used to produce this model:
+
+```yaml
+# slices:
+#   - sources:
+#       - model: psmathur/orca_mini_v3_13b
+#         layer_range: [0, 40]
+#       - model: garage-bAInd/Platypus2-13B
+#         layer_range: [0, 40]
+# # or, the equivalent models: syntax:
+# # models:
+# #   - model: psmathur/orca_mini_v3_13b
+# #   - model: garage-bAInd/Platypus2-13B
+# merge_method: slerp
+# base_model: psmathur/orca_mini_v3_13b
+# parameters:
+#   t:
+#     - filter: self_attn
+#       value: [0, 0.5, 0.3, 0.7, 1]
+#     - filter: mlp
+#       value: [1, 0.5, 0.7, 0.3, 0]
+#     - value: 0.5 # fallback for rest of tensors
+# dtype: float16
+# slices:
+models:
+      - model: /workspace/storage-shared/cnm/embedding/code/llama_1B_legal_retrieval_v2_2.3M_clean_150726_finetune2_07150726/checkpoint-8000 # Mô hình A (Base gốc)
+        # layer_range: [0, 36]
+      - model: /workspace/storage-shared/cnm/huyhq21_v2/mergekit/output_slerp_1B_data_v2_clean_v6   # Mô hình B (Đã Finetune)
+        # layer_range: [0, 36]
+merge_method: slerp
+base_model: /workspace/storage-shared/cnm/huyhq21_v2/mergekit/output_slerp_1B_data_v2_clean_v6 # Mô hình B (Đã Finetune)
+parameters:
+  t: 
+    #- filter: self_attn
+    - filter: layers
+      value: [0.2, 0.3, 0.5, 0.6, 0.7, 0.8]
+    #- filter: mlp
+     # value: [0.2, 0.3, 0.5, 0.6, 0.7, 0.8]
+    - value: 0.5 # fallback cho layernorm
+dtype: bfloat16
+
+# models:
+#   - model: /workspace/storage-shared/cnm/cuongnq23/project_embeding/models/Qwen/Qwen3-Embedding-8B
+#   - model: /workspace/storage-shared/cnm/embedding/code/qwen8b_SFT_kalm_vnlaw_3M/checkpoint-20000
+# merge_method: slerp
+# base_model: /workspace/storage-shared/cnm/cuongnq23/project_embeding/models/Qwen/Qwen3-Embedding-8B
+# parameters:
+#   t:
+#     # Trượt dốc từ 10% (đáy) lên 90% (đỉnh)
+#     - value: [0.1, 0.3, 0.5, 0.7, 0.9] 
+# dtype: bfloat16
+```

.ipynb_checkpoints/config-checkpoint.json

@@ -0,0 +1,43 @@
+{
+  "_name_or_path": "nvidia/llama-3.2-nv-embedqa-1b-v2",
+  "architectures": [
+    "LlamaBidirectionalModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "llama_bidirectional_model.LlamaBidirectionalConfig",
+    "AutoModel": "llama_bidirectional_model.LlamaBidirectionalModel"
+  },
+  "bos_token_id": 128000,
+  "eos_token_id": 128001,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intermediate_size": 8192,
+  "max_position_embeddings": 131072,
+  "mlp_bias": false,
+  "model_type": "llama_bidirec",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 16,
+  "num_key_value_heads": 8,
+  "pooling": "avg",
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": {
+    "factor": 32.0,
+    "high_freq_factor": 4.0,
+    "low_freq_factor": 1.0,
+    "original_max_position_embeddings": 8192,
+    "rope_type": "llama3",
+    "rope_theta": 500000.0
+  },
+  "rope_theta": 500000.0,
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.44.2",
+  "use_cache": true,
+  "vocab_size": 128256,
+  "use_bidirectional_attention": true
+}

.ipynb_checkpoints/llama_bidirectional_model-checkpoint.py

@@ -0,0 +1,226 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0.
+"""
+Bidirectional Llama model for embedding tasks.
+
+This module provides a modified LlamaModel that uses bidirectional (non-causal)
+attention, suitable for generating embeddings where each token should attend
+to all other tokens in the sequence.
+
+Supports transformers version 4.44 and above with a unified forward() implementation.
+
+Version compatibility notes:
+    - transformers 4.47: Setting _attn_implementation in __init__ had no effect due to
+      attention initialization order
+    - transformers 4.48+: Attention refactor (transformers#35235) activated the
+      _attn_implementation setting, which defaulted to "eager" instead of "sdpa"
+    - transformers < 4.53: LlamaModel has _update_causal_mask method that can be overridden
+    - transformers 4.53+: _update_causal_mask removed; masking moved to masking_utils module,
+      necessitating a full forward() override for custom attention masks
+    - transformers < 4.54: Decoder layer returns tuple, uses past_key_value (singular)
+    - transformers 4.54-4.55: Decoder layer returns tensor, uses past_key_value (singular)
+    - transformers 4.56+: Decoder layer returns tensor, uses past_key_values (plural),
+      DynamicCache accepts config parameter
+    - transformers 5.0+: Has native create_bidirectional_mask in masking_utils
+"""
+
+import inspect
+
+import torch
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from transformers.models.llama.configuration_llama import LlamaConfig
+from transformers.models.llama.modeling_llama import LlamaDecoderLayer, LlamaModel
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+# Check if native create_bidirectional_mask exists (transformers >= 5.0)
+try:
+    from transformers.masking_utils import create_bidirectional_mask
+
+    _HAS_NATIVE_BIDIRECTIONAL_MASK = True
+except ImportError:
+    from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
+
+    _HAS_NATIVE_BIDIRECTIONAL_MASK = False
+
+# Detect API differences via introspection
+_decoder_forward_params = inspect.signature(LlamaDecoderLayer.forward).parameters
+_dynamic_cache_init_params = inspect.signature(DynamicCache.__init__).parameters
+
+# past_key_value (singular) in < 4.56, past_key_values (plural) in >= 4.56
+_USE_PLURAL_CACHE_PARAM = "past_key_values" in _decoder_forward_params
+# DynamicCache accepts config parameter in >= 4.56
+_DYNAMIC_CACHE_ACCEPTS_CONFIG = "config" in _dynamic_cache_init_params
+
+
+class LlamaBidirectionalConfig(LlamaConfig):
+    """Configuration for LlamaBidirectionalModel with pooling and temperature settings."""
+
+    model_type = "llama_bidirec"
+
+    def __init__(
+        self, pooling: str = "avg", temperature: float = 1.0, **kwargs
+    ) -> None:
+        """
+        Initialize bidirectional Llama configuration.
+
+        Args:
+            pooling: Pooling strategy for embeddings ("avg", "cls", "last", etc.)
+            temperature: Temperature scaling for embeddings
+            **kwargs: Additional arguments passed to LlamaConfig
+        """
+        self.pooling = pooling
+        self.temperature = temperature
+        super().__init__(**kwargs)
+
+
+class LlamaBidirectionalModel(LlamaModel):
+    """
+    LlamaModel modified to use bidirectional (non-causal) attention.
+
+    In standard Llama, each token can only attend to previous tokens (causal attention).
+    This model removes that restriction, allowing each token to attend to all tokens
+    in the sequence, which is useful for embedding tasks.
+
+    The key modifications are:
+        1. Setting is_causal=False on all attention layers
+        2. Using a bidirectional attention mask instead of causal mask
+    """
+
+    config_class = LlamaBidirectionalConfig
+
+    def __init__(self, config: LlamaConfig) -> None:
+        super().__init__(config)
+        for layer in self.layers:
+            layer.self_attn.is_causal = False
+
+    def _create_bidirectional_mask(
+        self,
+        input_embeds: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+    ) -> torch.Tensor | None:
+        """
+        Create bidirectional attention mask.
+
+        Args:
+            input_embeds: Input embeddings tensor of shape (batch_size, seq_len, hidden_size)
+            attention_mask: Optional 2D attention mask of shape (batch_size, seq_len)
+                where 1 indicates tokens to attend to and 0 indicates masked tokens
+
+        Returns:
+            4D attention mask suitable for the attention implementation, or None
+            if no masking is needed
+        """
+        if attention_mask is None:
+            return None
+
+        if _HAS_NATIVE_BIDIRECTIONAL_MASK:
+            return create_bidirectional_mask(
+                config=self.config,
+                input_embeds=input_embeds,
+                attention_mask=attention_mask,
+            )
+
+        # Fallback for transformers < 5.0 without create_bidirectional_mask
+
+        # Flash attention handles 2D masks internally; only pass mask if there
+        # are actually masked tokens (zeros), otherwise return None for efficiency
+        if getattr(self.config, "_attn_implementation", None) == "flash_attention_2":
+            has_masked_tokens = (attention_mask == 0).any()
+            return attention_mask if has_masked_tokens else None
+
+        return _prepare_4d_attention_mask(attention_mask, input_embeds.dtype)
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        cache_position: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        **kwargs,
+    ) -> BaseModelOutputWithPast:
+        """
+        Forward pass with bidirectional attention.
+
+        Args:
+            input_ids: Input token IDs of shape (batch_size, seq_len)
+            attention_mask: Attention mask of shape (batch_size, seq_len)
+            position_ids: Position IDs for rotary embeddings
+            past_key_values: Cached key/value states for incremental decoding
+            inputs_embeds: Pre-computed input embeddings (alternative to input_ids)
+            cache_position: Position indices for cache updates
+            use_cache: Whether to return cached key/value states
+            **kwargs: Additional arguments passed to decoder layers
+
+        Returns:
+            BaseModelOutputWithPast containing last_hidden_state and past_key_values
+        """
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You must specify exactly one of input_ids or inputs_embeds"
+            )
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        # Initialize cache if needed
+        if use_cache and past_key_values is None:
+            if _DYNAMIC_CACHE_ACCEPTS_CONFIG:
+                past_key_values = DynamicCache(config=self.config)
+            else:
+                past_key_values = DynamicCache()
+
+        if cache_position is None:
+            past_seen_tokens = (
+                past_key_values.get_seq_length() if past_key_values is not None else 0
+            )
+            cache_position = torch.arange(
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        bidirectional_mask = self._create_bidirectional_mask(
+            inputs_embeds, attention_mask
+        )
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # Build decoder layer kwargs with correct cache parameter name
+        # (past_key_value in < 4.56, past_key_values in >= 4.56)
+        layer_kwargs = {
+            "attention_mask": bidirectional_mask,
+            "position_ids": position_ids,
+            "use_cache": use_cache,
+            "cache_position": cache_position,
+            "position_embeddings": position_embeddings,
+        }
+        if _USE_PLURAL_CACHE_PARAM:
+            layer_kwargs["past_key_values"] = past_key_values
+        else:
+            layer_kwargs["past_key_value"] = past_key_values
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            layer_outputs = decoder_layer(hidden_states, **layer_kwargs)
+
+            # Decoder returns tuple in < 4.54, tensor in >= 4.54
+            if isinstance(layer_outputs, tuple):
+                hidden_states = layer_outputs[0]
+            else:
+                hidden_states = layer_outputs
+
+        hidden_states = self.norm(hidden_states)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )

.ipynb_checkpoints/mergekit_config-checkpoint.yml

@@ -0,0 +1,48 @@
+# slices:
+#   - sources:
+#       - model: psmathur/orca_mini_v3_13b
+#         layer_range: [0, 40]
+#       - model: garage-bAInd/Platypus2-13B
+#         layer_range: [0, 40]
+# # or, the equivalent models: syntax:
+# # models:
+# #   - model: psmathur/orca_mini_v3_13b
+# #   - model: garage-bAInd/Platypus2-13B
+# merge_method: slerp
+# base_model: psmathur/orca_mini_v3_13b
+# parameters:
+#   t:
+#     - filter: self_attn
+#       value: [0, 0.5, 0.3, 0.7, 1]
+#     - filter: mlp
+#       value: [1, 0.5, 0.7, 0.3, 0]
+#     - value: 0.5 # fallback for rest of tensors
+# dtype: float16
+# slices:
+models:
+      - model: /workspace/storage-shared/cnm/embedding/code/llama_1B_legal_retrieval_v2_2.3M_clean_150726_finetune2_07150726/checkpoint-8000 # Mô hình A (Base gốc)
+        # layer_range: [0, 36]
+      - model: /workspace/storage-shared/cnm/huyhq21_v2/mergekit/output_slerp_1B_data_v2_clean_v6   # Mô hình B (Đã Finetune)
+        # layer_range: [0, 36]
+merge_method: slerp
+base_model: /workspace/storage-shared/cnm/huyhq21_v2/mergekit/output_slerp_1B_data_v2_clean_v6 # Mô hình B (Đã Finetune)
+parameters:
+  t: 
+    #- filter: self_attn
+    - filter: layers
+      value: [0.2, 0.3, 0.5, 0.6, 0.7, 0.8]
+    #- filter: mlp
+     # value: [0.2, 0.3, 0.5, 0.6, 0.7, 0.8]
+    - value: 0.5 # fallback cho layernorm
+dtype: bfloat16
+
+# models:
+#   - model: /workspace/storage-shared/cnm/cuongnq23/project_embeding/models/Qwen/Qwen3-Embedding-8B
+#   - model: /workspace/storage-shared/cnm/embedding/code/qwen8b_SFT_kalm_vnlaw_3M/checkpoint-20000
+# merge_method: slerp
+# base_model: /workspace/storage-shared/cnm/cuongnq23/project_embeding/models/Qwen/Qwen3-Embedding-8B
+# parameters:
+#   t:
+#     # Trượt dốc từ 10% (đáy) lên 90% (đỉnh)
+#     - value: [0.1, 0.3, 0.5, 0.7, 0.9] 
+# dtype: bfloat16

1_Pooling/.ipynb_checkpoints/config-checkpoint.json

@@ -0,0 +1,10 @@
+{
+    "word_embedding_dimension": 2048,
+    "pooling_mode_cls_token": false,
+    "pooling_mode_mean_tokens": true,
+    "pooling_mode_max_tokens": false,
+    "pooling_mode_mean_sqrt_len_tokens": false,
+    "pooling_mode_weightedmean_tokens": false,
+    "pooling_mode_lasttoken": false,
+    "include_prompt": true
+}

1_Pooling/config.json

@@ -0,0 +1,10 @@
+{
+    "word_embedding_dimension": 2048,
+    "pooling_mode_cls_token": false,
+    "pooling_mode_mean_tokens": true,
+    "pooling_mode_max_tokens": false,
+    "pooling_mode_mean_sqrt_len_tokens": false,
+    "pooling_mode_weightedmean_tokens": false,
+    "pooling_mode_lasttoken": false,
+    "include_prompt": true
+}

README.md

@@ -0,0 +1,77 @@
+---
+base_model: []
+library_name: transformers
+tags:
+- mergekit
+- merge
+
+---
+# output_merge_multiSFT_v6_1
+
+This is a merge of pre-trained language models created using [mergekit](https://github.com/cg123/mergekit).
+
+## Merge Details
+### Merge Method
+
+This model was merged using the [SLERP](https://en.wikipedia.org/wiki/Slerp) merge method.
+
+### Models Merged
+
+The following models were included in the merge:
+* /workspace/storage-shared/cnm/embedding/code/llama_1B_legal_retrieval_v2_2.3M_clean_150726_finetune2_07150726/checkpoint-8000
+* /workspace/storage-shared/cnm/huyhq21_v2/mergekit/output_slerp_1B_data_v2_clean_v6
+
+### Configuration
+
+The following YAML configuration was used to produce this model:
+
+```yaml
+# slices:
+#   - sources:
+#       - model: psmathur/orca_mini_v3_13b
+#         layer_range: [0, 40]
+#       - model: garage-bAInd/Platypus2-13B
+#         layer_range: [0, 40]
+# # or, the equivalent models: syntax:
+# # models:
+# #   - model: psmathur/orca_mini_v3_13b
+# #   - model: garage-bAInd/Platypus2-13B
+# merge_method: slerp
+# base_model: psmathur/orca_mini_v3_13b
+# parameters:
+#   t:
+#     - filter: self_attn
+#       value: [0, 0.5, 0.3, 0.7, 1]
+#     - filter: mlp
+#       value: [1, 0.5, 0.7, 0.3, 0]
+#     - value: 0.5 # fallback for rest of tensors
+# dtype: float16
+# slices:
+models:
+      - model: /workspace/storage-shared/cnm/embedding/code/llama_1B_legal_retrieval_v2_2.3M_clean_150726_finetune2_07150726/checkpoint-8000 # Mô hình A (Base gốc)
+        # layer_range: [0, 36]
+      - model: /workspace/storage-shared/cnm/huyhq21_v2/mergekit/output_slerp_1B_data_v2_clean_v6   # Mô hình B (Đã Finetune)
+        # layer_range: [0, 36]
+merge_method: slerp
+base_model: /workspace/storage-shared/cnm/huyhq21_v2/mergekit/output_slerp_1B_data_v2_clean_v6 # Mô hình B (Đã Finetune)
+parameters:
+  t: 
+    #- filter: self_attn
+    - filter: layers
+      value: [0.2, 0.3, 0.5, 0.6, 0.7, 0.8]
+    #- filter: mlp
+     # value: [0.2, 0.3, 0.5, 0.6, 0.7, 0.8]
+    - value: 0.5 # fallback cho layernorm
+dtype: bfloat16
+
+# models:
+#   - model: /workspace/storage-shared/cnm/cuongnq23/project_embeding/models/Qwen/Qwen3-Embedding-8B
+#   - model: /workspace/storage-shared/cnm/embedding/code/qwen8b_SFT_kalm_vnlaw_3M/checkpoint-20000
+# merge_method: slerp
+# base_model: /workspace/storage-shared/cnm/cuongnq23/project_embeding/models/Qwen/Qwen3-Embedding-8B
+# parameters:
+#   t:
+#     # Trượt dốc từ 10% (đáy) lên 90% (đỉnh)
+#     - value: [0.1, 0.3, 0.5, 0.7, 0.9] 
+# dtype: bfloat16
+```

config.json

@@ -0,0 +1,47 @@
+{
+  "architectures": [
+    "LlamaBidirectionalModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "llama_bidirectional_model.LlamaBidirectionalConfig",
+    "AutoModel": "llama_bidirectional_model.LlamaBidirectionalModel"
+  },
+  "bos_token_id": 128000,
+  "dtype": "bfloat16",
+  "eos_token_id": 128001,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intermediate_size": 8192,
+  "max_position_embeddings": 131072,
+  "mlp_bias": false,
+  "model_type": "llama_bidirec",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 16,
+  "num_key_value_heads": 8,
+  "pad_token_id": null,
+  "pooling": "avg",
+  "pretraining_tp": 1,
+  "quantization_config": {
+    "quant_method": "quanto",
+    "weights": "int8"
+  },
+  "rms_norm_eps": 1e-05,
+  "rope_parameters": {
+    "factor": 32.0,
+    "high_freq_factor": 4.0,
+    "low_freq_factor": 1.0,
+    "original_max_position_embeddings": 8192,
+    "rope_theta": 500000.0,
+    "rope_type": "llama3"
+  },
+  "temperature": 1.0,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.14.1",
+  "use_bidirectional_attention": true,
+  "use_cache": true,
+  "vocab_size": 128256
+}

config_1.json

@@ -0,0 +1,40 @@
+{
+  "architectures": [
+    "LlamaModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 128000,
+  "dtype": "bfloat16",
+  "eos_token_id": 128001,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intermediate_size": 8192,
+  "max_position_embeddings": 131072,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 16,
+  "num_key_value_heads": 8,
+  "pad_token_id": 128001,
+  "pooling": "avg",
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_parameters": {
+    "factor": 32.0,
+    "high_freq_factor": 4.0,
+    "low_freq_factor": 1.0,
+    "original_max_position_embeddings": 8192,
+    "rope_theta": 500000.0,
+    "rope_type": "llama3"
+  },
+  "rope_scaling": null,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": true,
+  "transformers_version": "4.56.0",
+  "use_bidirectional_attention": true,
+  "use_cache": false,
+  "vocab_size": 128256
+}

config_sentence_transformers.json

@@ -0,0 +1,14 @@
+{
+  "model_type": "SentenceTransformer",
+  "__version__": {
+    "sentence_transformers": "5.0.1",
+    "transformers": "4.47.1",
+    "pytorch": "2.9.1+cu126"
+  },
+  "prompts": {
+    "query": "query: ",
+    "document": "passage: "
+  },
+  "default_prompt_name": null,
+  "similarity_fn_name": "cosine"
+}

llama_bidirectional_model.py

@@ -0,0 +1,226 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0.
+"""
+Bidirectional Llama model for embedding tasks.
+
+This module provides a modified LlamaModel that uses bidirectional (non-causal)
+attention, suitable for generating embeddings where each token should attend
+to all other tokens in the sequence.
+
+Supports transformers version 4.44 and above with a unified forward() implementation.
+
+Version compatibility notes:
+    - transformers 4.47: Setting _attn_implementation in __init__ had no effect due to
+      attention initialization order
+    - transformers 4.48+: Attention refactor (transformers#35235) activated the
+      _attn_implementation setting, which defaulted to "eager" instead of "sdpa"
+    - transformers < 4.53: LlamaModel has _update_causal_mask method that can be overridden
+    - transformers 4.53+: _update_causal_mask removed; masking moved to masking_utils module,
+      necessitating a full forward() override for custom attention masks
+    - transformers < 4.54: Decoder layer returns tuple, uses past_key_value (singular)
+    - transformers 4.54-4.55: Decoder layer returns tensor, uses past_key_value (singular)
+    - transformers 4.56+: Decoder layer returns tensor, uses past_key_values (plural),
+      DynamicCache accepts config parameter
+    - transformers 5.0+: Has native create_bidirectional_mask in masking_utils
+"""
+
+import inspect
+
+import torch
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from transformers.models.llama.configuration_llama import LlamaConfig
+from transformers.models.llama.modeling_llama import LlamaDecoderLayer, LlamaModel
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+# Check if native create_bidirectional_mask exists (transformers >= 5.0)
+try:
+    from transformers.masking_utils import create_bidirectional_mask
+
+    _HAS_NATIVE_BIDIRECTIONAL_MASK = True
+except ImportError:
+    from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
+
+    _HAS_NATIVE_BIDIRECTIONAL_MASK = False
+
+# Detect API differences via introspection
+_decoder_forward_params = inspect.signature(LlamaDecoderLayer.forward).parameters
+_dynamic_cache_init_params = inspect.signature(DynamicCache.__init__).parameters
+
+# past_key_value (singular) in < 4.56, past_key_values (plural) in >= 4.56
+_USE_PLURAL_CACHE_PARAM = "past_key_values" in _decoder_forward_params
+# DynamicCache accepts config parameter in >= 4.56
+_DYNAMIC_CACHE_ACCEPTS_CONFIG = "config" in _dynamic_cache_init_params
+
+
+class LlamaBidirectionalConfig(LlamaConfig):
+    """Configuration for LlamaBidirectionalModel with pooling and temperature settings."""
+
+    model_type = "llama_bidirec"
+
+    def __init__(
+        self, pooling: str = "avg", temperature: float = 1.0, **kwargs
+    ) -> None:
+        """
+        Initialize bidirectional Llama configuration.
+
+        Args:
+            pooling: Pooling strategy for embeddings ("avg", "cls", "last", etc.)
+            temperature: Temperature scaling for embeddings
+            **kwargs: Additional arguments passed to LlamaConfig
+        """
+        self.pooling = pooling
+        self.temperature = temperature
+        super().__init__(**kwargs)
+
+
+class LlamaBidirectionalModel(LlamaModel):
+    """
+    LlamaModel modified to use bidirectional (non-causal) attention.
+
+    In standard Llama, each token can only attend to previous tokens (causal attention).
+    This model removes that restriction, allowing each token to attend to all tokens
+    in the sequence, which is useful for embedding tasks.
+
+    The key modifications are:
+        1. Setting is_causal=False on all attention layers
+        2. Using a bidirectional attention mask instead of causal mask
+    """
+
+    config_class = LlamaBidirectionalConfig
+
+    def __init__(self, config: LlamaConfig) -> None:
+        super().__init__(config)
+        for layer in self.layers:
+            layer.self_attn.is_causal = False
+
+    def _create_bidirectional_mask(
+        self,
+        input_embeds: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+    ) -> torch.Tensor | None:
+        """
+        Create bidirectional attention mask.
+
+        Args:
+            input_embeds: Input embeddings tensor of shape (batch_size, seq_len, hidden_size)
+            attention_mask: Optional 2D attention mask of shape (batch_size, seq_len)
+                where 1 indicates tokens to attend to and 0 indicates masked tokens
+
+        Returns:
+            4D attention mask suitable for the attention implementation, or None
+            if no masking is needed
+        """
+        if attention_mask is None:
+            return None
+
+        if _HAS_NATIVE_BIDIRECTIONAL_MASK:
+            return create_bidirectional_mask(
+                config=self.config,
+                input_embeds=input_embeds,
+                attention_mask=attention_mask,
+            )
+
+        # Fallback for transformers < 5.0 without create_bidirectional_mask
+
+        # Flash attention handles 2D masks internally; only pass mask if there
+        # are actually masked tokens (zeros), otherwise return None for efficiency
+        if getattr(self.config, "_attn_implementation", None) == "flash_attention_2":
+            has_masked_tokens = (attention_mask == 0).any()
+            return attention_mask if has_masked_tokens else None
+
+        return _prepare_4d_attention_mask(attention_mask, input_embeds.dtype)
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        cache_position: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        **kwargs,
+    ) -> BaseModelOutputWithPast:
+        """
+        Forward pass with bidirectional attention.
+
+        Args:
+            input_ids: Input token IDs of shape (batch_size, seq_len)
+            attention_mask: Attention mask of shape (batch_size, seq_len)
+            position_ids: Position IDs for rotary embeddings
+            past_key_values: Cached key/value states for incremental decoding
+            inputs_embeds: Pre-computed input embeddings (alternative to input_ids)
+            cache_position: Position indices for cache updates
+            use_cache: Whether to return cached key/value states
+            **kwargs: Additional arguments passed to decoder layers
+
+        Returns:
+            BaseModelOutputWithPast containing last_hidden_state and past_key_values
+        """
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You must specify exactly one of input_ids or inputs_embeds"
+            )
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        # Initialize cache if needed
+        if use_cache and past_key_values is None:
+            if _DYNAMIC_CACHE_ACCEPTS_CONFIG:
+                past_key_values = DynamicCache(config=self.config)
+            else:
+                past_key_values = DynamicCache()
+
+        if cache_position is None:
+            past_seen_tokens = (
+                past_key_values.get_seq_length() if past_key_values is not None else 0
+            )
+            cache_position = torch.arange(
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        bidirectional_mask = self._create_bidirectional_mask(
+            inputs_embeds, attention_mask
+        )
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # Build decoder layer kwargs with correct cache parameter name
+        # (past_key_value in < 4.56, past_key_values in >= 4.56)
+        layer_kwargs = {
+            "attention_mask": bidirectional_mask,
+            "position_ids": position_ids,
+            "use_cache": use_cache,
+            "cache_position": cache_position,
+            "position_embeddings": position_embeddings,
+        }
+        if _USE_PLURAL_CACHE_PARAM:
+            layer_kwargs["past_key_values"] = past_key_values
+        else:
+            layer_kwargs["past_key_value"] = past_key_values
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            layer_outputs = decoder_layer(hidden_states, **layer_kwargs)
+
+            # Decoder returns tuple in < 4.54, tensor in >= 4.54
+            if isinstance(layer_outputs, tuple):
+                hidden_states = layer_outputs[0]
+            else:
+                hidden_states = layer_outputs
+
+        hidden_states = self.norm(hidden_states)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )

mergekit_config.yml

@@ -0,0 +1,48 @@
+# slices:
+#   - sources:
+#       - model: psmathur/orca_mini_v3_13b
+#         layer_range: [0, 40]
+#       - model: garage-bAInd/Platypus2-13B
+#         layer_range: [0, 40]
+# # or, the equivalent models: syntax:
+# # models:
+# #   - model: psmathur/orca_mini_v3_13b
+# #   - model: garage-bAInd/Platypus2-13B
+# merge_method: slerp
+# base_model: psmathur/orca_mini_v3_13b
+# parameters:
+#   t:
+#     - filter: self_attn
+#       value: [0, 0.5, 0.3, 0.7, 1]
+#     - filter: mlp
+#       value: [1, 0.5, 0.7, 0.3, 0]
+#     - value: 0.5 # fallback for rest of tensors
+# dtype: float16
+# slices:
+models:
+      - model: /workspace/storage-shared/cnm/embedding/code/llama_1B_legal_retrieval_v2_2.3M_clean_150726_finetune2_07150726/checkpoint-8000 # Mô hình A (Base gốc)
+        # layer_range: [0, 36]
+      - model: /workspace/storage-shared/cnm/huyhq21_v2/mergekit/output_slerp_1B_data_v2_clean_v6   # Mô hình B (Đã Finetune)
+        # layer_range: [0, 36]
+merge_method: slerp
+base_model: /workspace/storage-shared/cnm/huyhq21_v2/mergekit/output_slerp_1B_data_v2_clean_v6 # Mô hình B (Đã Finetune)
+parameters:
+  t: 
+    #- filter: self_attn
+    - filter: layers
+      value: [0.2, 0.3, 0.5, 0.6, 0.7, 0.8]
+    #- filter: mlp
+     # value: [0.2, 0.3, 0.5, 0.6, 0.7, 0.8]
+    - value: 0.5 # fallback cho layernorm
+dtype: bfloat16
+
+# models:
+#   - model: /workspace/storage-shared/cnm/cuongnq23/project_embeding/models/Qwen/Qwen3-Embedding-8B
+#   - model: /workspace/storage-shared/cnm/embedding/code/qwen8b_SFT_kalm_vnlaw_3M/checkpoint-20000
+# merge_method: slerp
+# base_model: /workspace/storage-shared/cnm/cuongnq23/project_embeding/models/Qwen/Qwen3-Embedding-8B
+# parameters:
+#   t:
+#     # Trượt dốc từ 10% (đáy) lên 90% (đỉnh)
+#     - value: [0.1, 0.3, 0.5, 0.7, 0.9] 
+# dtype: bfloat16

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:79b68957119a1ba2974b5b2846956975657c6c375c9ba47ee70ee76c59a091cf
+size 1499356520

modules.json

@@ -0,0 +1,20 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "sentence_transformers.models.Transformer"
+  },
+  {
+    "idx": 1,
+    "name": "1",
+    "path": "1_Pooling",
+    "type": "sentence_transformers.models.Pooling"
+  },
+  {
+    "idx": 2,
+    "name": "2",
+    "path": "2_Normalize",
+    "type": "sentence_transformers.models.Normalize"
+  }
+]

pooling.py

@@ -0,0 +1,29 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0.
+from torch import Tensor
+import torch
+
+
+def pool(last_hidden_states: Tensor, attention_mask: Tensor, pool_type: str) -> Tensor:
+    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
+
+    if pool_type == "avg":
+        emb = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
+    elif pool_type == "weighted_avg":
+        emb = last_hidden.sum(dim=1)
+    elif pool_type == "cls":
+        emb = last_hidden[:, 0]
+    elif pool_type == "last":
+        left_padding = attention_mask[:, -1].sum() == attention_mask.shape[0]
+        if left_padding:
+            emb = last_hidden[:, -1]
+        else:
+            sequence_lengths = attention_mask.sum(dim=1) - 1
+            batch_size = last_hidden.shape[0]
+            emb = last_hidden[
+                torch.arange(batch_size, device=last_hidden.device), sequence_lengths
+            ]
+    else:
+        raise ValueError(f"pool_type {pool_type} not supported")
+
+    return emb

sentence_bert_config.json

@@ -0,0 +1,4 @@
+{
+    "max_seq_length": 8192,
+    "do_lower_case": false
+}

special_tokens_map.json

@@ -0,0 +1,23 @@
+{
+  "bos_token": {
+    "content": "<|begin_of_text|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|end_of_text|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|end_of_text|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb956fd3e742594d4410fdb86df5edf0c1b0f1f382117a2fbd877f2372ec158d
+size 17210019

tokenizer_config.json

@@ -0,0 +1,19 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|begin_of_text|>",
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "<|end_of_text|>",
+  "is_local": true,
+  "local_files_only": false,
+  "max_length": 8192,
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 131072,
+  "pad_token": "<|end_of_text|>",
+  "stride": 0,
+  "tokenizer_class": "TokenizersBackend",
+  "truncation_side": "right",
+  "truncation_strategy": "longest_first"
+}