Upload load_koochik_lite.py with huggingface_hub

load_koochik_lite.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+"""Load Shenava-Koochik-Lite-v1.0 — a LITEASR-compressed (-21.6% encoder) Koochik.
+
+Needs the base model Reza2kn/Shenava-Koochik-v1.0 (.nemo) plus this repo's
+`koochik_lite099_enc.pt` (compressed encoder state_dict) and `koochik_lite099_kmap.json`
+(per-layer low-rank sizes k). LITEASR replaces each targeted encoder Linear with a
+rank-k Sequential(Linear(D_in->k, bias=False), Linear(k->D_out, bias=True)); this loader
+rebuilds that skeleton from the kmap, then loads the compressed weights.
+
+    from huggingface_hub import hf_hub_download, snapshot_download
+    base = hf_hub_download("Reza2kn/Shenava-Koochik-v1.0", "shenava-koochik-v1.0.nemo")
+    repo = snapshot_download("Reza2kn/Shenava-Koochik-Lite-v1.0")
+    from load_koochik_lite import load_koochik_lite
+    m = load_koochik_lite(base, f"{repo}/koochik_lite099_enc.pt", f"{repo}/koochik_lite099_kmap.json")
+    print(m.transcribe(["clip.wav"])[0].text)
+"""
+import json, torch, torch.nn as nn
+import nemo.collections.asr as A
+
+
+def load_koochik_lite(base_nemo, enc_sd_path, kmap_path, device="cuda", ctc=True):
+    m = A.models.ASRModel.restore_from(base_nemo, map_location=device).eval()
+    if ctc:
+        try:
+            m.change_decoding_strategy(decoder_type="ctc")
+        except Exception:
+            pass
+
+    def parent(root, name):
+        parts = name.split("."); p = root
+        for x in parts[:-1]:
+            p = getattr(p, x)
+        return p, parts[-1]
+
+    for n, k in json.load(open(kmap_path)).items():
+        par, at = parent(m.encoder, n)
+        mod = getattr(par, at)
+        Dout, Din = mod.weight.shape
+        setattr(par, at, nn.Sequential(nn.Linear(Din, k, bias=False), nn.Linear(k, Dout, bias=True)))
+    m.encoder.load_state_dict(torch.load(enc_sd_path, map_location=device), strict=True)
+    return m.to(device).eval()
+
+
+if __name__ == "__main__":
+    import sys
+    m = load_koochik_lite(sys.argv[1], sys.argv[2], sys.argv[3], device="cuda" if torch.cuda.is_available() else "cpu")
+    for r in m.transcribe(sys.argv[4:], batch_size=8):
+        print(r.text if hasattr(r, "text") else r)