Update weights and modeling code to latest version

__init__.py

@@ -0,0 +1,56 @@
+#!/usr/bin/env python
+# encoding: utf-8
+'''
+@license: (C) Copyright 2025, Hey.
+@author: Hey
+@email: [email protected]
+@tel: 137****6540
+@datetime: 2025/12/30 11:32
+@project: lucaone
+@file: configuration_lucaone
+@desc: configuration_lucaone
+'''
+
+from .configuration_lucaone import LucaGPLMConfig
+from .tokenization_lucaone import LucaGPLMTokenizer, LucaGPLMTokenizerFast
+from .modeling_lucaone import (
+    LucaGPLMModel,
+    LucaGPLMPreTrainedModel,
+    LucaGPLMForMaskedLM,
+    LucaGPLMForSequenceClassification,
+    LucaGPLMForTokenClassification
+)
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoModelForMaskedLM,
+    AutoModelForSequenceClassification,
+    AutoModelForTokenClassification
+)
+
+__all__ = [
+    "LucaGPLMConfig",
+    "LucaGPLMModel",
+    "LucaGPLMPreTrainedModel",
+    "LucaGPLMTokenizer",
+    "LucaGPLMTokenizerFast",
+    "LucaGPLMForMaskedLM",
+    "LucaGPLMForSequenceClassification",
+    "LucaGPLMForTokenClassification"
+]
+
+
+# 1. 注册配置类 (必选)
+AutoConfig.register("lucaone", LucaGPLMConfig)
+
+# 2. 注册基础模型 (用于 AutoModel.from_pretrained)
+AutoModel.register(LucaGPLMConfig, LucaGPLMModel)
+
+# 3. 注册序列分类模型 (用于 AutoModelForSequenceClassification)
+AutoModelForSequenceClassification.register(LucaGPLMConfig, LucaGPLMForSequenceClassification)
+
+# 4. 注册 Token 分类模型 (用于 AutoModelForTokenClassification)
+AutoModelForTokenClassification.register(LucaGPLMConfig, LucaGPLMForTokenClassification)
+
+# 5. 注册掩码语言模型 (用于 AutoModelForMaskedLM)
+AutoModelForMaskedLM.register(LucaGPLMConfig, LucaGPLMForMaskedLM)

config.json

@@ -1,30 +1,54 @@
 {
   "alphabet": "gene_prot",
   "architectures": [
-    "LucaGPLMModel"
+    "LucaGPLMForMaskedLM"
   ],
   "attention_probs_dropout_prob": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_lucaone.LucaGPLMConfig",
+    "AutoModel": "modeling_lucaone.LucaGPLMModel",
+    "AutoModelForMaskedLM": "modeling_lucaone.LucaGPLMForMaskedLM",
+    "AutoModelForSequenceClassification": "modeling_lucaone.LucaGPLMForSequenceClassification",
+    "AutoModelForTokenClassification": "modeling_lucaone.LucaGPLMForTokenClassification",
+    "AutoTokenizer": [
+      "tokenization_lucaone.LucaGPLMTokenizer",
+      null
+    ]
+  },
+  "bos_token_id": 2,
   "classifier_dropout_prob": 0.0,
+  "classifier_loss_reduction": "mean",
+  "classifier_loss_type": "cross_entropy",
+  "classifier_num_labels": -1,
+  "classifier_pooling_type": "value_attention",
+  "classifier_pos_weight": 1.0,
+  "classifier_weight": null,
   "embed_scale": 1.0,
+  "eos_token_id": 3,
   "ffn_dim": 10240,
+  "hidden_act": "gelu",
   "hidden_dropout_prob": 0.0,
   "hidden_size": 2560,
   "ignore_index": -100,
   "initializer_range": 0.02,
   "layer_norm_eps": 1e-12,
-  "max_position_embeddings": 1280,
-  "model_type": "lucagplm",
+  "mask_token_id": 4,
+  "max_position_embeddings": 4096,
+  "model_type": "lucaone",
   "no_position_embeddings": true,
   "no_token_type_embeddings": false,
   "num_attention_heads": 40,
   "num_hidden_layers": 20,
   "pad_token_id": 0,
-  "pooling_type": "avg",
-  "pooling_units": null,
+  "sep_token_id": 3,
+  "task_level": "seq_level",
+  "task_type": "embedding",
+  "tie_word_embeddings": false,
   "token_dropout": false,
   "torch_dtype": "float32",
-  "transformers_version": "4.45.2",
+  "transformers_version": "4.41.2",
   "type_vocab_size": 2,
+  "unk_token_id": 1,
   "use_embed_layer_norm": false,
   "use_last_layer_norm": true,
   "vocab_size": 39

configuration_lucaone.py

@@ -0,0 +1,102 @@
+#!/usr/bin/env python
+# encoding: utf-8
+'''
+@license: (C) Copyright 2025, Hey.
+@author: Hey
+@email: [email protected]
+@tel: 137****6540
+@datetime: 2025/12/30 11:34
+@project: lucaone
+@file: tokenization_lucaone
+@desc: tokenization_lucaone
+'''
+
+from typing import Literal
+from transformers import PretrainedConfig
+
+class LucaGPLMConfig(PretrainedConfig):
+    model_type = "lucaone"
+    
+    def __init__(
+        self,
+        vocab_size: int = 39,
+        pad_token_id: int = 0,
+        unk_token_id: int = 1,
+        bos_token_id: int = 2,
+        eos_token_id: int = 3,
+        sep_token_id: int = 3,
+        mask_token_id: int = 4,
+        hidden_act: str = "gelu",
+        max_position_embeddings: int = 4096,
+        type_vocab_size: int = 2,
+        num_hidden_layers: int = 20,
+        num_attention_heads: int = 40,
+        hidden_size: int = 2560,
+        ffn_dim: int = 10240,
+        no_position_embeddings: bool = True,
+        no_token_type_embeddings: bool = False,
+        alphabet: str = "gene_prot",
+        token_dropout: bool = False,
+        attention_probs_dropout_prob: float = 0.0,
+        hidden_dropout_prob: float = 0.0,
+        use_embed_layer_norm: bool = False,
+        use_last_layer_norm: bool = True,
+        embed_scale: float = 1.0,
+        ignore_index: int = -100,
+        layer_norm_eps: float = 1e-12,
+        initializer_range: float = 0.02,
+        task_level: Literal["seq_level", "token_level"]  = "seq_level",
+        task_type: Literal["embedding", "mlm", "multi_class", "binary_class", "regression", "multi_label"]  = "embedding",
+        classifier_num_labels: int = -1,
+        classifier_dropout_prob: float = 0.1,
+        classifier_pooling_type: Literal["cls", "value_attention", "context_attention", "mean"]  = "value_attention",
+        classifier_loss_type: Literal["binary_cross_entropy", "cross_entropy", "mse", "mae"]  = "cross_entropy",
+        classifier_loss_reduction: Literal["mean", "sum", "none"]  = "mean",
+        classifier_pos_weight: float=1.0,
+        classifier_weight: list=None,
+        tie_word_embeddings: bool=True,
+        **kwargs
+    ):
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            pad_token_id=pad_token_id,
+            **kwargs
+        )
+        
+        self.alphabet = alphabet
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.type_vocab_size = type_vocab_size
+        self.no_token_type_embeddings = no_token_type_embeddings
+        self.no_position_embeddings = no_position_embeddings
+        self.num_hidden_layers = num_hidden_layers
+        self.hidden_size = hidden_size
+        self.num_attention_heads = num_attention_heads
+        self.ffn_dim = ffn_dim
+        self.token_dropout = token_dropout
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.classifier_dropout_prob = classifier_dropout_prob
+        self.ignore_index = ignore_index
+        self.use_embed_layer_norm = use_embed_layer_norm
+        self.use_last_layer_norm = use_last_layer_norm
+        self.embed_scale = embed_scale
+        self.layer_norm_eps = layer_norm_eps
+        self.initializer_range = initializer_range
+        self.unk_token_id = unk_token_id
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.sep_token_id = sep_token_id
+        self.mask_token_id = mask_token_id
+        self.hidden_act = hidden_act
+        self.classifier_num_labels = classifier_num_labels
+        self.classifier_pooling_type = classifier_pooling_type
+        self.task_level = task_level
+        self.task_type = task_type
+        self.classifier_loss_type = classifier_loss_type
+        self.classifier_loss_reduction = classifier_loss_reduction
+        self.classifier_pos_weight = classifier_pos_weight
+        self.classifier_weight = classifier_weight
+
+
+__all__ = ["LucaGPLMConfig"]

model-00001-of-00002.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:abba84d3e29bcafdd171cdbc587402c93694831bc41f276a7007ce23f8ca7aca
-size 4930878944
+oid sha256:f4f62f17030bbf8353cfabd8d8291688b9b7a6632d55710178651b4ea8a529c6
+size 4930881120

model-00002-of-00002.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9882d0048f5392ccb8752df13a963e49a7328de7cd40c0adcb54c56fc1fb39b8
-size 1363720440
+oid sha256:1b2b9ec5fa2159a09c671e5d0dde6c1ea9537219f5ceab51a53e189cb79c5fb5
+size 1390366196

model.safetensors.index.json

