Upload model_ts.py with huggingface_hub

model_ts.py

@@ -0,0 +1,106 @@
+from transformers import PretrainedConfig, PreTrainedModel, AutoModel, AutoConfig
+import torch
+import os
+import json
+from huggingface_hub import snapshot_download
+
+class IndicASRConfig(PretrainedConfig):
+    model_type = "iasr"
+    
+    def __init__(self, ts_folder: str = "path", BLANK_ID: int = 256, RNNT_MAX_SYMBOLS: int = 10,
+                 PRED_RNN_LAYERS: int = 2, PRED_RNN_HIDDEN_DIM: int = 640, SOS: int = 256, **kwargs):
+        super().__init__(**kwargs)
+        self.ts_folder = ts_folder
+        self.BLANK_ID = BLANK_ID
+        self.RNNT_MAX_SYMBOLS = RNNT_MAX_SYMBOLS
+        self.PRED_RNN_LAYERS = PRED_RNN_LAYERS
+        self.PRED_RNN_HIDDEN_DIM = PRED_RNN_HIDDEN_DIM
+        self.SOS = SOS
+
+class IndicASRModel(PreTrainedModel):
+    config_class = IndicASRConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        
+        # Load model components
+        self.models = {}
+        component_names = ['preprocessor', 'encoder', 'ctc_decoder', 'rnnt_decoder', 'joint_enc', 'joint_pred', 'joint_pre_net']
+        component_names += [f'joint_post_net_{lang}' for lang in ['as', 'bn', 'brx', 'doi', 'gu', 'hi', 'kn', 'kok', 'ks', 'mai', 'ml', 'mni', 'mr', 'ne', 'or', 'pa', 'sa', 'sat', 'sd', 'ta', 'te', 'ur']]
+
+        for name in component_names:
+            component_path = os.path.join(config.ts_folder, 'assets', f'{name}.ts')
+            if os.path.exists(component_path):
+                self.models[name] = torch.jit.load(component_path, map_location=torch.device(self.config.device))
+            else:
+                self.models[name] = None
+                print(f'Warning: {component_path} not found')
+
+        # Load vocab and language masks
+        with open(os.path.join(config.ts_folder, 'assets', 'vocab.json')) as f:
+            self.vocab = json.load(f)
+        
+        with open(os.path.join(config.ts_folder, 'assets', 'language_masks.json')) as f:
+            self.language_masks = json.load(f)
+    
+    def forward(self, wav, lang, decoding='ctc'):
+        encoder_outputs, encoded_lengths = self.encode(wav)
+        if decoding == 'ctc':
+            return self._ctc_decode(encoder_outputs, encoded_lengths, lang)
+        if decoding == 'rnnt':
+            return self._rnnt_decode(encoder_outputs, encoded_lengths, lang)
+
+    def encode(self, wav):
+        audio_signal, length = self.models['preprocessor'](wav.to(torch.device(self.config.device)), torch.tensor([wav.shape[-1]]).to(torch.device(self.config.device)))
+        outputs, encoded_lengths = self.models['encoder'](audio_signal, length)
+        return outputs, encoded_lengths
+
+    def _ctc_decode(self, encoder_outputs, encoded_lengths, lang):
+        logprobs = self.models['ctc_decoder'](encoder_outputs)
+        logprobs = logprobs[:, :, self.language_masks[lang]].log_softmax(dim=-1)
+        indices = torch.argmax(logprobs[0], dim=-1)
+        collapsed_indices = torch.unique_consecutive(indices, dim=-1)
+        return ''.join([self.vocab[lang][x] for x in collapsed_indices if x != self.config.BLANK_ID]).replace('▁', ' ').strip()
+    
+    def _rnnt_decode(self, encoder_outputs, encoded_lengths, lang):    
+        joint_enc = self.models['joint_enc'](encoder_outputs.transpose(1, 2))
+        hyp = [self.config.SOS]
+        prev_dec_state = (torch.zeros(self.config.PRED_RNN_LAYERS,1,self.config.PRED_RNN_HIDDEN_DIM).to(torch.device(self.config.device)), 
+                          torch.zeros(self.config.PRED_RNN_LAYERS,1,self.config.PRED_RNN_HIDDEN_DIM).to(torch.device(self.config.device)))
+        
+        for t in range(joint_enc.size(1)):
+            f = joint_enc[:, t, :].unsqueeze(1)
+            not_blank = True
+            symbols_added = 0
+            
+            while not_blank and (self.config.RNNT_MAX_SYMBOLS is None or symbols_added < self.config.RNNT_MAX_SYMBOLS):
+                g, _, dec_state = self.models['rnnt_decoder'](torch.Tensor([[hyp[-1]]]).long().to(torch.device(self.config.device)), torch.tensor([1]).to(torch.device(self.config.device)), states=prev_dec_state)
+                g = self.models['joint_pred'](g.transpose(1, 2))
+                joint_out = f + g  
+                joint_out = self.models['joint_pre_net'](joint_out)
+                logits = self.models[f'joint_post_net_{lang}'](joint_out)
+                log_probs = logits.log_softmax(dim=-1)
+                pred_token = log_probs.argmax(dim=-1).item()
+                
+                if pred_token == self.config.BLANK_ID:
+                    not_blank = False
+                else:
+                    hyp.append(pred_token)
+                    prev_dec_state = dec_state
+                symbols_added += 1
+        
+        return ''.join([self.vocab[lang][x] for x in hyp if x != self.config.SOS]).replace('▁', ' ').strip()
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        loc = snapshot_download(repo_id=pretrained_model_name_or_path)
+        return cls(IndicASRConfig(ts_folder=loc, **kwargs))
+
+# Register model classes
+AutoConfig.register("iasr", IndicASRConfig)
+AutoModel.register(IndicASRConfig, IndicASRModel)
+
+if __name__ == '__main__':
+    model_path = "ai4bharat/indic-conformer-600m-multilingual"
+    model = AutoModel.from_pretrained(model_path, trust_remote_code=True)
+    print(f"Loaded model: {type(model)}")