feat: U-net style transformer

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_prefix>"
-FIM_SUFFIX = "<fim_suffix>"
-FIM_MIDDLE = "<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,
     }
 
 

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)))

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()
+

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"]