@@ -1,351 +1,357 @@
 {
   "metadata": {
-    "total_size": 6294561280
+    "total_size": 6321205916
   },
   "weight_map": {
-    "embeddings.embed_tokens.weight": "model-00001-of-00002.safetensors",
-    "embeddings.embed_type.weight": "model-00001-of-00002.safetensors",
-    "encoder.last_layer_norm.bias": "model-00002-of-00002.safetensors",
-    "encoder.last_layer_norm.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.0.fc1.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.0.fc1.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.0.fc2.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.0.fc2.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.0.post_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.0.post_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.0.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.0.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.0.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.0.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.0.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.0.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.0.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.0.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.0.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
-    "encoder.layers.0.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.0.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.1.fc1.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.1.fc1.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.1.fc2.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.1.fc2.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.1.post_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.1.post_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.1.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.1.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.1.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.1.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.1.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.1.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.1.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.1.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.1.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
-    "encoder.layers.1.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.1.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.10.fc1.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.10.fc1.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.10.fc2.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.10.fc2.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.10.post_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.10.post_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.10.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.10.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.10.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.10.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.10.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.10.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.10.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.10.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.10.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
-    "encoder.layers.10.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.10.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.11.fc1.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.11.fc1.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.11.fc2.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.11.fc2.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.11.post_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.11.post_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.11.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.11.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.11.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.11.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.11.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.11.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.11.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.11.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.11.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
-    "encoder.layers.11.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.11.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.12.fc1.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.12.fc1.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.12.fc2.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.12.fc2.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.12.post_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.12.post_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.12.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.12.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.12.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.12.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.12.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.12.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.12.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.12.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.12.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
-    "encoder.layers.12.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.12.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.13.fc1.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.13.fc1.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.13.fc2.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.13.fc2.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.13.post_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.13.post_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.13.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.13.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.13.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.13.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.13.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.13.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.13.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.13.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.13.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
-    "encoder.layers.13.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.13.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.14.fc1.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.14.fc1.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.14.fc2.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.14.fc2.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.14.post_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.14.post_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.14.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.14.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.14.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.14.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.14.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.14.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.14.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.14.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.14.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
-    "encoder.layers.14.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.14.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.15.fc1.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.15.fc1.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.15.fc2.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.15.fc2.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.15.post_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.15.post_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.15.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.15.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.15.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.15.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.15.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.15.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.15.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.15.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.15.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
-    "encoder.layers.15.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.15.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.16.fc1.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.16.fc1.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.16.fc2.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.16.fc2.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.16.post_layer_norm.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.16.post_layer_norm.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.16.pre_layer_norm.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.16.pre_layer_norm.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.16.self_attn.k_proj.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.16.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.16.self_attn.out_proj.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.16.self_attn.out_proj.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.16.self_attn.q_proj.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.16.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.16.self_attn.rot_emb.inv_freq": "model-00002-of-00002.safetensors",
-    "encoder.layers.16.self_attn.v_proj.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.16.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.17.fc1.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.17.fc1.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.17.fc2.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.17.fc2.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.17.post_layer_norm.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.17.post_layer_norm.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.17.pre_layer_norm.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.17.pre_layer_norm.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.17.self_attn.k_proj.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.17.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.17.self_attn.out_proj.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.17.self_attn.out_proj.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.17.self_attn.q_proj.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.17.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.17.self_attn.rot_emb.inv_freq": "model-00002-of-00002.safetensors",
-    "encoder.layers.17.self_attn.v_proj.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.17.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.18.fc1.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.18.fc1.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.18.fc2.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.18.fc2.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.18.post_layer_norm.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.18.post_layer_norm.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.18.pre_layer_norm.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.18.pre_layer_norm.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.18.self_attn.k_proj.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.18.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.18.self_attn.out_proj.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.18.self_attn.out_proj.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.18.self_attn.q_proj.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.18.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.18.self_attn.rot_emb.inv_freq": "model-00002-of-00002.safetensors",
-    "encoder.layers.18.self_attn.v_proj.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.18.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.19.fc1.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.19.fc1.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.19.fc2.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.19.fc2.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.19.post_layer_norm.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.19.post_layer_norm.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.19.pre_layer_norm.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.19.pre_layer_norm.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.19.self_attn.k_proj.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.19.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.19.self_attn.out_proj.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.19.self_attn.out_proj.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.19.self_attn.q_proj.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.19.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.19.self_attn.rot_emb.inv_freq": "model-00002-of-00002.safetensors",
-    "encoder.layers.19.self_attn.v_proj.bias": "model-00002-of-00002.safetensors",
-    "encoder.layers.19.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
-    "encoder.layers.2.fc1.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.2.fc1.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.2.fc2.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.2.fc2.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.2.post_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.2.post_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.2.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.2.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.2.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.2.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.2.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.2.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.2.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.2.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.2.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
-    "encoder.layers.2.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.2.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.3.fc1.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.3.fc1.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.3.fc2.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.3.fc2.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.3.post_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.3.post_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.3.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.3.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.3.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.3.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.3.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.3.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.3.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.3.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.3.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
-    "encoder.layers.3.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.3.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.4.fc1.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.4.fc1.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.4.fc2.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.4.fc2.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.4.post_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.4.post_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.4.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.4.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.4.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.4.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.4.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.4.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.4.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.4.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.4.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
-    "encoder.layers.4.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.4.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.5.fc1.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.5.fc1.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.5.fc2.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.5.fc2.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.5.post_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.5.post_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.5.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.5.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.5.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.5.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.5.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.5.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.5.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.5.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.5.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
-    "encoder.layers.5.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.5.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.6.fc1.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.6.fc1.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.6.fc2.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.6.fc2.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.6.post_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.6.post_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.6.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.6.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.6.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.6.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.6.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.6.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.6.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.6.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.6.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
-    "encoder.layers.6.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.6.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.7.fc1.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.7.fc1.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.7.fc2.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.7.fc2.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.7.post_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.7.post_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.7.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.7.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.7.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.7.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.7.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.7.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.7.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.7.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.7.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
-    "encoder.layers.7.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.7.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.8.fc1.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.8.fc1.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.8.fc2.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.8.fc2.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.8.post_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.8.post_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.8.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.8.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.8.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.8.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.8.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.8.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.8.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.8.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.8.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
-    "encoder.layers.8.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.8.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.9.fc1.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.9.fc1.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.9.fc2.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.9.fc2.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.9.post_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.9.post_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.9.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.9.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.9.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.9.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.9.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.9.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.9.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.9.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
-    "encoder.layers.9.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
-    "encoder.layers.9.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
-    "encoder.layers.9.self_attn.v_proj.weight": "model-00001-of-00002.safetensors"
+    "lm_head.bias": "model-00002-of-00002.safetensors",
+    "lm_head.decoder.weight": "model-00002-of-00002.safetensors",
+    "lm_head.dense.bias": "model-00002-of-00002.safetensors",
+    "lm_head.dense.weight": "model-00002-of-00002.safetensors",
+    "lm_head.layer_norm.bias": "model-00002-of-00002.safetensors",
+    "lm_head.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "lucaone.embeddings.embed_tokens.weight": "model-00001-of-00002.safetensors",
+    "lucaone.embeddings.embed_type.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.last_layer_norm.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.last_layer_norm.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.0.fc1.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.0.fc1.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.0.fc2.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.0.fc2.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.0.post_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.0.post_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.0.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.0.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.0.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.0.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.0.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.0.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.0.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.0.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.0.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.0.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.0.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.1.fc1.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.1.fc1.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.1.fc2.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.1.fc2.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.1.post_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.1.post_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.1.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.1.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.1.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.1.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.1.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.1.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.1.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.1.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.1.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.1.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.1.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.10.fc1.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.10.fc1.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.10.fc2.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.10.fc2.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.10.post_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.10.post_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.10.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.10.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.10.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.10.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.10.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.10.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.10.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.10.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.10.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.10.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.10.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.11.fc1.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.11.fc1.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.11.fc2.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.11.fc2.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.11.post_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.11.post_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.11.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.11.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.11.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.11.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.11.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.11.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.11.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.11.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.11.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.11.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.11.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.12.fc1.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.12.fc1.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.12.fc2.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.12.fc2.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.12.post_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.12.post_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.12.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.12.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.12.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.12.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.12.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.12.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.12.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.12.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.12.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.12.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.12.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.13.fc1.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.13.fc1.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.13.fc2.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.13.fc2.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.13.post_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.13.post_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.13.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.13.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.13.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.13.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.13.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.13.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.13.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.13.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.13.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.13.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.13.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.14.fc1.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.14.fc1.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.14.fc2.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.14.fc2.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.14.post_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.14.post_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.14.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.14.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.14.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.14.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.14.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.14.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.14.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.14.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.14.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.14.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.14.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.15.fc1.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.15.fc1.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.15.fc2.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.15.fc2.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.15.post_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.15.post_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.15.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.15.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.15.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.15.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.15.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.15.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.15.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.15.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.15.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.15.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.15.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.16.fc1.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.16.fc1.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.16.fc2.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.16.fc2.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.16.post_layer_norm.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.16.post_layer_norm.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.16.pre_layer_norm.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.16.pre_layer_norm.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.16.self_attn.k_proj.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.16.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.16.self_attn.out_proj.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.16.self_attn.out_proj.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.16.self_attn.q_proj.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.16.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.16.self_attn.rot_emb.inv_freq": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.16.self_attn.v_proj.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.16.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.17.fc1.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.17.fc1.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.17.fc2.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.17.fc2.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.17.post_layer_norm.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.17.post_layer_norm.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.17.pre_layer_norm.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.17.pre_layer_norm.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.17.self_attn.k_proj.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.17.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.17.self_attn.out_proj.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.17.self_attn.out_proj.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.17.self_attn.q_proj.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.17.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.17.self_attn.rot_emb.inv_freq": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.17.self_attn.v_proj.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.17.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.18.fc1.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.18.fc1.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.18.fc2.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.18.fc2.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.18.post_layer_norm.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.18.post_layer_norm.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.18.pre_layer_norm.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.18.pre_layer_norm.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.18.self_attn.k_proj.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.18.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.18.self_attn.out_proj.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.18.self_attn.out_proj.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.18.self_attn.q_proj.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.18.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.18.self_attn.rot_emb.inv_freq": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.18.self_attn.v_proj.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.18.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.19.fc1.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.19.fc1.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.19.fc2.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.19.fc2.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.19.post_layer_norm.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.19.post_layer_norm.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.19.pre_layer_norm.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.19.pre_layer_norm.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.19.self_attn.k_proj.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.19.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.19.self_attn.out_proj.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.19.self_attn.out_proj.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.19.self_attn.q_proj.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.19.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.19.self_attn.rot_emb.inv_freq": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.19.self_attn.v_proj.bias": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.19.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "lucaone.encoder.layers.2.fc1.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.2.fc1.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.2.fc2.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.2.fc2.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.2.post_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.2.post_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.2.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.2.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.2.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.2.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.2.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.2.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.2.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.2.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.2.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.2.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.2.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.3.fc1.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.3.fc1.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.3.fc2.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.3.fc2.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.3.post_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.3.post_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.3.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.3.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.3.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.3.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.3.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.3.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.3.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.3.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.3.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.3.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.3.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.4.fc1.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.4.fc1.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.4.fc2.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.4.fc2.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.4.post_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.4.post_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.4.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.4.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.4.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.4.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.4.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.4.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.4.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.4.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.4.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.4.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.4.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.5.fc1.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.5.fc1.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.5.fc2.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.5.fc2.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.5.post_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.5.post_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.5.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.5.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.5.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.5.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.5.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.5.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.5.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.5.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.5.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.5.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.5.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.6.fc1.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.6.fc1.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.6.fc2.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.6.fc2.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.6.post_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.6.post_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.6.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.6.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.6.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.6.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.6.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.6.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.6.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.6.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.6.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.6.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.6.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.7.fc1.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.7.fc1.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.7.fc2.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.7.fc2.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.7.post_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.7.post_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.7.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.7.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.7.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.7.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.7.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.7.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.7.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.7.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.7.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.7.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.7.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.8.fc1.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.8.fc1.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.8.fc2.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.8.fc2.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.8.post_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.8.post_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.8.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.8.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.8.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.8.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.8.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.8.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.8.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.8.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.8.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.8.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.8.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.9.fc1.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.9.fc1.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.9.fc2.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.9.fc2.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.9.post_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.9.post_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.9.pre_layer_norm.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.9.pre_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.9.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.9.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.9.self_attn.out_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.9.self_attn.out_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.9.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.9.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.9.self_attn.rot_emb.inv_freq": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.9.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
+    "lucaone.encoder.layers.9.self_attn.v_proj.weight": "model-00001-of-00002.safetensors"
   }
 }

