Update model files

config.json

@@ -0,0 +1,51 @@
+{"model_type": "hourglass_transformer",
+  "auto_map": {
+    "AutoConfig": "hourglass_transformer.HourglassTransformerConfig",
+    "AutoModel": "hourglass_transformer.HourglassTransformerForMaskedLM"
+  },
+  "activation_function": "gelu",
+  "architectures": [
+    "HourglassTransformerForMaskedLM"
+  ],
+  "attn_resampling": false,
+  "bias": false,
+  "depth": [
+    4,
+    [
+      4,
+      4,
+      4
+    ],
+    4
+  ],
+  "dim": 768,
+  "dim_head": 64,
+  "heads": 8,
+  "inference": false,
+  "metadata_dim": 3072,
+  "model_type": "hourglass_transformer",
+  "norm_out": false,
+  "predict_expression_mode": false,
+  "predict_seq": true,
+  "predict_taxonomy": false,
+  "predict_tracks": true,
+  "rotary_emb_dim": 32,
+  "seq_vocab_size": 11,
+  "shorten_factor": [
+    8,
+    8
+  ],
+  "sliding_window": [
+    512,
+    512,
+    -1
+  ],
+  "taxonomy_vocab_size": 2604,
+  "torch_dtype": "float32",
+  "track_activation_fn": null,
+  "track_output_dim": 4,
+  "transformers_version": "4.44.2",
+  "updown_sample_type": "linear",
+  "use_metadata": true,
+  "use_taxonomy": false
+}

hourglass_transformer.py

