File size: 4,195 Bytes
cb634e7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | """HuggingFace wrapper for the NucEngram genomic language model.
A char-level DNA masked language model (A/C/G/T/N, 9-token vocab) built on a
ModernBERT encoder with an n-gram memory module injected into two early layers.
`AutoModel.from_pretrained(..., trust_remote_code=True)` returns the encoder;
`last_hidden_state` is a per-nucleotide embedding you can pool + probe/fine-tune.
"""
from __future__ import annotations
import numpy as np
import torch
import torch.nn as nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import BaseModelOutput, MaskedLMOutput
from .configuration_nucengram import NucEngramConfig
from .engram import EngramConfig
from .model_modernbert import ModernBertGenomicConfig
from .model_modernbert_engram import ModernBertEngramWrapper
from .tokenizer import encode as _encode, PAD_ID
_VOCAB = {i: c for i, c in zip(range(9), "_XXXACGTN")} # display only
class NucEngramPreTrainedModel(PreTrainedModel):
config_class = NucEngramConfig
base_model_prefix = "net"
_no_split_modules = ["ModernBertEncoderLayer", "Engram"]
def _init_weights(self, module):
pass # weights come from the pretrained checkpoint
# --- tokenisation helper (the tokenizer is a fixed char->id map) ---
@staticmethod
def encode(sequences, max_length=8192, device="cpu"):
"""str | list[str] -> (input_ids, attention_mask) LongTensors."""
if isinstance(sequences, str):
sequences = [sequences]
B = len(sequences)
L = min(max(len(s) for s in sequences), max_length)
ids = np.full((B, L), PAD_ID, dtype=np.int64)
for i, s in enumerate(sequences):
e = _encode(s).astype(np.int64)[:L]
ids[i, :len(e)] = e
ids = torch.from_numpy(ids)
return ids.to(device), (ids != PAD_ID).long().to(device)
class NucEngramModel(NucEngramPreTrainedModel):
"""Encoder — returns per-token `last_hidden_state` (use for embeddings / FT)."""
def __init__(self, config: NucEngramConfig):
super().__init__(config)
mb = ModernBertGenomicConfig(**config.backbone)
eng = EngramConfig(**config.engram)
self.net = ModernBertEngramWrapper(mb, eng)
self.post_init()
def forward(self, input_ids, attention_mask=None, **kwargs):
if attention_mask is None:
attention_mask = (input_ids != PAD_ID).long()
# the engram hooks read the raw ids off the wrapper
self.net._current_input_ids = input_ids
out = self.net.model.model(input_ids=input_ids, attention_mask=attention_mask)
return BaseModelOutput(last_hidden_state=out.last_hidden_state)
@torch.no_grad()
def embed(self, sequences, pooling="mean", max_length=8192):
"""Convenience: sequence(s) -> pooled embedding [B, hidden]."""
self.eval()
dev = next(self.parameters()).device
ids, am = self.encode(sequences, max_length=max_length, device=dev)
h = self.forward(ids, am).last_hidden_state.float()
if pooling == "cls":
return h[:, 0]
m = am.float().unsqueeze(-1)
return (h * m).sum(1) / m.sum(1).clamp_min(1.0)
class NucEngramForMaskedLM(NucEngramPreTrainedModel):
"""Same weights, exposing the MLM head (returns `logits` over the 9-token vocab)."""
def __init__(self, config: NucEngramConfig):
super().__init__(config)
mb = ModernBertGenomicConfig(**config.backbone)
eng = EngramConfig(**config.engram)
self.net = ModernBertEngramWrapper(mb, eng)
self.post_init()
def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
if attention_mask is None:
attention_mask = (input_ids != PAD_ID).long()
self.net._current_input_ids = input_ids
out = self.net.model(input_ids=input_ids, attention_mask=attention_mask)
logits = out.logits if hasattr(out, "logits") else out
loss = None
if labels is not None:
loss = nn.functional.cross_entropy(
logits.view(-1, logits.size(-1)), labels.view(-1), ignore_index=-100)
return MaskedLMOutput(loss=loss, logits=logits)
|