modeling_lucaone.py

@@ -0,0 +1,1344 @@
+#!/usr/bin/env python
+# encoding: utf-8
+'''
+@license: (C) Copyright 2025, Hey.
+@author: Hey
+@email: [email protected]
+@tel: 137****6540
+@datetime: 2025/12/30 11:35
+@project: lucaone
+@file: modeling_lucaone
+@desc: modeling_lucaone
+'''
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutput
+from transformers.modeling_outputs import MaskedLMOutput
+from transformers.modeling_outputs import SequenceClassifierOutput
+from transformers.modeling_outputs import TokenClassifierOutput
+from typing import Optional, List, Union, Tuple
+from .configuration_lucaone import LucaGPLMConfig
+try:
+    from apex.normalization import FusedLayerNorm as _FusedLayerNorm
+    class LucaGPLM1bLayerNorm(_FusedLayerNorm):
+        @torch.jit.unused
+        def forward(self, x):
+            if not x.is_cuda:
+                return super().forward(x)
+            else:
+                with torch.cuda.device(x.device):
+                    return super().forward(x)
+except ImportError:
+    from torch.nn import LayerNorm as LucaGPLM1bLayerNorm
+
+def gelu(x):
+    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))
+
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+def apply_rotary_pos_emb(x, cos, sin):
+    cos = cos[:, : x.shape[-2], :]
+    sin = sin[:, : x.shape[-2], :]
+    return (x * cos) + (rotate_half(x) * sin)
+
+class LucaGPLMRotaryEmbedding(torch.nn.Module):
+    def __init__(self, dim: int, *_, **__):
+        super().__init__()
+        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+
+        self._seq_len_cached = None
+        self._cos_cached = None
+        self._sin_cached = None
+
+    def _update_cos_sin_tables(self, x, seq_dimension=1):
+        seq_len = x.shape[seq_dimension]
+
+        if (seq_len != self._seq_len_cached or 
+            self._cos_cached is None or 
+            self._sin_cached is None or
+            self._cos_cached.device != x.device):
+            self._seq_len_cached = seq_len
+            t = torch.arange(x.shape[seq_dimension], device=x.device).type_as(self.inv_freq)
+            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
+
+            self._cos_cached = emb.cos()[None, :, :]
+            self._sin_cached = emb.sin()[None, :, :]
+
+        return self._cos_cached, self._sin_cached
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        self._cos_cached, self._sin_cached = self._update_cos_sin_tables(k, seq_dimension=-2)
+
+        return (
+            apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached),
+            apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached),
+        )
+
+class LucaGPLMGlobalMaskWeightedAttentionPooling1D(nn.Module):
+    def __init__(self, embed_size, use_bias=False):
+        super(LucaGPLMGlobalMaskWeightedAttentionPooling1D, self).__init__()
+        self.embed_size = embed_size
+        self.use_bias = use_bias
+
+        self.W = nn.Parameter(torch.Tensor(self.embed_size))
+        nn.init.trunc_normal_(self.W, std=0.01)
+        if self.use_bias:
+            self.b = nn.Parameter(torch.Tensor(1))
+            nn.init.trunc_normal_(self.b, std=0.01)
+
+    def forward(self, x, mask=None):
+        # (B, Len, Embed) x (Embed,) = (B, Len)
+        logits = torch.matmul(x, self.W)
+        if self.use_bias:
+            logits += self.b
+
+        if mask is not None:
+            attention_probs = nn.Softmax(dim=-1)(logits + (1.0 - mask) * -10000)
+        else:
+            attention_probs = nn.Softmax(dim=-1)(logits)
+        x = torch.sum(torch.unsqueeze(attention_probs, dim=-1) * x, dim=1)
+        return x
+
+    def __repr__(self):
+        return self.__class__.__name__ + ' (' + str(self.embed_size) + (', bias=%r)' % self.use_bias)
+
+class LucaGPLMGlobalMaskContextAttentionPooling1D(nn.Module):
+    def __init__(self, embed_size, units=None, use_additive_bias=False, use_attention_bias=False):
+        super(LucaGPLMGlobalMaskContextAttentionPooling1D, self).__init__()
+        self.embed_size = embed_size
+        self.use_additive_bias = use_additive_bias
+        self.use_attention_bias = use_attention_bias
+        self.units = units if units else embed_size
+
+        self.U = nn.Parameter(torch.Tensor(self.embed_size, self.units))
+        self.V = nn.Parameter(torch.Tensor(self.embed_size, self.units))
+        if self.use_additive_bias:
+            self.b1 = nn.Parameter(torch.Tensor(self.units))
+            nn.init.trunc_normal_(self.b1, std=0.01)
+        if self.use_attention_bias:
+            self.b2 = nn.Parameter(torch.Tensor(1))
+            nn.init.trunc_normal_(self.b2, std=0.01)
+
+        self.c = nn.Parameter(torch.Tensor(self.units))
+
+        nn.init.trunc_normal_(self.U, std=0.01)
+        nn.init.trunc_normal_(self.V, std=0.01)
+        nn.init.trunc_normal_(self.c, std=0.01)
+
+    def forward(self, x, mask=None):
+        # (B, Len, Embed) x (Embed, Units) = (B, Len, Units)
+        q = torch.matmul(x, self.U)
+        k = torch.matmul(x, self.V)
+        if self.use_additive_bias:
+            h = torch.tanh(q + k + self.b1)
+        else:
+            h = torch.tanh(q + k)
+
+        if self.use_attention_bias:
+            e = torch.matmul(h, self.c) + self.b2
+        else:
+            e = torch.matmul(h, self.c)
+        if mask is not None:
+            attention_probs = nn.Softmax(dim=-1)(e + (1.0 - mask) * -10000)
+        else:
+            attention_probs = nn.Softmax(dim=-1)(e)
+        x = torch.sum(torch.unsqueeze(attention_probs, dim=-1) * x, dim=1)
+        return x
+
+    def __repr__(self):
+        return self.__class__.__name__ + ' (' + str(self.embed_size) + ' -> ' + str(self.units) + ', bias=(%r, %r))' % (self.use_additive_bias, self.use_attention_bias)
+
+class LucaGPLMGlobalMaskValueAttentionPooling1D(nn.Module):
+    def __init__(self, embed_size, units=None, use_additive_bias=False, use_attention_bias=False):
+        super(LucaGPLMGlobalMaskValueAttentionPooling1D, self).__init__()
+        self.embed_size = embed_size
+        self.use_additive_bias = use_additive_bias
+        self.use_attention_bias = use_attention_bias
+        self.units = units if units else embed_size
+
+        self.U = nn.Parameter(torch.Tensor(self.embed_size, self.units))
+        self.V = nn.Parameter(torch.Tensor(self.embed_size, self.units))
+        if self.use_additive_bias:
+            self.b1 = nn.Parameter(torch.Tensor(self.units))
+            nn.init.trunc_normal_(self.b1, std=0.01)
+        if self.use_attention_bias:
+            self.b2 = nn.Parameter(torch.Tensor(self.embed_size))
+            nn.init.trunc_normal_(self.b2, std=0.01)
+
+        self.W = nn.Parameter(torch.Tensor(self.units, self.embed_size))
+
+        nn.init.trunc_normal_(self.U, std=0.01)
+        nn.init.trunc_normal_(self.V, std=0.01)
+        nn.init.trunc_normal_(self.W, std=0.01)
+
+    def forward(self, x, mask=None):
+        # (B, Len, Embed) x (Embed, Units) = (B, Len, Units)
+        q = torch.matmul(x, self.U)
+        k = torch.matmul(x, self.V)
+        if self.use_additive_bias:
+            h = torch.tanh(q + k + self.b1)
+        else:
+            h = torch.tanh(q + k)
+
+        # (B, Len, Units) x (Units, Embed) = (B, Len, Embed)
+        if self.use_attention_bias:
+            e = torch.matmul(h, self.W) + self.b2
+        else:
+            e = torch.matmul(h, self.W)
+        if mask is not None:
+            attention_probs = nn.Softmax(dim=1)(e + torch.unsqueeze((1.0 - mask) * -10000, dim=-1))
+        else:
+            attention_probs = nn.Softmax(dim=1)(e)
+        x = torch.sum(attention_probs * x, dim=1)
+        return x
+
+    def __repr__(self):
+        return self.__class__.__name__ + ' (' + str(self.embed_size) + ' -> ' + str(self.units) + ', bias=(%r, %r))' % (self.use_additive_bias, self.use_attention_bias)
+
+class LucaGPLM1LayerNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-12, affine=True):
+        super().__init__()
+        self.hidden_size = (hidden_size,) if isinstance(hidden_size, int) else tuple(hidden_size)
+        self.eps = eps
+        self.affine = bool(affine)
+        if self.affine:
+            self.weight = nn.Parameter(torch.ones(hidden_size))
+            self.bias = nn.Parameter(torch.zeros(hidden_size))
+        else:
+            self.weight, self.bias = None, None
+
+    def forward(self, x):
+        dims = tuple(-(i + 1) for i in range(len(self.hidden_size)))
+        means = x.mean(dims, keepdim=True)
+        x_zeromean = x - means
+        variances = x_zeromean.pow(2).mean(dims, keepdim=True)
+        x = x_zeromean / torch.sqrt(variances + self.eps)
+        if self.affine:
+            x = (self.weight * x) + self.bias
+        return x
+
+class LucaGPLMMultiheadAttention(nn.Module):
+    def __init__(
+        self,
+        embed_dim,
+        num_heads,
+        kdim=None,
+        vdim=None,
+        dropout=0.0,
+        bias=True,
+        add_bias_kv: bool = False,
+        add_zero_attn: bool = False,
+        self_attention: bool = False,
+        encoder_decoder_attention: bool = False,
+        use_rotary_embeddings: bool = False,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.kdim = kdim if kdim is not None else embed_dim
+        self.vdim = vdim if vdim is not None else embed_dim
+        self.qkv_same_dim = self.kdim == embed_dim and self.vdim == embed_dim
+
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        assert (
+            self.head_dim * num_heads == self.embed_dim
+        ), "embed_dim must be divisible by num_heads"
+        self.scaling = self.head_dim**-0.5
+
+        self.self_attention = self_attention
+        self.encoder_decoder_attention = encoder_decoder_attention
+
+        assert not self.self_attention or self.qkv_same_dim, (
+            "Self-attention requires query, key and " "value to be of the same size"
+        )
+
+        self.k_proj = nn.Linear(self.kdim, embed_dim, bias=bias)
+        self.v_proj = nn.Linear(self.vdim, embed_dim, bias=bias)
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+        if add_bias_kv:
+            self.bias_k = nn.Parameter(torch.Tensor(1, 1, embed_dim))
+            self.bias_v = nn.Parameter(torch.Tensor(1, 1, embed_dim))
+        else:
+            self.bias_k = self.bias_v = None
+
+        self.add_zero_attn = add_zero_attn
+
+        self.reset_parameters()
+
+        self.rot_emb = None
+        if use_rotary_embeddings:
+            self.rot_emb = LucaGPLMRotaryEmbedding(dim=self.head_dim)
+
+    def reset_parameters(self):
+        nn.init.xavier_uniform_(self.k_proj.weight, gain=nn.init.calculate_gain("relu"))
+        nn.init.xavier_uniform_(self.v_proj.weight, gain=nn.init.calculate_gain("relu"))
+        nn.init.xavier_uniform_(self.q_proj.weight, gain=nn.init.calculate_gain("relu"))
+        nn.init.xavier_uniform_(self.out_proj.weight, gain=nn.init.calculate_gain("relu"))
+        
+        if self.out_proj.bias is not None:
+            nn.init.constant_(self.out_proj.bias, 0.0)
+        if self.bias_k is not None:
+            nn.init.xavier_normal_(self.bias_k)
+        if self.bias_v is not None:
+            nn.init.xavier_normal_(self.bias_v)
+
+    def forward(
+        self,
+        query,
+        key: Optional[torch.Tensor] = None,
+        value: Optional[torch.Tensor] = None,
+        key_padding_mask: Optional[torch.Tensor] = None,
+        need_weights: bool = True,
+        attn_mask: Optional[torch.Tensor] = None,
+        need_head_weights: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        if need_head_weights:
+            need_weights = True
+
+        tgt_len, bsz, embed_dim = query.size()
+        assert embed_dim == self.embed_dim
+
+        if self.self_attention:
+            q = self.q_proj(query)
+            k = self.k_proj(query)
+            v = self.v_proj(query)
+        else:
+            assert key is not None and value is not None
+            q = self.q_proj(query)
+            k = self.k_proj(key)
+            v = self.v_proj(value)
+
+        q *= self.scaling
+
+        if self.bias_k is not None:
+            assert self.bias_v is not None
+            k = torch.cat([k, self.bias_k.repeat(1, bsz, 1)])
+            v = torch.cat([v, self.bias_v.repeat(1, bsz, 1)])
+            if attn_mask is not None:
+                attn_mask = torch.cat(
+                    [attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1
+                )
+            if key_padding_mask is not None:
+                key_padding_mask = torch.cat(
+                    [
+                        key_padding_mask,
+                        key_padding_mask.new_zeros(key_padding_mask.size(0), 1),
+                    ],
+                    dim=1,
+                )
+
+        q = q.contiguous().view(tgt_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)
+        if k is not None:
+            k = k.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)
+        if v is not None:
+            v = v.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)
+
+        assert k is not None
+        src_len = k.size(1)
+
+        if self.rot_emb:
+            q, k = self.rot_emb(q, k)
+
+        attn_weights = torch.bmm(q, k.transpose(1, 2))
+        assert list(attn_weights.size()) == [bsz * self.num_heads, tgt_len, src_len]
+
+        if attn_mask is not None:
+            attn_mask = attn_mask.unsqueeze(0)
+            attn_weights += attn_mask
+
+        if key_padding_mask is not None:
+            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
+            attn_weights = attn_weights.masked_fill(
+                key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool), float("-inf")
+            )
+            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
+
+        attn_weights_float = F.softmax(attn_weights, dim=-1)
+        attn_weights = attn_weights_float.type_as(attn_weights)
+        attn_probs = F.dropout(
+            attn_weights_float.type_as(attn_weights),
+            p=self.dropout,
+            training=self.training,
+        )
+
+        assert v is not None
+        attn = torch.bmm(attn_probs, v)
+        assert list(attn.size()) == [bsz * self.num_heads, tgt_len, self.head_dim]
+        attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
+        attn = self.out_proj(attn)
+
+        attn_weights_output: Optional[torch.Tensor] = None
+        if need_weights:
+            attn_weights_output = attn_weights_float.view(
+                bsz, self.num_heads, tgt_len, src_len
+            ).type_as(attn).transpose(1, 0)
+            if not need_head_weights:
+                # average attention weights over heads
+                attn_weights_output = attn_weights_output.mean(dim=0)
+
+        return attn, attn_weights_output
+
+class LucaGPLMMultiheadAttentionWithSDPA(nn.Module):
+    def __init__(
+        self,
+        embed_dim,
+        num_heads,
+        kdim=None,
+        vdim=None,
+        dropout=0.0,
+        bias=True,
+        add_bias_kv: bool = False,
+        add_zero_attn: bool = False,
+        self_attention: bool = False,
+        encoder_decoder_attention: bool = False,
+        use_rotary_embeddings: bool = True,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.kdim = kdim if kdim is not None else embed_dim
+        self.vdim = vdim if vdim is not None else embed_dim
+        self.qkv_same_dim = self.kdim == embed_dim and self.vdim == embed_dim
+
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        assert (
+            self.head_dim * num_heads == self.embed_dim
+        ), "embed_dim must be divisible by num_heads"
+        self.scaling = self.head_dim**-0.5
+
+        self.self_attention = self_attention
+        self.encoder_decoder_attention = encoder_decoder_attention
+
+        assert not self.self_attention or self.qkv_same_dim, (
+            "Self-attention requires query, key and " "value to be of the same size"
+        )
+
+        self.k_proj = nn.Linear(self.kdim, embed_dim, bias=bias)
+        self.v_proj = nn.Linear(self.vdim, embed_dim, bias=bias)
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+        if add_bias_kv:
+            self.bias_k = nn.Parameter(torch.Tensor(1, 1, embed_dim))
+            self.bias_v = nn.Parameter(torch.Tensor(1, 1, embed_dim))
+        else:
+            self.bias_k = self.bias_v = None
+
+        self.add_zero_attn = add_zero_attn
+
+        self.reset_parameters()
+
+        self.rot_emb = None
+        if use_rotary_embeddings:
+            self.rot_emb = LucaGPLMRotaryEmbedding(dim=self.head_dim)
+
+    def reset_parameters(self):
+        nn.init.xavier_uniform_(self.k_proj.weight, gain=nn.init.calculate_gain("relu"))
+        nn.init.xavier_uniform_(self.v_proj.weight, gain=nn.init.calculate_gain("relu"))
+        nn.init.xavier_uniform_(self.q_proj.weight, gain=nn.init.calculate_gain("relu"))
+        nn.init.xavier_uniform_(self.out_proj.weight, gain=nn.init.calculate_gain("relu"))
+        
+        if self.out_proj.bias is not None:
+            nn.init.constant_(self.out_proj.bias, 0.0)
+        if self.bias_k is not None:
+            nn.init.xavier_normal_(self.bias_k)
+        if self.bias_v is not None:
+            nn.init.xavier_normal_(self.bias_v)
+
+    def forward(
+        self,
+        query,
+        key: Optional[torch.Tensor] = None,
+        value: Optional[torch.Tensor] = None,
+        key_padding_mask: Optional[torch.Tensor] = None,
+        need_weights: bool = True,
+        attn_mask: Optional[torch.Tensor] = None,
+        need_head_weights: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        
+        tgt_len, bsz, embed_dim = query.size()
+        assert embed_dim == self.embed_dim
+
+        if self.self_attention:
+            q = self.q_proj(query)
+            k = self.k_proj(query)
+            v = self.v_proj(query)
+        else:
+            assert key is not None and value is not None
+            q = self.q_proj(query)
+            k = self.k_proj(key)
+            v = self.v_proj(value)
+
+        if self.bias_k is not None:
+            assert self.bias_v is not None
+            k = torch.cat([k, self.bias_k.repeat(1, bsz, 1)])
+            v = torch.cat([v, self.bias_v.repeat(1, bsz, 1)])
+            if attn_mask is not None:
+                attn_mask = torch.cat(
+                    [attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1
+                )
+            if key_padding_mask is not None:
+                key_padding_mask = torch.cat(
+                    [
+                        key_padding_mask,
+                        key_padding_mask.new_zeros(key_padding_mask.size(0), 1),
+                    ],
+                    dim=1,
+                )
+
+        # ----------------------------------------------------------------------
+        # Flash Attention Optimization
+        # ----------------------------------------------------------------------
+        # 如果不需要返回 head weights 且 PyTorch 版本支持，则使用 Flash Attention
+        if not need_head_weights and hasattr(F, "scaled_dot_product_attention"):
+            # Reshape inputs to (Batch, Head, Seq_Len, Dim) for SDPA
+            # q, k, v input shape: (Seq_Len, Batch, Embed_Dim)
+            q_sdpa = q.view(tgt_len, bsz, self.num_heads, self.head_dim).permute(1, 2, 0, 3)
+            k_sdpa = k.view(-1, bsz, self.num_heads, self.head_dim).permute(1, 2, 0, 3)
+            v_sdpa = v.view(-1, bsz, self.num_heads, self.head_dim).permute(1, 2, 0, 3)
+
+            # Apply Rotary Embedding if needed
+            if self.rot_emb:
+                # Rotary expects inputs (..., Seq_Len, Dim)
+                # It handles broadcasting over Batch and Head
+                q_sdpa, k_sdpa = self.rot_emb(q_sdpa, k_sdpa)
+
+            # Prepare Mask
+            # SDPA accepts a broadcastable boolean mask or float mask
+            # key_padding_mask is (Batch, Seq_Len), True where padding
+            sdpa_mask = None
+            if attn_mask is not None or key_padding_mask is not None:
+                # Start with a float mask suitable for SDPA
+                target_shape = (bsz, 1, tgt_len, k_sdpa.size(2))
+                sdpa_mask = torch.zeros(target_shape, device=q.device, dtype=q.dtype)
+                
+                if key_padding_mask is not None:
+                    # key_padding_mask is (Batch, Seq_Len) -> (Batch, 1, 1, Seq_Len)
+                    sdpa_mask = sdpa_mask.masked_fill(
+                        key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool), 
+                        float("-inf")
+                    )
+                
+                if attn_mask is not None:
+                    if attn_mask.dim() == 2:
+                        sdpa_mask = sdpa_mask + attn_mask.unsqueeze(0).unsqueeze(0)
+                    elif attn_mask.dim() == 3:
+                         pass
+                    else:
+                        sdpa_mask = sdpa_mask + attn_mask
+
+            # Call Flash Attention
+            # 【关键修改】：添加 scale=1.0，因为 q 已经被手动缩放过了
+            attn_output = F.scaled_dot_product_attention(
+                q_sdpa,
+                k_sdpa,
+                v_sdpa,
+                attn_mask=sdpa_mask, 
+                dropout_p=self.dropout if self.training else 0.0,
+                is_causal=False
+            )
+
+            # Reshape back to (Seq_Len, Batch, Embed_Dim)
+            # (B, H, L, D) -> (L, B, H, D) -> (L, B, E)
+            attn_output = attn_output.permute(2, 0, 1, 3).contiguous().view(tgt_len, bsz, self.embed_dim)
+            
+            # Linear projection
+            attn_output = self.out_proj(attn_output)
+
+            # Return None for weights (optimization trade-off)
+            return attn_output, None
+
+        q = q * self.scaling
+        # ----------------------------------------------------------------------
+        # Original Implementation (Fallback)
+        # ----------------------------------------------------------------------
+        # print('Fall back to slow implementation.')
+        q = q.contiguous().view(tgt_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)
+        if k is not None:
+            k = k.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)
+        if v is not None:
+            v = v.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)
+
+        assert k is not None
+        src_len = k.size(1)
+
+        if self.rot_emb:
+            q, k = self.rot_emb(q, k)
+
+        attn_weights = torch.bmm(q, k.transpose(1, 2))
+        assert list(attn_weights.size()) == [bsz * self.num_heads, tgt_len, src_len]
+
+        if attn_mask is not None:
+            attn_mask = attn_mask.unsqueeze(0)
+            attn_weights += attn_mask
+
+        if key_padding_mask is not None:
+            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
+            attn_weights = attn_weights.masked_fill(
+                key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool), float("-inf")
+            )
+            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
+
+        attn_weights_float = F.softmax(attn_weights, dim=-1)
+        attn_weights = attn_weights_float.type_as(attn_weights)
+        attn_probs = F.dropout(
+            attn_weights_float.type_as(attn_weights),
+            p=self.dropout,
+            training=self.training,
+        )
+
+        assert v is not None
+        attn = torch.bmm(attn_probs, v)
+        assert list(attn.size()) == [bsz * self.num_heads, tgt_len, self.head_dim]
+        attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
+        attn = self.out_proj(attn)
+
+        attn_weights_output: Optional[torch.Tensor] = None
+        if need_weights:
+            attn_weights_output = attn_weights_float.view(
+                bsz, self.num_heads, tgt_len, src_len
+            ).type_as(attn).transpose(1, 0)
+            if not need_head_weights:
+                # average attention weights over heads
+                attn_weights_output = attn_weights_output.mean(dim=0)
+
+        return attn, attn_weights_output
+
+class LucaGPLMRobertaLMHead(nn.Module):
+    def __init__(self, embed_dim, output_dim):
+        super().__init__()
+        self.dense = nn.Linear(embed_dim, embed_dim)
+        self.layer_norm = LucaGPLM1bLayerNorm(embed_dim)
+        # 使用标准的 nn.Linear
+        self.decoder = nn.Linear(embed_dim, output_dim, bias=False)
+        self.bias = nn.Parameter(torch.zeros(output_dim))
+
+    def forward(self, features):
+        x = self.dense(features)
+        x = gelu(x)
+        x = self.layer_norm(x)
+        # project back to size of vocabulary with bias
+        # x = F.linear(x, self.weight) + self.bias
+        x = self.decoder(x) + self.bias
+        return x
+
+class LucaGPLMTransformerLayer(nn.Module):
+    def __init__(
+        self,
+        embed_dim,
+        ffn_embed_dim,
+        attention_heads,
+        add_bias_kv=True,
+        use_lucagplm1b_layer_norm=False,
+        use_rotary_embeddings: bool=True,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.ffn_embed_dim = ffn_embed_dim
+        self.attention_heads = attention_heads
+        self.use_rotary_embeddings = use_rotary_embeddings
+        
+        LucaGPLMLayerNorm = LucaGPLM1bLayerNorm if use_lucagplm1b_layer_norm else LucaGPLM1LayerNorm
+
+        self.pre_layer_norm = LucaGPLMLayerNorm(self.embed_dim)
+
+        self.self_attn = LucaGPLMMultiheadAttentionWithSDPA(
+            self.embed_dim,
+            self.attention_heads,
+            add_bias_kv=add_bias_kv,
+            add_zero_attn=False,
+            self_attention=True,
+            use_rotary_embeddings=self.use_rotary_embeddings,
+        )
+
+        # post layer norm
+        self.post_layer_norm = LucaGPLMLayerNorm(self.embed_dim)
+
+        # dimension increase by the fully connected layer
+        self.fc1 = nn.Linear(self.embed_dim, self.ffn_embed_dim)
+
+        # dimension reduction by the fully connected layer
+        self.fc2 = nn.Linear(self.ffn_embed_dim, self.embed_dim)
+
+    def forward(
+        self,
+        x,
+        self_attn_mask=None,
+        self_attn_padding_mask=None,
+        need_head_weights=False
+    ):
+        residual = x
+        x = self.pre_layer_norm(x)
+        x, attn = self.self_attn(
+            query=x,
+            key=x,
+            value=x,
+            key_padding_mask=self_attn_padding_mask,
+            need_weights=True,
+            need_head_weights=need_head_weights,
+            attn_mask=self_attn_mask,
+        )
+        x = residual + x
+
+        residual = x
+        x = self.post_layer_norm(x)
+        x = gelu(self.fc1(x))
+        x = self.fc2(x)
+        x = residual + x
+
+        return x, attn
+
+class LucaGPLMEmbeddings(nn.Module):
+    def __init__(self, config: LucaGPLMConfig):
+        super().__init__()
+        
+        # Store config flags for forward pass
+        self.no_position_embeddings = getattr(config, 'no_position_embeddings', False)
+        self.no_token_type_embeddings = getattr(config, 'no_token_type_embeddings', False)
+        self.use_embed_layer_norm = getattr(config, 'use_embed_layer_norm', True)
+        self.embed_scale = getattr(config, 'embed_scale', 1.0)
+        self.token_dropout = getattr(config, 'token_dropout', False)
+        
+        # Token ids for special tokens (matching old model)
+        self.mask_idx = getattr(config, 'mask_token_id', 4)
+        self.padding_idx = getattr(config, 'pad_token_id', 0)
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+        
+        # Only create position embeddings if not disabled
+        if not self.no_position_embeddings:
+            self.embed_pos = nn.Embedding(config.max_position_embeddings, config.hidden_size)
+        else:
+            self.embed_pos = None
+            
+        # Only create token type embeddings if not disabled    
+        if not self.no_token_type_embeddings:
+            self.embed_type = nn.Embedding(config.type_vocab_size, config.hidden_size)
+        else:
+            self.embed_type = None
+            
+        # Only create layer norm if enabled
+        if self.use_embed_layer_norm:
+            self.embed_layer_norm = LucaGPLM1bLayerNorm(config.hidden_size)
+        else:
+            self.embed_layer_norm = None
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        token_type_ids: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        input_shape = input_ids.size()
+        seq_length = input_shape[1]
+
+        # Start with token embeddings and apply embed_scale
+        inputs_embeds = self.embed_scale * self.embed_tokens(input_ids)
+        
+        # Add position embeddings if enabled
+        if not self.no_position_embeddings and self.embed_pos is not None:
+            if position_ids is None:
+                position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device)
+                position_ids = position_ids.unsqueeze(0).expand(input_shape)
+            position_embeddings = self.embed_scale * self.embed_pos(position_ids)
+            inputs_embeds = inputs_embeds + position_embeddings
+
+        # Add token type embeddings if enabled
+        if not self.no_token_type_embeddings and self.embed_type is not None:
+            if token_type_ids is None:
+                token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=input_ids.device)
+            token_type_embeddings = self.embed_scale * self.embed_type(token_type_ids)
+            inputs_embeds = inputs_embeds + token_type_embeddings
+        
+        # Apply layer norm if enabled
+        if self.use_embed_layer_norm and self.embed_layer_norm is not None:
+            embeddings = self.embed_layer_norm(inputs_embeds)
+        else:
+            embeddings = inputs_embeds
+
+        # Apply token dropout (matching old model behavior)
+        if self.token_dropout and self.training:
+            # Zero out masked token embeddings
+            embeddings = embeddings.masked_fill((input_ids == self.mask_idx).unsqueeze(-1), 0.0)
+            
+            # Apply token dropout scaling
+            mask_ratio_train = 0.15 * 0.8
+            padding_mask = input_ids.eq(self.padding_idx)
+            src_lengths = (~padding_mask).sum(-1)
+            mask_ratio_observed = (input_ids == self.mask_idx).sum(-1).to(embeddings.dtype) / src_lengths
+            embeddings = embeddings * (1 - mask_ratio_train) / (1 - mask_ratio_observed)[:, None, None]
+
+        # Apply padding mask to embeddings
+        padding_mask = input_ids.eq(self.padding_idx)
+        if padding_mask.any():
+            embeddings = embeddings * (1 - padding_mask.unsqueeze(-1).type_as(embeddings))
+
+        return embeddings
+
+class LucaGPLMEncoder(nn.Module):
+    def __init__(self, config: LucaGPLMConfig):
+        super().__init__()
+
+        self.layers = nn.ModuleList([
+            LucaGPLMTransformerLayer(
+                config.hidden_size,
+                4 * config.hidden_size,  # ffn_embed_dim = 4 * embed_dim
+                config.num_attention_heads,
+                add_bias_kv=False,
+                use_lucagplm1b_layer_norm=True,
+                use_rotary_embeddings=True,
+            )
+            for _ in range(config.num_hidden_layers)
+        ])
+        
+        self.use_last_layer_norm = getattr(config, 'use_last_layer_norm', True)
+        if self.use_last_layer_norm:
+            self.last_layer_norm = LucaGPLM1bLayerNorm(config.hidden_size)
+        else:
+            self.last_layer_norm = None
+
+        self.padding_idx = config.pad_token_id
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        output_hidden_states: bool = False,
+        return_dict: bool = True,
+        need_head_weights: bool = False,
+        repr_layers: Optional[List[int]] = None,
+        use_last_layer_norm: bool = True,
+    ) -> Union[Tuple[torch.Tensor], BaseModelOutput]:
+        all_hidden_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+        
+        if repr_layers is None:
+            repr_layers = [-1]
+        
+        # 转换为原始模型的索引系统
+        layer_size = len(self.layers)
+        repr_layers = [(i + layer_size + 1) % (layer_size + 1) for i in repr_layers]
+        repr_layers = set(repr_layers)
+        hidden_representations = {}
+
+        # Process attention mask - 原始模型期望的是padding mask
+        if attention_mask is None:
+            padding_mask = hidden_states.new_zeros(hidden_states.shape[:2]).eq(self.padding_idx)
+        else:
+            # 原始模型中 padding_mask 是 True 表示 padding位置
+            padding_mask = attention_mask.eq(0)
+
+        # 0: embedding layer
+        if 0 in repr_layers:
+            hidden_representations[0] = hidden_states
+
+        # 转换为 (seq_len, batch_size, hidden_size) 格式，与原始模型一致
+        hidden_states = hidden_states.transpose(0, 1)
+        
+        if not padding_mask.any():
+            padding_mask = None
+
+        # 是否需要返回head weights
+        if need_head_weights or output_attentions:
+            attn_weights = []
+
+        for layer_idx, layer_module in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states.transpose(0, 1),)
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    layer_module.__call__,
+                    hidden_states,
+                    None,  # self_attn_mask
+                    padding_mask,
+                    need_head_weights or output_attentions,
+                )
+            else:
+                layer_outputs = layer_module(
+                    hidden_states,
+                    self_attn_mask=None,
+                    self_attn_padding_mask=padding_mask,
+                    need_head_weights=need_head_weights or output_attentions,
+                )
+
+            hidden_states, attn = layer_outputs
+
+            if (layer_idx + 1) in repr_layers:
+                hidden_representations[layer_idx + 1] = hidden_states.transpose(0, 1)
+
+            if need_head_weights or output_attentions:
+                # (H, B, L, L) => (B, H, L, L)
+                attn_weights.append(attn.transpose(1, 0))
+
+        # 应用最后的layer norm
+        if self.last_layer_norm is not None and use_last_layer_norm:
+            hidden_states = self.last_layer_norm(hidden_states)
+
+        # 转换回 (batch_size, seq_len, hidden_size) 格式
+        hidden_states = hidden_states.transpose(0, 1)
+
+        # last hidden representation should have layer norm applied
+        if (layer_idx + 1) in repr_layers:
+            hidden_representations[layer_idx + 1] = hidden_states
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if need_head_weights or output_attentions:
+            # 将attention weights转换为正确格式
+            if attn_weights:
+                # B x Layers x H x L x L
+                all_attentions = torch.stack(attn_weights, 1)
+                if padding_mask is not None:
+                    attention_mask_expanded = 1 - padding_mask.type_as(all_attentions)
+                    attention_mask_expanded = attention_mask_expanded.unsqueeze(1) * attention_mask_expanded.unsqueeze(2)
+                    all_attentions = all_attentions * attention_mask_expanded[:, None, None, :, :]
+            
+            if not output_attentions:
+                all_attentions = None
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, all_hidden_states, all_attentions] if v is not None)
+
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+            attentions=all_attentions,
+        )
+
+class LucaGPLMPreTrainedModel(PreTrainedModel):
+    config_class = LucaGPLMConfig
+    base_model_prefix = "lucaone"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["LucaGPLMTransformerLayer"]
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, (LucaGPLM1LayerNorm, LucaGPLM1bLayerNorm)):
+            if hasattr(module, 'weight') and module.weight is not None:
+                module.weight.data.fill_(1.0)
+            if hasattr(module, 'bias') and module.bias is not None:
+                module.bias.data.zero_()
+
+class LucaGPLMModel(LucaGPLMPreTrainedModel):
+    """
+    The LucaGPLM model for extracting sequence representations and optionally predicting contacts.
+    Based on the original LucaGPLM implementation but restructured to use modern transformers architecture.
+    """
+    
+    def __init__(self, config: LucaGPLMConfig):
+        super().__init__(config)
+        self.config = config
+        self.embeddings = LucaGPLMEmbeddings(self.config)
+        self.encoder = LucaGPLMEncoder(self.config)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embeddings.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embeddings.embed_tokens = value
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_contacts: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        need_head_weights: Optional[bool] = None,
+        repr_layers: Optional[List[int]] = None,
+        use_last_layer_norm: Optional[bool] = True,
+    ) -> Union[Tuple[torch.Tensor], BaseModelOutput]:
+        
+        output_attentions = output_attentions if output_attentions is not None else getattr(self.config, 'output_attentions', False)
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else getattr(self.config, 'output_hidden_states', False)
+        return_contacts = return_contacts if return_contacts is not None else False
+        return_dict = return_dict if return_dict is not None else getattr(self.config, 'use_return_dict', True)
+        need_head_weights = need_head_weights if need_head_weights is not None else return_contacts  # Need attention weights for contacts
+        use_last_layer_norm = use_last_layer_norm if use_last_layer_norm is not None else True
+
+        # Force output_attentions=True when return_contacts=True since we need attention weights
+        if return_contacts:
+            output_attentions = True
+            need_head_weights = True
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            input_shape = input_ids.size()
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        batch_size, seq_length = input_shape
+        device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+        # Create attention mask if not provided
+        if attention_mask is None:
+            attention_mask = torch.ones(input_shape, device=device)
+
+        # Get embeddings
+        if inputs_embeds is None:
+            embedding_output = self.embeddings(
+                input_ids=input_ids,
+                position_ids=position_ids,
+                token_type_ids=token_type_ids,
+            )
+        else:
+            embedding_output = inputs_embeds
+
+        # Pass through encoder
+        encoder_outputs = self.encoder(
+            embedding_output,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            need_head_weights=need_head_weights,
+            repr_layers=repr_layers,
+            use_last_layer_norm=use_last_layer_norm,
+        )
+        
+        sequence_output = encoder_outputs[0]
+        
+        # Handle contact prediction
+        contacts = None
+        if return_contacts and encoder_outputs.attentions is not None:
+            # Simple contact prediction using attention weights
+            # This is a simplified implementation - you can enhance this later
+            attentions = encoder_outputs.attentions
+            # Average over layers and heads, then symmetrize
+            averaged_attention = attentions.mean(dim=(1, 2))  # Average over layers and heads
+            contacts = (averaged_attention + averaged_attention.transpose(-1, -2)) / 2
+            
+            # Remove special tokens (BOS/EOS) if present
+            if attention_mask is not None:
+                # Find actual sequence positions (non-padding)
+                seq_lens = attention_mask.sum(dim=1)
+                # For now, keep the full contact map - you can trim special tokens later if needed
+
+        if not return_dict:
+            outputs = (sequence_output, ) + encoder_outputs[1:]
+            if contacts is not None:
+                outputs = outputs + (contacts,)
+            return outputs
+
+        # Create output object with contacts
+        output = BaseModelOutput(
+            last_hidden_state=sequence_output,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )
+        
+        # Add contacts as an attribute if computed
+        if contacts is not None:
+            output.contacts = contacts
+            
+        return output
+
+class LucaGPLMForMaskedLM(LucaGPLMPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        # 基础编码器
+        self.lucaone = LucaGPLMModel(config)
+
+        # MLM 预测头
+        self.lm_head = LucaGPLMRobertaLMHead(
+            embed_dim=config.hidden_size,
+            output_dim=config.vocab_size
+        )
+        self._tied_weights_keys = [
+            "lucaone.embeddings.embed_tokens.weight",
+            "lm_head.decoder.weight"
+        ]
+        # 初始化权重并进行权重绑定
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.lucaone.get_input_embeddings()
+
+    def get_output_embeddings(self):
+        return self.lm_head.decoder
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head.decoder = new_embeddings
+
+    def forward(
+            self,
+            input_ids: Optional[torch.Tensor] = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            token_type_ids: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.Tensor] = None,
+            labels: Optional[torch.Tensor] = None, # MLM 训练时的标签
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, MaskedLMOutput]:
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # 1. 获取基础模型的输出 (Hidden States)
+        outputs = self.lucaone(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        sequence_output = outputs[0] # (batch_size, seq_len, hidden_size)
+
+        # 2. 通过 MLM Head 得到预测结果 (Logits)
+        prediction_scores = self.lm_head(sequence_output)
+
+        masked_lm_loss = None
+        if labels is not None:
+            # 3. 计算 MLM Loss
+            loss_fct = nn.CrossEntropyLoss(ignore_index=-100) # 默认 ignore_index=-100
+            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
+
+        if not return_dict:
+            output = (prediction_scores,) + outputs[2:]
+            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
+
+        return MaskedLMOutput(
+            loss=masked_lm_loss,
+            logits=prediction_scores,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+class LucaGPLMForSequenceClassification(LucaGPLMPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.classifier_num_labels
+        self.task_level = config.task_level
+        self.task_type = config.task_type
+        assert self.task_level == "seq_level"
+        self.classifier_pooling_type = config.classifier_pooling_type
+        self.classifier_loss_type = config.classifier_loss_type
+        self.classifier_loss_reduction = config.classifier_loss_reduction
+        self.classifier_pos_weight = config.classifier_pos_weight
+        self.classifier_weight = config.classifier_weight
+        self.lucaone = LucaGPLMModel(config) # 基础模型
+        if self.classifier_pooling_type == "value_attention":
+            self.pooler = LucaGPLMGlobalMaskValueAttentionPooling1D(config.hidden_size)
+        elif self.classifier_pooling_type == "context_attention":
+            self.pooler = LucaGPLMGlobalMaskContextAttentionPooling1D(embed_size=config.hidden_size)
+        elif self.classifier_pooling_type == "weighted_attention":
+            self.pooler = LucaGPLMGlobalMaskWeightedAttentionPooling1D(embed_size=config.hidden_size)
+        else:
+            self.pooler = None
+        self.dropout = nn.Dropout(config.classifier_dropout_prob)
+
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+        if self.task_type == "multi_class":
+            weight = None
+            if self.classifier_weight:
+                if isinstance(self.classifier_weight, str) or isinstance(self.classifier_weight, int):
+                    weight = torch.tensor([float(self.classifier_weight)] * self.num_labels, dtype=torch.float32)
+                elif isinstance(self.classifier_weight, float):
+                    weight = torch.tensor([self.classifier_weight] * self.num_labels, dtype=torch.float32)
+                elif isinstance(self.classifier_weight, list):
+                    weight = torch.tensor(self.classifier_weight, dtype=torch.float32)
+            self.loss_fct = nn.CrossEntropyLoss(weight=weight, reduction="mean")
+        elif self.task_type == "binary_class":
+            pos_weight = None
+            if self.classifier_pos_weight:
+                if isinstance(self.classifier_pos_weight, str) or isinstance(self.classifier_pos_weight, int):
+                    pos_weight = torch.tensor([float(self.classifier_pos_weight)], dtype=torch.float32)
+                elif isinstance(self.classifier_pos_weight, float):
+                    pos_weight = torch.tensor([self.classifier_pos_weight], dtype=torch.float32)
+            self.loss_fct = nn.BCEWithLogitsLoss(pos_weight=pos_weight, reduction="mean")
+        elif self.task_type == "regression":
+            if self.classifier_loss_type == "mae":
+                self.loss_fct = nn.L1Loss(reduction="mean")
+            else:
+                self.loss_fct = nn.MSELoss(reduction="mean")
+        elif self.task_type == "multi_label":
+            pos_weight = None
+            if self.classifier_pos_weight:
+                if isinstance(self.classifier_pos_weight, str) or isinstance(self.classifier_pos_weight, int):
+                   pos_weight = torch.tensor([float(self.classifier_pos_weight)] * self.num_labels, dtype=torch.float32)
+                elif isinstance(self.classifier_pos_weight, float):
+                    pos_weight = torch.tensor([self.classifier_pos_weight] * self.num_labels, dtype=torch.float32)
+            self.loss_fct = nn.BCEWithLogitsLoss(pos_weight=pos_weight, reduction=self.classifier_loss_reduction)
+        else:
+            raise ValueError("Invalid task type: %s" % self.task_type)
+        self.post_init()
+
+    def forward(
+            self,
+            input_ids=None,
+            token_type_ids=None,
+            attention_mask=None,
+            labels=None,
+            return_dict=None
+    ):
+        return_dict = return_dict if return_dict is not None else getattr(self.config, 'use_return_dict', True)
+        outputs = self.lucaone(
+            input_ids,
+            token_type_ids=token_type_ids,
+            attention_mask=attention_mask,
+            return_dict=return_dict
+        )
+        if self.pooler is not None:
+            pooled_output = self.pooler(outputs[0])
+        elif self.classifier_pooling_type == "cls":
+            # 取 CLS token
+            pooled_output = outputs[0][:, 0, :]
+        elif self.classifier_pooling_type == "mean":
+            pooled_output = outputs[0].mean(dim=1)
+        else:
+            raise ValueError("Invalid classifier pooling type: %s" % self.classifier_pooling_type)
+
+        pooled_output = self.dropout(pooled_output)
+        logits = self.classifier(pooled_output)
+
+        loss = None
+        if labels is not None:
+            if self.task_type == "multi_class":
+                loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.task_type == "binary_class":
+                loss = self.loss_fct(logits.view(-1), labels.view(-1).float())
+            elif self.task_type == "regression":
+                loss = self.loss_fct(logits.view(-1), labels.view(-1))
+            elif self.task_type == "multi_label":
+                loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1, self.num_labels).float())
+            else:
+                raise ValueError("Invalid task type: %s" % self.task_type)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(loss=loss, logits=logits)
+
+class LucaGPLMForTokenClassification(LucaGPLMPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.classifier_num_labels
+        self.task_level = config.task_level
+        self.task_type = config.task_type
+        assert self.task_level == "token_level"
+        self.classifier_pooling_type = config.classifier_pooling_type
+        self.classifier_loss_type = config.classifier_loss_type
+        self.classifier_loss_reduction = config.classifier_loss_reduction
+        self.classifier_pos_weight = config.classifier_pos_weight
+        self.classifier_weight = config.classifier_weight
+        self.lucaone = LucaGPLMModel(config) # 基础模型
+        self.dropout = nn.Dropout(config.classifier_dropout_prob)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+        if self.task_type == "multi_class":
+            weight = None
+            if self.classifier_weight:
+                # [1, 1, 1, ,1, 1...] length: num_labels
+                if isinstance(self.classifier_weight, str) or isinstance(self.classifier_weight, int):
+                    weight = torch.tensor([float(self.classifier_weight)] * self.num_labels, dtype=torch.float32)
+                elif isinstance(self.classifier_weight, float):
+                    weight = torch.tensor([self.classifier_weight] * self.num_labels, dtype=torch.float32)
+                elif isinstance(self.classifier_weight, list):
+                    weight = torch.tensor(self.classifier_weight, dtype=torch.float32)
+            self.loss_fct = nn.CrossEntropyLoss(weight=weight, reduction="mean")
+        elif self.task_type == "binary_class":
+            pos_weight = None
+            if self.classifier_pos_weight:
+                if isinstance(self.classifier_pos_weight, str) or isinstance(self.classifier_pos_weight, int):
+                    pos_weight = torch.tensor([float(self.classifier_pos_weight)], dtype=torch.float32)
+                elif isinstance(self.classifier_pos_weight, float):
+                    pos_weight = torch.tensor([float(self.classifier_pos_weight)], dtype=torch.float32)
+            self.loss_fct = nn.BCEWithLogitsLoss(pos_weight=pos_weight, reduction="mean")
+        elif self.task_type == "regression":
+            if self.classifier_loss_type == "mae":
+                self.loss_fct = nn.L1Loss(reduction="mean")
+            else:
+                self.loss_fct = nn.MSELoss(reduction="mean")
+        elif self.task_type == "multi_label":
+            pos_weight = None
+            if self.classifier_pos_weight:
+                if isinstance(self.classifier_pos_weight, str) or isinstance(self.classifier_pos_weight, int):
+                    pos_weight = torch.tensor([float(self.classifier_pos_weight)] * self.num_labels, dtype=torch.float32)
+                elif isinstance(self.classifier_pos_weight, float):
+                    pos_weight = torch.tensor([self.classifier_pos_weight] * self.num_labels, dtype=torch.float32)
+            self.loss_fct = nn.BCEWithLogitsLoss(pos_weight=pos_weight, reduction=self.classifier_loss_reduction)
+        else:
+            raise ValueError("Invalid task type: %s" % self.task_type)
+        self.post_init()
+
+    def forward(
+            self,
+            input_ids=None,
+            token_type_ids=None,
+            attention_mask=None,
+            labels=None,
+            return_dict=None
+    ):
+        return_dict = return_dict if return_dict is not None else getattr(self.config, 'use_return_dict', True)
+        outputs = self.lucaone(
+            input_ids,
+            token_type_ids=token_type_ids,
+            attention_mask=attention_mask,
+            return_dict=return_dict
+        )
+        sequence_output = outputs[0][:, 1:-1, :] # (B, L, H)
+
+        sequence_output = self.dropout(sequence_output)
+        logits = self.classifier(sequence_output)
+
+        loss = None
+        if labels is not None:
+            if self.task_type == "multi_class":
+                loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.task_type == "binary_class":
+                loss = self.loss_fct(logits.view(-1), labels.view(-1).float())
+            elif self.task_type == "regression":
+                loss = self.loss_fct(logits.view(-1), labels.view(-1))
+            elif self.task_type == "multi_label":
+                loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1, self.num_labels).float())
+            else:
+                raise ValueError("Invalid task type: %s" % self.task_type)
+
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+        return TokenClassifierOutput(loss=loss, logits=logits)
+
+__all__ = [
+    "LucaGPLMModel",
+    "LucaGPLMPreTrainedModel",
+    "LucaGPLMForMaskedLM",
+    "LucaGPLMForSequenceClassification",
+    "LucaGPLMForTokenClassification"
+]