@@ -0,0 +1,201 @@
+"""
+HuggingFace model wrapper for HourglassTransformerLM.
+This allows the model to be saved and loaded in HuggingFace format.
+"""
+
+from typing import Optional, Union
+from dataclasses import dataclass
+import torch
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import MaskedLMOutput
+from rnalm.utils.hydra_utils import to_tuple_recursive
+from rnalm.models.networks.hourglass_transformer import HourglassTransformerLM
+
+
+@dataclass
+class HourglassTransformerOutput(MaskedLMOutput):
+    # Standard MaskedLMOutput fields (inherited)
+    # loss: Optional[torch.FloatTensor] = None
+    # logits: torch.FloatTensor = None
+    # hidden_states: Optional[tuple] = None
+    # attentions: Optional[tuple] = None
+
+    # Custom multi-task fields
+    seq_logits: Optional[torch.FloatTensor] = None
+    tax_logits: Optional[torch.FloatTensor] = None
+    track_yhat: Optional[torch.FloatTensor] = None
+    expression_mode: Optional[torch.FloatTensor] = None
+    last_hidden_state: Optional[torch.FloatTensor] = None
+    last_hidden_state_track: Optional[torch.FloatTensor] = None
+
+    def __post_init__(self):
+        """Sync standard and custom field names for compatibility."""
+        # Call parent __post_init__ if it exists
+        if hasattr(super(), "__post_init__"):
+            super().__post_init__()
+
+        # Map seq_logits to logits if logits is None
+        if self.logits is None and self.seq_logits is not None:
+            object.__setattr__(self, "logits", self.seq_logits)
+
+
+class HourglassTransformerConfig(PretrainedConfig):
+    model_type = "hourglass_transformer"
+
+    def __init__(
+        self,
+        seq_vocab_size: int = 11,
+        taxonomy_vocab_size: int = 2604,
+        dim: int = 128,
+        depth: tuple = (2, 2, 2),
+        shorten_factor: Union[int, tuple] = 4,
+        sliding_window: tuple = (512, 512),
+        attn_resampling: bool = False,
+        updown_sample_type: str = "linear",
+        heads: int = 8,
+        dim_head: int = 64,
+        norm_out: bool = False,
+        bias: bool = True,
+        activation_function: str = "gelu",
+        rotary_emb_dim: int = 32,
+        use_taxonomy: bool = False,
+        use_metadata: bool = False,
+        predict_taxonomy: bool = False,
+        predict_tracks: bool = False,
+        predict_seq: bool = True,
+        track_activation_fn: Optional[str] = None,
+        track_output_dim: int = 4,
+        predict_expression_mode: bool = False,
+        inference: bool = False,
+        metadata_dim: int = 3072,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.seq_vocab_size = seq_vocab_size
+        self.taxonomy_vocab_size = taxonomy_vocab_size
+        self.dim = dim
+        if isinstance(depth, tuple):
+            self.depth = depth
+        elif isinstance(depth, list):
+            self.depth = tuple(depth)
+        else:
+            self.depth = depth
+        if isinstance(sliding_window, tuple):
+            self.sliding_window = sliding_window
+        elif isinstance(sliding_window, list):
+            self.sliding_window = tuple(sliding_window)
+        else:
+            self.sliding_window = sliding_window
+        self.rotary_emb_dim = rotary_emb_dim
+        self.shorten_factor = shorten_factor
+        self.attn_resampling = attn_resampling
+        self.updown_sample_type = updown_sample_type
+        self.heads = heads
+        self.dim_head = dim_head
+        self.norm_out = norm_out
+        self.bias = bias
+        self.activation_function = activation_function
+        self.use_taxonomy = use_taxonomy
+        self.use_metadata = use_metadata
+        self.metadata_dim = metadata_dim
+        self.predict_taxonomy = predict_taxonomy
+        self.predict_tracks = predict_tracks
+        self.predict_seq = predict_seq
+        self.track_activation_fn = track_activation_fn
+        self.track_output_dim = track_output_dim
+        self.predict_expression_mode = predict_expression_mode
+        self.inference = inference
+
+
+class HourglassTransformerForMaskedLM(PreTrainedModel):
+
+    config_class = HourglassTransformerConfig
+
+    def __init__(self, config: HourglassTransformerConfig):
+        super().__init__(config)
+
+        # Convert config to dict for model initialization
+        model_kwargs = {
+            "seq_vocab_size": config.seq_vocab_size,
+            "taxonomy_vocab_size": config.taxonomy_vocab_size,
+            "dim": config.dim,
+            "depth": to_tuple_recursive(config.depth),
+            "sliding_window": config.sliding_window,
+            "rotary_emb_dim": config.rotary_emb_dim,
+            "shorten_factor": config.shorten_factor,
+            "attn_resampling": config.attn_resampling,
+            "updown_sample_type": config.updown_sample_type,
+            "heads": config.heads,
+            "dim_head": config.dim_head,
+            "norm_out": config.norm_out,
+            "bias": config.bias,
+            "activation_function": config.activation_function,
+            "use_taxonomy": config.use_taxonomy,
+            "use_metadata": config.use_metadata,
+            "metadata_dim": config.metadata_dim,
+            "predict_taxonomy": config.predict_taxonomy,
+            "predict_tracks": config.predict_tracks,
+            "predict_seq": config.predict_seq,
+            "track_activation_fn": config.track_activation_fn,
+            "track_output_dim": config.track_output_dim,
+            "predict_expression_mode": config.predict_expression_mode,
+            "inference": config.inference,
+        }
+
+        self.model = HourglassTransformerLM(**model_kwargs)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        masked_taxonomy: Optional[torch.Tensor] = None,
+        metadata: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        **kwargs,
+    ) -> HourglassTransformerOutput:
+        """
+        Forward pass of the model.
+
+        Args:
+            input_ids: Tokenized input sequences (batch_size, seq_len)
+            masked_taxonomy: Optional taxonomy tokens (batch_size, 8)
+            metadata: Optional metadata embeddings
+            attention_mask: Optional attention mask (batch_size, seq_len)
+            labels: Optional labels for computing loss (batch_size, seq_len)
+            output_attentions: Whether to return attentions (not supported)
+            output_hidden_states: Whether to return hidden states
+
+        Returns:
+            HourglassTransformerOutput containing all model outputs
+        """
+        # Get the base model output
+        outputs = self.model(
+            masked_seq=input_ids,
+            masked_taxonomy=masked_taxonomy,
+            metadata=metadata,
+            mask=attention_mask,
+        )
+
+        # Convert to HourglassTransformerOutput
+        # This extends MaskedLMOutput for HuggingFace compatibility
+        hf_output = HourglassTransformerOutput(
+            loss=None,  # Loss should be computed externally if labels provided
+            logits=outputs.seq_logits,  # Standard HuggingFace field
+            hidden_states=(
+                (outputs.last_hidden_state,)
+                if (output_hidden_states and outputs.last_hidden_state is not None)
+                else None
+            ),
+            attentions=None,  # Not currently supported
+            # Custom fields
+            seq_logits=outputs.seq_logits,
+            tax_logits=outputs.tax_logits,
+            track_yhat=outputs.track_yhat,
+            expression_mode=outputs.expression_mode,
+            last_hidden_state=outputs.last_hidden_state,
+            last_hidden_state_track=outputs.last_hidden_state_track,
+        )
+
+        return hf_output

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:423ac098fcb8ee69c96e99c310f19eb55d225804478abf3cd650e66471043301
+size 593330620