Add OgmaTokenizerFast + model.embed() high-level API

- tokenization_ogma.py: PreTrainedTokenizerFast subclass that shifts
content token ids by N_SPECIAL=7 (training-matched tokenization,
no manual offset needed by callers)
- tokenizer_config.json: wires AutoTokenizer to OgmaTokenizerFast
- ogma_model.py: adds TaskToken class attr + embed(texts, task=) method
so callers need only AutoModel + AutoTokenizer, no sys.modules digging

ogma_model.py

@@ -138,6 +138,56 @@ class OgmaModel(PreTrainedModel):
         )
         return self.forward(input_ids=token_ids, attention_mask=attention_mask, task_token_ids=task_ids)
 
+    # TaskToken re-exported as a class attribute for clean external access
+    TaskToken = TaskToken  # noqa: F821 (imported at module top)
+
+    @torch.no_grad()
+    def embed(
+        self,
+        texts,
+        task: str = "sym",
+        tokenizer=None,
+        batch_size: int = 32,
+        max_length: int = 1024,
+    ) -> "torch.Tensor":
+        """High-level text → L2-normalized embeddings.
+
+        Args:
+            texts: str or list[str] to encode.
+            task: "qry" / "doc" / "sym" (or a TaskToken enum member).
+            tokenizer: OgmaTokenizerFast instance. If None, loaded
+                automatically from self.name_or_path (requires the model to
+                have been loaded via AutoModel.from_pretrained).
+            batch_size: Texts per forward pass.
+            max_length: Token cap per text (default 1024).
+
+        Returns:
+            (len(texts), d_output) tensor of L2-normalized embeddings on the
+            same device as the model.
+        """
+        if isinstance(texts, str):
+            texts = [texts]
+        if isinstance(task, str):
+            task = self.TaskToken[task.upper()]
+        if tokenizer is None:
+            from transformers import AutoTokenizer
+            tokenizer = AutoTokenizer.from_pretrained(
+                self.name_or_path, trust_remote_code=True
+            )
+        device = next(self.parameters()).device
+        outs = []
+        for i in range(0, len(texts), batch_size):
+            enc = tokenizer(
+                texts[i : i + batch_size],
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=max_length,
+            ).to(device)
+            outs.append(self.encode(enc["input_ids"], enc["attention_mask"], task=task))
+        return torch.cat(outs, dim=0)
+
+
     def param_count(self) -> int:
         """Count total trainable parameters."""
         return sum(p.numel() for p in self.parameters() if p.requires_grad)

tokenization_ogma.py

@@ -0,0 +1,43 @@
+"""OgmaTokenizerFast — wraps PreTrainedTokenizerFast, shifts token ids by
+N_SPECIAL so they align with Ogma's embedding table.
+
+Ogma reserved vocab ids (0-6):
+  0 <pad>  1 <unk>  2 [CLS]  3 [SEP]  4 [MASK]  5 [DOC]  6 [SYM]
+Regular SentencePiece tokens start at 7.
+
+The tokenizer post-processor already adds [CLS] / [SEP] around every input.
+This wrapper shifts ALL content positions (attention_mask == 1) up by
+N_SPECIAL so that [CLS]->9, [SEP]->10, and content tokens land where the
+model was trained to see them.  Padding positions (attention_mask == 0) stay
+at 0 (Ogma pad id).
+"""
+from __future__ import annotations
+
+import torch
+from transformers import PreTrainedTokenizerFast
+from transformers.tokenization_utils_base import BatchEncoding
+
+__all__ = ["OgmaTokenizerFast"]
+
+N_SPECIAL = 7
+
+
+class OgmaTokenizerFast(PreTrainedTokenizerFast):
+    N_SPECIAL = N_SPECIAL
+
+    def _shift(self, ids, mask):
+        if isinstance(ids, torch.Tensor):
+            return ids + self.N_SPECIAL * mask.long()
+        return [
+            [i + self.N_SPECIAL if m else i for i, m in zip(row_i, row_m)]
+            for row_i, row_m in zip(ids, mask)
+        ]
+
+    def __call__(self, *args, **kwargs) -> BatchEncoding:
+        kwargs.setdefault("padding", True)
+        kwargs.setdefault("truncation", True)
+        kwargs.setdefault("max_length", self.model_max_length or 1024)
+        enc = super().__call__(*args, **kwargs)
+        if "input_ids" in enc and "attention_mask" in enc:
+            enc["input_ids"] = self._shift(enc["input_ids"], enc["attention_mask"])
+        return enc

tokenizer_config.json

@@ -1,17 +1,20 @@
 {
-  "add_prefix_space": true,
-  "backend": "tokenizers",
-  "bos_token": "[CLS]",
+  "tokenizer_class": "OgmaTokenizerFast",
+  "auto_map": {
+    "AutoTokenizer": [
+      null,
+      "tokenization_ogma.OgmaTokenizerFast"
+    ]
+  },
+  "model_max_length": 1024,
+  "padding_side": "right",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>",
   "cls_token": "[CLS]",
-  "do_lower_case": true,
+  "sep_token": "[SEP]",
+  "bos_token": "[CLS]",
   "eos_token": "[SEP]",
-  "is_local": false,
-  "keep_accents": false,
   "mask_token": "[MASK]",
-  "model_max_length": 512,
-  "pad_token": "<pad>",
-  "sep_token": "[SEP]",
-  "tokenizer_class": "AlbertTokenizer",
-  "trim_offsets": true,
-  "unk_token": "<unk>"
+  "do_lower_case": true,
+  "backend": "tokenizers"
 }