diff --git a/src/proxyz/data/dataset.py b/src/proxyz/data/dataset.py index c64f597..a623dd7 100644 --- a/src/proxyz/data/dataset.py +++ b/src/proxyz/data/dataset.py @@ -3,18 +3,13 @@ from typing import Literal, Sequence, Union from Bio.Data.PDBData import protein_letters_3to1 +from biotite.structure import alphabet from datasets import Dataset import torch from proxyz.data.utils import lines, opener -FIM_PREFIX = "" -FIM_SUFFIX = "" -FIM_MIDDLE = "" -FIM_TOKENS = [FIM_PREFIX, FIM_SUFFIX, FIM_MIDDLE] - - def line_iterator(file_paths: Sequence[str], batch_size=64): batch = [] @@ -75,18 +70,54 @@ def pdb_transform(data_dir: Union[pathlib.Path, str], examples: dict): assert "id" in examples coord, coord_mask, residue_idx, seq = [], [], [], [] + cle, pseudo_beta, pseudo_beta_mask = [], [], [] for pid in examples["id"]: graph = torch.load(processed_dir / f"{pid}.pt", weights_only=False) coord.append(graph.coords) coord_mask.append(graph.coord_mask) residue_idx.append(graph.residue_pdb_idx - graph.residue_pdb_idx[0]) - seq.append("".join(protein_letters_3to1[r] for r in graph.residues)) + seq.append("".join(protein_letters_3to1.get(r, "A") for r in graph.residues)) + + # atom indices + n_idx, ca_idx, c_idx, cb_idx = 0, 1, 2, 3 + + # pseudo_beta + is_gly = (graph.residue_type == 7) + pseudo_beta.append( + torch.where( + is_gly[:, None], graph.coords[:, ca_idx, :], graph.coords[:, cb_idx, :] + ) + ) + pseudo_beta_mask.append( + torch.where( + is_gly, graph.coord_mask[:, ca_idx], graph.coord_mask[:, cb_idx] + ) + ) + + # 3di + nan = torch.full((3, ), torch.nan) + bbxyz = torch.stack( + ( + torch.where(graph.coord_mask[:, ca_idx, None], graph.coords[:, ca_idx, :], nan), + torch.where(graph.coord_mask[:, cb_idx, None], graph.coords[:, cb_idx, :], nan), + torch.where(graph.coord_mask[:, n_idx, None], graph.coords[:, n_idx, :], nan), + torch.where(graph.coord_mask[:, c_idx, None], graph.coords[:, c_idx, :], nan), + ) + ) + cle.append( + torch.from_numpy( + alphabet.i3d.Encoder().encode(*bbxyz.numpy()).filled() + ).long() + ) return { "text": seq, "residue_idx": residue_idx, "coord": coord, - "coord_mask": coord_mask + "coord_mask": coord_mask, + "pseudo_beta": pseudo_beta, + "pseudo_beta_mask": pseudo_beta_mask, + "cle_labels": cle, } diff --git a/src/proxyz/data/sampler.py b/src/proxyz/data/sampler.py index c88326c..eaf9b54 100644 --- a/src/proxyz/data/sampler.py +++ b/src/proxyz/data/sampler.py @@ -9,104 +9,9 @@ from proxyz.data.utils import lines, opener -class DistributedSamplerWrapper(Sampler): - """ - Wraps a non-distributed Sampler (like WeightedRandomSampler) - to make it compatible with Distributed Data Parallel (DDP). - """ - def __init__( - self, - sampler: Sampler, - num_replicas: int = None, - rank: int = None, - shuffle: bool = True, - seed: int = None, - drop_last: bool = False, - ) -> None: - super().__init__() - - # Automatically detect DDP environment details if not provided - if num_replicas is None: - if not torch.distributed.is_available() or not torch.distributed.is_initialized(): - raise RuntimeError("Requires distributed package to be available") - num_replicas = torch.distributed.get_world_size() - if rank is None: - if not torch.distributed.is_available() or not torch.distributed.is_initialized(): - raise RuntimeError("Requires distributed package to be available") - rank = torch.distributed.get_rank() - if rank >= num_replicas or rank < 0: - raise ValueError( - f"Invalid rank {rank}, rank should be in the interval [0, {num_replicas - 1}]" - ) - - self.sampler = sampler - self.num_replicas = num_replicas - self.rank = rank - self.epoch = 0 - # If the dataset length is evenly divisible by # of replicas, then there - # is no need to drop any data, since the dataset will be split equally. - if self.drop_last and len(self.sampler) % self.num_replicas != 0: - # Split to nearest available length that is evenly divisible. - # This is to ensure each rank receives the same amount of data when - # using this Sampler. - self.num_samples = math.ceil( - (len(self.sampler) - self.num_replicas) / self.num_replicas - ) - else: - self.num_samples = math.ceil(len(self.sampler) / self.num_replicas) - self.total_size = self.num_samples * self.num_replicas - self.shuffle = shuffle - self.seed = seed - - def __iter__(self): - # Get the underlying sample order from the WeightedRandomSampler - indices = list(self.sampler) - - if self.shuffle: - # deterministically shuffle based on epoch and seed - g = torch.Generator() - g.manual_seed(self.seed + self.epoch) - - # Shuffle the underlying indices globally before splitting - indices = [indices[i] for i in torch.randperm(len(indices), generator=g).tolist()] - - if not self.drop_last: - # add extra samples to make it evenly divisible - padding_size = self.total_size - len(indices) - if padding_size <= len(indices): - indices += indices[:padding_size] - else: - indices += (indices * math.ceil(padding_size / len(indices)))[ - :padding_size - ] - else: - # remove tail of data to make it evenly divisible. - indices = indices[: self.total_size] - if len(indices) != self.total_size: - raise AssertionError( - f"Number of indices ({len(indices)}) does not match total_size ({self.total_size})" - ) - - # subsample - indices = indices[self.rank : self.total_size : self.num_replicas] - if len(indices) != self.num_samples: - raise AssertionError( - f"Number of subsampled indices ({len(indices)}) does not match num_samples ({self.num_samples})" - ) - - # pyrefly: ignore [bad-return] - return iter(indices) - - def __len__(self): - return self.num_samples - - def set_epoch(self, epoch): - self.epoch = epoch - - def from_cluster_files( dataset: Dataset, file_paths: Sequence[str], generator: torch.Generator = None -): +) -> Sampler: seqeuence_to_idx = {} for idx, example in enumerate(dataset): seqeuence_to_idx[example["id"]] = idx @@ -120,7 +25,7 @@ def from_cluster_files( cluster_map[cluster_id] += data_row_ids # Assign weights to each sample - sample_weights = [1.0] * len(dataset) + sample_weights = [0.0] * len(dataset) # disabled by default for cluster_id, data_row_ids in cluster_map.items(): # Compute sampling weights: n / (1 + log(n)) for each cluster weight = 1 / (1 + math.log(len(data_row_ids))) diff --git a/src/proxyz/models/__init__.py b/src/proxyz/models/__init__.py new file mode 100644 index 0000000..865214a --- /dev/null +++ b/src/proxyz/models/__init__.py @@ -0,0 +1,6 @@ +from .configuration_xyz import * +from .modeling_xyz import * +from .processing_xyz import * + +XYZProcessor.register_for_auto_class() + diff --git a/src/proxyz/models/configuration_xyz.py b/src/proxyz/models/configuration_xyz.py new file mode 100644 index 0000000..2952a0b --- /dev/null +++ b/src/proxyz/models/configuration_xyz.py @@ -0,0 +1,170 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/xyz/modular_xyz.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_xyz.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +import contextlib + +from huggingface_hub.dataclasses import strict +from proxyz.utils import attr + +from transformers.configuration_utils import PreTrainedConfig +from transformers.modeling_rope_utils import RopeParameters +from transformers.utils import auto_docstring +from transformers.utils.type_validators import interval + + +@auto_docstring(checkpoint="bigict/ProXYZ") +@strict +class XYZConfig(PreTrainedConfig): + """ + Configuration for the XYZ U-Net style protein language model. + + A three-stage architecture: + + 1. **Char encoder** — transformer at character (residue) granularity + 2. **Token trunk** — transformer at BPE-token granularity + 3. **Char decoder** — transformer at character granularity with U-Net skip + connections from the encoder + + The char encoder/decoder share the same ``head_dim`` as the trunk so that + a single ``RotaryEmbedding`` instance can be used across all three stacks. + + This configuration extends [`LlamaConfig`]. All token-level fields + (``hidden_size``, ``intermediate_size``, ``num_hidden_layers``, etc.) are + inherited and control the *trunk* transformer. The char-level fields + below control the encoder and decoder. + + Constraint: ``char_hidden_size == char_num_attention_heads × head_dim``. + + char_hidden_size (`int`, *optional*, defaults to 2048): + Hidden size of the char encoder / decoder transformer. + char_intermediate_size (`int`, *optional*, defaults to 5504): + FFN intermediate size for char-level layers. + char_num_hidden_layers (`int`, *optional*, defaults to 4): + Number of transformer layers in the char encoder and char decoder. + char_num_attention_heads (`int`, *optional*, defaults to 16): + Number of attention heads for char-level self-attention. + char_num_key_value_heads (`int`, *optional*): + Number of KV heads for GQA at char level. Defaults to + ``char_num_attention_heads`` (i.e. multi-head attention) when + not specified. + + Note: + ``char_head_dim`` is intentionally omitted — the char stacks reuse the + token-level ``head_dim`` so that RoPE can be shared. + + Example: + ```python + >>> from proxyz.models import XYZConfig, XYZForCausalLM + >>> config = XYZConfig() + >>> model = XYZForCausalLM(config) + ``` + """ + + model_type = "xyz" + keys_to_ignore_at_inference = ["past_key_values", "char_past_key_values"] + + # ---- Tensor-parallel / pipeline-parallel plans (inherited from Llama) ---- + base_model_tp_plan = { + ".*.layers.*.self_attn.q_proj": "colwise", + ".*.layers.*.self_attn.k_proj": "colwise", + ".*.layers.*.self_attn.v_proj": "colwise", + ".*.layers.*.self_attn.o_proj": "rowwise", + ".*.layers.*.mlp.gate_proj": "colwise", + ".*.layers.*.mlp.up_proj": "colwise", + ".*.layers.*.mlp.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + ".*.layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + ".*.norm": (["hidden_states"], ["hidden_states"]), + } + + vocab_size: int = 32000 + hidden_size: int = 4096 + intermediate_size: int = 11008 + num_hidden_layers: int = 32 + num_attention_heads: int = 32 + num_key_value_heads: int | None = None + hidden_act: str = "silu" + max_position_embeddings: int = 2048 + initializer_range: float = interval(min=0.0, max=1.0)(default=0.02) + rms_norm_eps: float = 1e-6 + use_cache: bool = True + pad_token_id: int | None = None + bos_token_id: int | None = 1 + eos_token_id: int | list[int] | None = 2 + pretraining_tp: int | None = 1 + tie_word_embeddings: bool = False + rope_parameters: RopeParameters | dict | None = None + attention_bias: bool = False + attention_dropout: int | float | None = 0.0 + mlp_bias: bool = False + head_dim: int | None = None + + char_hidden_size: int = 2048 + char_intermediate_size: int = 5504 + char_num_hidden_layers: int = 4 + char_num_attention_heads: int = 16 + char_num_key_value_heads: int | None = None + + has_char_lm_head: bool = False + has_cle_lm_head: bool = False + has_distogram_lm_head: bool = False + + def __post_init__(self, **kwargs): + # Default char KV heads to char query heads (MHA) when unspecified. + if self.char_num_key_value_heads is None: + self.char_num_key_value_heads = self.char_num_attention_heads + if self.head_dim is None: + self.head_dim = self.hidden_size // self.num_attention_heads + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + + super().__post_init__(**kwargs) + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + # Ensure char_hidden_size is compatible with (char_num_attention_heads, head_dim). + if self.char_hidden_size != self.char_num_attention_heads * self.head_dim: + raise ValueError( + f"The char hidden size ({self.char_hidden_size}) must equal " + f"char_num_attention_heads ({self.char_num_attention_heads}) × head_dim ({self.head_dim})." + ) + if self.hidden_size % self.num_attention_heads != 0: + raise ValueError( + f"The hidden size ({self.hidden_size}) is not a multiple of the number of attention " + f"heads ({self.num_attention_heads})." + ) + + @contextlib.contextmanager + def tokenization(self): + """Context manager that yields *self* with token-level config active. + + This is a no-op passthrough — the inherited Llama fields already hold + the token-level values. Provided for symmetry with ``characterization``. + """ + yield self + + @contextlib.contextmanager + def characterization(self): + """Context manager that temporarily swaps token-level fields with + char-level equivalents so that ``XYZDecoderLayer`` / ``XYZDecoderLayers`` + can be constructed or invoked with char-granularity dimensions. + + On exit, all fields are restored to their original (token-level) values. + """ + with attr( + self, + hidden_size=self.char_hidden_size, + intermediate_size=self.char_intermediate_size, + num_hidden_layers=self.char_num_hidden_layers, + num_attention_heads=self.char_num_attention_heads, + num_key_value_heads=self.char_num_key_value_heads, + ): + yield self + + +__all__ = ["XYZConfig"] diff --git a/src/proxyz/models/modeling_xyz.py b/src/proxyz/models/modeling_xyz.py new file mode 100644 index 0000000..896ce35 --- /dev/null +++ b/src/proxyz/models/modeling_xyz.py @@ -0,0 +1,837 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/xyz/modular_xyz.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_xyz.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +from collections.abc import Callable +from dataclasses import dataclass +from typing import Optional + +import torch +from torch import nn + +from transformers.activations import ACT2FN +from transformers.cache_utils import Cache, DynamicCache +from transformers.generation import GenerationMixin +from transformers.integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +from transformers.masking_utils import create_causal_mask +from transformers.modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from transformers.processing_utils import Unpack +from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple +from transformers.utils.generic import maybe_autocast, merge_with_config_defaults +from transformers.utils.output_capturing import capture_outputs + +from .configuration_xyz import XYZConfig + + +@use_kernel_forward_from_hub("RMSNorm") +class XYZRMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + XYZRMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +class XYZRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: XYZConfig, device=None): + super().__init__() + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + self.config = config + + self.rope_type = self.config.rope_parameters["rope_type"] + rope_init_fn: Callable = self.compute_default_rope_parameters + if self.rope_type != "default": + rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + inv_freq, self.attention_scaling = rope_init_fn(self.config, device) + + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False) + + @staticmethod + def compute_default_rope_parameters( + config: XYZConfig | None = None, + device: Optional["torch.device"] = None, + seq_len: int | None = None, + ) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies according to the original RoPE implementation + Args: + config ([`~transformers.PreTrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length. Unused for this type of RoPE. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). + """ + base = config.rope_parameters["rope_theta"] + dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + + attention_factor = 1.0 # Unused in this type of RoPE + + # Compute the inverse frequencies + inv_freq = 1.0 / ( + base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) + ) + return inv_freq, attention_factor + + @torch.no_grad() + @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) + def forward(self, x, position_ids): + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) + position_ids_expanded = position_ids[:, None, :].float() + + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with maybe_autocast(device_type=device_type, enabled=False): # Force float32 + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class XYZMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +@use_kernel_func_from_hub("rotary_pos_emb") +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +@use_kernelized_func(apply_rotary_pos_emb) +class XYZAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: XYZConfig, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = True + + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class XYZDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: XYZConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = XYZAttention(config=config, layer_idx=layer_idx) + + self.mlp = XYZMLP(config) + self.input_layernorm = XYZRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = XYZRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + # Self Attention + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +class XYZDecoderLayers(nn.Module): + def __init__(self, config: XYZConfig): + super().__init__() + self.config = config # FIX: AttributeError: 'XYZDecoderLayers' object has no attribute 'config' + + self.layers = nn.ModuleList( + [XYZDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = XYZRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + inputs_embeds: torch.FloatTensor, + causal_mask: torch.Tensor, + position_embeddings: torch.Tensor, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPast: + """ + Runs a stack of [`XYZDecoderLayer`] modules followed by RMS normalization. + + Args: + inputs_embeds (`torch.FloatTensor` of shape `(batch, seq_len, hidden_size)`): + Input embeddings (output of a projection or embedding layer). + causal_mask (`torch.Tensor`): + Pre-computed causal attention mask. + position_embeddings (`torch.Tensor`): + RoPE cos/sin embeddings for the current positions. + position_ids (`torch.LongTensor`, *optional*): + Position indices. Inferred from *past_key_values* when *None*. + past_key_values (`Cache`, *optional*): + Key-value cache for incremental decoding. + use_cache (`bool`, *optional*): + Whether to populate and return the KV cache. + + Returns: + [`BaseModelOutputWithPast`] with the normalized hidden states and + (optionally) the updated KV cache. + """ + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + hidden_states = inputs_embeds + for decoder_layer in self.layers: + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_embeddings=position_embeddings, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +@auto_docstring +class XYZPreTrainedModel(PreTrainedModel): + config: XYZConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["XYZDecoderLayer"] + _skip_keys_device_placement = ["past_key_values", "char_past_key_values"] + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = True + + _can_compile_fullgraph = True + _supports_attention_backend = True + _can_record_outputs = { + "hidden_states": XYZDecoderLayer, + "attentions": XYZAttention, + } + + +@auto_docstring +@dataclass +class XYZModelOutputWithPast(BaseModelOutputWithPast): + """ + Output of [`XYZModel`]. + + Extends [`BaseModelOutputWithPast`] with character-granularity outputs from + the U-Net decoder branch. + + char_last_hidden_state (`torch.Tensor` of shape `(batch, char_seq_len, char_hidden_size)`, *optional*): + Character-level hidden states from the char decoder. + char_past_key_values (`Tuple[Cache, Cache]`, *optional*): + Pair of KV caches — one for the char encoder, one for the char + decoder — used during incremental generation. + """ + + char_last_hidden_state: torch.FloatTensor | None = None + char_past_key_values: tuple[Cache, Cache] | None = None + char_hidden_states: tuple[torch.FloatTensor, ...] | None = None + char_attentions: tuple[torch.FloatTensor, ...] | None = None + + +class XYZModel(XYZPreTrainedModel): + def __init__(self, config: XYZConfig): + super().__init__(config) + + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.rotary_emb = XYZRotaryEmbedding(config=config) + + # Encoder (character granularity) + self.to_char_encoder = nn.Sequential(nn.Linear(config.hidden_size, config.char_hidden_size, bias=False)) + with self.config.characterization() as config: + self.char_encoder = XYZDecoderLayers(config) + self.from_char_encoder = nn.Sequential( + nn.Linear(config.char_hidden_size, config.hidden_size, bias=False), + ) + + # Trunk + with self.config.tokenization() as config: + self.trunk = XYZDecoderLayers(config) + + # Decoder (character granularity) + self.to_char_decoder = nn.Sequential(nn.Linear(config.hidden_size, config.char_hidden_size, bias=False)) + with self.config.characterization() as config: + self.char_decoder = XYZDecoderLayers(config) + + self.gradient_checkpointing = False + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + char_input_ids: torch.LongTensor | None = None, + char_attention_mask: torch.Tensor | None = None, + char_position_ids: torch.LongTensor | None = None, + char_past_key_values: tuple[Cache, Cache] | None = None, + char_inputs_embeds: torch.FloatTensor | None = None, + repr_char_idx: torch.LongTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> XYZModelOutputWithPast: + """ + U-Net forward pass: char encoder → token trunk → char decoder. + + **Stage 1 — Char encoder.** Token embeddings are projected to + *char_hidden_size* and processed by the char encoder stack. + + **Stage 2 — Token trunk.** Char hidden states are gathered at + representative positions (``repr_char_idx``), projected back to + *hidden_size*, and added to the token embeddings. The result is + processed by the trunk transformer. + + **Stage 3 — Char decoder.** Trunk output is projected to + *char_hidden_size* and scattered back to char positions via + ``scatter_add`` (skip connection from the trunk). The char decoder + stack refines the representation with U-Net skip connections from the + encoder. + + Args: + input_ids (`torch.LongTensor` of shape `(batch, token_seq_len)`, *optional*): + BPE token indices. Mutually exclusive with *inputs_embeds*. + attention_mask (`torch.Tensor` of shape `(batch, token_seq_len)`, *optional*): + Token-level attention mask (1 = attend, 0 = mask). + position_ids (`torch.LongTensor`, *optional*): + Token position indices. Inferred from *past_key_values* when *None*. + past_key_values (`Cache`, *optional*): + Trunk KV cache for incremental decoding. + inputs_embeds (`torch.FloatTensor`, *optional*): + Pre-computed token embeddings. Mutually exclusive with *input_ids*. + char_input_ids (`torch.LongTensor` of shape `(batch, char_seq_len)`, *optional*): + Character-level token indices. Mutually exclusive with + *char_inputs_embeds*. + char_attention_mask (`torch.Tensor` of shape `(batch, char_seq_len)`, *optional*): + Character-level attention mask. + char_position_ids (`torch.LongTensor`, *optional*): + Character position indices. Inferred from *char_past_key_values* + when *None*. + char_past_key_values (`Tuple[Cache, Cache]`, *optional*): + Pair of KV caches for the char encoder and char decoder. + char_inputs_embeds (`torch.FloatTensor`, *optional*): + Pre-computed char embeddings. Mutually exclusive with *char_input_ids*. + repr_char_idx (`torch.LongTensor` of shape `(batch, token_seq_len)`, *optional*): + Maps each BPE token to its representative character position in + the char sequence. Used to gather encoder output into the trunk + and scatter trunk output back to the decoder. + use_cache (`bool`, *optional*): + Whether to populate and return KV caches. + + Returns: + [`XYZModelOutputWithPast`]: + - **last_hidden_state** — token-level output from the trunk. + - **char_last_hidden_state** — char-level output from the char + decoder. + - **past_key_values** / **char_past_key_values** — updated caches. + """ + # character level. + if (char_input_ids is None) ^ (char_inputs_embeds is not None): + raise ValueError("You must specify exactly one of char_input_ids or char_inputs_embeds") + + if char_inputs_embeds is None: + char_inputs_embeds: torch.Tensor = self.to_char_encoder(self.embed_tokens(char_input_ids)) + + if use_cache and char_past_key_values is None: + with self.config.characterization() as config: + char_past_key_values = (DynamicCache(config=config), DynamicCache(config=config)) + + if char_position_ids is None: + char_past_seen_tokens = char_past_key_values[0].get_seq_length() if char_past_key_values is not None else 0 + char_position_ids = ( + torch.arange(char_inputs_embeds.shape[1], device=char_inputs_embeds.device) + char_past_seen_tokens + ) + char_position_ids = char_position_ids.unsqueeze(0) + + char_position_embeddings = self.rotary_emb(char_inputs_embeds, position_ids=char_position_ids) + + # token level + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + with self.config.tokenization() as config: + past_key_values = DynamicCache(config=self.config) + + if position_ids is None: + # position_ids = char_position_ids.expand( + # repr_char_idx.shape[0], char_position_ids.shape[1] + # ).gather(1, repr_char_idx) + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + position_embeddings = self.rotary_emb(inputs_embeds, position_ids=position_ids) + + # encode + with self.config.characterization() as config: + char_causal_mask = create_causal_mask( + config=config, + inputs_embeds=char_inputs_embeds, + attention_mask=char_attention_mask, + past_key_values=char_past_key_values[0] if char_past_key_values else None, + position_ids=char_position_ids, + ) + encoder_outputs: BaseModelOutputWithPast = self.char_encoder( + inputs_embeds=char_inputs_embeds, + causal_mask=char_causal_mask, + position_embeddings=char_position_embeddings, + position_ids=char_position_ids, + past_key_values=char_past_key_values[0] if char_past_key_values else None, + use_cache=use_cache, + **kwargs, + ) + + # trunk + inputs_embeds = inputs_embeds + self.from_char_encoder( + encoder_outputs.last_hidden_state.gather( + 1, repr_char_idx[..., None].expand(*repr_char_idx.shape, encoder_outputs.last_hidden_state.shape[2]) + ) + ) + with self.config.tokenization() as config: + causal_mask = create_causal_mask( + config=config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + trunk_outputs: BaseModelOutputWithPast = self.trunk( + inputs_embeds=inputs_embeds, + causal_mask=causal_mask, + position_embeddings=position_embeddings, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + # decode + trunk_last_hidden_state = self.to_char_decoder(trunk_outputs.last_hidden_state) + char_inputs_embeds = encoder_outputs.last_hidden_state.scatter_add( # skip connection + 1, repr_char_idx[..., None].expand_as(trunk_last_hidden_state), trunk_last_hidden_state + ) + with self.config.characterization() as config: + decoder_outputs: BaseModelOutputWithPast = self.char_decoder( + inputs_embeds=char_inputs_embeds, + causal_mask=char_causal_mask, + position_embeddings=char_position_embeddings, + position_ids=char_position_ids, + past_key_values=char_past_key_values[1] if char_past_key_values else None, + use_cache=use_cache, + **kwargs, + ) + + return XYZModelOutputWithPast( + last_hidden_state=trunk_outputs.last_hidden_state, + past_key_values=trunk_outputs.past_key_values, + char_last_hidden_state=decoder_outputs.last_hidden_state, + char_past_key_values=(encoder_outputs.past_key_values, decoder_outputs.past_key_values), + ) + + +@auto_docstring +@dataclass +class XYZCausalLMOutputWithPast(CausalLMOutputWithPast): + """ + Output of [`XYZCausalLMOutputWithPast`]. + + Extends [`CausalLMOutputWithPast`] with character-granularity logits and + hidden states for the auxiliary next-character prediction head. + + char_logits (`torch.FloatTensor` of shape `(batch, char_seq_len, vocab_size)`, *optional*): + Next-character prediction logits from the char decoder LM head. + char_last_hidden_state (`torch.Tensor` of shape `(batch, char_seq_len, char_hidden_size)`, *optional*): + Character-level hidden states from the char decoder. + char_past_key_values (`Tuple[Cache, Cache]`, *optional*): + Pair of KV caches — one for the char encoder, one for the char + decoder — used during incremental generation. + """ + + char_logits: torch.FloatTensor | None = None + char_past_key_values: tuple[Cache, Cache] | None = None + char_hidden_states: tuple[torch.FloatTensor, ...] | None = None + char_attentions: tuple[torch.FloatTensor, ...] | None = None + cle_logits: torch.FloatTensor | None = None + + +@auto_docstring +class XYZForCausalLM(XYZPreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + _tp_plan = {"lm_head": "colwise_gather_output"} + _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + + def __init__(self, config: XYZConfig): + super().__init__(config) + self.model = XYZModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + if config.has_char_lm_head: + self.char_lm_head = nn.Linear(config.char_hidden_size, config.vocab_size, bias=False) + if config.has_cle_lm_head: + self.cle_lm_head = nn.Linear(config.char_hidden_size, 26, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + logits_to_keep: int | torch.Tensor = 0, + char_input_ids: torch.LongTensor | None = None, + char_attention_mask: torch.Tensor | None = None, + char_position_ids: torch.LongTensor | None = None, + char_past_key_values: tuple[Cache, Cache] | None = None, + char_inputs_embeds: torch.FloatTensor | None = None, + char_labels: torch.LongTensor | None = None, + char_logits_to_keep: int | torch.Tensor = 0, + cle_labels: torch.LongTensor | None = None, + repr_char_idx: torch.LongTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> XYZCausalLMOutputWithPast: + """ + Causal language-modeling forward pass with dual prediction heads. + + Runs the U-Net backbone ([`XYZModel`]) and applies two LM heads: + + - **Next-token** — ``lm_head`` on trunk output → token logits + - **Next-character** — ``char_lm_head`` on char-decoder output → char logits + + Losses are summed when both *labels* and *char_labels* are provided. + + Args: + input_ids (`torch.LongTensor` of shape `(batch, token_seq_len)`, *optional*): + BPE token indices. + attention_mask (`torch.Tensor`, *optional*): + Token-level attention mask. + position_ids (`torch.LongTensor`, *optional*): + Token position indices. + past_key_values (`Cache`, *optional*): + Trunk KV cache. + inputs_embeds (`torch.FloatTensor`, *optional*): + Pre-computed token embeddings. + labels (`torch.LongTensor` of shape `(batch, token_seq_len)`, *optional*): + Ground-truth token ids for next-token loss. Shifted internally. + char_input_ids (`torch.LongTensor` of shape `(batch, char_seq_len)`, *optional*): + Character-level token indices. + char_attention_mask (`torch.Tensor`, *optional*): + Character-level attention mask. + char_position_ids (`torch.LongTensor`, *optional*): + Character position indices. + char_past_key_values (`Tuple[Cache, Cache]`, *optional*): + Char encoder / decoder KV caches. + char_inputs_embeds (`torch.FloatTensor`, *optional*): + Pre-computed char embeddings. + char_labels (`torch.LongTensor` of shape `(batch, char_seq_len)`, *optional*): + Ground-truth char ids for next-character loss. Shifted internally. + repr_char_idx (`torch.LongTensor`, *optional*): + Token → representative-char mapping. + use_cache (`bool`, *optional*): + Whether to populate and return KV caches. + + Returns: + [`XYZForCausalLMOutput`]: + - **loss** — combined next-token + next-character CE loss. + - **logits** — next-token logits (used by `generate()`). + - **char_logits** — next-character logits. + """ + outputs: XYZModelOutputWithPast = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + char_input_ids=char_input_ids, + char_attention_mask=char_attention_mask, + char_position_ids=char_position_ids, + char_past_key_values=char_past_key_values, + char_inputs_embeds=char_inputs_embeds, + repr_char_idx=repr_char_idx, + use_cache=use_cache, + **kwargs, + ) + + # trunk + hidden_states = outputs.last_hidden_state + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) + + # decoder + char_hidden_states = outputs.char_last_hidden_state + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = ( + slice(-char_logits_to_keep, None) if isinstance(char_logits_to_keep, int) else char_logits_to_keep + ) + + aux_loss = [] + + char_logits = None + if self.config.has_char_lm_head: + char_logits = self.char_lm_head(char_hidden_states[:, slice_indices, :]) + if char_labels is not None: + aux_loss.append( + self.loss_function( + logits=char_logits, labels=char_labels, vocab_size=self.config.vocab_size, **kwargs + ) + ) + cle_logits = None + if self.config.has_cle_lm_head: + cle_logits = self.cle_lm_head(char_hidden_states[:, slice_indices, :]) + if cle_labels is not None: + aux_loss.append(self.loss_function(logits=cle_logits, labels=cle_labels, vocab_size=26, **kwargs)) + + if loss is not None and aux_loss: + loss = loss + sum(aux_loss) + elif aux_loss: + loss = sum(aux_loss) + + return XYZCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + char_logits=char_logits, + char_past_key_values=outputs.char_past_key_values, + char_hidden_states=outputs.char_hidden_states, + char_attentions=outputs.char_attentions, + cle_logits=cle_logits, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + attention_mask=None, + inputs_embeds=None, + **kwargs, + ): + return None + + +class XYZForSequenceClassification(GenericForSequenceClassification, XYZPreTrainedModel): + pass + + +__all__ = [ + "XYZPreTrainedModel", + "XYZModelOutputWithPast", + "XYZModel", + "XYZCausalLMOutputWithPast", + "XYZForCausalLM", + "XYZForSequenceClassification", +] diff --git a/src/proxyz/models/modular_xyz.py b/src/proxyz/models/modular_xyz.py new file mode 100644 index 0000000..54a62fb --- /dev/null +++ b/src/proxyz/models/modular_xyz.py @@ -0,0 +1,907 @@ +"""U-Net style language model. + +Three-stage architecture: + -> char encoder (character granularity) + -> token transformmer (token granularity) + -> char decoder (character granularity) + => predict next token +""" +import contextlib +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple, Union +import functools +import random + +import torch +from torch import nn + +from transformers.models.llama.configuration_llama import LlamaConfig +from transformers.models.llama.modeling_llama import ( + LlamaDecoderLayer, + LlamaForCausalLM, + LlamaForSequenceClassification, + LlamaMLP, + LlamaModel, + LlamaPreTrainedModel, + LlamaRMSNorm, + LlamaRotaryEmbedding, + eager_attention_forward, +) +from transformers.activations import ACT2FN +from transformers.cache_utils import Cache, DynamicCache +from transformers.configuration_utils import PreTrainedConfig +from transformers.generation import GenerationMixin +from transformers.integrations import ( + use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func +) +from transformers.masking_utils import create_causal_mask +from transformers.modeling_layers import ( + GenericForSequenceClassification, GradientCheckpointingLayer +) +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_rope_utils import ( + ROPE_INIT_FUNCTIONS, dynamic_rope_update, RopeParameters +) +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from transformers.processing_utils import ProcessorMixin, Unpack +from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, logging +from transformers.utils.generic import maybe_autocast, merge_with_config_defaults +from transformers.utils.output_capturing import capture_outputs +from transformers.utils.type_validators import interval + +from proxyz.utils import attr + + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="bigict/ProXYZ") +@strict +class XYZConfig(LlamaConfig): + """ + Configuration for the XYZ U-Net style protein language model. + + A three-stage architecture: + + 1. **Char encoder** — transformer at character (residue) granularity + 2. **Token trunk** — transformer at BPE-token granularity + 3. **Char decoder** — transformer at character granularity with U-Net skip + connections from the encoder + + The char encoder/decoder share the same ``head_dim`` as the trunk so that + a single ``RotaryEmbedding`` instance can be used across all three stacks. + + This configuration extends [`LlamaConfig`]. All token-level fields + (``hidden_size``, ``intermediate_size``, ``num_hidden_layers``, etc.) are + inherited and control the *trunk* transformer. The char-level fields + below control the encoder and decoder. + + Constraint: ``char_hidden_size == char_num_attention_heads × head_dim``. + + char_hidden_size (`int`, *optional*, defaults to 2048): + Hidden size of the char encoder / decoder transformer. + char_intermediate_size (`int`, *optional*, defaults to 5504): + FFN intermediate size for char-level layers. + char_num_hidden_layers (`int`, *optional*, defaults to 4): + Number of transformer layers in the char encoder and char decoder. + char_num_attention_heads (`int`, *optional*, defaults to 16): + Number of attention heads for char-level self-attention. + char_num_key_value_heads (`int`, *optional*): + Number of KV heads for GQA at char level. Defaults to + ``char_num_attention_heads`` (i.e. multi-head attention) when + not specified. + + Note: + ``char_head_dim`` is intentionally omitted — the char stacks reuse the + token-level ``head_dim`` so that RoPE can be shared. + + Example: + ```python + >>> from proxyz.models import XYZConfig, XYZForCausalLM + >>> config = XYZConfig() + >>> model = XYZForCausalLM(config) + ``` + """ + + model_type = "xyz" + keys_to_ignore_at_inference = ["past_key_values", "char_past_key_values"] + + # ---- Tensor-parallel / pipeline-parallel plans (inherited from Llama) ---- + base_model_tp_plan = { + ".*.layers.*.self_attn.q_proj": "colwise", + ".*.layers.*.self_attn.k_proj": "colwise", + ".*.layers.*.self_attn.v_proj": "colwise", + ".*.layers.*.self_attn.o_proj": "rowwise", + ".*.layers.*.mlp.gate_proj": "colwise", + ".*.layers.*.mlp.up_proj": "colwise", + ".*.layers.*.mlp.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + ".*.layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + ".*.norm": (["hidden_states"], ["hidden_states"]), + } + + char_hidden_size: int = 2048 + char_intermediate_size: int = 5504 + char_num_hidden_layers: int = 4 + char_num_attention_heads: int = 16 + char_num_key_value_heads: int | None = None + + has_char_lm_head: bool = False + has_cle_lm_head: bool = False + has_distogram_lm_head: bool = False + + def __post_init__(self, **kwargs): + # Default char KV heads to char query heads (MHA) when unspecified. + if self.char_num_key_value_heads is None: + self.char_num_key_value_heads = self.char_num_attention_heads + super().__post_init__(**kwargs) + + def validate_architecture(self): + # Ensure char_hidden_size is compatible with (char_num_attention_heads, head_dim). + if self.char_hidden_size != self.char_num_attention_heads * self.head_dim: + raise ValueError( + f"The char hidden size ({self.char_hidden_size}) must equal " + f"char_num_attention_heads ({self.char_num_attention_heads}) × head_dim ({self.head_dim})." + ) + super().validate_architecture() + + @contextlib.contextmanager + def tokenization(self): + """Context manager that yields *self* with token-level config active. + + This is a no-op passthrough — the inherited Llama fields already hold + the token-level values. Provided for symmetry with ``characterization``. + """ + yield self + + @contextlib.contextmanager + def characterization(self): + """Context manager that temporarily swaps token-level fields with + char-level equivalents so that ``XYZDecoderLayer`` / ``XYZDecoderLayers`` + can be constructed or invoked with char-granularity dimensions. + + On exit, all fields are restored to their original (token-level) values. + """ + with attr( + self, + hidden_size=self.char_hidden_size, + intermediate_size=self.char_intermediate_size, + num_hidden_layers=self.char_num_hidden_layers, + num_attention_heads=self.char_num_attention_heads, + num_key_value_heads=self.char_num_key_value_heads, + ): + yield self + + +class XYZRMSNorm(LlamaRMSNorm): + pass + + +class XYZRotaryEmbedding(LlamaRotaryEmbedding): + pass + + +class XYZMLP(LlamaMLP): + pass + + +class XYZDecoderLayer(LlamaDecoderLayer): + pass + + +class XYZDecoderLayers(nn.Module): + def __init__(self, config: XYZConfig): + super().__init__() + self.config = config # FIX: AttributeError: 'XYZDecoderLayers' object has no attribute 'config' + + self.layers = nn.ModuleList( + [XYZDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = XYZRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + inputs_embeds: torch.FloatTensor, + causal_mask: torch.Tensor, + position_embeddings: torch.Tensor, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPast: + """ + Runs a stack of [`XYZDecoderLayer`] modules followed by RMS normalization. + + Args: + inputs_embeds (`torch.FloatTensor` of shape `(batch, seq_len, hidden_size)`): + Input embeddings (output of a projection or embedding layer). + causal_mask (`torch.Tensor`): + Pre-computed causal attention mask. + position_embeddings (`torch.Tensor`): + RoPE cos/sin embeddings for the current positions. + position_ids (`torch.LongTensor`, *optional*): + Position indices. Inferred from *past_key_values* when *None*. + past_key_values (`Cache`, *optional*): + Key-value cache for incremental decoding. + use_cache (`bool`, *optional*): + Whether to populate and return the KV cache. + + Returns: + [`BaseModelOutputWithPast`] with the normalized hidden states and + (optionally) the updated KV cache. + """ + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + hidden_states = inputs_embeds + for decoder_layer in self.layers: + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_embeddings=position_embeddings, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +class XYZPreTrainedModel(LlamaPreTrainedModel): + _skip_keys_device_placement = ["past_key_values", "char_past_key_values"] + + +@auto_docstring +@dataclass +class XYZModelOutputWithPast(BaseModelOutputWithPast): + """ + Output of [`XYZModel`]. + + Extends [`BaseModelOutputWithPast`] with character-granularity outputs from + the U-Net decoder branch. + + char_last_hidden_state (`torch.Tensor` of shape `(batch, char_seq_len, char_hidden_size)`, *optional*): + Character-level hidden states from the char decoder. + char_past_key_values (`Tuple[Cache, Cache]`, *optional*): + Pair of KV caches — one for the char encoder, one for the char + decoder — used during incremental generation. + """ + + char_last_hidden_state: torch.FloatTensor | None = None + char_past_key_values: Tuple[Cache, Cache] | None = None + char_hidden_states: tuple[torch.FloatTensor, ...] | None = None + char_attentions: tuple[torch.FloatTensor, ...] | None = None + + +class XYZModel(XYZPreTrainedModel): + def __init__(self, config: XYZConfig): + super().__init__(config) + + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.rotary_emb = XYZRotaryEmbedding(config=config) + + # Encoder (character granularity) + self.to_char_encoder = nn.Sequential( + nn.Linear(config.hidden_size, config.char_hidden_size, bias=False) + ) + with self.config.characterization() as config: + self.char_encoder = XYZDecoderLayers(config) + self.from_char_encoder = nn.Sequential( + nn.Linear(config.char_hidden_size, config.hidden_size, bias=False), + ) + + # Trunk + with self.config.tokenization() as config: + self.trunk = XYZDecoderLayers(config) + + # Decoder (character granularity) + self.to_char_decoder = nn.Sequential( + nn.Linear(config.hidden_size, config.char_hidden_size, bias=False) + ) + with self.config.characterization() as config: + self.char_decoder = XYZDecoderLayers(config) + + self.gradient_checkpointing = False + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + char_input_ids: torch.LongTensor | None = None, + char_attention_mask: torch.Tensor | None = None, + char_position_ids: torch.LongTensor | None = None, + char_past_key_values: Tuple[Cache, Cache] | None = None, + char_inputs_embeds: torch.FloatTensor | None = None, + repr_char_idx: torch.LongTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> XYZModelOutputWithPast: + """ + U-Net forward pass: char encoder → token trunk → char decoder. + + **Stage 1 — Char encoder.** Token embeddings are projected to + *char_hidden_size* and processed by the char encoder stack. + + **Stage 2 — Token trunk.** Char hidden states are gathered at + representative positions (``repr_char_idx``), projected back to + *hidden_size*, and added to the token embeddings. The result is + processed by the trunk transformer. + + **Stage 3 — Char decoder.** Trunk output is projected to + *char_hidden_size* and scattered back to char positions via + ``scatter_add`` (skip connection from the trunk). The char decoder + stack refines the representation with U-Net skip connections from the + encoder. + + Args: + input_ids (`torch.LongTensor` of shape `(batch, token_seq_len)`, *optional*): + BPE token indices. Mutually exclusive with *inputs_embeds*. + attention_mask (`torch.Tensor` of shape `(batch, token_seq_len)`, *optional*): + Token-level attention mask (1 = attend, 0 = mask). + position_ids (`torch.LongTensor`, *optional*): + Token position indices. Inferred from *past_key_values* when *None*. + past_key_values (`Cache`, *optional*): + Trunk KV cache for incremental decoding. + inputs_embeds (`torch.FloatTensor`, *optional*): + Pre-computed token embeddings. Mutually exclusive with *input_ids*. + char_input_ids (`torch.LongTensor` of shape `(batch, char_seq_len)`, *optional*): + Character-level token indices. Mutually exclusive with + *char_inputs_embeds*. + char_attention_mask (`torch.Tensor` of shape `(batch, char_seq_len)`, *optional*): + Character-level attention mask. + char_position_ids (`torch.LongTensor`, *optional*): + Character position indices. Inferred from *char_past_key_values* + when *None*. + char_past_key_values (`Tuple[Cache, Cache]`, *optional*): + Pair of KV caches for the char encoder and char decoder. + char_inputs_embeds (`torch.FloatTensor`, *optional*): + Pre-computed char embeddings. Mutually exclusive with *char_input_ids*. + repr_char_idx (`torch.LongTensor` of shape `(batch, token_seq_len)`, *optional*): + Maps each BPE token to its representative character position in + the char sequence. Used to gather encoder output into the trunk + and scatter trunk output back to the decoder. + use_cache (`bool`, *optional*): + Whether to populate and return KV caches. + + Returns: + [`XYZModelOutputWithPast`]: + - **last_hidden_state** — token-level output from the trunk. + - **char_last_hidden_state** — char-level output from the char + decoder. + - **past_key_values** / **char_past_key_values** — updated caches. + """ + # character level. + if (char_input_ids is None) ^ (char_inputs_embeds is not None): + raise ValueError("You must specify exactly one of char_input_ids or char_inputs_embeds") + + if char_inputs_embeds is None: + char_inputs_embeds: torch.Tensor = self.to_char_encoder(self.embed_tokens(char_input_ids)) + + if use_cache and char_past_key_values is None: + with self.config.characterization() as config: + char_past_key_values = ( + DynamicCache(config=config), DynamicCache(config=config) + ) + + if char_position_ids is None: + char_past_seen_tokens = char_past_key_values[0].get_seq_length() if char_past_key_values is not None else 0 + char_position_ids = torch.arange(char_inputs_embeds.shape[1], device=char_inputs_embeds.device) + char_past_seen_tokens + char_position_ids = char_position_ids.unsqueeze(0) + + char_position_embeddings = self.rotary_emb(char_inputs_embeds, position_ids=char_position_ids) + + # token level + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + with self.config.tokenization() as config: + past_key_values = DynamicCache(config=self.config) + + if position_ids is None: + # position_ids = char_position_ids.expand( + # repr_char_idx.shape[0], char_position_ids.shape[1] + # ).gather(1, repr_char_idx) + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + position_embeddings = self.rotary_emb(inputs_embeds, position_ids=position_ids) + + # encode + with self.config.characterization() as config: + char_causal_mask = create_causal_mask( + config=config, + inputs_embeds=char_inputs_embeds, + attention_mask=char_attention_mask, + past_key_values=char_past_key_values[0] if char_past_key_values else None, + position_ids=char_position_ids, + ) + encoder_outputs: BaseModelOutputWithPast = self.char_encoder( + inputs_embeds=char_inputs_embeds, + causal_mask=char_causal_mask, + position_embeddings=char_position_embeddings, + position_ids=char_position_ids, + past_key_values=char_past_key_values[0] if char_past_key_values else None, + use_cache=use_cache, + **kwargs, + ) + + # trunk + inputs_embeds = inputs_embeds + self.from_char_encoder( + encoder_outputs.last_hidden_state.gather( + 1, repr_char_idx[..., None].expand( + *repr_char_idx.shape, encoder_outputs.last_hidden_state.shape[2] + ) + ) + ) + with self.config.tokenization() as config: + causal_mask = create_causal_mask( + config=config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + trunk_outputs: BaseModelOutputWithPast = self.trunk( + inputs_embeds=inputs_embeds, + causal_mask=causal_mask, + position_embeddings=position_embeddings, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + # decode + trunk_last_hidden_state = self.to_char_decoder(trunk_outputs.last_hidden_state) + char_inputs_embeds = encoder_outputs.last_hidden_state.scatter_add( # skip connection + 1, repr_char_idx[..., None].expand_as(trunk_last_hidden_state), trunk_last_hidden_state + ) + with self.config.characterization() as config: + decoder_outputs: BaseModelOutputWithPast = self.char_decoder( + inputs_embeds=char_inputs_embeds, + causal_mask=char_causal_mask, + position_embeddings=char_position_embeddings, + position_ids=char_position_ids, + past_key_values=char_past_key_values[1] if char_past_key_values else None, + use_cache=use_cache, + **kwargs, + ) + + return XYZModelOutputWithPast( + last_hidden_state=trunk_outputs.last_hidden_state, + past_key_values=trunk_outputs.past_key_values, + char_last_hidden_state=decoder_outputs.last_hidden_state, + char_past_key_values=( + encoder_outputs.past_key_values, decoder_outputs.past_key_values + ) + ) + + +@auto_docstring +@dataclass +class XYZCausalLMOutputWithPast(CausalLMOutputWithPast): + """ + Output of [`XYZCausalLMOutputWithPast`]. + + Extends [`CausalLMOutputWithPast`] with character-granularity logits and + hidden states for the auxiliary next-character prediction head. + + char_logits (`torch.FloatTensor` of shape `(batch, char_seq_len, vocab_size)`, *optional*): + Next-character prediction logits from the char decoder LM head. + char_last_hidden_state (`torch.Tensor` of shape `(batch, char_seq_len, char_hidden_size)`, *optional*): + Character-level hidden states from the char decoder. + char_past_key_values (`Tuple[Cache, Cache]`, *optional*): + Pair of KV caches — one for the char encoder, one for the char + decoder — used during incremental generation. + """ + + char_logits: torch.FloatTensor | None = None + char_past_key_values: Tuple[Cache, Cache] | None = None + char_hidden_states: tuple[torch.FloatTensor, ...] | None = None + char_attentions: tuple[torch.FloatTensor, ...] | None = None + cle_logits: torch.FloatTensor | None = None + + +class XYZForCausalLM(LlamaForCausalLM): + def __init__(self, config: XYZConfig): + super().__init__(config) + + if config.has_char_lm_head: + self.char_lm_head = nn.Linear( + config.char_hidden_size, config.vocab_size, bias=False + ) + if config.has_cle_lm_head: + self.cle_lm_head = nn.Linear(config.char_hidden_size, 26, bias=False) + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + logits_to_keep: int | torch.Tensor = 0, + char_input_ids: torch.LongTensor | None = None, + char_attention_mask: torch.Tensor | None = None, + char_position_ids: torch.LongTensor | None = None, + char_past_key_values: Tuple[Cache, Cache] | None = None, + char_inputs_embeds: torch.FloatTensor | None = None, + char_labels: torch.LongTensor | None = None, + char_logits_to_keep: int | torch.Tensor = 0, + cle_labels: torch.LongTensor | None = None, + repr_char_idx: torch.LongTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> XYZCausalLMOutputWithPast: + """ + Causal language-modeling forward pass with dual prediction heads. + + Runs the U-Net backbone ([`XYZModel`]) and applies two LM heads: + + - **Next-token** — ``lm_head`` on trunk output → token logits + - **Next-character** — ``char_lm_head`` on char-decoder output → char logits + + Losses are summed when both *labels* and *char_labels* are provided. + + Args: + input_ids (`torch.LongTensor` of shape `(batch, token_seq_len)`, *optional*): + BPE token indices. + attention_mask (`torch.Tensor`, *optional*): + Token-level attention mask. + position_ids (`torch.LongTensor`, *optional*): + Token position indices. + past_key_values (`Cache`, *optional*): + Trunk KV cache. + inputs_embeds (`torch.FloatTensor`, *optional*): + Pre-computed token embeddings. + labels (`torch.LongTensor` of shape `(batch, token_seq_len)`, *optional*): + Ground-truth token ids for next-token loss. Shifted internally. + char_input_ids (`torch.LongTensor` of shape `(batch, char_seq_len)`, *optional*): + Character-level token indices. + char_attention_mask (`torch.Tensor`, *optional*): + Character-level attention mask. + char_position_ids (`torch.LongTensor`, *optional*): + Character position indices. + char_past_key_values (`Tuple[Cache, Cache]`, *optional*): + Char encoder / decoder KV caches. + char_inputs_embeds (`torch.FloatTensor`, *optional*): + Pre-computed char embeddings. + char_labels (`torch.LongTensor` of shape `(batch, char_seq_len)`, *optional*): + Ground-truth char ids for next-character loss. Shifted internally. + repr_char_idx (`torch.LongTensor`, *optional*): + Token → representative-char mapping. + use_cache (`bool`, *optional*): + Whether to populate and return KV caches. + + Returns: + [`XYZForCausalLMOutput`]: + - **loss** — combined next-token + next-character CE loss. + - **logits** — next-token logits (used by `generate()`). + - **char_logits** — next-character logits. + """ + outputs: XYZModelOutputWithPast = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + char_input_ids=char_input_ids, + char_attention_mask=char_attention_mask, + char_position_ids=char_position_ids, + char_past_key_values=char_past_key_values, + char_inputs_embeds=char_inputs_embeds, + repr_char_idx=repr_char_idx, + use_cache=use_cache, + **kwargs, + ) + + # trunk + hidden_states = outputs.last_hidden_state + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function( + logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs + ) + + # decoder + char_hidden_states = outputs.char_last_hidden_state + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-char_logits_to_keep, None) if isinstance(char_logits_to_keep, int) else char_logits_to_keep + + aux_loss = [] + + char_logits = None + if self.config.has_char_lm_head: + char_logits = self.char_lm_head(char_hidden_states[:, slice_indices, :]) + if char_labels is not None: + aux_loss.append( + self.loss_function( + logits=char_logits, labels=char_labels, vocab_size=self.config.vocab_size, **kwargs + ) + ) + cle_logits = None + if self.config.has_cle_lm_head: + cle_logits = self.cle_lm_head(char_hidden_states[:, slice_indices, :]) + if cle_labels is not None: + aux_loss.append( + self.loss_function( + logits=cle_logits, labels=cle_labels, vocab_size=26, **kwargs + ) + ) + + if loss is not None and aux_loss: + loss = loss + sum(aux_loss) + elif aux_loss: + loss = sum(aux_loss) + + return XYZCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + char_logits=char_logits, + char_past_key_values=outputs.char_past_key_values, + char_hidden_states=outputs.char_hidden_states, + char_attentions=outputs.char_attentions, + cle_logits=cle_logits, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + attention_mask=None, + inputs_embeds=None, + **kwargs, + ): + return None + + +class XYZForSequenceClassification(LlamaForSequenceClassification): + pass + + +class XYZProcessor(ProcessorMixin): + """Processes protein sequences at both amino-acid and BPE-token granularity. + + Inherits ``ProcessorMixin`` so it supports ``save_pretrained`` / + ``from_pretrained`` out of the box. The wrapped BPE tokenizer is + persisted alongside the processor config. + + ``AutoProcessor.from_pretrained(dir)`` works as long as *proxyz* is + importable (the ``auto_map`` entry in *preprocessor_config.json* + points to ``proxyz.processor.XYZProcessor``). + """ + + attributes = ["tokenizer"] + tokenizer_class = "AutoTokenizer" + + FIM_PREFIX = "" + FIM_SUFFIX = "" + FIM_MIDDLE = "" + FIM_TOKENS = [FIM_PREFIX, FIM_SUFFIX, FIM_MIDDLE] + + def __init__( + self, tokenizer, text_column: str = "text", features: List[str] = None, **kwargs + ) -> None: + tokenizer.add_special_tokens({"additional_special_tokens": self.FIM_TOKENS}) + logger.info( + f"Added FIM tokens: {self.FIM_TOKENS} to tokenizer. " + f"(vocabu size: {len(tokenizer)})" + ) + + super().__init__(tokenizer, **kwargs) + + self.text_column = text_column + self.features = features + + def __call__( + self, + examples: Dict, + *, + bpe_dropout: Optional[float] = None, + char_apply: bool = False, + fim_apply: bool = False, + fim_spm_rate: float = 0.5, + fim_sft_style: bool = False, + max_length: Optional[int] = None, + **kwargs + ) -> Dict: + examples = self.apply_crop(examples, max_length=max_length) + if fim_apply: + examples = self.apply_fim(examples, spm_rate=fim_spm_rate) + examples = self.apply_wrap(examples) + + def tokenize_with_dropout(dropout, prefix=""): + with attr(self.tokenizer.backend_tokenizer.model, dropout=dropout): + tokenized = self.tokenizer( + examples[self.text_column], + truncation=True, + return_tensors="pt", + padding=True, + return_offsets_mapping=char_apply, + ) + tokenized["labels"] = tokenized["input_ids"].where( + tokenized["attention_mask"] > 0, -100 + ) + if fim_apply and fim_sft_style: + middle_pos = ( + tokenized["labels"] == self.tokenizer.convert_tokens_to_ids(self.FIM_MIDDLE) + ).cumsum(1) + tokenized["labels"] = tokenized["labels"].where( + middle_pos.cumsum(1) > 1, -100 # the is excluded + ) + return {f"{prefix}{k}": v for k, v in tokenized.items()} + + tokenized = tokenize_with_dropout(bpe_dropout, prefix="") + if char_apply: + tokenized.update(tokenize_with_dropout(1.0, prefix="char_")) + + # Align characters with tokenized *text* + tokenized["char_to_token_idx"] = tokenized["char_attention_mask"].new_zeros( + tokenized["char_attention_mask"].size() + ) + tokenized["repr_char_idx"] = tokenized["attention_mask"].new_zeros( + tokenized["attention_mask"].size() + ) + # FIX: use torch.searchsorted instead + for k in range(len(examples[self.text_column])): + i, j = 0, 0 + while i < len(tokenized["offset_mapping"][k]) and j < len(tokenized["char_offset_mapping"][k]): + if tokenized["attention_mask"][k, i] == 0: + i += 1 + elif tokenized["char_attention_mask"][k, j] == 0: + j += 1 + else: + si, ei = tokenized["offset_mapping"][k, i] + sj, ej = tokenized["char_offset_mapping"][k, j] + if si <= sj and ej <= ei: + tokenized["char_to_token_idx"][k, j] = i + if ej == ei: + tokenized["repr_char_idx"][k, i] = j + j += 1 + elif ei <= sj: + i += 1 + else: + assert False, (text[k], i, j) + + del tokenized["char_offset_mapping"] + del tokenized["offset_mapping"] + + # copy features + if self.features is not None: + for column in self.features: + if column in examples: + max_len = max( + examples[column][k].shape[0] for k in range(len(examples[column])) + ) + for k in range(len(examples[column])): + pad = examples[column][k].new_full( + (max_len - examples[column][k].shape[0], *examples[column][k].shape[1:]), + -100 if "labels" in column else 0, + ) + examples[column][k] = torch.cat((examples[column][k], pad)) + tokenized[column] = torch.stack(examples[column]) + + return tokenized + + def apply_fim(self, examples: Dict, spm_rate: float = 0.5) -> Dict: + """Split content into prefix/middle/suffix and rearrange for FIM training. + Prefix or suffix may be empty, but middle is always non-empty.""" + for idx, text in enumerate(examples[self.text_column]): + n = len(text) + cut1 = random.randint(0, n - 1) + cut2 = random.randint(cut1 + 1, n) + prefix, middle, suffix = text[:cut1], text[cut1: cut2], text[cut2:] + + is_spm = random.random() < spm_rate + if is_spm: + # SPM: + first_tag, second_tag = self.FIM_SUFFIX, self.FIM_PREFIX + first, second = suffix, prefix + else: + # PSM: + first_tag, second_tag = self.FIM_PREFIX, self.FIM_SUFFIX + first, second = prefix, suffix + + examples[self.text_column][idx] = ( + first_tag + first + second_tag + second + self.FIM_MIDDLE + middle + ) + if self.features is not None: + for column in self.features: + if column in examples: + pad = examples[column][idx].new_full( + (1, *examples[column][idx].shape[1:]), + -100 if "labels" in column else 0, + ) + if is_spm: + examples[column][idx] = torch.cat( + ( + pad, examples[column][idx][cut2:, ...], + pad, examples[column][idx][:cut1, ...], + pad, examples[column][idx][cut1:cut2, ...], + ) + ) + else: + examples[column][idx] = torch.cat( + ( + pad, examples[column][idx][:cut1, ...], + pad, examples[column][idx][cut2:, ...], + pad, examples[column][idx][cut1:cut2, ...], + ) + ) + return examples + + def apply_crop(self, examples: Dict, max_length: Optional[int] = None) -> Dict: + for idx, text in enumerate(examples[self.text_column]): + n = len(text) + if max_length and max_length < n: + cut = random.randint(0, n - max_length) + text = text[cut: cut + max_length] + examples[self.text_column][idx] = text + if self.features is not None: + for column in self.features: + if column in examples: + examples[column][idx] = examples[column][idx][ + cut: cut + max_length, ... + ] + return examples + + def apply_wrap(self, examples: Dict) -> Dict: + for idx, text in enumerate(examples[self.text_column]): + examples[self.text_column][idx] = ( + f"{self.tokenizer.bos_token}{text}{self.tokenizer.eos_token}" + ) + if self.features is not None: + for column in self.features: + if column in examples: + pad = examples[column][idx].new_full( + (1, *examples[column][idx].shape[1:]), + -100 if "labels" in column else 0, + ) + examples[column][idx] = torch.cat((pad, examples[column][idx], pad)) + return examples + + +__all__ = [ + "XYZPreTrainedModel", + "XYZModelOutputWithPast", + "XYZModel", + "XYZCausalLMOutputWithPast", + "XYZForCausalLM", + "XYZForSequenceClassification", + "XYZConfig", + "XYZProcessor", +] diff --git a/src/proxyz/models/processing_xyz.py b/src/proxyz/models/processing_xyz.py new file mode 100644 index 0000000..b2c4777 --- /dev/null +++ b/src/proxyz/models/processing_xyz.py @@ -0,0 +1,209 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/xyz/modular_xyz.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_xyz.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +import random + +import torch +from proxyz.utils import attr + +from transformers.processing_utils import ProcessorMixin +from transformers.utils import logging + + +logger = logging.get_logger(__name__) + + +class XYZProcessor(ProcessorMixin): + """Processes protein sequences at both amino-acid and BPE-token granularity. + + Inherits ``ProcessorMixin`` so it supports ``save_pretrained`` / + ``from_pretrained`` out of the box. The wrapped BPE tokenizer is + persisted alongside the processor config. + + ``AutoProcessor.from_pretrained(dir)`` works as long as *proxyz* is + importable (the ``auto_map`` entry in *preprocessor_config.json* + points to ``proxyz.processor.XYZProcessor``). + """ + + attributes = ["tokenizer"] + tokenizer_class = "AutoTokenizer" + + FIM_PREFIX = "" + FIM_SUFFIX = "" + FIM_MIDDLE = "" + FIM_TOKENS = [FIM_PREFIX, FIM_SUFFIX, FIM_MIDDLE] + + def __init__(self, tokenizer, text_column: str = "text", features: list[str] = None, **kwargs) -> None: + tokenizer.add_special_tokens({"additional_special_tokens": self.FIM_TOKENS}) + logger.info(f"Added FIM tokens: {self.FIM_TOKENS} to tokenizer. (vocabu size: {len(tokenizer)})") + + super().__init__(tokenizer, **kwargs) + + self.text_column = text_column + self.features = features + + def __call__( + self, + examples: dict, + *, + bpe_dropout: float | None = None, + char_apply: bool = False, + fim_apply: bool = False, + fim_spm_rate: float = 0.5, + fim_sft_style: bool = False, + max_length: int | None = None, + **kwargs, + ) -> dict: + examples = self.apply_crop(examples, max_length=max_length) + if fim_apply: + examples = self.apply_fim(examples, spm_rate=fim_spm_rate) + examples = self.apply_wrap(examples) + + def tokenize_with_dropout(dropout, prefix=""): + with attr(self.tokenizer.backend_tokenizer.model, dropout=dropout): + tokenized = self.tokenizer( + examples[self.text_column], + truncation=True, + return_tensors="pt", + padding=True, + return_offsets_mapping=char_apply, + ) + tokenized["labels"] = tokenized["input_ids"].where(tokenized["attention_mask"] > 0, -100) + if fim_apply and fim_sft_style: + middle_pos = (tokenized["labels"] == self.tokenizer.convert_tokens_to_ids(self.FIM_MIDDLE)).cumsum(1) + tokenized["labels"] = tokenized["labels"].where( + middle_pos.cumsum(1) > 1, + -100, # the is excluded + ) + return {f"{prefix}{k}": v for k, v in tokenized.items()} + + tokenized = tokenize_with_dropout(bpe_dropout, prefix="") + if char_apply: + tokenized.update(tokenize_with_dropout(1.0, prefix="char_")) + + # Align characters with tokenized *text* + tokenized["char_to_token_idx"] = tokenized["char_attention_mask"].new_zeros( + tokenized["char_attention_mask"].size() + ) + tokenized["repr_char_idx"] = tokenized["attention_mask"].new_zeros(tokenized["attention_mask"].size()) + # FIX: use torch.searchsorted instead + for k in range(len(examples[self.text_column])): + i, j = 0, 0 + while i < len(tokenized["offset_mapping"][k]) and j < len(tokenized["char_offset_mapping"][k]): + if tokenized["attention_mask"][k, i] == 0: + i += 1 + elif tokenized["char_attention_mask"][k, j] == 0: + j += 1 + else: + si, ei = tokenized["offset_mapping"][k, i] + sj, ej = tokenized["char_offset_mapping"][k, j] + if si <= sj and ej <= ei: + tokenized["char_to_token_idx"][k, j] = i + if ej == ei: + tokenized["repr_char_idx"][k, i] = j + j += 1 + elif ei <= sj: + i += 1 + else: + assert False, (text[k], i, j) + + del tokenized["char_offset_mapping"] + del tokenized["offset_mapping"] + + # copy features + if self.features is not None: + for column in self.features: + if column in examples: + max_len = max(examples[column][k].shape[0] for k in range(len(examples[column]))) + for k in range(len(examples[column])): + pad = examples[column][k].new_full( + (max_len - examples[column][k].shape[0], *examples[column][k].shape[1:]), + -100 if "labels" in column else 0, + ) + examples[column][k] = torch.cat((examples[column][k], pad)) + tokenized[column] = torch.stack(examples[column]) + + return tokenized + + def apply_fim(self, examples: dict, spm_rate: float = 0.5) -> dict: + """Split content into prefix/middle/suffix and rearrange for FIM training. + Prefix or suffix may be empty, but middle is always non-empty.""" + for idx, text in enumerate(examples[self.text_column]): + n = len(text) + cut1 = random.randint(0, n - 1) + cut2 = random.randint(cut1 + 1, n) + prefix, middle, suffix = text[:cut1], text[cut1:cut2], text[cut2:] + + is_spm = random.random() < spm_rate + if is_spm: + # SPM: + first_tag, second_tag = self.FIM_SUFFIX, self.FIM_PREFIX + first, second = suffix, prefix + else: + # PSM: + first_tag, second_tag = self.FIM_PREFIX, self.FIM_SUFFIX + first, second = prefix, suffix + + examples[self.text_column][idx] = first_tag + first + second_tag + second + self.FIM_MIDDLE + middle + if self.features is not None: + for column in self.features: + if column in examples: + pad = examples[column][idx].new_full( + (1, *examples[column][idx].shape[1:]), + -100 if "labels" in column else 0, + ) + if is_spm: + examples[column][idx] = torch.cat( + ( + pad, + examples[column][idx][cut2:, ...], + pad, + examples[column][idx][:cut1, ...], + pad, + examples[column][idx][cut1:cut2, ...], + ) + ) + else: + examples[column][idx] = torch.cat( + ( + pad, + examples[column][idx][:cut1, ...], + pad, + examples[column][idx][cut2:, ...], + pad, + examples[column][idx][cut1:cut2, ...], + ) + ) + return examples + + def apply_crop(self, examples: dict, max_length: int | None = None) -> dict: + for idx, text in enumerate(examples[self.text_column]): + n = len(text) + if max_length and max_length < n: + cut = random.randint(0, n - max_length) + text = text[cut : cut + max_length] + examples[self.text_column][idx] = text + if self.features is not None: + for column in self.features: + if column in examples: + examples[column][idx] = examples[column][idx][cut : cut + max_length, ...] + return examples + + def apply_wrap(self, examples: dict) -> dict: + for idx, text in enumerate(examples[self.text_column]): + examples[self.text_column][idx] = f"{self.tokenizer.bos_token}{text}{self.tokenizer.eos_token}" + if self.features is not None: + for column in self.features: + if column in examples: + pad = examples[column][idx].new_full( + (1, *examples[column][idx].shape[1:]), + -100 if "labels" in column else 0, + ) + examples[column][idx] = torch.cat((pad, examples[column][idx], pad)) + return examples + + +__all__ = ["XYZProcessor"] diff --git a/src/proxyz/models/test_xyz.py b/src/proxyz/models/test_xyz.py new file mode 100644 index 0000000..e6b9f71 --- /dev/null +++ b/src/proxyz/models/test_xyz.py @@ -0,0 +1,60 @@ +import click + +from transformers import ( + AutoProcessor, + AutoTokenizer, + PreTrainedTokenizerFast, +) +from proxyz.models.processing_xyz import XYZProcessor +from proxyz.data import dataset + +XYZProcessor.register_for_auto_class() + +class dict2object(object): + def __init__(self, **args): + self.__dict__.update(args) + +@click.command(context_settings={'show_default': True}) +@click.option( + "--output_dir", + type=click.Path(), + default="./deepseek_style_model", + help="Where checkpoints and the final model are saved.", +) +@click.option( + "--tokenizer_file", + type=click.Path(), + default="my_tokenizer.json", + help="Path to the tokenizer json file.", +) +def main(**args): + args = dict2object(**args) + + fim_tokens = dataset.FIM_TOKENS + + tokenizer = PreTrainedTokenizerFast( + tokenizer_file=args.tokenizer_file, + unk_token="[UNK]", + pad_token="[PAD]", + bos_token="[BOS]", + eos_token="[EOS]", + ) + + processor = XYZProcessor(tokenizer=tokenizer) + print(tokenizer.all_special_tokens, tokenizer.all_special_ids) + print(processor.attributes) + print(dir(processor)) + print(processor(["ABC", "ACCGGGWGCG"], fim_apply=False, char_apply=True)) + + processor.save_pretrained(args.output_dir) + + processor = AutoProcessor.from_pretrained(args.output_dir, trust_remote_code=True) + print(processor.attributes) + print(dir(processor)) + + + print(processor(["ABC", "ACCGGGWGCG"], fim_apply=False, char_apply=True)) + +if __name__ == "__main__": + main() + diff --git a/src/proxyz/train.py b/src/proxyz/train.py index 0787342..73fd0ee 100644 --- a/src/proxyz/train.py +++ b/src/proxyz/train.py @@ -10,16 +10,15 @@ from transformers import ( LlamaConfig, LlamaForCausalLM, - DeepseekV2Config, - DeepseekV2ForCausalLM, PreTrainedTokenizerFast, Trainer, TrainingArguments, ) from datasets import Dataset, load_dataset -from proxyz.utils import dict2object, compose from proxyz.data import dataset, sampler +from proxyz.models import XYZConfig, XYZForCausalLM, XYZProcessor +from proxyz.utils import dict2object, compose @click.command(context_settings={'show_default': True}) @@ -97,39 +96,43 @@ help="Model Grouped-Query Attention (GQA) for speed.", ) @click.option( - "--use_mla", + "--use_unet", is_flag=True, - help="Use Multi-head Latent Attention (MLA) from DeepSeek-V2 instead of standard Llama attention.", + help="Use U-net style XYZForCausalLM instead of standard Llama attention.", ) @click.option( - "--kv_lora_rank", + "--model_char_hidden_size", type=int, - default=512, - help="MLA: Rank for KV low-rank compression.", + default=768, + help="Character: Model width.", ) @click.option( - "--q_lora_rank", + "--model_char_intermediate_size", type=int, - default=0, - help="MLA: Rank for Q low-rank compression (0 to disable).", + default=2064, + help="Character: Model SwiGLU hidden dimension (usually ~8/3 of hidden_size).", ) @click.option( - "--qk_nope_head_dim", + "--model_char_num_hidden_layers", type=int, - default=128, - help="MLA: Dimension of non-RoPE part in Q/K heads.", + default=2, + help="Character: Model depth.", ) @click.option( - "--qk_rope_head_dim", + "--model_char_num_attention_heads", type=int, - default=64, - help="MLA: Dimension of RoPE part in Q/K heads.", + default=6, + help="Character: Model attention heads.", ) @click.option( - "--v_head_dim", - type=int, - default=128, - help="MLA: Dimension of V heads.", + "--model_has_char_lm_head", + is_flag=True, + help="Character: Model has char_lm_head", +) +@click.option( + "--model_has_cle_lm_head", + is_flag=True, + help="Character: Model has cle_lm_head", ) @click.option( "--max_position_embeddings", type=int, default=4096, help="Context window length." @@ -272,15 +275,13 @@ def main(**args): bos_token="[BOS]", eos_token="[EOS]", ) - if args.tokenizer_bpe_dropout > 0: - # FIXME: Turn it Off (0.0) for evaluation - tokenizer.backend_tokenizer.model.dropout = args.tokenizer_bpe_dropout - - # Add FIM special tokens if FIM training is enabled - if args.fim_rate > 0: - fim_tokens = dataset.FIM_TOKENS - tokenizer.add_special_tokens({"additional_special_tokens": fim_tokens}) - print(f"Added FIM tokens: {fim_tokens} (vocab size: {len(tokenizer)})") + processor = XYZProcessor( + tokenizer=tokenizer, + text_column=args.text_column, + features=[ + "coord", "coord_mask", "pseudo_beta", "pseudo_beta_mask", "cle_labels", + ], + ) # Ensure the embedding layer matches this size exactly vocab_size = len(tokenizer) @@ -290,6 +291,17 @@ def main(**args): # ========================================== use_cuda = torch.cuda.is_available() + label_names = ["labels"] + keys_to_ignore_at_inference = ["past_key_values", "char_past_key_values"] + if args.model_has_char_lm_head: + label_names += ["char_labels"] + else: + keys_to_ignore_at_inference += ["char_logits"] + if args.model_has_cle_lm_head: + label_names += ["cle_labels"] + else: + keys_to_ignore_at_inference += ["cle_logits"] + # Shared config parameters common_config = dict( vocab_size=vocab_size, @@ -297,7 +309,9 @@ def main(**args): intermediate_size=args.model_intermediate_size, num_hidden_layers=args.model_num_hidden_layers, num_attention_heads=args.model_num_attention_heads, + num_key_value_heads=args.model_num_key_value_heads, max_position_embeddings=args.max_position_embeddings, + hidden_act="silu", initializer_range=0.02, rms_norm_eps=1e-6, pad_token_id=tokenizer.pad_token_id, @@ -306,27 +320,25 @@ def main(**args): attn_implementation=args.attn_implementation, torch_dtype=torch.bfloat16, tie_word_embeddings=False, + keys_to_ignore_at_inference=keys_to_ignore_at_inference ) - if args.use_mla: - config = DeepseekV2Config( + if args.use_unet: + config = XYZConfig( **common_config, - kv_lora_rank=args.kv_lora_rank, - q_lora_rank=args.q_lora_rank, - qk_nope_head_dim=args.qk_nope_head_dim, - qk_rope_head_dim=args.qk_rope_head_dim, - v_head_dim=args.v_head_dim, + char_hidden_size=args.model_char_hidden_size, + char_intermediate_size=args.model_char_intermediate_size, + char_num_hidden_layers=args.model_char_num_hidden_layers, + char_num_attention_heads=args.model_char_num_attention_heads, + has_char_lm_head=args.model_has_char_lm_head, + has_cle_lm_head=args.model_has_cle_lm_head, ) - model = DeepseekV2ForCausalLM(config) - attn_type = "MLA (DeepSeek-V2)" + model = XYZForCausalLM(config) else: config = LlamaConfig( **common_config, - num_key_value_heads=args.model_num_key_value_heads, - hidden_act="silu", ) model = LlamaForCausalLM(config) - attn_type = "Llama GQA" # Ensure all parameters are bf16 — FlashAttention requires fp16 or bf16 use_cuda = torch.cuda.is_available() @@ -337,7 +349,6 @@ def main(**args): total_params = sum(p.numel() for p in model.parameters()) trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) print(f"--- Dense DeepSeek-Style Model ---") - print(f"Attention: {attn_type}") print(f"Attention backend: {args.attn_implementation}") print(f"Total Parameters: {total_params:,}") print(f"Trainable Parameters: {trainable_params:,}") @@ -345,74 +356,24 @@ def main(**args): # ========================================== # 3. PREPARE YOUR DATASET # ========================================== - def apply_crop(text, max_length): - if len(text) > max_length: - cut = random.randint(0, len(text) - max_length) - text = text[cut: cut + max_length] - return text - - def apply_fim(content): - """Split content into prefix/middle/suffix and rearrange for FIM training. - Prefix or suffix may be empty, but middle is always non-empty.""" - n = len(content) - cut1 = random.randint(0, n - 1) - cut2 = random.randint(cut1 + 1, n) - prefix, middle, suffix = content[:cut1], content[cut1:cut2], content[cut2:] - - is_spm = random.random() < args.fim_spm_rate - if is_spm: - # SPM: - first_tag, second_tag = dataset.FIM_SUFFIX, dataset.FIM_PREFIX - first, second = suffix, prefix - else: - # PSM: - first_tag, second_tag = dataset.FIM_PREFIX, dataset.FIM_SUFFIX - first, second = prefix, suffix - - return first_tag + first + second_tag + second + dataset.FIM_MIDDLE + middle - max_sequence_length = args.max_sequence_length if max_sequence_length is None: max_sequence_length = args.max_position_embeddings def tokenize_function(examples): - do_fim = random.random() < args.fim_rate - + fim_apply = random.random() < args.fim_rate if args.data_format == "pdb": - examples = dataset.pdb_transform(os.environ["DATA_DIR"], examples) - - if do_fim: - text_process_fn = compose( - apply_fim, - functools.partial(apply_crop, max_length=max_sequence_length - 5) - ) - else: - text_process_fn = compose( - functools.partial(apply_crop, max_length=max_sequence_length - 2) - ) - - wrapped = [ - f"{tokenizer.bos_token}{text_process_fn(text)}{tokenizer.eos_token}" - for text in examples[args.text_column] - ] - tokenized = tokenizer( - wrapped, - truncation=True, - return_tensors="pt", - padding=True, - ) - tokenized["labels"] = tokenized["input_ids"].where( - tokenized["attention_mask"] > 0, -100 + examples = dataset.pdb_transform(os.environ["DATA_PATH"], examples) + examples = processor( + examples, + bpe_dropout=args.tokenizer_bpe_dropout, + char_apply=args.use_unet, + fim_apply=fim_apply, + fim_spm_rate=args.fim_spm_rate, + fim_sft_style=args.fim_sft_style, + max_length=max_sequence_length - (5 if fim_apply else 2), ) - if do_fim and args.fim_sft_style: - middle_pos = ( - tokenized["labels"] == tokenizer.convert_tokens_to_ids(dataset.FIM_MIDDLE) - ).cumsum(1) - tokenized["labels"] = tokenized["labels"].where( - middle_pos.cumsum(1) > 1, -100 # the is excluded - ) - - return tokenized + return examples # Load dataset from HuggingFace or local files if args.dataset_name: @@ -494,7 +455,9 @@ def compute_loss( ): # Always cache data for FIM loss tracking (training and eval) # Cache data BEFORE calling super (which may modify inputs) - labels = inputs["labels"].clone() + labels = tuple(inputs[k].clone() for k in self.args.label_names) + if len(labels) == 1: + labels = labels[0] # Call parent compute_loss (handles label smoothing, loss scaling, etc.) loss, outputs = super().compute_loss( @@ -502,12 +465,24 @@ def compute_loss( ) # ONLY track training metrics if the model is actively training - if model.training: - for key, val in self.aux_metric_calculator( - ( - self.aux_preprocess_logits_for_metrics(outputs.logits, labels), - labels, - ), prefix="" + if model.training and self.compute_metrics is not None: + ignore_keys = [] + + module = model + if not hasattr(module, "config") and hasattr(module, "module"): # DistributedDataParallel + module = module.module + if hasattr(module, "config"): + ignore_keys = getattr( + module.config, "keys_to_ignore_at_inference", ["past_key_values"] + ) + + logits = tuple(v for k, v in outputs.items() if k not in ignore_keys + ["loss"]) + if len(logits) == 1: + logits = logits[0] + if self.preprocess_logits_for_metrics is not None: + logits = self.preprocess_logits_for_metrics(logits, labels) + for key, val in self.compute_metrics( + (logits, labels), prefix="" ).items(): if key in self._logs: self._logs[key].append(val) @@ -526,9 +501,15 @@ def log(self, logs, start_time=None): super().log(logs, start_time=start_time) - @staticmethod - def aux_preprocess_logits_for_metrics(logits, labels): + @classmethod + def aux_preprocess_logits_for_metrics(cls, logits, labels): loss_fct = torch.nn.CrossEntropyLoss(reduction="none") + if isinstance(logits, tuple) and isinstance(labels, tuple): + return tuple( + cls.aux_preprocess_logits_for_metrics(p, l) for p, l in zip(logits, labels) + ) + assert not isinstance(logits, tuple), len(logits) + assert not isinstance(labels, tuple), len(labels) shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() loss_per_token = loss_fct( @@ -536,17 +517,31 @@ def aux_preprocess_logits_for_metrics(logits, labels): ) return loss_per_token.view(-1, shift_labels.size(-1)) - @staticmethod - def aux_metric_calculator(preds, prefix="eval_"): + @classmethod + def aux_metric_calculator(cls, preds, prefix="eval_", label_names=None): """Computes metrics: n_fim / n_std, loss_fim / loss_std etc""" loss_per_token, labels = preds logs = {} + + if isinstance(loss_per_token, tuple) and isinstance(labels, tuple): + assert len(label_names) == len(loss_per_token) + for p, l, label in zip(loss_per_token, labels, label_names): + if label.endswith("labels"): + label = label[:-len("labels")] + logs.update( + cls.aux_metric_calculator((p, l), prefix=f"{prefix}{label}") + ) + return logs + with torch.no_grad(): # Detect FIM examples: labels start with -100 - is_fim = ( - labels == tokenizer.convert_tokens_to_ids(dataset.FIM_MIDDLE) - ).any(1) + if args.fim_sft_style: + is_fim = (labels[..., 0] == -100) + else: + is_fim = ( + labels == processor.tokenizer.convert_tokens_to_ids(processor.FIM_MIDDLE) + ).any(1) for tag, mask in [("fim", is_fim), ("std", ~is_fim)]: # Add FIM/standard counts to logs @@ -583,6 +578,7 @@ def aux_metric_calculator(preds, prefix="eval_"): eval_strategy=args.eval_strategy if (args.eval_files or args.dataset_eval_split) else "no", eval_steps=args.eval_steps if args.eval_strategy == "steps" else None, eval_accumulation_steps=args.gradient_accumulation_steps, + label_names=label_names, bf16=use_cuda, # bf16 is preferred over fp16 on modern GPUs ddp_find_unused_parameters=False, # disabled warning num_train_epochs=args.num_train_epochs, @@ -598,10 +594,12 @@ def aux_metric_calculator(preds, prefix="eval_"): args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset, - processing_class=tokenizer, # transformers >=5 renamed `tokenizer` + processing_class=processor, # transformers >=5 renamed `tokenizer` train_sampler=train_sampler, preprocess_logits_for_metrics=FIMTrainer.aux_preprocess_logits_for_metrics, - compute_metrics=FIMTrainer.aux_metric_calculator, + compute_metrics=functools.partial( + FIMTrainer.aux_metric_calculator, label_names=training_args.label_names + ), ) @@ -637,7 +635,7 @@ def aux_metric_calculator(preds, prefix="eval_"): # Save final weights and configuration trainer.save_model(args.output_dir) - tokenizer.save_pretrained(args.output_dir) + processor.save_pretrained(args.output_dir) # Clean up distributed process group to avoid resource leaks if torch.distributed.is_initialized(): diff --git a/src/proxyz/utils/__init__.py b/src/proxyz/utils/__init__.py index 241516b..4f6faeb 100644 --- a/src/proxyz/utils/__init__.py +++ b/src/proxyz/utils/__init__.py @@ -1,3 +1,4 @@ +import contextlib import functools @@ -8,3 +9,14 @@ def __init__(self, **args): def compose(*funcs): return functools.reduce(lambda g, f: lambda x: f(g(x)), funcs) + + +@contextlib.contextmanager +def attr(obj, **kwags): + t = {key: getattr(obj, key) for key in kwags} + + for key in kwags: + setattr(obj, key, kwags[key]) + yield obj + for key in kwags: + setattr(obj, key, t[key])