Upload 10 files

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+decoder_joint-asr.fp16.ts filter=lfs diff=lfs merge=lfs -text
+encoder-asr.fp16.ts filter=lfs diff=lfs merge=lfs -text

config.json

@@ -0,0 +1,28 @@
+{
+  "model_type": "fastconformer_tdt",
+  "architectures": [
+    "FastConformerTDTForCTC"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_fastconformer.FastConformerTDTConfig",
+    "AutoModel": "modeling_fastconformer.FastConformerTDTForCTC",
+    "AutoModelForCTC": "modeling_fastconformer.FastConformerTDTForCTC",
+    "AutoProcessor": "processing_fastconformer.FastConformerProcessor"
+  },
+  "vocab_size": 5000,
+  "blank_id": 5000,
+  "durations": [
+    0,
+    1,
+    2,
+    3,
+    4
+  ],
+  "num_durations": 5,
+  "pred_hidden": 640,
+  "pred_rnn_layers": 1,
+  "max_symbols": 10,
+  "enc_d_model": 1024,
+  "feat_in": 128,
+  "sample_rate": 16000
+}

configuration_fastconformer.py

@@ -0,0 +1,33 @@
+from transformers import PretrainedConfig
+
+
+class FastConformerTDTConfig(PretrainedConfig):
+    """Config for the FastConformer TDT (RNNT) ASR model wrapped for HF transformers."""
+
+    model_type = "fastconformer_tdt"
+
+    def __init__(
+        self,
+        vocab_size=3000,
+        blank_id=3000,
+        durations=(0, 1, 2, 3, 4),
+        num_durations=5,
+        pred_hidden=640,
+        pred_rnn_layers=1,
+        max_symbols=10,
+        enc_d_model=1024,
+        feat_in=128,
+        sample_rate=16000,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.blank_id = blank_id
+        self.durations = list(durations)
+        self.num_durations = num_durations
+        self.pred_hidden = pred_hidden
+        self.pred_rnn_layers = pred_rnn_layers
+        self.max_symbols = max_symbols
+        self.enc_d_model = enc_d_model
+        self.feat_in = feat_in
+        self.sample_rate = sample_rate
+        super().__init__(**kwargs)

decoder_joint-asr.fp16.ts

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6396aa30a8c252a7cf981ce3f4c63b1dc8783c83da7e9390acb04e4e1ec1a5e0
+size 21539480

encoder-asr.fp16.ts

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3b2c3789d91b06c0f19cd4b55156688fd291e325990361f7024b7f225444fb0d
+size 887312617

example_inference.py

@@ -0,0 +1,13 @@
+import sys, torch, torchaudio
+from transformers import AutoModel, AutoProcessor
+
+REPO = "."
+DEV = "cuda" if torch.cuda.is_available() else "cpu"
+model = AutoModel.from_pretrained(REPO, trust_remote_code=True).to(DEV).eval()
+proc = AutoProcessor.from_pretrained(REPO, trust_remote_code=True)
+
+for path in sys.argv[1:]:
+    wav, sr = torchaudio.load(path)
+    wav = wav.mean(0) if wav.shape[0] > 1 else wav.squeeze(0)
+    inputs = proc(wav, sampling_rate=sr, return_tensors="pt").to(DEV)
+    print(f"{path}\t{proc.batch_decode(model.generate(**inputs))[0]}")

model_meta.json

@@ -0,0 +1,18 @@
+{
+  "vocab_size": 5000,
+  "blank_id": 5000,
+  "durations": [
+    0,
+    1,
+    2,
+    3,
+    4
+  ],
+  "num_durations": 5,
+  "subsampling_factor": 8,
+  "pred_hidden": 640,
+  "pred_rnn_layers": 1,
+  "enc_d_model": 1024,
+  "max_symbols": 10,
+  "feat_in": 128
+}

modeling_fastconformer.py

@@ -0,0 +1,168 @@
+"""
+HF `transformers`-compatible inference wrapper for the FastConformer TDT (RNNT)
+ASR model. Wraps the exported TorchScript encoder + decoder_joint graphs.
+
+Runtime deps: torch, transformers, sentencepiece. **No nemo_toolkit.**
+"""
+import os
+import torch
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+from dataclasses import dataclass
+from typing import Optional, List
+
+from .configuration_fastconformer import FastConformerTDTConfig
+
+
+def _normalize_per_feature(x, seq_len, constant):
+    B, _, max_time = x.shape
+    steps = torch.arange(max_time, device=x.device).unsqueeze(0).expand(B, max_time)
+    valid = steps < seq_len.unsqueeze(1)
+    denom = valid.sum(dim=1)
+    mean = torch.where(valid.unsqueeze(1), x, torch.zeros_like(x)).sum(dim=2) / denom.unsqueeze(1)
+    var = torch.sum(torch.where(valid.unsqueeze(1), x - mean.unsqueeze(2), torch.zeros_like(x)) ** 2,
+                    dim=2) / (denom.unsqueeze(1) - 1.0)
+    std = torch.sqrt(var).masked_fill(torch.sqrt(var).isnan(), 0.0) + constant
+    return (x - mean.unsqueeze(2)) / std.unsqueeze(2)
+
+
+@dataclass
+class ASRGreedyOutput(ModelOutput):
+    sequences: Optional[torch.LongTensor] = None
+    token_lists: Optional[List[List[int]]] = None
+
+
+class FastConformerTDTForCTC(PreTrainedModel):
+    """Named *ForCTC for AutoModel discoverability, but decoding is TDT/RNNT greedy."""
+
+    config_class = FastConformerTDTConfig
+    main_input_name = "input_features"
+
+    def __init__(self, config: FastConformerTDTConfig):
+        super().__init__(config)
+        self._anchor = torch.nn.Parameter(torch.zeros(1), requires_grad=False)
+        self._artifacts_dir = getattr(config, "_name_or_path", ".") or "."
+        self.encoder = None
+        self.decoder_joint = None
+        self._loaded = False
+        self._io_dtype = torch.float32
+        self._req_dtype = None  # user-requested compute dtype (torch_dtype / .half())
+        self.post_init()
+
+    def _ensure_loaded(self, device=None):
+        if self._loaded:
+            return
+        d = self._artifacts_dir
+        dev = device or self._anchor.device
+        on_cpu = (torch.device(dev).type == "cpu")
+        # The shipped graphs are fp16 on disk (half the size) but exported *unfrozen*,
+        # so weights are real parameters. CPU has no fp16 conv kernel, so on CPU we
+        # upcast to fp32 (lossless widening). On GPU we keep fp16 only if the user asked
+        # for it (torch_dtype=float16 or .half()); otherwise we upcast for exact fp32
+        # numerics. A legacy frozen fp32 graph (no parameters) is used as-is.
+        want_fp16 = (self._req_dtype == torch.float16) or (self._anchor.dtype == torch.float16)
+
+        def load(fp16_name, fp32_name):
+            p16 = os.path.join(d, fp16_name)
+            path = p16 if os.path.exists(p16) else os.path.join(d, fp32_name)
+            mod = torch.jit.load(path, map_location=dev).eval()
+            params = list(mod.parameters())
+            pdtype = params[0].dtype if params else torch.float32
+            target = torch.float16 if (want_fp16 and not on_cpu and params) else torch.float32
+            if params and pdtype != target:
+                mod = mod.half() if target == torch.float16 else mod.float()
+            return mod, target
+
+        self.encoder, io = load("encoder-asr.fp16.ts", "encoder-asr.ts")
+        self.decoder_joint, _ = load("decoder_joint-asr.fp16.ts", "decoder_joint-asr.ts")
+        self._io_dtype = io
+        pp = torch.load(os.path.join(d, "preproc.pt"), map_location="cpu")
+        self._p = pp["params"]
+        self.register_buffer("_window", pp["window"].to(dev), persistent=False)
+        self.register_buffer("_fb", pp["fb"].to(dev), persistent=False)
+        self._loaded = True
+
+    @classmethod
+    def from_pretrained(cls, path, *args, **kwargs):
+        config = kwargs.pop("config", None) or FastConformerTDTConfig.from_pretrained(path)
+        config._name_or_path = path
+        model = cls(config)
+        model._artifacts_dir = path
+        dtype = kwargs.get("torch_dtype", None)
+        if isinstance(dtype, str):
+            dtype = None if dtype == "auto" else getattr(torch, dtype, None)
+        if dtype in (torch.float16, torch.float32):
+            model._req_dtype = dtype
+            if dtype == torch.float16:
+                model = model.half()
+        return model
+
+    @torch.no_grad()
+    def extract_features(self, wav, wav_len):
+        self._ensure_loaded()
+        p = self._p
+        dev = self._anchor.device
+        n_fft, hop, win = p["n_fft"], p["hop_length"], p["win_length"]
+        wav = wav.to(dev).float()
+        wav_len = wav_len.to(dev)
+        seq_len = torch.div(wav_len + n_fft - n_fft, hop, rounding_mode="floor")
+        tmask = torch.arange(wav.shape[1], device=dev).unsqueeze(0) < wav_len.unsqueeze(1)
+        x = torch.cat((wav[:, :1], wav[:, 1:] - p["preemph"] * wav[:, :-1]), dim=1).masked_fill(~tmask, 0.0)
+        spec = torch.stft(x, n_fft=n_fft, hop_length=hop, win_length=win, window=self._window.float(),
+                          center=True, pad_mode="constant", return_complex=True)
+        x = torch.view_as_real(spec)
+        x = torch.sqrt(x.pow(2).sum(-1)).pow(p["mag_power"])
+        x = torch.matmul(self._fb.float(), x)
+        x = torch.log(x + p["log_zero_guard_value"])
+        x = _normalize_per_feature(x, seq_len, p["CONSTANT"])
+        max_len = x.size(-1)
+        m = torch.arange(max_len, device=dev).repeat(x.size(0), 1) >= seq_len.unsqueeze(1)
+        return x.masked_fill(m.unsqueeze(1), p["pad_value"]), seq_len
+
+    @torch.no_grad()
+    def forward(self, input_features, feature_lengths=None, **kwargs):
+        self._ensure_loaded()
+        if feature_lengths is None:
+            feature_lengths = torch.full((input_features.size(0),), input_features.size(-1),
+                                         dtype=torch.long, device=input_features.device)
+        feats = input_features.to(self._anchor.device).to(self._io_dtype)
+        enc, enc_len = self.encoder(feats, feature_lengths.to(self._anchor.device))
+        return ModelOutput(last_hidden_state=enc, encoder_lengths=enc_len)
+
+    @torch.no_grad()
+    def _greedy_one(self, enc_out, T):
+        cfg, dev = self.config, self._anchor.device
+        nd, blank, durs = cfg.num_durations, cfg.blank_id, cfg.durations
+        h = torch.zeros(cfg.pred_rnn_layers, 1, cfg.pred_hidden, device=dev, dtype=self._io_dtype)
+        c = torch.zeros(cfg.pred_rnn_layers, 1, cfg.pred_hidden, device=dev, dtype=self._io_dtype)
+        last, toks = blank, []
+        tlen = torch.ones(1, dtype=torch.int32, device=dev)
+        t = 0
+        while t < T:
+            f = enc_out.narrow(2, t, 1)
+            added, need = 0, True
+            while need and added < cfg.max_symbols:
+                tgt = torch.tensor([[last]], dtype=torch.int32, device=dev)
+                logits, _, h2, c2 = self.decoder_joint(f, tgt, tlen, h, c)
+                logits = logits[0, 0, 0]
+                k = int(logits[:-nd].argmax().item())
+                skip = durs[int(logits[-nd:].argmax().item())]
+                if k != blank:
+                    toks.append(k); h, c, last = h2, c2, k
+                added += 1; t += skip; need = (skip == 0)
+            if added == cfg.max_symbols:
+                t += 1
+        return toks
+
+    @torch.no_grad()
+    def generate(self, input_features=None, feature_lengths=None, **kwargs):
+        out = self.forward(input_features, feature_lengths)
+        enc, enc_len = out.last_hidden_state, out.encoder_lengths
+        token_lists = [self._greedy_one(enc[i:i + 1], int(enc_len[i].item())) for i in range(enc.size(0))]
+        maxlen = max((len(t) for t in token_lists), default=0)
+        pad = self.config.blank_id
+        seqs = torch.full((len(token_lists), max(maxlen, 1)), pad, dtype=torch.long)
+        for i, t in enumerate(token_lists):
+            if t:
+                seqs[i, :len(t)] = torch.tensor(t, dtype=torch.long)
+        return ASRGreedyOutput(sequences=seqs, token_lists=token_lists)

preproc.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8119ff424215f42c11eaf4e745c12bbe3076d5a037593ab523c1f15ef32f5b2f
+size 135269

processing_fastconformer.py

@@ -0,0 +1,92 @@
+"""HF processor: raw audio -> mel input_features, and token ids -> text (SentencePiece)."""
+import os
+import torch
+import sentencepiece as spm
+from transformers.feature_extraction_utils import BatchFeature
+
+
+def _normalize_per_feature(x, seq_len, constant):
+    B, _, max_time = x.shape
+    steps = torch.arange(max_time, device=x.device).unsqueeze(0).expand(B, max_time)
+    valid = steps < seq_len.unsqueeze(1)
+    denom = valid.sum(dim=1)
+    mean = torch.where(valid.unsqueeze(1), x, torch.zeros_like(x)).sum(dim=2) / denom.unsqueeze(1)
+    var = torch.sum(torch.where(valid.unsqueeze(1), x - mean.unsqueeze(2), torch.zeros_like(x)) ** 2,
+                    dim=2) / (denom.unsqueeze(1) - 1.0)
+    std = torch.sqrt(var).masked_fill(torch.sqrt(var).isnan(), 0.0) + constant
+    return (x - mean.unsqueeze(2)) / std.unsqueeze(2)
+
+
+class FastConformerProcessor:
+    def __init__(self, sp, window, fb, params, blank_id=3000, sample_rate=16000):
+        self.sp = sp
+        self.window = window
+        self.fb = fb
+        self.p = params
+        self.blank_id = blank_id
+        self.sample_rate = sample_rate
+
+    @classmethod
+    def register_for_auto_class(cls, auto_class="AutoProcessor"):
+        return None
+
+    @classmethod
+    def from_pretrained(cls, path, **kwargs):
+        sp = spm.SentencePieceProcessor(model_file=os.path.join(path, "tokenizer.model"))
+        pp = torch.load(os.path.join(path, "preproc.pt"), map_location="cpu")
+        blank = pp.get("blank_id", 3000)
+        return cls(sp, pp["window"], pp["fb"], pp["params"], blank_id=blank)
+
+    def save_pretrained(self, path, **kwargs):
+        os.makedirs(path, exist_ok=True)
+
+    @staticmethod
+    def _to_2d(audio):
+        if isinstance(audio, torch.Tensor):
+            a = audio
+        else:
+            import numpy as np
+            a = torch.as_tensor(np.asarray(audio), dtype=torch.float32)
+        if a.dim() == 1:
+            a = a.unsqueeze(0)
+        return a.float()
+
+    @torch.no_grad()
+    def __call__(self, audio, sampling_rate=16000, return_tensors="pt"):
+        wav = self._to_2d(audio)
+        if sampling_rate != self.sample_rate:
+            import torchaudio
+            wav = torchaudio.functional.resample(wav, sampling_rate, self.sample_rate)
+        p = self.p
+        n_fft, hop, win = p["n_fft"], p["hop_length"], p["win_length"]
+        wav_len = torch.tensor([wav.shape[1]] * wav.shape[0])
+        seq_len = torch.div(wav_len + n_fft - n_fft, hop, rounding_mode="floor")
+        tmask = torch.arange(wav.shape[1]).unsqueeze(0) < wav_len.unsqueeze(1)
+        x = torch.cat((wav[:, :1], wav[:, 1:] - p["preemph"] * wav[:, :-1]), dim=1).masked_fill(~tmask, 0.0)
+        spec = torch.stft(x, n_fft=n_fft, hop_length=hop, win_length=win, window=self.window,
+                          center=True, pad_mode="constant", return_complex=True)
+        x = torch.view_as_real(spec)
+        x = torch.sqrt(x.pow(2).sum(-1)).pow(p["mag_power"])
+        x = torch.matmul(self.fb.to(x.dtype), x)
+        x = torch.log(x + p["log_zero_guard_value"])
+        x = _normalize_per_feature(x, seq_len, p["CONSTANT"])
+        max_len = x.size(-1)
+        m = torch.arange(max_len).repeat(x.size(0), 1) >= seq_len.unsqueeze(1)
+        x = x.masked_fill(m.unsqueeze(1), p["pad_value"])
+        return BatchFeature({"input_features": x, "feature_lengths": seq_len}, tensor_type=return_tensors)
+
+    def _clean(self, ids):
+        return [int(i) for i in ids if int(i) != self.blank_id]
+
+    def batch_decode(self, sequences, **kwargs):
+        token_lists = getattr(sequences, "token_lists", None)
+        if token_lists is None:
+            seqs = getattr(sequences, "sequences", sequences)
+            if isinstance(seqs, torch.Tensor):
+                token_lists = [self._clean(row.tolist()) for row in seqs]
+            else:
+                token_lists = [self._clean(row) for row in seqs]
+        return [self.sp.decode(t) for t in token_lists]
+
+    def decode(self, sequence, **kwargs):
+        return self.batch_decode([sequence], **kwargs)[0]

tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0af4086ec53482ac0cef0375369cf4ce7bafaf8b0a7203e97d126d0599ab90a6
+size 325287