import numpy as np import torch import torch.nn.functional as F from einops import rearrange from torch import nn from tqdm import tqdm from transformers import PretrainedConfig, PreTrainedModel, ModernBertConfig, ModernBertModel from transformers.modeling_outputs import ModelOutput, SequenceClassifierOutput class GenoJEPAConfig(PretrainedConfig): model_type = "GenoJEPA" def __init__(self, **kwargs): super().__init__(**kwargs) self.vocab_num = kwargs.get("vocab_num", 5) self.pad_value = kwargs.get("pad_value", 4) self.global_num = kwargs.get("global_num", 2) self.global_scale = kwargs.get("global_scale", [0.65, 0.8]) self.local_num = kwargs.get("local_num", 6) self.local_scale = kwargs.get("local_scale", [0.35, 0.4]) self.patch_size = kwargs.get("patch_size", 16) self.embed_size = kwargs.get("embed_size", 16) self.head_size = kwargs.get("head_size", 16) self.hidden_size = kwargs.get("hidden_size", 512) self.intermediate_size = kwargs.get("intermediate_size", 2048) self.layer_num = kwargs.get("layer_num", 12) self.head_num = kwargs.get("head_num", 16) self.slice_num = kwargs.get("slice_num", 256) self.point_num = kwargs.get("point_num", 17) self.max_t = kwargs.get("max_t", 3.0) self.alpha = kwargs.get("alpha", 0.05) self.auto_map = { "AutoConfig": "model.GenoJEPAConfig", "AutoModel": "model.GenoJEPAForSequenceEmbedding", "AutoModelForSequenceClassification": "model.GenoJEPAForSequenceClassification", } def to_dict(self): output = super().to_dict() output["auto_map"] = self.auto_map return output class Patch(nn.Module): def __init__(self, config): super().__init__() self.pad_value = config.pad_value self.patch_size = config.patch_size @torch.no_grad() def forward(self, data): batch_size, seq_len = data.shape pad_len = (self.patch_size - seq_len % self.patch_size) % self.patch_size pad_data = F.pad(data, (0, pad_len), value=self.pad_value) patch_data = rearrange(pad_data, "B (N P) -> B N P", P=self.patch_size) patch_mask = (patch_data != self.pad_value).any(dim=-1).long() return patch_data.detach(), patch_mask.detach() class Augment(nn.Module): def __init__(self, config): super().__init__() @torch.no_grad() def _crop(self, data, scale): batch_size, seq_len = data.shape sample_scale = np.random.uniform(scale[0], scale[1]) sample_len = max(int(seq_len * sample_scale), 1) start_idx = torch.randint(0, seq_len - sample_len + 1, (batch_size, 1), device=data.device) gather_idx = start_idx + torch.arange(sample_len, device=data.device).unsqueeze(0) aug_data = torch.gather(data, 1, gather_idx) return aug_data.detach() @torch.no_grad() def forward(self, data, scale): aug_data = self._crop(data, scale) return aug_data.detach() class SIGReg(nn.Module): def __init__(self, config): super().__init__() self.slice_num = config.slice_num t = torch.linspace(0, config.max_t, config.point_num) window = torch.exp(-t.square() / 2.0) dt = config.max_t / (config.point_num - 1) weight = torch.full((config.point_num,), 2 * dt) weight[[0, -1]] = dt self.register_buffer("t", t) self.register_buffer("phi", window) self.register_buffer("weight", weight * window) def forward(self, embed): view_num, batch_size, hidden_size = embed.shape A = torch.randn((hidden_size, self.slice_num)).to(embed.device) A = A.div_(A.norm(p=2, dim=0)) x_t = (embed @ A).unsqueeze(-1) * self.t err = (x_t.cos().mean(dim=-3) - self.phi).square() + x_t.sin().mean(dim=-3).square() loss = (err @ self.weight) * batch_size loss = loss.mean() return loss class Embedder(nn.Module): def __init__(self, config): super().__init__() self.embedding = nn.Embedding(config.vocab_num, config.embed_size) self.embedder = nn.Conv1d( in_channels=config.embed_size, out_channels=config.hidden_size, kernel_size=config.patch_size, stride=config.patch_size, bias=True, ) def forward(self, data): batch_size, patch_num, patch_size = data.shape embed = self.embedding(data) embed = rearrange(embed, "B N P D -> (B N) D P") embed = self.embedder(embed) embed = rearrange(embed, "(B N) D 1 -> B N D", B=batch_size) return embed class Model(nn.Module): def __init__(self, config): super().__init__() self.cls_token = nn.Parameter(torch.randn(1, 1, config.hidden_size), requires_grad=True) self.patch = Patch(config) self.embedder = Embedder(config) self.encoder = ModernBertModel( ModernBertConfig( vocab_size=config.vocab_num, pad_token_id=config.pad_value, global_attn_every_n_layers=1, hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, num_hidden_layers=config.layer_num, num_attention_heads=config.head_num, ) ) def forward(self, data): batch_size, seq_len = data.shape patch_data, patch_mask = self.patch(data) init_patch_embed = self.embedder(patch_data) init_cls_embed = self.cls_token.expand(batch_size, -1, -1) cls_mask = torch.ones(batch_size, 1, device=data.device).long() input_embed = torch.cat((init_cls_embed, init_patch_embed), dim=1) input_mask = torch.cat((cls_mask, patch_mask), dim=1) output = self.encoder(inputs_embeds=input_embed, attention_mask=input_mask).last_hidden_state cls_embed, patch_embed = output[:, 0, :], output[:, 1:, :] patch_mask_expanded = patch_mask.unsqueeze(-1) mask_sum = patch_mask_expanded.sum(dim=1) sum_embed = (patch_embed * patch_mask_expanded).sum(dim=1) avg_embed = sum_embed / mask_sum.clamp(min=1) raw_max = patch_embed.masked_fill(patch_mask_expanded == 0, float("-inf")).max(dim=1)[0] max_embed = torch.where(mask_sum > 0, raw_max, torch.zeros_like(raw_max)) return ModelOutput(patch_embed=patch_embed, cls_embed=cls_embed, avg_embed=avg_embed, max_embed=max_embed) class Head(nn.Module): def __init__(self, config): super().__init__() self.head = nn.Sequential( nn.Conv1d(config.hidden_size, config.intermediate_size, kernel_size=1, stride=1, bias=True), nn.BatchNorm1d(config.intermediate_size), nn.GELU(), nn.Conv1d(config.intermediate_size, config.head_size, kernel_size=1, stride=1, bias=False), ) def forward(self, embed): embed = rearrange(embed, "V B D -> B D V") embed = self.head(embed) embed = rearrange(embed, "B D V -> V B D") return embed class GenoJEPAForPreTraining(PreTrainedModel): config_class = GenoJEPAConfig def __init__(self, config): super().__init__(config) self.alpha = config.alpha self.pred_loss = 0.0 self.sigreg_loss = 0.0 self.global_num = int(config.global_num) self.global_scale = config.global_scale self.local_num = int(config.local_num) self.local_scale = config.local_scale self.augment = Augment(config) self.model = Model(config) self.head = Head(config) self.pred_criterion = nn.MSELoss() self.sigreg_criterion = SIGReg(config) self.post_init() @torch.no_grad() def get_detailed_metric(self): return ModelOutput(pred_loss=self.pred_loss, sigreg_loss=self.sigreg_loss) def forward(self, input_ids, attention_mask=None, labels=None): if not self.training: sequence_embed = self.model(input_ids.long()).avg_embed return SequenceClassifierOutput(loss=sequence_embed.std(dim=0).mean(), logits=sequence_embed) embed_list = [] for _ in range(self.global_num): global_input_ids = self.augment(input_ids.long(), self.global_scale) embed_list.append(self.model(global_input_ids).avg_embed) for _ in range(self.local_num): local_input_ids = self.augment(input_ids.long(), self.local_scale) embed_list.append(self.model(local_input_ids).avg_embed) embed = self.head(torch.stack(embed_list, dim=0)) global_embed, local_embed = torch.split(embed, [self.global_num, self.local_num], dim=0) pred_loss = self.pred_criterion(embed, global_embed.mean(dim=0, keepdim=True)) sigreg_loss = self.sigreg_criterion(embed) self.pred_loss, self.sigreg_loss = pred_loss.item(), sigreg_loss.item() loss = (1 - self.alpha) * pred_loss + self.alpha * sigreg_loss return ModelOutput(loss=loss) class GenoJEPAForSequenceEmbedding(PreTrainedModel): config_class = GenoJEPAConfig def __init__(self, config): super().__init__(config) self.model = Model(config) self.post_init() def forward(self, input_ids, attention_mask=None, labels=None): embed = self.model(input_ids.long()).avg_embed return embed def encode(self, sequence, tokenizer, batch_size=64): original_training = self.training self.eval() embed_list = [] with torch.no_grad(): for i in tqdm(range(0, len(sequence), batch_size)): batch_sequence = sequence[i:i + batch_size] batch_input_dict = tokenizer(batch_sequence, return_tensors="pt", padding=True) batch_input_ids = batch_input_dict["input_ids"].to(self.model.cls_token.device) batch_embed = self.forward(input_ids=batch_input_ids) embed_list.append(batch_embed.detach().cpu()) embed = torch.concat(embed_list, dim=0).numpy() if original_training: self.train() return embed class GenoJEPAForSequenceClassification(PreTrainedModel): config_class = GenoJEPAConfig def __init__(self, config): super().__init__(config) self.model = Model(config) self.classifier = nn.Sequential( nn.Linear(config.hidden_size, config.intermediate_size), nn.LayerNorm(config.intermediate_size), nn.GELU(), nn.Linear(config.intermediate_size, config.num_labels), ) self.criterion = nn.CrossEntropyLoss() self.post_init() def forward(self, input_ids, attention_mask=None, labels=None): embeds = self.model(input_ids.long()).cls_embed logits = self.classifier(embeds) loss = None if labels is not None: loss = self.criterion(logits, labels) return SequenceClassifierOutput(loss=loss, logits=logits)