Enable AutoModel loading

ogma_model.py

@@ -5,17 +5,12 @@ from __future__ import annotations
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from transformers import PreTrainedModel
 
 from .config import OgmaConfig, TaskToken, VariantType
 from .embeddings import TokenEmbedding
 from .pooling import create_pooling
-from .variants.conv import ConvVariant
-from .variants.deep_narrow import DeepNarrowVariant
-from .variants.linear_attention import LinearAttentionVariant
-from .variants.mlp_mixer import MLPMixerVariant
-from .variants.transformer import TransformerVariant
-from .variants.transformer_resa import TransformerReSAVariant
-from .variants.gla import GLAVariant
+from .transformer import TransformerVariant
 
 __all__ = ["OgmaModel"]
 
@@ -23,25 +18,13 @@ MAX_PARAMS = 10_000_000
 
 
 def _build_variant(config: OgmaConfig) -> nn.Module:
-    """Instantiate the appropriate architecture variant."""
-    if config.variant == VariantType.TRANSFORMER:
-        return TransformerVariant(config)
-    elif config.variant == VariantType.DEEP_NARROW:
-        return DeepNarrowVariant(config)
-    elif config.variant == VariantType.CONV:
-        return ConvVariant(config)
-    elif config.variant == VariantType.LINEAR_ATTENTION:
-        return LinearAttentionVariant(config)
-    elif config.variant == VariantType.MLP_MIXER:
-        return MLPMixerVariant(config)
-    elif config.variant == VariantType.TRANSFORMER_RESA:
-        return TransformerReSAVariant(config)
-    elif config.variant == VariantType.GLA:
-        return GLAVariant(config)
-    raise ValueError(f"Unknown variant: {config.variant}")
-
-
-class OgmaModel(nn.Module):
+    """Instantiate the released Ogma architecture variant."""
+    if config.variant != VariantType.TRANSFORMER:
+        raise ValueError(f"This HF release supports transformer checkpoints, got {config.variant}")
+    return TransformerVariant(config)
+
+
+class OgmaModel(PreTrainedModel):
     """Ogma embedding model.
 
     Wraps any architecture variant with shared embedding, pooling, and
@@ -49,8 +32,14 @@ class OgmaModel(nn.Module):
     Matryoshka-compatible at configured sub-dimensions.
     """
 
+    config_class = OgmaConfig
+    base_model_prefix = "ogma"
+    supports_gradient_checkpointing = False
+    _tied_weights_keys: list[str] = []
+    all_tied_weights_keys: dict[str, str] = {}
+
     def __init__(self, config: OgmaConfig) -> None:
-        super().__init__()
+        super().__init__(config)
         self.config = config
         self.embedding = TokenEmbedding(config)
         self.variant = _build_variant(config)
@@ -71,20 +60,37 @@ class OgmaModel(nn.Module):
 
     def forward(
         self,
-        token_ids: torch.Tensor,
-        attention_mask: torch.Tensor,
-        task_token_ids: torch.Tensor,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        task_token_ids: torch.Tensor | None = None,
+        token_ids: torch.Tensor | None = None,
+        **_: object,
     ) -> torch.Tensor:
         """Forward pass producing L2-normalized embeddings.
 
         Args:
-            token_ids: (B, S) token IDs.
+            input_ids: (B, S) token IDs, Hugging Face style.
             attention_mask: (B, S) attention mask (1=valid, 0=pad).
             task_token_ids: (B,) task token IDs (4=QRY, 5=DOC, 6=SYM).
+            token_ids: Backward-compatible alias for input_ids.
 
         Returns:
             (B, d_output) L2-normalized embeddings.
         """
+        if input_ids is None:
+            input_ids = token_ids
+        if input_ids is None:
+            raise ValueError("input_ids or token_ids must be provided")
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+        if task_token_ids is None:
+            task_token_ids = torch.full(
+                (input_ids.shape[0],),
+                self.config.sym_id,
+                device=input_ids.device,
+                dtype=torch.long,
+            )
+        token_ids = input_ids
         # Embed tokens with task token prepended -> (B, S+1, d_model)
         x = self.embedding(token_ids, task_token_ids)
 
@@ -130,7 +136,7 @@ class OgmaModel(nn.Module):
             device=token_ids.device,
             dtype=torch.long,
         )
-        return self.forward(token_ids, attention_mask, task_ids)
+        return self.forward(input_ids=token_ids, attention_mask=attention_mask, task_token_ids=task_ids)
 
     def param_count(self) -> int:
         """Count total trainable parameters."""
@@ -147,7 +153,8 @@ class OgmaModel(nn.Module):
     def from_config(cls, config: OgmaConfig) -> OgmaModel:
         """Factory method to build a model from config."""
         model = cls(config)
-        model.assert_param_budget()
+        if model.param_count() < MAX_PARAMS:
+            model.assert_param_budget()
         return model
 
     @classmethod