Upload folder using huggingface_hub

README.md

@@ -1,3 +1,14 @@
----
-license: mit
----
+# TransHLA2.0
+
+**TransHLA2.0** 是一个基于 ESM2 + LoRA + CNN-Transformer 融合的蛋白质序列建模模型。  
+支持 Hugging Face Hub 上的 `trust_remote_code` 加载和自定义微调。
+
+## 快速加载
+
+```python
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained(
+    "your-username/TransHLA2.0",
+    trust_remote_code=True
+)

TransHLA2.0.pt

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

TransHLA2.0.py

@@ -0,0 +1,163 @@
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel, PretrainedConfig
+
+from peft import LoraConfig, get_peft_model, TaskType
+from transformers import EsmModel
+
+class TransHLA2Config(PretrainedConfig):
+    model_type = "transhla2"
+    def __init__(self, d_model=480, **kwargs):
+        super().__init__(**kwargs)
+        self.d_model = d_model
+        # 可加入其它自定义参数
+
+class LoraESM(nn.Module):
+    def __init__(self, d_model=480):
+        super().__init__()
+        self.model_name_or_path = "facebook/esm2_t12_35M_UR50D"
+        self.tokenizer_name_or_path = "facebook/esm2_t12_35M_UR50D"
+        self.peft_config = LoraConfig(
+            target_modules=['query', 'out_proj', 'value', 'key', 'dense', 'regression'],
+            task_type=TaskType.FEATURE_EXTRACTION, 
+            inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1
+        )
+        self.esm = EsmModel.from_pretrained(self.model_name_or_path)
+        self.lora_esm = get_peft_model(self.esm, self.peft_config)
+        self.fc_task = nn.Sequential(
+            nn.Linear(d_model, d_model // 4),
+            nn.BatchNorm1d(d_model // 4),
+            nn.Dropout(0.2),
+            nn.SiLU(),
+            nn.Linear(d_model // 4, 32),
+            nn.BatchNorm1d(32),
+        )
+        self.classifier = nn.Linear(32, 2)
+
+    def forward(self, x_in):
+        lora_outputs = self.lora_esm(x_in)
+        last_hidden_state = lora_outputs.last_hidden_state
+        out_linear = last_hidden_state.mean(dim=1)
+        H = self.fc_task(out_linear)
+        output = self.classifier(H)
+        return output, last_hidden_state
+
+class TransHLA2(PreTrainedModel):
+    config_class = TransHLA2Config
+
+    def __init__(self, config):
+        super().__init__(config)
+        n_layers = 4
+        n_head = 8
+        d_model = config.d_model
+        d_ff = 64
+        cnn_num_channel = 256
+        region_embedding_size = 3
+        cnn_kernel_size = 3
+        cnn_padding_size = 1
+        cnn_stride = 1
+        pooling_size = 2
+
+        self.lora_esm = LoraESM(d_model=d_model)
+
+        self.region_cnn1 = nn.Conv1d(d_model, cnn_num_channel, region_embedding_size)
+        self.region_cnn2 = nn.Conv1d(d_model, cnn_num_channel, region_embedding_size)
+        self.padding1 = nn.ConstantPad1d((1, 1), 0)
+        self.padding2 = nn.ConstantPad1d((0, 1), 0)
+        self.relu = nn.SiLU()
+        self.cnn1 = nn.Conv1d(cnn_num_channel, cnn_num_channel, kernel_size=cnn_kernel_size,
+                              padding=cnn_padding_size, stride=cnn_stride)
+        self.cnn2 = nn.Conv1d(cnn_num_channel, cnn_num_channel, kernel_size=cnn_kernel_size,
+                              padding=cnn_padding_size, stride=cnn_stride)
+        self.maxpooling = nn.MaxPool1d(kernel_size=pooling_size)
+        self.epitope_transformer_layers = nn.TransformerEncoderLayer(
+            d_model=d_model, nhead=n_head, dim_feedforward=d_ff, dropout=0.2)
+        self.epitope_transformer_encoder = nn.TransformerEncoder(
+            self.epitope_transformer_layers, num_layers=n_layers)
+        self.hla_transformer_layers = nn.TransformerEncoderLayer(
+            d_model=d_model, nhead=n_head, dim_feedforward=d_ff, dropout=0.2)
+        self.hla_transformer_encoder = nn.TransformerEncoder(
+            self.hla_transformer_layers, num_layers=n_layers)
+        
+        # Cross Attention layers
+        self.cross_attention_epitope_layers = nn.ModuleList(
+            [nn.MultiheadAttention(d_model, n_head, dropout=0.2) for _ in range(4)])
+        self.cross_attention_hla_layers = nn.ModuleList(
+            [nn.MultiheadAttention(d_model, n_head, dropout=0.2) for _ in range(4)])
+
+        self.bn1 = nn.BatchNorm1d(cnn_num_channel)
+        self.bn2 = nn.BatchNorm1d(cnn_num_channel)
+        self.fc_task = nn.Sequential(
+            nn.Linear(2*d_model + 2*cnn_num_channel, 2 * (d_model + cnn_num_channel) // 4),
+            nn.BatchNorm1d(2 * (d_model + cnn_num_channel) // 4),
+            nn.Dropout(0.2),
+            nn.SiLU(),
+            nn.Linear(2 * (d_model + cnn_num_channel) // 4, 96),
+            nn.BatchNorm1d(96),
+        )
+        self.classifier = nn.Linear(96, 2)
+
+    def cnn_block1(self, x):
+        return self.cnn1(self.relu(x))
+
+    def cnn_block2(self, x):
+        x = self.padding2(x)
+        px = self.maxpooling(x)
+        x = self.relu(px)
+        x = self.cnn1(x)
+        x = self.relu(x)
+        x = self.cnn1(x)
+        x = px + x
+        return x
+
+    def structure_block1(self, x):
+        return self.cnn2(self.relu(x))
+
+    def structure_block2(self, x):
+        x = self.padding2(x)
+        px = self.maxpooling(x)
+        x = self.relu(px)
+        x = self.cnn2(x)
+        x = self.relu(x)
+        x = self.cnn2(x)
+        x = px + x
+        return x
+
+    def forward(self, epitope_in, hla_in):
+        _, epitope_emb = self.lora_esm(epitope_in)
+        _, hla_emb = self.lora_esm(hla_in)
+
+        epitope_trans = self.epitope_transformer_encoder(epitope_emb.transpose(0, 1))
+        hla_trans = self.hla_transformer_encoder(hla_emb.transpose(0, 1))
+
+        # Cross Attention layers
+        for cross_attention_epitope, cross_attention_hla in zip(self.cross_attention_epitope_layers, self.cross_attention_hla_layers):
+            epitope_trans, _ = cross_attention_epitope(epitope_trans, hla_trans, hla_trans)
+            hla_trans, _ = cross_attention_hla(hla_trans, epitope_trans, epitope_trans)
+
+        # Mean Pooling
+        epitope_mean = epitope_trans.mean(dim=0)
+        hla_mean = hla_trans.mean(dim=0)
+        
+        epitope_cnn_emb = self.region_cnn1(epitope_emb.transpose(1, 2))
+        epitope_cnn_emb = self.padding1(epitope_cnn_emb)
+        conv = epitope_cnn_emb + self.cnn_block1(self.cnn_block1(epitope_cnn_emb))
+        while conv.size(-1) >= 2:
+            conv = self.cnn_block2(conv)
+        epitope_cnn_out = torch.squeeze(conv, dim=-1)
+        epitope_cnn_out = self.bn1(epitope_cnn_out)
+
+        hla_cnn_emb = self.region_cnn2(hla_emb.transpose(1, 2))
+        hla_cnn_emb = self.padding1(hla_cnn_emb)
+        hla_conv = hla_cnn_emb + self.structure_block1(self.structure_block1(hla_cnn_emb))
+        while hla_conv.size(-1) >= 2:
+            hla_conv = self.structure_block2(hla_conv)
+
+        hla_cnn_out = torch.squeeze(hla_conv, dim=-1)
+        hla_cnn_out = self.bn2(hla_cnn_out)
+
+        representation = torch.cat((epitope_mean, hla_mean, epitope_cnn_out, hla_cnn_out), dim=1)
+        reduction_feature = self.fc_task(representation)
+        logits_clsf = self.classifier(reduction_feature)
+        logits_clsf = torch.nn.functional.softmax(logits_clsf, dim=1)
+        return logits_clsf, reduction_feature

config.json

@@ -0,0 +1,5 @@
+{
+    "architectures": ["TransHLA2"],
+    "model_type": "transhla2",
+    "d_model": 480
+  }