tokenization_lucaone.py

@@ -0,0 +1,432 @@
+#!/usr/bin/env python
+# encoding: utf-8
+'''
+@license: (C) Copyright 2025, Hey.
+@author: Hey
+@email: [email protected]
+@tel: 137****6540
+@datetime: 2025/12/30 11:33
+@project: lucaone
+@file: tokenization_lucaone
+@desc: tokenization_lucaone
+'''
+
+import os
+import json
+import itertools
+from typing import List, Optional, Dict, Any, Tuple, Union
+from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
+
+def gene_seq_replace(seq):
+    """
+    Gene sequence preprocessing: A->1, U/T->2, C->3, G->4, N->5
+    Optimized for performance.
+    """
+    # 使用字典映射比 if-else 判断快
+    mapping = {
+        'A': '1', 'a': '1',
+        'T': '2', 't': '2', 'U': '2', 'u': '2',
+        'C': '3', 'c': '3',
+        'G': '4', 'g': '4'
+    }
+    # 对于不在字典中的字符（如 N），默认返回 '5'
+    return "".join([mapping.get(ch, '5') for ch in seq])
+
+class LucaGPLMTokenizer(PreTrainedTokenizer):
+    """
+    HuggingFace-compatible tokenizer that performs identical tokenization 
+    to the old model's Alphabet class.
+    """
+    
+    # Vocabulary definitions matching the old model
+    gene_prepend_toks = ['[PAD]', '[UNK]']
+    gene_append_toks = ['[CLS]', '[SEP]', '[MASK]']
+    gene_standard_toks = ['1', '2', '3', '4', '5', '.', '-', '*']
+    
+    prot_prepend_toks = ['[PAD]', '[UNK]']
+    prot_append_toks = ['[CLS]', '[SEP]', '[MASK]']
+    prot_standard_toks = ['L', 'A', 'G', 'V', 'S', 'E', 'R', 'T', 'I', 'D', 'P', 'K', 'Q', 'N', 'F', 'Y', 'M', 'H', 'W', 'C', 'X', 'B', 'U', 'Z', 'O', 'J', '.', '-', '*']
+    
+    gene_prot_prepend_toks = ['[PAD]', '[UNK]']
+    gene_prot_append_toks = ['[CLS]', '[SEP]', '[MASK]']
+    # EXACT VOCABULARY ORDER FROM ORIGINAL ALPHABET CLASS
+
+    gene_prot_standard_toks = [
+        '1',      # 5 - gene A (after gene_seq_replace)
+        '2',      # 6 - gene T/U (after gene_seq_replace) 
+        '3',      # 7 - gene C (after gene_seq_replace)
+        '4',      # 8 - gene G (after gene_seq_replace)
+        '5',      # 9 - gene N/unknown
+        'L',      # 10 - protein
+        'A',      # 11 - protein
+        'G',      # 12 - protein
+        'V',      # 13 - protein
+        'S',      # 14 - protein
+        'E',      # 15 - protein
+        'R',      # 16 - protein
+        'T',      # 17 - protein
+        'I',      # 18 - protein
+        'D',      # 19 - protein
+        'P',      # 20 - protein
+        'K',      # 21 - protein
+        'Q',      # 22 - protein
+        'N',      # 23 - protein
+        'F',      # 24 - protein
+        'Y',      # 25 - protein
+        'M',      # 26 - protein
+        'H',      # 27 - protein
+        'W',      # 28 - protein
+        'C',      # 29 - protein
+        'X',      # 30 - protein unknown
+        'B',      # 31 - protein
+        'U',      # 32 - protein
+        'Z',      # 33 - protein
+        'O',      # 34 - protein
+        'J',      # 35 - protein
+        '.',      # 36 - special
+        '-',      # 37 - special
+        '*'       # 38 - special
+    ]
+
+    def __init__(
+        self,
+        vocab_type: str = "gene_prot",
+        prepend_bos: bool = True,
+        append_eos: bool = True,
+        unk_token="[UNK]",
+        pad_token="[PAD]",
+        cls_token="[CLS]",
+        sep_token="[SEP]",
+        mask_token="[MASK]",
+        **kwargs
+    ):
+        # Set vocabulary based on type
+        if vocab_type.lower() == "prot":
+            prepend_toks = self.prot_prepend_toks
+            append_toks = self.prot_append_toks
+            standard_toks = self.prot_standard_toks
+        elif vocab_type.lower() == "gene":
+            prepend_toks = self.gene_prepend_toks
+            append_toks = self.gene_append_toks
+            standard_toks = self.gene_standard_toks
+        elif vocab_type.lower() in ["gene_prot", "prot_gene"]:
+            prepend_toks = self.gene_prot_prepend_toks
+            append_toks = self.gene_prot_append_toks
+            standard_toks = self.gene_prot_standard_toks
+        else:
+            raise ValueError(f"Not support tokenizer vocab_type: {vocab_type}")
+        
+        # Build vocabulary
+        self.all_toks = list(prepend_toks) + list(append_toks) + list(standard_toks)
+        self.tok_to_idx = {tok: i for i, tok in enumerate(self.all_toks)}
+        self.idx_to_tok = {i: tok for i, tok in enumerate(self.all_toks)}
+        
+        # Store configuration
+        self.vocab_type = vocab_type
+        self.prepend_bos = prepend_bos
+        self.append_eos = append_eos
+        self.unique_no_split_tokens = self.all_toks.copy()
+        
+        # Special token indices
+        self.unk_idx = self.tok_to_idx.get("[UNK]", 1)
+        self.padding_idx = self.tok_to_idx.get("[PAD]", 0)
+        self.cls_idx = self.tok_to_idx.get("[CLS]", 2)
+        self.mask_idx = self.tok_to_idx.get("[MASK]", 4)
+        self.eos_idx = self.tok_to_idx.get("[SEP]", 3)
+        
+        super().__init__(
+            unk_token=unk_token,
+            pad_token=pad_token,
+            cls_token=cls_token,
+            sep_token=sep_token,
+            mask_token=mask_token,
+            **kwargs
+        )
+
+    def get_vocab(self) -> Dict[str, int]:
+        return self.tok_to_idx.copy()
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.all_toks)
+
+    def get_idx(self, tok):
+        return self.tok_to_idx.get(tok, self.unk_idx)
+
+    def get_tok(self, idx):
+        return self.idx_to_tok.get(idx, "[UNK]")
+
+    def _tokenize_char_level(self, text: str) -> List[str]:
+        """Simple character-level tokenization (fallback)"""
+        return list(text)
+
+    def _tokenize(self, text: str) -> List[str]:
+        """
+        Tokenize text using the same logic as the old Alphabet.tokenize() method
+        """
+        text = text.strip()
+        if not text:
+            return []
+            
+        return list(text)
+
+    def _convert_token_to_id(self, token: str) -> int:
+        return self.get_idx(token)
+
+    def _convert_id_to_token(self, index: int) -> str:
+        return self.get_tok(index)
+
+    def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        return "".join(tokens)
+
+    def _convert_text_to_ids(self, text: str, seq_type: str) -> List[int]:
+        """Internal helper to convert text to IDs without special tokens."""
+        if seq_type == "gene":
+            text = gene_seq_replace(text)
+        tokens = self._tokenize(text)
+        return [self._convert_token_to_id(token) for token in tokens]
+
+    def build_inputs_with_special_tokens(self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None) -> List[int]:
+        """
+        Build model inputs from a sequence by adding special tokens.
+        This mimics the old model's prepend_bos and append_eos behavior.
+        """
+        result = token_ids_0.copy()
+        
+        if self.prepend_bos:
+            result = [self.cls_idx] + result
+        if self.append_eos:
+            result = result + [self.eos_idx]
+            
+        return result
+
+    def get_special_tokens_mask(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
+    ) -> List[int]:
+        """
+        Retrieve sequence ids from a token list.
+        """
+        if already_has_special_tokens:
+            return super().get_special_tokens_mask(
+                token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
+            )
+
+        result = [0] * len(token_ids_0)
+        if self.prepend_bos:
+            result = [1] + result
+        if self.append_eos:
+            result = result + [1]
+        return result
+
+    def encode(
+        self,
+        text: str,
+        seq_type: str = "gene",
+        add_special_tokens: bool = True,
+        padding: Union[bool, str] = False, # 虽然 encode 通常不处理 padding，但保持 API 兼容性
+        truncation: bool = False,          # <--- 关键参数
+        max_length: Optional[int] = None,  # <--- 关键参数
+        **kwargs
+    ) -> List[int]:
+        
+        # 1. 基础转换
+        token_ids = self._convert_text_to_ids(text, seq_type)
+        
+        # 2. 添加特殊 token
+        if add_special_tokens:
+            token_ids = self.build_inputs_with_special_tokens(token_ids)
+            
+        # 3. 执行截断 (修复点：之前这里缺失逻辑)
+        if truncation and max_length is not None and len(token_ids) > max_length:
+            token_ids = token_ids[:max_length]
+            # 如果启用了 append_eos，强行把截断后的最后一位改回 SEP
+            if add_special_tokens and self.append_eos:
+                token_ids[-1] = self.eos_idx
+                
+        return token_ids
+
+    def __call__(
+        self,
+        text: Union[str, List[str]],
+        text_pair: Optional[Union[str, List[str]]] = None,
+        seq_type: str = "gene",
+        add_special_tokens: bool = True,
+        padding: Union[bool, str] = False,
+        max_length: Optional[int] = None,
+        return_attention_mask: bool = True,
+        return_token_type_ids: bool = True,
+        return_tensors: Optional[str] = None,
+        truncation: bool = False,
+        **kwargs
+    ) -> Union[Dict[str, Any], List[Dict[str, Any]]]:
+        """
+        Main callable method for tokenization - HuggingFace standard interface
+        """
+        if isinstance(text, list):
+            # Handle batch processing
+            return self.batch_encode_plus(
+                text,
+                text_pair=text_pair,
+                seq_type=seq_type,
+                add_special_tokens=add_special_tokens,
+                padding=padding,
+                max_length=max_length,
+                return_attention_mask=return_attention_mask,
+                return_token_type_ids=return_token_type_ids,
+                return_tensors=return_tensors,
+                truncation=truncation,
+                **kwargs
+            )
+        else:
+            # Handle single text
+            return self.encode_plus(
+                text,
+                text_pair=text_pair,
+                seq_type=seq_type,
+                add_special_tokens=add_special_tokens,
+                padding=padding,
+                max_length=max_length,
+                return_attention_mask=return_attention_mask,
+                return_token_type_ids=return_token_type_ids,
+                return_tensors=return_tensors,
+                truncation=truncation,
+                **kwargs
+            )
+
+    def batch_encode_plus(self, *args, **kwargs):
+        # 显式调用父类，或者保留你原有的实现，只要确保内部调用的是修复后的 encode_plus 即可
+        return super().batch_encode_plus(*args, **kwargs)
+
+    def encode_plus(
+        self,
+        text: str,
+        text_pair: Optional[str] = None,
+        seq_type: str = "gene",
+        add_special_tokens: bool = True,
+        padding: Union[bool, str] = False,
+        max_length: Optional[int] = None,
+        return_attention_mask: bool = True,
+        return_token_type_ids: bool = True,
+        return_tensors: Optional[str] = None,
+        truncation: bool = False,
+        **kwargs
+    ) -> Dict[str, Any]:
+        
+        # 调用修复后的 encode，它现在会正确处理截断
+        token_ids = self.encode(
+            text, 
+            seq_type=seq_type, 
+            add_special_tokens=add_special_tokens,
+            truncation=truncation,
+            max_length=max_length
+        )
+        
+        # 处理 Padding
+        attention_mask = [1] * len(token_ids)
+        if padding == "max_length" and max_length is not None:
+            if len(token_ids) < max_length:
+                pad_length = max_length - len(token_ids)
+                token_ids.extend([self.padding_idx] * pad_length)
+                attention_mask.extend([0] * pad_length)
+        # 注意：padding=True (dynamic padding) 通常由 batch_encode_plus 处理，这里单条通常不处理
+        
+        result = {"input_ids": token_ids}
+        
+        if return_attention_mask:
+            result["attention_mask"] = attention_mask
+        
+        if return_token_type_ids:
+            # 0 for gene, 1 for protein
+            type_value = 0 if seq_type == "gene" else 1
+            result["token_type_ids"] = [type_value] * len(token_ids)
+        
+        if return_tensors == "pt":
+            import torch
+            for key, value in result.items():
+                result[key] = torch.tensor(value, dtype=torch.long).unsqueeze(0)
+        
+        return result
+
+    def encode_old_model_style(
+        self,
+        text: str,
+        seq_type: str = "gene", 
+        max_length: int = None
+    ) -> List[int]:
+        """
+        Encode using the EXACT same process as the old model's encoder function.
+        This replicates the logic from src/llm/lucaone_virus/get_embedding.py:encoder()
+        """
+        # Preprocess gene sequences (done in get_embedding function BEFORE calling encoder)
+        if seq_type == "gene":
+            text = gene_seq_replace(text)
+        
+        # Call tokenizer.encode (which does NOT include BOS/EOS in old model)
+        seq_encoded = self.encode(text, seq_type=seq_type, add_special_tokens=False)
+        
+        # Apply max_length truncation if specified  
+        if max_length and len(seq_encoded) > max_length:
+            seq_encoded = seq_encoded[:max_length]
+        
+        # Calculate processed_seq_len (as done in old model)
+        processed_seq_len = len(seq_encoded) + int(self.prepend_bos) + int(self.append_eos)
+        
+        # Create input_ids tensor (as done in old model encoder function)
+        input_ids = [self.padding_idx] * processed_seq_len
+        
+        # Add BOS token if enabled
+        if self.prepend_bos:
+            input_ids[0] = self.cls_idx
+            
+        # Place the encoded sequence
+        start_idx = int(self.prepend_bos)
+        for i, token_id in enumerate(seq_encoded):
+            input_ids[start_idx + i] = token_id
+            
+        # Add EOS token if enabled  
+        if self.append_eos:
+            input_ids[len(seq_encoded) + int(self.prepend_bos)] = self.eos_idx
+            
+        return input_ids
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        """
+        Save the tokenizer vocabulary to a JSON file.
+        Required by HuggingFace tokenizer interface.
+        """
+        if filename_prefix is None:
+            filename_prefix = ""
+        else:
+            filename_prefix = filename_prefix + "-"
+        
+        vocab_file = os.path.join(save_directory, f"{filename_prefix}vocab.json")
+        vocab_dict = self.get_vocab()
+        with open(vocab_file, "w", encoding="utf-8") as f:
+            json.dump(vocab_dict, f, ensure_ascii=False, indent=2)
+        
+        return (vocab_file,)
+
+    @classmethod 
+    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
+        """
+        Load tokenizer from pretrained model path (standard HuggingFace interface)
+        """
+        vocab_file = os.path.join(pretrained_model_name_or_path, "vocab.json")
+        if os.path.exists(vocab_file):
+            print("Load from saved vocabulary (not implemented yet, use default)")
+            return cls(vocab_type="gene_prot", **kwargs)
+        else:
+            return cls(vocab_type="gene_prot", **kwargs)
+
+class LucaGPLMTokenizerFast(PreTrainedTokenizerFast):
+    """
+    Fast tokenizer version - currently just delegates to slow tokenizer
+    """
+    slow_tokenizer_class = LucaGPLMTokenizer
+    
+    def __init__(self, **kwargs):
+        # For now, this is just a placeholder
+        # In a full implementation, you would use the tokenizers library
+        super().__init__(**kwargs)
+
+__all__ = ["LucaGPLMTokenizer", "LucaGPLMTokenizerFast", "gene_seq_replace"]

tokenizer_config.json

@@ -41,12 +41,18 @@
       "special": true
     }
   },
-  "clean_up_tokenization_spaces": false,
+  "clean_up_tokenization_spaces": true,
   "cls_token": "[CLS]",
   "mask_token": "[MASK]",
   "model_max_length": 1000000000000000019884624838656,
   "pad_token": "[PAD]",
   "sep_token": "[SEP]",
   "tokenizer_class": "LucaGPLMTokenizer",
-  "unk_token": "[UNK]"
-}
+  "unk_token": "[UNK]",
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_lucaone.LucaGPLMTokenizer",
+      null
+    ]
+  }
+}