Update configuration and tokenization

modeling_rotary_indictrans.py

@@ -43,7 +43,7 @@ try:
     )
 except ImportError:
     logger.warning(
-        "It is highly recommended to use `flash_attention_2` for better performance with RotaryIndicTrans." 
+        "It is highly recommended to use `flash_attention_2` for better performance with RotaryIndicTrans."
         "Falling back to the default `eager` implementation."
     )
 
@@ -96,25 +96,23 @@ def rotate_half(x):
 @autocast("cuda", enabled=False)
 def apply_rotary_emb(cos, sin, t):
     rot_dim = cos.shape[-1]
-    assert rot_dim <= t.shape[-1] and cos.shape == sin.shape
     t_left, t_right = t[..., :rot_dim], t[..., rot_dim:]
     t_transformed = (t_left * cos) + (rotate_half(t_left) * sin)
     return torch.cat((t_transformed, t_right), dim=-1).type(t.dtype)
 
 
 class RotaryEmbedding(torch.nn.Module):
-    def __init__(
-        self, dim, theta=10000, interpolate_factor=1.0, cache_max_seq_len=8192
-    ):
+    def __init__(self, dim, theta=10000, scaling_factor=1.0, max_seq_len=8192):
         super().__init__()
 
-        freqs_ = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
-        self.cache_max_seq_len = cache_max_seq_len
-        self.interpolate_factor = interpolate_factor
+        self.max_seq_len = max_seq_len
+        self.scaling_factor = scaling_factor
+
+        inv_freq_ = 1.0 / (theta ** (torch.arange(0, dim, 2, device=device).float() / dim))
 
-        self.freqs = torch.nn.Parameter(freqs_, requires_grad=False).to(device)
+        self.register_buffer("inv_freq", inv_freq_, persistent=False)
+        self.precompute_freqs(max_seq_len)
         self.apply_rotary_emb = staticmethod(apply_rotary_emb)
-        self.precompute_freqs(cache_max_seq_len)
 
     def precompute_freqs(self, max_seq_len):
         thetas = self.forward(max_seq_len, device=device)
@@ -124,9 +122,9 @@ class RotaryEmbedding(torch.nn.Module):
     def rotate_queries_or_keys(self, t, seq_dim=-2, offset=0):
         seq_len = t.shape[seq_dim]
 
-        if seq_len > self.cache_max_seq_len:
-            self.cache_max_seq_len = seq_len * 2
-            self.precompute_freqs(self.cache_max_seq_len)
+        if seq_len > self.max_seq_len:
+            self.max_seq_len = seq_len * 2
+            self.precompute_freqs(self.max_seq_len)
 
         cos, sin = (
             self.cached_cos[offset : (offset + seq_len)],
@@ -136,8 +134,8 @@ class RotaryEmbedding(torch.nn.Module):
 
     @autocast("cuda", enabled=False)
     def forward(self, seq_len, device):
-        seq = torch.arange(seq_len, device=device) / self.interpolate_factor
-        thetas = einsum("..., f -> ... f", seq, self.freqs)
+        seq = torch.arange(seq_len, device=device) / self.scaling_factor
+        thetas = einsum("..., f -> ... f", seq, self.inv_freq)
         thetas = repeat(thetas, "... n -> ... (n r)", r=2)
         return thetas
 
@@ -176,7 +174,7 @@ class RotaryIndicTransAttention(nn.Module):
             RotaryEmbedding(
                 dim=self.head_dim // 2,
                 theta=config.rope_args.get("theta", 10000),
-                interpolate_factor=config.rope_args.get("interpolate_factor", 1.0),
+                scaling_factor=config.rope_args.get("scaling_factor", 1.0),
             )
             if not is_cross_attention
             else None
@@ -1653,4 +1651,4 @@ class RotaryIndicTransForConditionalGeneration(
                     past_state.index_select(0, beam_idx) for past_state in layer_past
                 ),
             )
-        return reordered_past
+        return reordered_past

tokenization_indictrans.py

@@ -1,6 +1,5 @@
 import os
 import json
-from functools import lru_cache
 
 from transformers.utils import logging
 from typing import Dict, List, Optional, Union, Tuple
@@ -11,10 +10,8 @@ from transformers.tokenization_utils import PreTrainedTokenizer
 
 logger = logging.get_logger(__name__)
 
-SPIECE_UNDERLINE = "▁"
-
-# Convert SPECIAL_TAGS to a frozen set for faster lookups
-SPECIAL_TAGS = frozenset(
+# Convert LANGUAGE_TAGS to a frozen set for faster lookups
+LANGUAGE_TAGS = frozenset(
     {
         "asm_Beng",
         "awa_Deva",
@@ -137,9 +134,9 @@ class IndicTransTokenizer(PreTrainedTokenizer):
             **kwargs,
         )
 
-    def add_new_special_tags(self, new_tags: List[str]) -> None:
-        global SPECIAL_TAGS
-        SPECIAL_TAGS = frozenset(SPECIAL_TAGS | set(new_tags))
+    def add_new_language_tags(self, new_tags: List[str]) -> None:
+        global LANGUAGE_TAGS
+        LANGUAGE_TAGS = frozenset(LANGUAGE_TAGS | set(new_tags))
 
     def _switch_to_input_mode(self) -> None:
         self.spm = self.src_spm
@@ -197,10 +194,12 @@ class IndicTransTokenizer(PreTrainedTokenizer):
         return self.decoder.get(index, self.unk_token)
 
     def convert_tokens_to_string(self, tokens: List[str]) -> str:
-        return "".join(tokens).replace(SPIECE_UNDERLINE, " ").strip()
+        return "".join(tokens).replace("▁", " ").strip()
 
     def _src_tokenize(self, text: str) -> List[str]:
         src_lang, tgt_lang, text = text.split(" ", 2)
+        assert src_lang in LANGUAGE_TAGS, f"Invalid source language tag: {src_lang}"
+        assert tgt_lang in LANGUAGE_TAGS, f"Invalid target language tag: {tgt_lang}"
         return [src_lang, tgt_lang] + self.spm.EncodeAsPieces(text)
 
     def _tgt_tokenize(self, text: str) -> List[str]: