update

README.md

@@ -0,0 +1,25 @@
+# ESM2 T36
+
+```python
+from transformers import AutoTokenizer, AutoModel
+import torch
+
+tokenizer = AutoTokenizer.from_pretrained('facebook/esm2_t36_3B_UR50D')
+encoder = AutoModel.from_pretrained("facebook/esm2_t36_3B_UR50D")
+protenrich = AutoModel.from_pretrained("SaeedLab/ProtEnrich-ESM2-T36", trust_remote_code=True)
+
+seqs = ["MKTFFVLLL"]
+seqs = [" ".join(i) for i in seqs]
+inputs = tokenizer(seqs, return_tensors="pt", padding=True)
+
+with torch.no_grad():
+  outputs = encoder(**inputs)
+  pooled = outputs.last_hidden_state[0, 1:-1].mean(axis=0)
+  enriched = protenrich(pooled)
+
+print('H enrich:', enriched.h_enrich)
+print('H anchor:', enriched.h_anchor)
+print('H algn:', enriched.h_algn)
+print('Structure:', enriched.struct)
+print('Dynamics:', enriched.dyn)
+```

config.json

@@ -0,0 +1,19 @@
+{
+  "alpha_max": 0.3,
+  "auto_map": {
+    "AutoConfig": "configuration_protenrich.ProtEnrichConfig",
+    "AutoModel": "modeling_protenrich.ProtEnrichModel",
+    "AutoModelForSequenceClassification": "modeling_protenrich.ProtEnrichForSequenceClassification"
+  },
+  "architectures": [
+    "ProtEnrichModel"
+  ],
+  "dtype": "float32",
+  "dyn_dim": 20,
+  "embed_dim": 1024,
+  "model_type": "protenrich",
+  "project_dim": 256,
+  "seq_dim": 2560,
+  "struct_dim": 1024,
+  "transformers_version": "4.57.3"
+}

configuration_protenrich.py

@@ -0,0 +1,23 @@
+from transformers import PretrainedConfig
+
+class ProtEnrichConfig(PretrainedConfig):
+  model_type = "protenrich"
+
+  def __init__(
+    self,
+    seq_dim : int = 2560, 
+    struct_dim: int = 1024, 
+    dyn_dim: int = 20, 
+    embed_dim: int = 1024, 
+    project_dim: int = 256, 
+    alpha_max: float = 0.3,
+    num_labels: int = 2,
+    **kwargs):
+    super().__init__(**kwargs)
+    self.seq_dim = seq_dim
+    self.struct_dim = struct_dim
+    self.dyn_dim = dyn_dim
+    self.embed_dim = embed_dim
+    self.project_dim = project_dim
+    self.alpha_max = alpha_max
+    self.num_labels = num_labels

model.safetensors

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

modeling_protenrich.py

@@ -0,0 +1,122 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dataclasses import dataclass
+from typing import Optional, Tuple
+import torch
+from transformers.utils import ModelOutput
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutputWithPooling, SequenceClassifierOutput
+from transformers import AutoModel
+
+from .configuration_protenrich import ProtEnrichConfig
+
+@dataclass
+class ProtEnrichModelOutput(ModelOutput):
+  h_enrich: torch.FloatTensor = None
+  h_anchor: Optional[torch.FloatTensor] = None
+  h_algn: Optional[torch.FloatTensor] = None
+  struct: Optional[torch.FloatTensor] = None
+  dyn: Optional[torch.FloatTensor] = None
+
+class MLPEncoder(nn.Module):
+  def __init__(self, in_dim, out_dim, hidden_dim=1024, n_layers=2, dropout=0.1):
+    super().__init__()
+    layers = []
+    d = in_dim
+    for _ in range(n_layers - 1):
+      layers += [
+        nn.Linear(d, hidden_dim),
+        nn.LayerNorm(hidden_dim),
+        nn.GELU(),
+        nn.Dropout(dropout),
+      ]
+      d = hidden_dim
+    layers.append(nn.Linear(d, out_dim))
+    self.net = nn.Sequential(*layers)
+
+  def forward(self, x):
+    return self.net(x)
+
+class ProtEnrichModel(PreTrainedModel):
+  config_class = ProtEnrichConfig
+  base_model_prefix = "protenrich"
+
+  def __init__(self, config: ProtEnrichConfig):
+    super().__init__(config)
+
+    self.seq_anchor = MLPEncoder(config.seq_dim, config.embed_dim)
+    self.seq_algn = MLPEncoder(config.seq_dim, config.embed_dim)
+    self.struct_encoder = MLPEncoder(config.struct_dim, config.embed_dim)
+    self.dyn_encoder    = MLPEncoder(config.dyn_dim, config.embed_dim)
+
+    for p in self.struct_encoder.parameters():
+      p.requires_grad = False
+    for p in self.dyn_encoder.parameters():
+      p.requires_grad = False
+
+    self.seq_projector = nn.Linear(config.embed_dim, config.project_dim)
+    self.struct_projector = nn.Linear(config.embed_dim, config.project_dim)
+    self.dyn_projector = nn.Linear(config.embed_dim, config.project_dim)
+
+    self.seq_decoder = MLPEncoder(config.embed_dim, config.seq_dim)
+    self.struct_decoder = MLPEncoder(config.embed_dim, config.struct_dim)
+    self.dyn_decoder = MLPEncoder(config.embed_dim, config.dyn_dim)
+
+    self.alpha_logit = nn.Parameter(torch.tensor(-2.0))
+    self.alpha_max = config.alpha_max
+
+    self.norm_anchor = nn.LayerNorm(config.embed_dim)
+    self.norm_algn = nn.LayerNorm(config.embed_dim)
+
+    self.post_init()
+
+  def forward(self, seq: torch.Tensor, return_dict: Optional[bool] = None):
+
+    h_anchor = self.norm_anchor(self.seq_anchor(seq))
+    h_algn = self.norm_algn(self.seq_algn(seq))
+
+    struct = self.struct_decoder(h_algn)
+    dyn = self.dyn_decoder(h_algn)
+
+    alpha = torch.sigmoid(self.alpha_logit) * self.alpha_max
+    h_enrich = h_anchor + alpha * h_algn
+
+    return ProtEnrichModelOutput(
+      h_enrich=h_enrich,
+      h_anchor=h_anchor,
+      h_algn=h_algn,
+      struct=struct,
+      dyn=dyn,
+    )
+
+class ProtEnrichForSequenceClassification(PreTrainedModel):
+  config_class = ProtEnrichConfig
+
+  def __init__(self, config: ProtEnrichConfig):
+    super().__init__(config)
+
+    self.num_labels = config.num_labels
+
+    self.protenrich = ProtEnrichModel(config)
+    self.classifier = nn.Linear(config.embed_dim, config.num_labels)
+
+    self.post_init()
+
+  def forward(self, seq: torch.Tensor, labels: Optional[torch.Tensor] = None, return_dict: Optional[bool] = None):
+
+    outputs = self.protenrich(seq=seq, return_dict=return_dict)
+    pooled = outputs.h_enrich
+
+    logits = self.classifier(pooled)
+
+    loss = None
+    if labels is not None:
+      loss_fct = nn.CrossEntropyLoss()
+      loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+    return SequenceClassifierOutput(
+      loss=loss,
+      logits=logits,
+      hidden_states=pooled,
+    )