Upload folder using huggingface_hub

config.json

@@ -0,0 +1,34 @@
+{
+  "architectures": [
+    "HELMBertForMaskedLM"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "bos_token_id": 1,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 256,
+  "intermediate_size": 3072,
+  "mask_token_id": 4,
+  "max_position_embeddings": 512,
+  "max_relative_positions": 512,
+  "model_type": "helmbert",
+  "ngie_dropout": 0.1,
+  "ngie_kernel_size": 3,
+  "num_attention_heads": 4,
+  "num_hidden_layers": 2,
+  "pad_token_id": 0,
+  "pos_att_type": "c2p|p2c",
+  "position_buckets": 256,
+  "sep_token_id": 2,
+  "share_att_key": false,
+  "transformers_version": "4.57.3",
+  "vocab_size": 78,
+  "auto_map": {
+    "AutoConfig": "configuration_helmbert.HELMBertConfig",
+    "AutoModel": "modeling_helmbert.HELMBertModel",
+    "AutoModelForMaskedLM": "modeling_helmbert.HELMBertForMaskedLM",
+    "AutoModelForSequenceClassification": "modeling_helmbert.HELMBertForSequenceClassification",
+    "AutoTokenizer": "tokenization_helmbert.HELMBertTokenizer"
+  }
+}

configuration_helmbert.py

@@ -0,0 +1,104 @@
+"""HELM-BERT configuration."""
+
+from transformers import PretrainedConfig
+
+
+class HELMBertConfig(PretrainedConfig):
+    """Configuration class for HELM-BERT model.
+
+    This configuration class stores all the parameters needed to instantiate a HELM-BERT model.
+    It inherits from PretrainedConfig and can be used with HuggingFace's from_pretrained and
+    save_pretrained methods.
+
+    Args:
+        vocab_size: Size of the vocabulary (default: 78 for HELM character vocabulary)
+        hidden_size: Dimensionality of the encoder layers (default: 768)
+        num_hidden_layers: Number of transformer layers (default: 6)
+        num_attention_heads: Number of attention heads (default: 12)
+        intermediate_size: Dimensionality of the feed-forward layer (default: 3072)
+        hidden_dropout_prob: Dropout probability for hidden layers (default: 0.1)
+        attention_probs_dropout_prob: Dropout probability for attention (default: 0.1)
+        max_position_embeddings: Maximum sequence length (default: 512)
+        max_relative_positions: Maximum relative position distance (default: 512)
+        position_buckets: Number of position buckets for log-bucketing (default: 256)
+        pos_att_type: Position attention types, pipe-separated (default: "c2p|p2c")
+        share_att_key: Whether to share attention key projections (default: False)
+        ngie_kernel_size: Kernel size for nGiE convolution (default: 3)
+        ngie_dropout: Dropout for nGiE layer (default: 0.1)
+        pad_token_id: ID of padding token (default: 0)
+        bos_token_id: ID of beginning-of-sequence token (default: 1)
+        eos_token_id: ID of end-of-sequence token (default: 2)
+        sep_token_id: ID of separator token (default: 2)
+        mask_token_id: ID of mask token (default: 4)
+
+    Example:
+        >>> from helmbert import HELMBertConfig, HELMBertModel
+        >>> config = HELMBertConfig(hidden_size=768, num_hidden_layers=6)
+        >>> model = HELMBertModel(config)
+    """
+
+    model_type = "helmbert"
+
+    def __init__(
+        self,
+        vocab_size: int = 78,
+        hidden_size: int = 768,
+        num_hidden_layers: int = 6,
+        num_attention_heads: int = 12,
+        intermediate_size: int = 3072,
+        hidden_dropout_prob: float = 0.1,
+        attention_probs_dropout_prob: float = 0.1,
+        max_position_embeddings: int = 512,
+        # Disentangled attention parameters
+        max_relative_positions: int = 512,
+        position_buckets: int = 256,
+        pos_att_type: str = "c2p|p2c",
+        share_att_key: bool = False,
+        # nGiE parameters
+        ngie_kernel_size: int = 3,
+        ngie_dropout: float = 0.1,
+        # Special token IDs
+        pad_token_id: int = 0,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+        sep_token_id: int = 2,
+        mask_token_id: int = 4,
+        # Classification/regression
+        num_labels: int = 2,
+        problem_type: str = None,
+        **kwargs,
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            **kwargs,
+        )
+
+        # Core transformer parameters
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.max_position_embeddings = max_position_embeddings
+
+        # Disentangled attention parameters
+        self.max_relative_positions = max_relative_positions
+        self.position_buckets = position_buckets
+        self.pos_att_type = pos_att_type
+        self.share_att_key = share_att_key
+
+        # nGiE parameters
+        self.ngie_kernel_size = ngie_kernel_size
+        self.ngie_dropout = ngie_dropout
+
+        # Special token IDs
+        self.sep_token_id = sep_token_id
+        self.mask_token_id = mask_token_id
+
+        # Classification/regression
+        self.num_labels = num_labels
+        self.problem_type = problem_type

model.safetensors

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

modeling_helmbert.py

@@ -0,0 +1,988 @@
+"""HELM-BERT model implementation.
+
+This module implements the HELM-BERT model with:
+- Disentangled attention (DeBERTa-style)
+- Enhanced Mask Decoder (EMD) for MLM
+- n-gram Induced Encoding (nGiE) layer
+"""
+
+import math
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from packaging import version
+from torch import _softmax_backward_data
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import (
+    BaseModelOutput,
+    BaseModelOutputWithPooling,
+    MaskedLMOutput,
+    SequenceClassifierOutput,
+)
+
+from .configuration_helmbert import HELMBertConfig
+
+
+# -----------------------------------------------------------------------------
+# Utility Functions
+# -----------------------------------------------------------------------------
+
+
+def masked_layer_norm(
+    layer_norm: nn.LayerNorm, x: torch.Tensor, mask: Optional[torch.Tensor] = None
+) -> torch.Tensor:
+    """Apply LayerNorm with masking to avoid updates on padding tokens.
+
+    Args:
+        layer_norm: LayerNorm module
+        x: Input tensor (batch_size, seq_len, hidden_size)
+        mask: Mask tensor where 0 = padding (ignored), 1 = valid token
+
+    Returns:
+        Normalized tensor with padding positions zeroed out
+    """
+    output = layer_norm(x).to(x)
+    if mask is None:
+        return output
+    if mask.dim() != x.dim():
+        if mask.dim() == 4:
+            mask = mask.squeeze(1).squeeze(1)
+        mask = mask.unsqueeze(2)
+    mask = mask.to(output.dtype)
+    return output * mask
+
+
+class XSoftmax(torch.autograd.Function):
+    """Masked Softmax optimized for memory efficiency."""
+
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, mask: Optional[torch.Tensor], dim: int) -> torch.Tensor:
+        ctx.dim = dim
+        if mask is not None:
+            rmask = ~(mask.bool())
+            if rmask.dim() == 2:
+                rmask = rmask.unsqueeze(1).unsqueeze(2)
+            elif rmask.dim() == 3:
+                rmask = rmask.unsqueeze(2)
+            output = input.masked_fill(rmask, float("-inf"))
+        else:
+            output = input
+        output = torch.softmax(output, ctx.dim)
+        if mask is not None:
+            output.masked_fill_(rmask, 0)
+        ctx.save_for_backward(output)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None, None]:
+        (output,) = ctx.saved_tensors
+        if version.Version(torch.__version__) >= version.Version("1.11.0"):
+            input_grad = _softmax_backward_data(grad_output, output, ctx.dim, output.dtype)
+        else:
+            input_grad = _softmax_backward_data(grad_output, output, ctx.dim, output)
+        return input_grad, None, None
+
+
+def build_relative_position(
+    query_size: int,
+    key_size: int,
+    bucket_size: int = -1,
+    max_position: int = 512,
+    device: Optional[torch.device] = None,
+) -> torch.Tensor:
+    """Build relative position matrix with optional log-bucketing."""
+    q_ids = torch.arange(query_size, dtype=torch.long, device=device)
+    k_ids = torch.arange(key_size, dtype=torch.long, device=device)
+    rel_pos = q_ids.unsqueeze(1) - k_ids.unsqueeze(0)
+
+    if bucket_size > 0:
+        rel_buckets = 0
+        num_buckets = bucket_size
+        rel_buckets += (rel_pos > 0).long() * (num_buckets // 2)
+        rel_pos = torch.abs(rel_pos)
+
+        max_exact = num_buckets // 4
+        is_small = rel_pos < max_exact
+
+        rel_pos_if_large = max_exact + (
+            torch.log(rel_pos.float() / max_exact)
+            / math.log(max_position / max_exact)
+            * (num_buckets // 4 - 1)
+        ).long()
+        rel_pos_if_large = torch.min(
+            rel_pos_if_large, torch.full_like(rel_pos_if_large, num_buckets // 2 - 1)
+        )
+
+        rel_buckets += torch.where(is_small, rel_pos, rel_pos_if_large)
+        return rel_buckets
+    else:
+        rel_pos = torch.clamp(rel_pos, -max_position, max_position)
+        return rel_pos + max_position
+
+
+# -----------------------------------------------------------------------------
+# Attention Modules
+# -----------------------------------------------------------------------------
+
+
+class DisentangledSelfAttention(nn.Module):
+    """Disentangled self-attention with content and position separation.
+
+    Implements content-to-content, content-to-position, and position-to-content
+    attention as described in DeBERTa.
+    """
+
+    def __init__(self, config: HELMBertConfig):
+        super().__init__()
+
+        if config.hidden_size % config.num_attention_heads != 0:
+            raise ValueError(
+                f"hidden_size ({config.hidden_size}) must be divisible by "
+                f"num_attention_heads ({config.num_attention_heads})"
+            )
+
+        self.num_heads = config.num_attention_heads
+        self.head_size = config.hidden_size // config.num_attention_heads
+        self.all_head_size = self.num_heads * self.head_size
+
+        # Content projections
+        self.query_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
+        self.key_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
+        self.value_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
+
+        # Position attention configuration
+        self.pos_att_type = [x.strip() for x in config.pos_att_type.lower().split("|")]
+        self.max_relative_positions = config.max_relative_positions
+        self.position_buckets = config.position_buckets
+        self.share_att_key = config.share_att_key
+
+        # Position embedding size
+        self.pos_ebd_size = config.max_relative_positions
+        if config.position_buckets > 0:
+            self.pos_ebd_size = config.position_buckets
+
+        # Position embeddings
+        self.rel_embeddings = nn.Embedding(self.pos_ebd_size * 2, config.hidden_size)
+
+        # Position projections
+        if not self.share_att_key:
+            if "c2p" in self.pos_att_type or "p2p" in self.pos_att_type:
+                self.pos_key_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
+            if "p2c" in self.pos_att_type or "p2p" in self.pos_att_type:
+                self.pos_query_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=False)
+
+        # Dropout
+        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+        self.pos_dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+    def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
+        """Reshape tensor for attention computation."""
+        new_shape = x.size()[:-1] + (self.num_heads, self.head_size)
+        x = x.view(*new_shape)
+        return x.permute(0, 2, 1, 3).contiguous().view(-1, x.size(1), x.size(-1))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        query_states: Optional[torch.Tensor] = None,
+        relative_pos: Optional[torch.Tensor] = None,
+        rel_embeddings: Optional[torch.Tensor] = None,
+    ) -> Dict[str, Any]:
+        """Forward pass of disentangled attention."""
+        if query_states is None:
+            query_states = hidden_states
+
+        # Compute Q, K, V
+        query_layer = self.transpose_for_scores(self.query_proj(query_states)).float()
+        key_layer = self.transpose_for_scores(self.key_proj(hidden_states)).float()
+        value_layer = self.transpose_for_scores(self.value_proj(hidden_states))
+
+        # Calculate scale factor
+        scale_factor = 1
+        if "c2p" in self.pos_att_type:
+            scale_factor += 1
+        if "p2c" in self.pos_att_type:
+            scale_factor += 1
+        if "p2p" in self.pos_att_type:
+            scale_factor += 1
+
+        scale = 1.0 / math.sqrt(self.head_size * scale_factor)
+
+        # Content-to-content attention (c2c)
+        c2c_scores = torch.bmm(query_layer, key_layer.transpose(-1, -2) * scale)
+        attention_scores = c2c_scores
+
+        # Add relative position bias if enabled
+        if len(self.pos_att_type) > 0 and self.pos_att_type[0]:
+            rel_att = self._disentangled_attention_bias(
+                query_layer, key_layer, relative_pos, rel_embeddings, scale_factor
+            )
+            if rel_att is not None:
+                attention_scores = attention_scores + rel_att
+
+        # Normalize scores for numerical stability
+        attention_scores = attention_scores - attention_scores.max(dim=-1, keepdim=True)[0].detach()
+        attention_scores = attention_scores.to(hidden_states.dtype)
+
+        # Reshape for XSoftmax
+        attention_scores = attention_scores.view(
+            -1, self.num_heads, attention_scores.size(-2), attention_scores.size(-1)
+        )
+
+        # Apply XSoftmax
+        attention_probs = XSoftmax.apply(attention_scores, attention_mask, -1)
+        attention_probs = self.dropout(attention_probs)
+
+        # Apply attention to values
+        attention_probs_flat = attention_probs.view(-1, attention_probs.size(-2), attention_probs.size(-1))
+        context_layer = torch.bmm(attention_probs_flat, value_layer)
+
+        # Reshape output
+        context_layer = context_layer.view(-1, self.num_heads, context_layer.size(-2), context_layer.size(-1))
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(*new_shape)
+
+        result = {"hidden_states": context_layer, "attention_probs": attention_probs}
+        return result
+
+    def _disentangled_attention_bias(
+        self,
+        query_layer: torch.Tensor,
+        key_layer: torch.Tensor,
+        relative_pos: Optional[torch.Tensor],
+        rel_embeddings: Optional[torch.Tensor],
+        scale_factor: int,
+    ) -> Optional[torch.Tensor]:
+        """Compute disentangled attention bias."""
+        if relative_pos is None:
+            q_size = query_layer.size(-2)
+            k_size = key_layer.size(-2)
+            relative_pos = build_relative_position(
+                q_size,
+                k_size,
+                bucket_size=self.position_buckets,
+                max_position=self.max_relative_positions,
+                device=query_layer.device,
+            )
+
+        if relative_pos.dim() == 2:
+            relative_pos = relative_pos.unsqueeze(0).unsqueeze(0)
+        elif relative_pos.dim() == 3:
+            relative_pos = relative_pos.unsqueeze(1)
+
+        batch_size = query_layer.size(0) // self.num_heads
+
+        # Get position embeddings
+        if rel_embeddings is None:
+            rel_embeddings = self.rel_embeddings.weight
+
+        att_span = self.pos_ebd_size
+        rel_embeddings = rel_embeddings[
+            self.pos_ebd_size - att_span : self.pos_ebd_size + att_span, :
+        ].unsqueeze(0)
+        rel_embeddings = self.pos_dropout(rel_embeddings)
+
+        score = torch.zeros_like(query_layer[:, :, :1]).expand(-1, -1, key_layer.size(-2))
+
+        # Prepare position indices
+        c2p_pos = torch.clamp(relative_pos + att_span, 0, att_span * 2 - 1)
+        c2p_pos = c2p_pos.squeeze(0).expand(query_layer.size(0), query_layer.size(1), relative_pos.size(-1))
+
+        # Content-to-position (c2p)
+        if "c2p" in self.pos_att_type:
+            pos_key_layer = (
+                self.pos_key_proj(rel_embeddings) if not self.share_att_key else self.key_proj(rel_embeddings)
+            )
+            pos_key_layer = self.transpose_for_scores(pos_key_layer).repeat(batch_size, 1, 1)
+
+            c2p_scale = 1.0 / math.sqrt(self.head_size * scale_factor)
+            c2p_att = torch.bmm(query_layer, pos_key_layer.transpose(-1, -2) * c2p_scale)
+            c2p_att = torch.gather(c2p_att, dim=-1, index=c2p_pos)
+            score = score + c2p_att
+
+        # Position-to-content (p2c)
+        if "p2c" in self.pos_att_type:
+            pos_query_layer = (
+                self.pos_query_proj(rel_embeddings) if not self.share_att_key else self.query_proj(rel_embeddings)
+            )
+            pos_query_layer = self.transpose_for_scores(pos_query_layer).repeat(batch_size, 1, 1)
+
+            p2c_scale = 1.0 / math.sqrt(self.head_size * scale_factor)
+            p2c_att = torch.bmm(pos_query_layer * p2c_scale, key_layer.transpose(-1, -2))
+            p2c_att = torch.gather(p2c_att, dim=-2, index=c2p_pos)
+            score = score + p2c_att
+
+        return score
+
+
+# -----------------------------------------------------------------------------
+# Transformer Components
+# -----------------------------------------------------------------------------
+
+
+class HELMBertEmbeddings(nn.Module):
+    """Token and position embeddings for HELM-BERT."""
+
+    def __init__(self, config: HELMBertConfig):
+        super().__init__()
+        self.word_embeddings = nn.Embedding(
+            config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id
+        )
+        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
+        self.layer_norm = nn.LayerNorm(config.hidden_size)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass.
+
+        Returns:
+            Tuple of (token_embeddings, position_embeddings)
+        """
+        batch_size, seq_len = input_ids.shape
+
+        # Token embeddings
+        embeddings = self.word_embeddings(input_ids)
+
+        # Position embeddings
+        position_ids = torch.arange(seq_len, dtype=torch.long, device=input_ids.device)
+        position_ids = position_ids.unsqueeze(0).expand(batch_size, -1)
+        position_embeds = self.position_embeddings(position_ids)
+
+        # Layer norm and dropout
+        embeddings = masked_layer_norm(self.layer_norm, embeddings, attention_mask)
+        embeddings = self.dropout(embeddings)
+
+        return embeddings, position_embeds
+
+
+class NgieLayer(nn.Module):
+    """n-gram Induced Input Encoding (nGiE) layer.
+
+    Captures local n-gram patterns using 1D convolution.
+    """
+
+    def __init__(self, config: HELMBertConfig):
+        super().__init__()
+
+        self.conv = nn.Conv1d(
+            in_channels=config.hidden_size,
+            out_channels=config.hidden_size,
+            kernel_size=config.ngie_kernel_size,
+            padding=(config.ngie_kernel_size - 1) // 2,
+            groups=1,
+        )
+        self.activation = nn.Tanh()
+        self.layer_norm = nn.LayerNorm(config.hidden_size)
+        self.dropout = nn.Dropout(config.ngie_dropout)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        residual_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """Forward pass.
+
+        Args:
+            hidden_states: Input to convolution (batch, seq, hidden)
+            residual_states: States for residual connection (batch, seq, hidden)
+            attention_mask: Mask where 1 = valid, 0 = padding
+
+        Returns:
+            Output with n-gram information incorporated
+        """
+        # Apply 1D convolution
+        out = self.conv(hidden_states.permute(0, 2, 1).contiguous()).permute(0, 2, 1).contiguous()
+
+        # Create reverse mask for padding
+        if version.Version(torch.__version__) >= version.Version("1.2.0a"):
+            rmask = (1 - attention_mask).bool()
+        else:
+            rmask = (1 - attention_mask).byte()
+
+        # Zero out padding positions
+        out.masked_fill_(rmask.unsqueeze(-1).expand(out.size()), 0)
+
+        # Apply activation and dropout
+        out = self.activation(self.dropout(out))
+
+        # Residual connection with LayerNorm
+        output_states = masked_layer_norm(self.layer_norm, residual_states + out, attention_mask)
+
+        return output_states
+
+
+class TransformerBlock(nn.Module):
+    """Transformer block with disentangled attention and GELU FFN."""
+
+    def __init__(self, config: HELMBertConfig):
+        super().__init__()
+
+        self.self_attn = DisentangledSelfAttention(config)
+        self.attn_output_dense = nn.Linear(config.hidden_size, config.hidden_size)
+
+        # FFN with GELU
+        self.linear1 = nn.Sequential(
+            nn.Linear(config.hidden_size, config.intermediate_size), nn.GELU()
+        )
+        self.linear2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+        # Normalization and dropout
+        self.norm1 = nn.LayerNorm(config.hidden_size)
+        self.norm2 = nn.LayerNorm(config.hidden_size)
+        self.dropout1 = nn.Dropout(config.hidden_dropout_prob)
+        self.dropout2 = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(
+        self,
+        src: torch.Tensor,
+        src_key_padding_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        query_states: Optional[torch.Tensor] = None,
+        relative_pos: Optional[torch.Tensor] = None,
+        rel_embeddings: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Forward pass.
+
+        Args:
+            src: Input embeddings [seq_len, batch, hidden]
+            src_key_padding_mask: Padding mask [batch, seq_len]
+            output_attentions: Whether to return attention weights
+            query_states: Optional query for EMD
+            relative_pos: Relative position indices
+            rel_embeddings: Relative position embeddings
+
+        Returns:
+            Tuple of (output, optional attention weights)
+        """
+        # Transpose for attention [seq, batch, hidden] -> [batch, seq, hidden]
+        src_transposed = src.transpose(0, 1)
+
+        # Convert padding mask to attention mask (1=valid, 0=padding)
+        attention_mask = None
+        if src_key_padding_mask is not None:
+            attention_mask = (~src_key_padding_mask).float()
+
+        query_states_transposed = None
+        if query_states is not None:
+            query_states_transposed = query_states.transpose(0, 1)
+
+        # Self-attention
+        attn_result = self.self_attn(
+            src_transposed,
+            attention_mask,
+            output_attentions=output_attentions,
+            query_states=query_states_transposed,
+            relative_pos=relative_pos,
+            rel_embeddings=rel_embeddings,
+        )
+        attn_output = attn_result["hidden_states"].transpose(0, 1)
+        attn_weights = attn_result.get("attention_probs") if output_attentions else None
+
+        # Dense projection
+        attn_output = self.attn_output_dense(attn_output)
+
+        # Residual connection
+        residual_input = query_states if query_states is not None else src
+        src = residual_input + self.dropout1(attn_output)
+
+        # LayerNorm
+        src = src.transpose(0, 1)
+        src = masked_layer_norm(self.norm1, src)
+        src = src.transpose(0, 1)
+
+        # FFN
+        ff_output = self.linear1(src)
+        ff_output = self.linear2(ff_output)
+        ff_output = self.dropout2(ff_output)
+        src = src + ff_output
+
+        # LayerNorm
+        src = src.transpose(0, 1)
+        src = masked_layer_norm(self.norm2, src)
+        src = src.transpose(0, 1)
+
+        return src, attn_weights
+
+
+class HELMBertEncoder(nn.Module):
+    """Stack of transformer blocks with nGiE layer."""
+
+    def __init__(self, config: HELMBertConfig):
+        super().__init__()
+        self.config = config
+
+        # nGiE layer (applied after first transformer block)
+        self.ngie_layer = NgieLayer(config)
+
+        # Transformer blocks
+        self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.num_hidden_layers)])
+
+    def get_rel_embedding(self) -> Optional[torch.Tensor]:
+        """Get relative position embeddings from first layer."""
+        if len(self.layers) > 0:
+            first_layer = self.layers[0]
+            if hasattr(first_layer, "self_attn") and hasattr(first_layer.self_attn, "rel_embeddings"):
+                return first_layer.self_attn.rel_embeddings.weight
+        return None
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_embeddings: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        output_hidden_states: bool = False,
+        use_emd: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple], Optional[Tuple]]:
+        """Forward pass.
+
+        Args:
+            hidden_states: Input embeddings [batch, seq, hidden]
+            attention_mask: Attention mask [batch, seq]
+            position_embeddings: Position embeddings for EMD
+            output_attentions: Whether to return attention weights
+            output_hidden_states: Whether to return all hidden states
+            use_emd: Whether to use Enhanced Mask Decoder
+
+        Returns:
+            Tuple of (last_hidden_state, emd_output, all_hidden_states, all_attentions)
+        """
+        all_hidden_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
+        # Store for nGiE
+        ngie_input_states = hidden_states.clone()
+
+        # [batch, seq, hidden] -> [seq, batch, hidden]
+        hidden_states = hidden_states.transpose(0, 1)
+
+        # Key padding mask (True = padding)
+        key_padding_mask = None
+        if attention_mask is not None:
+            key_padding_mask = ~attention_mask.bool()
+
+        # Store layer[-2] for EMD
+        layer_minus_2 = None
+        num_layers = len(self.layers)
+
+        for layer_idx, layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states.transpose(0, 1),)
+
+            hidden_states, attn_weights = layer(
+                hidden_states,
+                src_key_padding_mask=key_padding_mask,
+                output_attentions=output_attentions,
+            )
+
+            if output_attentions and attn_weights is not None:
+                all_attentions = all_attentions + (attn_weights,)
+
+            # Apply nGiE after first layer
+            if layer_idx == 0:
+                hidden_states_batch = hidden_states.transpose(0, 1)
+                hidden_states_batch = self.ngie_layer(ngie_input_states, hidden_states_batch, attention_mask)
+                hidden_states = hidden_states_batch.transpose(0, 1)
+
+            # Store layer[-2] for EMD
+            if use_emd and layer_idx == num_layers - 2:
+                layer_minus_2 = hidden_states.clone()
+
+        # Convert back to [batch, seq, hidden]
+        hidden_states = hidden_states.transpose(0, 1)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        # Enhanced Mask Decoder (EMD) for MLM
+        emd_output = None
+        if use_emd and layer_minus_2 is not None and position_embeddings is not None:
+            emd_keys_values = layer_minus_2
+            emd_query = layer_minus_2.transpose(0, 1)
+            emd_query = position_embeddings + emd_query
+            emd_query = emd_query.transpose(0, 1)
+
+            rel_embeddings = self.get_rel_embedding()
+            last_layer = self.layers[-1]
+
+            for _ in range(2):
+                emd_query, _ = last_layer(
+                    emd_keys_values,
+                    src_key_padding_mask=key_padding_mask,
+                    query_states=emd_query,
+                    relative_pos=None,
+                    rel_embeddings=rel_embeddings,
+                )
+
+            emd_output = emd_query.transpose(0, 1)
+
+        return hidden_states, emd_output, all_hidden_states, all_attentions
+
+
+class HELMBertPooler(nn.Module):
+    """Mean pooling over sequence."""
+
+    def __init__(self, config: HELMBertConfig):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+    def forward(
+        self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """Apply mean pooling.
+
+        Args:
+            hidden_states: [batch, seq, hidden]
+            attention_mask: [batch, seq]
+
+        Returns:
+            Pooled output [batch, hidden]
+        """
+        if attention_mask is not None:
+            mask_expanded = attention_mask.unsqueeze(-1).expand(hidden_states.size()).float()
+            sum_embeddings = torch.sum(hidden_states * mask_expanded, 1)
+            eps = torch.finfo(hidden_states.dtype).eps
+            sum_mask = torch.clamp(mask_expanded.sum(1), min=eps)
+            return sum_embeddings / sum_mask
+        else:
+            return hidden_states.mean(dim=1)
+
+
+# -----------------------------------------------------------------------------
+# Pre-trained Model Base
+# -----------------------------------------------------------------------------
+
+
+class HELMBertPreTrainedModel(PreTrainedModel):
+    """Base class for HELM-BERT models."""
+
+    config_class = HELMBertConfig
+    base_model_prefix = "helmbert"
+    supports_gradient_checkpointing = True
+
+    def _init_weights(self, module: nn.Module) -> None:
+        """Initialize weights with BERT-style initialization."""
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.ones_(module.weight)
+            nn.init.zeros_(module.bias)
+
+
+# -----------------------------------------------------------------------------
+# Model Classes
+# -----------------------------------------------------------------------------
+
+
+class HELMBertModel(HELMBertPreTrainedModel):
+    """HELM-BERT base model.
+
+    This model outputs the last hidden states and optionally pooled output.
+
+    Example:
+        >>> from helmbert import HELMBertModel, HELMBertTokenizer
+        >>> tokenizer = HELMBertTokenizer()
+        >>> model = HELMBertModel.from_pretrained("./checkpoints/helmbert-base")
+        >>> inputs = tokenizer("PEPTIDE1{A.C.D.E}$$$$", return_tensors="pt")
+        >>> outputs = model(**inputs)
+        >>> last_hidden_state = outputs.last_hidden_state
+        >>> pooler_output = outputs.pooler_output
+    """
+
+    def __init__(self, config: HELMBertConfig):
+        super().__init__(config)
+        self.config = config
+
+        self.embeddings = HELMBertEmbeddings(config)
+        self.encoder = HELMBertEncoder(config)
+        self.pooler = HELMBertPooler(config)
+
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.embeddings.word_embeddings
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.embeddings.word_embeddings = value
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        output_hidden_states: bool = False,
+        return_dict: bool = True,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        """Forward pass.
+
+        Args:
+            input_ids: Token IDs [batch, seq]
+            attention_mask: Attention mask [batch, seq]
+            output_attentions: Whether to return attention weights
+            output_hidden_states: Whether to return all hidden states
+            return_dict: Whether to return a ModelOutput
+
+        Returns:
+            BaseModelOutputWithPooling or tuple
+        """
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+
+        # Embeddings
+        embeddings, position_embeddings = self.embeddings(input_ids, attention_mask)
+
+        # Encoder
+        encoder_outputs = self.encoder(
+            embeddings,
+            attention_mask=attention_mask,
+            position_embeddings=position_embeddings,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            use_emd=False,
+        )
+
+        last_hidden_state = encoder_outputs[0]
+        hidden_states = encoder_outputs[2]
+        attentions = encoder_outputs[3]
+
+        # Pooling
+        pooler_output = self.pooler(last_hidden_state, attention_mask)
+
+        if not return_dict:
+            return (last_hidden_state, pooler_output, hidden_states, attentions)
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=last_hidden_state,
+            pooler_output=pooler_output,
+            hidden_states=hidden_states,
+            attentions=attentions,
+        )
+
+
+class HELMBertLMHead(nn.Module):
+    """MLM head with weight tying."""
+
+    def __init__(self, config: HELMBertConfig):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.layer_norm = nn.LayerNorm(config.hidden_size)
+        self.activation = nn.GELU()
+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=True)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Forward pass.
+
+        Args:
+            hidden_states: [batch, seq, hidden]
+
+        Returns:
+            Logits [batch, seq, vocab]
+        """
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.activation(hidden_states)
+        hidden_states = self.layer_norm(hidden_states)
+        logits = self.decoder(hidden_states)
+        return logits
+
+
+class HELMBertForMaskedLM(HELMBertPreTrainedModel):
+    """HELM-BERT for Masked Language Modeling with Enhanced Mask Decoder (EMD).
+
+    Example:
+        >>> from helmbert import HELMBertForMaskedLM, HELMBertTokenizer
+        >>> tokenizer = HELMBertTokenizer()
+        >>> model = HELMBertForMaskedLM.from_pretrained("./checkpoints/helmbert-base")
+        >>> inputs = tokenizer("PEPTIDE1{A.¶.D.E}$$$$", return_tensors="pt")  # ¶ is mask
+        >>> outputs = model(**inputs)
+        >>> predictions = outputs.logits.argmax(dim=-1)
+    """
+
+    _tied_weights_keys = ["lm_head.decoder.weight", "lm_head.decoder.bias"]
+
+    def __init__(self, config: HELMBertConfig):
+        super().__init__(config)
+        self.helmbert = HELMBertModel(config)
+        self.lm_head = HELMBertLMHead(config)
+
+        self.post_init()
+
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.lm_head.decoder
+
+    def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
+        self.lm_head.decoder = new_embeddings
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        output_hidden_states: bool = False,
+        return_dict: bool = True,
+        use_emd: bool = True,
+    ) -> Union[Tuple, MaskedLMOutput]:
+        """Forward pass.
+
+        Args:
+            input_ids: Token IDs [batch, seq]
+            attention_mask: Attention mask [batch, seq]
+            labels: Labels for MLM (-100 for non-masked tokens)
+            output_attentions: Whether to return attention weights
+            output_hidden_states: Whether to return all hidden states
+            return_dict: Whether to return a ModelOutput
+            use_emd: Whether to use Enhanced Mask Decoder
+
+        Returns:
+            MaskedLMOutput or tuple
+        """
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+
+        # Embeddings
+        embeddings, position_embeddings = self.helmbert.embeddings(input_ids, attention_mask)
+
+        # Encoder with optional EMD
+        encoder_outputs = self.helmbert.encoder(
+            embeddings,
+            attention_mask=attention_mask,
+            position_embeddings=position_embeddings,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            use_emd=use_emd,
+        )
+
+        # Use EMD output if available, otherwise use last hidden state
+        if use_emd and encoder_outputs[1] is not None:
+            sequence_output = encoder_outputs[1]
+        else:
+            sequence_output = encoder_outputs[0]
+
+        hidden_states = encoder_outputs[2]
+        attentions = encoder_outputs[3]
+
+        # MLM head
+        prediction_scores = self.lm_head(sequence_output)
+
+        # Calculate loss if labels provided
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
+            loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
+
+        if not return_dict:
+            output = (prediction_scores, hidden_states, attentions)
+            return ((loss,) + output) if loss is not None else output
+
+        return MaskedLMOutput(
+            loss=loss,
+            logits=prediction_scores,
+            hidden_states=hidden_states,
+            attentions=attentions,
+        )
+
+
+class HELMBertForSequenceClassification(HELMBertPreTrainedModel):
+    """HELM-BERT for sequence classification/regression.
+
+    Example:
+        >>> from helmbert import HELMBertForSequenceClassification, HELMBertConfig
+        >>> config = HELMBertConfig(num_labels=1)  # Regression
+        >>> model = HELMBertForSequenceClassification(config)
+    """
+
+    def __init__(self, config: HELMBertConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+
+        self.helmbert = HELMBertModel(config)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        output_hidden_states: bool = False,
+        return_dict: bool = True,
+    ) -> Union[Tuple, SequenceClassifierOutput]:
+        """Forward pass.
+
+        Args:
+            input_ids: Token IDs [batch, seq]
+            attention_mask: Attention mask [batch, seq]
+            labels: Labels for classification/regression
+            output_attentions: Whether to return attention weights
+            output_hidden_states: Whether to return all hidden states
+            return_dict: Whether to return a ModelOutput
+
+        Returns:
+            SequenceClassifierOutput or tuple
+        """
+        outputs = self.helmbert(
+            input_ids,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+        )
+
+        pooled_output = outputs.pooler_output
+        pooled_output = self.dropout(pooled_output)
+        logits = self.classifier(pooled_output)
+
+        loss = None
+        if labels is not None:
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = nn.MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = nn.CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = nn.BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

special_tokens_map.json

@@ -0,0 +1,9 @@
+{
+  "bos_token": "@",
+  "cls_token": "@",
+  "eos_token": "\n",
+  "mask_token": "¶",
+  "pad_token": " ",
+  "sep_token": "\n",
+  "unk_token": "§"
+}

tokenization_helmbert.py

@@ -0,0 +1,285 @@
+"""HELM-BERT tokenizer."""
+
+import json
+import os
+from typing import Dict, List, Optional, Tuple, Union
+
+from transformers import PreTrainedTokenizer
+
+
+# Default vocabulary for HELM notation
+HELM_VOCAB = {
+    # Special tokens (0-4)
+    " ": 0,   # PAD
+    "@": 1,   # BOS/CLS
+    "\n": 2,  # EOS/SEP
+    "§": 3,   # UNK
+    "¶": 4,   # MASK
+
+    # Natural amino acids (5-25)
+    "A": 5, "R": 6, "N": 7, "D": 8, "C": 9,
+    "E": 10, "Q": 11, "G": 12, "H": 13, "I": 14,
+    "L": 15, "K": 16, "M": 17, "F": 18, "P": 19,
+    "S": 20, "T": 21, "W": 22, "Y": 23, "V": 24,
+    "X": 25,  # Unknown amino acid
+
+    # Structure symbols (26-37)
+    "[": 26, "]": 27, "{": 28, "}": 29, "(": 30, ")": 31,
+    "$": 32, ",": 33, ":": 34, "|": 35, "-": 36, ".": 37,
+
+    # Numbers (38-47)
+    "0": 38, "1": 39, "2": 40, "3": 41, "4": 42,
+    "5": 43, "6": 44, "7": 45, "8": 46, "9": 47,
+
+    # Uppercase non-amino acids (48-50)
+    "B": 48, "O": 49, ">": 50,
+
+    # Lowercase letters (51-72)
+    "a": 51, "b": 52, "c": 53, "d": 54, "e": 55,
+    "f": 56, "g": 57, "h": 58, "i": 59, "l": 60,
+    "m": 61, "n": 62, "o": 63, "p": 64, "r": 65,
+    "s": 66, "t": 67, "u": 68, "v": 69, "x": 70,
+    "y": 71, "z": 72,
+
+    # Encoded polymer markers (73-76)
+    "/": 73,   # PEPTIDE
+    "*": 74,   # me
+    "\t": 75,  # am
+    "&": 76,   # ac
+
+    # Miscellaneous (77)
+    "_": 77,
+}
+
+# Multi-character to single-character encoding
+HELM_ENCODE_MAP = {"PEPTIDE": "/", "me": "*", "am": "\t", "ac": "&"}
+HELM_DECODE_MAP = {v: k for k, v in HELM_ENCODE_MAP.items()}
+
+
+class HELMBertTokenizer(PreTrainedTokenizer):
+    """Tokenizer for HELM-BERT.
+
+    This tokenizer handles HELM (Hierarchical Editing Language for Macromolecules)
+    notation, converting peptide sequences into token IDs for the HELM-BERT model.
+
+    The tokenizer uses character-level tokenization with special handling for
+    multi-character HELM tokens like "PEPTIDE", "me", "am", "ac".
+
+    Example:
+        >>> from helmbert import HELMBertTokenizer
+        >>> tokenizer = HELMBertTokenizer()
+        >>> inputs = tokenizer("PEPTIDE1{A.C.D.E}$$$$", return_tensors="pt")
+        >>> inputs.input_ids
+        tensor([[ 1, 73, 39, 28, 5, 37, 9, 37, 8, 37, 10, 29, 32, 32, 32, 32,  2]])
+    """
+
+    vocab_files_names = {"vocab_file": "vocab.json"}
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file: Optional[str] = None,
+        unk_token: str = "§",
+        sep_token: str = "\n",
+        pad_token: str = " ",
+        cls_token: str = "@",
+        mask_token: str = "¶",
+        bos_token: str = "@",
+        eos_token: str = "\n",
+        model_max_length: int = 512,
+        **kwargs,
+    ):
+        # Load or create vocabulary
+        if vocab_file is not None and os.path.isfile(vocab_file):
+            with open(vocab_file, encoding="utf-8") as f:
+                self.vocab = json.load(f)
+        else:
+            self.vocab = HELM_VOCAB.copy()
+
+        self.ids_to_tokens = {v: k for k, v in self.vocab.items()}
+
+        # HELM encoding/decoding maps
+        self.encode_map = HELM_ENCODE_MAP.copy()
+        self.decode_map = HELM_DECODE_MAP.copy()
+
+        super().__init__(
+            unk_token=unk_token,
+            sep_token=sep_token,
+            pad_token=pad_token,
+            cls_token=cls_token,
+            mask_token=mask_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            model_max_length=model_max_length,
+            **kwargs,
+        )
+
+    @property
+    def vocab_size(self) -> int:
+        """Return the vocabulary size."""
+        return len(self.vocab)
+
+    def get_vocab(self) -> Dict[str, int]:
+        """Return the vocabulary as a dictionary."""
+        return self.vocab.copy()
+
+    def _encode_helm(self, text: str) -> str:
+        """Encode multi-character HELM tokens to single characters.
+
+        Args:
+            text: Raw HELM notation string
+
+        Returns:
+            Encoded string with single-character tokens
+        """
+        if not text:
+            return ""
+        result = text
+        for seq, tok in self.encode_map.items():
+            result = result.replace(seq, tok)
+        return result
+
+    def _decode_helm(self, text: str) -> str:
+        """Decode single-character tokens back to multi-character HELM tokens.
+
+        Args:
+            text: Encoded string with single-character tokens
+
+        Returns:
+            Decoded HELM notation string
+        """
+        if not text:
+            return ""
+        result = text
+        for tok, seq in self.decode_map.items():
+            result = result.replace(tok, seq)
+        return result
+
+    def _tokenize(self, text: str) -> List[str]:
+        """Tokenize a HELM string into a list of tokens.
+
+        Args:
+            text: HELM notation string
+
+        Returns:
+            List of single-character tokens
+        """
+        # First encode multi-character tokens to single characters
+        encoded = self._encode_helm(text)
+        # Return as list of characters
+        return list(encoded)
+
+    def _convert_token_to_id(self, token: str) -> int:
+        """Convert a token to its ID."""
+        return self.vocab.get(token, self.vocab.get(self.unk_token, 3))
+
+    def _convert_id_to_token(self, index: int) -> str:
+        """Convert an ID to its token."""
+        return self.ids_to_tokens.get(index, self.unk_token)
+
+    def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        """Convert a list of tokens to a HELM string.
+
+        Args:
+            tokens: List of tokens
+
+        Returns:
+            Decoded HELM notation string
+        """
+        # Join tokens and decode back to HELM notation
+        joined = "".join(tokens)
+        return self._decode_helm(joined)
+
+    def build_inputs_with_special_tokens(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """Build model inputs by adding special tokens.
+
+        Args:
+            token_ids_0: First sequence of token IDs
+            token_ids_1: Optional second sequence of token IDs
+
+        Returns:
+            List of token IDs with special tokens added
+        """
+        cls_id = [self.cls_token_id]
+        sep_id = [self.sep_token_id]
+
+        if token_ids_1 is None:
+            return cls_id + token_ids_0 + sep_id
+
+        return cls_id + token_ids_0 + sep_id + token_ids_1 + sep_id
+
+    def get_special_tokens_mask(
+        self,
+        token_ids_0: List[int],
+        token_ids_1: Optional[List[int]] = None,
+        already_has_special_tokens: bool = False,
+    ) -> List[int]:
+        """Get a mask identifying special tokens.
+
+        Args:
+            token_ids_0: First sequence of token IDs
+            token_ids_1: Optional second sequence of token IDs
+            already_has_special_tokens: Whether the sequences already have special tokens
+
+        Returns:
+            List of 0s and 1s (1 = special token)
+        """
+        if already_has_special_tokens:
+            return [1 if x in [self.cls_token_id, self.sep_token_id, self.pad_token_id] else 0 for x in token_ids_0]
+
+        if token_ids_1 is None:
+            return [1] + [0] * len(token_ids_0) + [1]
+
+        return [1] + [0] * len(token_ids_0) + [1] + [0] * len(token_ids_1) + [1]
+
+    def create_token_type_ids_from_sequences(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """Create token type IDs for sequence pairs.
+
+        Args:
+            token_ids_0: First sequence of token IDs
+            token_ids_1: Optional second sequence of token IDs
+
+        Returns:
+            List of token type IDs
+        """
+        sep = [self.sep_token_id]
+        cls = [self.cls_token_id]
+
+        if token_ids_1 is None:
+            return [0] * len(cls + token_ids_0 + sep)
+
+        return [0] * len(cls + token_ids_0 + sep) + [1] * len(token_ids_1 + sep)
+
+    def save_vocabulary(
+        self, save_directory: str, filename_prefix: Optional[str] = None
+    ) -> Tuple[str]:
+        """Save the vocabulary to a file.
+
+        Args:
+            save_directory: Directory to save the vocabulary
+            filename_prefix: Optional prefix for the filename
+
+        Returns:
+            Tuple containing the path to the saved vocabulary file
+        """
+        if not os.path.isdir(save_directory):
+            os.makedirs(save_directory, exist_ok=True)
+
+        vocab_file = os.path.join(
+            save_directory,
+            (filename_prefix + "-" if filename_prefix else "") + "vocab.json",
+        )
+
+        with open(vocab_file, "w", encoding="utf-8") as f:
+            json.dump(self.vocab, f, ensure_ascii=False, indent=2)
+
+        return (vocab_file,)
+
+    @property
+    def mask_token_id(self) -> int:
+        """Return the mask token ID."""
+        return self.vocab.get(self.mask_token, 4)

tokenizer_config.json

@@ -0,0 +1,58 @@
+{
+  "auto_map": {
+    "AutoTokenizer": ["tokenization_helmbert.HELMBertTokenizer", null]
+  },
+  "added_tokens_decoder": {
+    "0": {
+      "content": " ",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "@",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "\n",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "§",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "¶",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "@",
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "@",
+  "eos_token": "\n",
+  "extra_special_tokens": {},
+  "mask_token": "¶",
+  "model_max_length": 512,
+  "pad_token": " ",
+  "sep_token": "\n",
+  "tokenizer_class": "HELMBertTokenizer",
+  "unk_token": "§"
+}

vocab.json

@@ -0,0 +1,80 @@
+{
+  " ": 0,
+  "@": 1,
+  "\n": 2,
+  "§": 3,
+  "¶": 4,
+  "A": 5,
+  "R": 6,
+  "N": 7,
+  "D": 8,
+  "C": 9,
+  "E": 10,
+  "Q": 11,
+  "G": 12,
+  "H": 13,
+  "I": 14,
+  "L": 15,
+  "K": 16,
+  "M": 17,
+  "F": 18,
+  "P": 19,
+  "S": 20,
+  "T": 21,
+  "W": 22,
+  "Y": 23,
+  "V": 24,
+  "X": 25,
+  "[": 26,
+  "]": 27,
+  "{": 28,
+  "}": 29,
+  "(": 30,
+  ")": 31,
+  "$": 32,
+  ",": 33,
+  ":": 34,
+  "|": 35,
+  "-": 36,
+  ".": 37,
+  "0": 38,
+  "1": 39,
+  "2": 40,
+  "3": 41,
+  "4": 42,
+  "5": 43,
+  "6": 44,
+  "7": 45,
+  "8": 46,
+  "9": 47,
+  "B": 48,
+  "O": 49,
+  ">": 50,
+  "a": 51,
+  "b": 52,
+  "c": 53,
+  "d": 54,
+  "e": 55,
+  "f": 56,
+  "g": 57,
+  "h": 58,
+  "i": 59,
+  "l": 60,
+  "m": 61,
+  "n": 62,
+  "o": 63,
+  "p": 64,
+  "r": 65,
+  "s": 66,
+  "t": 67,
+  "u": 68,
+  "v": 69,
+  "x": 70,
+  "y": 71,
+  "z": 72,
+  "/": 73,
+  "*": 74,
+  "\t": 75,
+  "&": 76,
+  "_": 77
+}