Export TinyIndianLM checkpoint in HF format

config.json

@@ -17,6 +17,7 @@
   "dropout": 0.0,
   "rope_theta": 10000.0,
   "pad_token_id": 0,
+  "unk_token_id": 1,
   "bos_token_id": 2,
   "eos_token_id": 3,
   "torch_dtype": "float32",

configuration_tinyindianlm.py

@@ -16,6 +16,7 @@ class TinyIndianLMConfig(PretrainedConfig):
         dropout=0.0,
         rope_theta=10000.0,
         pad_token_id=0,
+        unk_token_id=1,
         bos_token_id=2,
         eos_token_id=3,
         **kwargs,
@@ -31,6 +32,7 @@ class TinyIndianLMConfig(PretrainedConfig):
         self.rope_theta = rope_theta
         super().__init__(
             pad_token_id=pad_token_id,
+            unk_token_id=unk_token_id,
             bos_token_id=bos_token_id,
             eos_token_id=eos_token_id,
             **kwargs,

modeling_tinyindianlm.py

@@ -48,6 +48,17 @@ class TinyIndianLMForCausalLM(PreTrainedModel):
         **kwargs,
     ):
         del attention_mask, kwargs
+
+        # Safety guard for exported repos where tokenizer vocab and
+        # model embedding size diverge. Any out-of-range token id is
+        # mapped to unk to avoid embedding lookup crashes.
+        if input_ids is not None:
+            vocab_size = self.model.embedding.num_embeddings
+            if (input_ids >= vocab_size).any():
+                unk_id = getattr(self.config, "unk_token_id", 1)
+                input_ids = input_ids.clone()
+                input_ids[input_ids >= vocab_size] = unk_id
+
         logits, loss = self.model(
             input_ids,
             targets=labels,