canary: support Indic-Transcribe (Bodhan AI) output modes

scripts/convert-indic-canary.py

@@ -0,0 +1,197 @@
+#!/usr/bin/env python3
+"""
+convert-indic-canary.py - convert Bodhan AI's Indic-Transcribe (Flex / Core)
+NeMo checkpoints into transcribe.cpp canary GGUFs, without a NeMo install.
+
+    python scripts/convert-indic-canary.py \
+        path/to/indic_transcribe_flex.nemo out.gguf \
+        --variant indic-transcribe-flex --default-mode mixed
+
+The .nemo archive is a tar of model_config.yaml, model_weights.ckpt and the
+tokenizer's SentencePiece models. This script rebuilds the three things
+convert-canary.py reads from a NeMo model -- cfg, state_dict() and the
+aggregate CanaryTokenizer -- from those files directly, then hands off to
+convert-canary.py's convert(). The tokenizer is Bodhan's
+CanaryMultilingualTokenizer: an aggregate of `spl_tokens` (prompt tokens) and
+one shared `multilingual` sub-tokenizer, routed exactly like NeMo's
+CanaryTokenizer, so the aggregate extraction path applies unchanged.
+
+On top of stock canary2 it writes:
+    stt.canary.special.romanized_id / noromanized_id   10-token prompt slot
+    stt.canary.default_mode    native | mixed | romanized (mode used when
+                               the caller leaves itn at DEFAULT)
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib.util
+import io
+import re
+import sys
+import tarfile
+import tempfile
+from pathlib import Path
+
+import sentencepiece as spm
+import torch
+from omegaconf import OmegaConf
+
+HERE = Path(__file__).resolve().parent
+_spec = importlib.util.spec_from_file_location("convert_canary", HERE / "convert-canary.py")
+cc = importlib.util.module_from_spec(_spec)
+sys.modules["convert_canary"] = cc
+_spec.loader.exec_module(cc)
+
+LANGUAGES = [
+    "as", "bgc", "bhb", "bho", "bn", "brx", "doi", "en", "gu", "hi", "hne", "kn", "kok",
+    "ks", "mai", "ml", "mni", "mr", "ne", "or", "pa", "sa", "sat", "sd", "ta", "te", "ur",
+]
+LICENSE = {
+    "license": "other",
+    "license_name": "Indic Open Model License v1.0",
+    "license_link": "https://github.com/Bodhan-AI/bodhan-model-info/blob/main/licenses/"
+                    "indic-open-model-license/v1/Indic_Open_Model_License.md",
+}
+PROFILES = {
+    "indic-transcribe-flex": {"display_name": "Indic-Transcribe-Flex", "version": "flex",
+                              "repo": "bodhan-ai/indic-transcribe-flex"},
+    "indic-transcribe-core": {"display_name": "Indic-Transcribe-Core", "version": "core",
+                              "repo": "bodhan-ai/indic-transcribe-core"},
+}
+SPECIAL_RE = re.compile(r"^<\|.+\|>$")
+
+
+class _SubTokenizer:
+    def __init__(self, path: Path):
+        self.sp = spm.SentencePieceProcessor(model_file=str(path))
+        self.vocab_size = self.sp.GetPieceSize()
+        self.unk_id = self.sp.unk_id()
+
+
+class CanaryTokenizer:
+    """Just the surface convert-canary.py's _extract_aggregate() reads.
+    The class name matters: extract_tokenizer() dispatches on it."""
+
+    def __init__(self, subs: dict[str, _SubTokenizer]):
+        self.tokenizers_dict = subs
+        self.token_id_offset, off = {}, 0
+        for name, sub in subs.items():
+            self.token_id_offset[name] = off
+            off += sub.vocab_size
+        self.vocab_size = off
+        spl = subs["spl_tokens"]
+        self.special_tokens = {
+            p: i for i in range(spl.vocab_size)
+            if SPECIAL_RE.match(p := spl.sp.IdToPiece(i)) or p == "<pad>"
+        }
+        self.bos_id = self.special_tokens["<|startoftranscript|>"]
+        self.eos_id = self.special_tokens["<|endoftext|>"]
+        self.pad_id = self.special_tokens["<pad>"]
+
+    def ids_to_tokens(self, ids):
+        out = []
+        for i in ids:
+            for name, off in reversed(list(self.token_id_offset.items())):
+                if i >= off:
+                    out.append(self.tokenizers_dict[name].sp.IdToPiece(i - off))
+                    break
+        return out
+
+
+class NemoArchive:
+    """cfg / state_dict() / tokenizer from an extracted .nemo, NeMo-free."""
+
+    def __init__(self, nemo_path: Path, workdir: Path):
+        print(f"Extracting {nemo_path} ...")
+        with tarfile.open(nemo_path) as tar:
+            tar.extractall(workdir, filter="data")
+        files = {p.name: p for p in workdir.rglob("*") if p.is_file()}
+        self.cfg = OmegaConf.load(files["model_config.yaml"])
+        print("Loading model_weights.ckpt ...")
+        self._sd = torch.load(files["model_weights.ckpt"], map_location="cpu", weights_only=True)
+        langs = self.cfg.tokenizer.langs
+        subs = {}
+        for name in langs:  # config order == NeMo's id-offset order
+            ref = str(langs[name].get("model_path") or langs[name].dir)
+            ref = ref.removeprefix("nemo:")
+            cand = [p for n, p in files.items() if n.endswith(Path(ref).name) and n.endswith(".model")]
+            if not cand:
+                cand = [p for n, p in files.items() if n.endswith(".model") and name in n]
+            if len(cand) != 1:
+                raise RuntimeError(f"cannot locate SentencePiece model for {name!r} ({ref}); have {sorted(files)}")
+            subs[name] = _SubTokenizer(cand[0])
+        if list(subs)[0] != "spl_tokens":
+            raise RuntimeError(f"expected spl_tokens first, got {list(subs)}")
+        self.tokenizer = CanaryTokenizer(subs)
+
+    def state_dict(self):
+        return self._sd
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("nemo", type=Path)
+    ap.add_argument("out", type=Path)
+    ap.add_argument("--variant", default="indic-transcribe-flex", choices=sorted(PROFILES))
+    ap.add_argument("--default-mode", default="native", choices=["native", "mixed", "romanized"])
+    args = ap.parse_args()
+
+    prof = PROFILES[args.variant]
+    with tempfile.TemporaryDirectory(dir=args.out.parent) as tmp:
+        archive = NemoArchive(args.nemo, Path(tmp))
+        cc.load_canary_model = lambda _spec: archive
+
+        tok = archive.tokenizer
+        extra = {name: tok.special_tokens.get(name) for name in ("<|romanized|>", "<|noromanized|>")}
+        if None in extra.values():
+            raise RuntimeError("checkpoint has no <|romanized|>/<|noromanized|> tokens; not an Indic-Transcribe model")
+
+        real_writer = cc.gguf_writer
+
+        def writer_with_indic_kvs(path, arch):
+            w = real_writer(path, arch)
+            # Flex/Core are ASR-only and their timestamp slot is untested in
+            # this port: advertise neither, whatever convert() writes later.
+            real_add_bool = w.add_bool
+
+            def add_bool(key, val):
+                if key in ("stt.capability.translate", "stt.capability.timestamps"):
+                    val = False
+                return real_add_bool(key, val)
+
+            w.add_bool = add_bool
+            w.add_uint32("stt.canary.special.romanized_id", int(extra["<|romanized|>"]))
+            w.add_uint32("stt.canary.special.noromanized_id", int(extra["<|noromanized|>"]))
+            w.add_string("stt.canary.default_mode", args.default_mode)
+            w.add_string("general.author", "Bodhan AI / AI4Bharat")
+            w.add_string("general.organization", "bodhan-ai")
+            w.add_string("general.base_model.0.name", "Canary 1B v2")
+            w.add_string("general.base_model.0.repo_url", "https://huggingface.co/nvidia/canary-1b-v2")
+            w.add_string("general.description",
+                         f"Built with {prof['display_name']} from Bodhan AI / AI4Bharat. "
+                         "Licensed under the Indic Open Model License v1.0.")
+            return w
+
+        cc.gguf_writer = writer_with_indic_kvs
+        # add_general_identity hardcodes author/organization from the profile
+        # call site ("NVIDIA"); drop those two so ours above stand.
+        real_identity = cc.add_general_identity
+
+        def identity(writer, **kw):
+            kw.pop("author", None)
+            kw.pop("organization", None)
+            return real_identity(writer, **kw)
+
+        cc.add_general_identity = identity
+        profile = {"size_label": "1B", "prompt_format": "canary2", **LICENSE, **prof}
+        # "<code>-Latn" (BCP-47 script subtag) requests romanized output; listing
+        # them in general.languages lets transcribe.cpp's language gate accept
+        # them. They get no token-id KV, so the runtime maps them to the base code.
+        languages = LANGUAGES + [f"{code}-Latn" for code in LANGUAGES]
+        cc.convert(str(args.nemo), args.out, args.variant, profile, languages, repo_id=prof["repo"])
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

src/arch/canary/model.cpp

@@ -388,6 +388,12 @@ transcribe_status load(Loader & loader, const transcribe_model_load_params * par
                             m->hparams.prompt_format == "canary2" &&
                                 (m->hparams.tokenizer_single_sp || !m->hparams.tok_lang_codes.empty()));
 
+    // Indic-Transcribe GGUFs select mixed-script output through the itn slot,
+    // so transcribe_run_params::itn is a real control there.
+    if (m->hparams.noromanized_id >= 0 && m->hparams.itn_id >= 0) {
+        transcribe::set_feature(m.get(), TRANSCRIBE_FEATURE_ITN, true);
+    }
+
     // Publish the input-length ceiling now that the encoder positional span
     // and frontend rate are known (apply_family_invariants ran before the
     // hparams were read). The encoder is the binding INPUT limit for canary.
@@ -592,7 +598,48 @@ std::vector<int32_t> build_prompt_canary(const CanaryHParams & hp,
 // build_prompt_canary2): nine task slots, plus the empty decoder-context
 // marker on single-SP tokenizers.
 int canary2_base_prompt_tokens(const CanaryHParams & hp) {
-    return hp.tokenizer_single_sp ? 10 : 9;
+    return (hp.tokenizer_single_sp ? 10 : 9) + (hp.noromanized_id >= 0 ? 1 : 0);
+}
+
+// "te-Latn" asks an Indic-Transcribe model for romanized output; the
+// language token is still the base code's.
+bool is_latn_request(const char * lang) {
+    if (lang == nullptr) {
+        return false;
+    }
+    const size_t n = std::strlen(lang);
+    return n > 5 && std::strcmp(lang + n - 5, "-Latn") == 0;
+}
+
+// Output-mode slots for this run. Stock canary: always <|noitn|>, no
+// romanized slot. Indic-Transcribe: the GGUF's default_mode, overridden by
+// transcribe_run_params::itn (ON = mixed, OFF = not mixed) and by a
+// "-Latn" language suffix (romanized).
+struct CanaryOutputMode {
+    bool itn       = false;
+    bool romanized = false;
+};
+
+CanaryOutputMode canary_output_mode(const CanaryHParams & hp, const transcribe_run_params * params) {
+    CanaryOutputMode m;
+    if (hp.noromanized_id < 0) {
+        return m;
+    }
+    m.itn       = hp.default_mode == "mixed";
+    m.romanized = hp.default_mode == "romanized";
+    if (params != nullptr) {
+        if (params->itn == TRANSCRIBE_ITN_MODE_ON) {
+            m.itn       = true;
+            m.romanized = false;
+        } else if (params->itn == TRANSCRIBE_ITN_MODE_OFF) {
+            m.itn = false;
+        }
+        if (is_latn_request(params->language)) {
+            m.itn       = false;
+            m.romanized = true;
+        }
+    }
+    return m;
 }
 
 // The decoder prompt is prefilled into the self-KV cache, so it must fit the
@@ -628,7 +675,8 @@ std::vector<int32_t> build_prompt_canary2(const CanaryModel &   cm,
                                           int                   tgt_lang_id,
                                           const char * /*task*/,
                                           bool                         pnc,
-                                          const std::vector<int32_t> & prefix_ids) {
+                                          const std::vector<int32_t> & prefix_ids,
+                                          CanaryOutputMode             mode = {}) {
     // canary2 prompt template:
     //   <|startofcontext|> [decodercontext] <|startoftranscript|>
     //   <|emo:?|> <|src_lang|> <|tgt_lang|> <|pnc|> <|itn|> <|timestamp|> <|diarize|>
@@ -672,10 +720,19 @@ std::vector<int32_t> build_prompt_canary2(const CanaryModel &   cm,
     }
     ids.push_back(pnc_slot);
 
-    if (hp.noitn_id < 0) {
+    const int itn_slot = mode.itn ? hp.itn_id : hp.noitn_id;
+    if (itn_slot < 0) {
         return {};
     }
-    ids.push_back(hp.noitn_id);
+    ids.push_back(itn_slot);
+
+    if (hp.noromanized_id >= 0) {
+        const int rom_slot = mode.romanized ? hp.romanized_id : hp.noromanized_id;
+        if (rom_slot < 0) {
+            return {};
+        }
+        ids.push_back(rom_slot);
+    }
 
     if (hp.notimestamp_id < 0) {
         return {};
@@ -731,6 +788,10 @@ int find_language_id(const CanaryHParams & hp, const char * lang) {
     if (lang == nullptr) {
         return -1;
     }
+    if (hp.noromanized_id >= 0 && is_latn_request(lang)) {
+        const std::string base(lang, std::strlen(lang) - 5);
+        return find_language_id(hp, base.c_str());
+    }
     for (size_t i = 0; i < hp.languages.size() && i < hp.language_ids.size(); ++i) {
         if (hp.languages[i] == lang) {
             return hp.language_ids[i];
@@ -994,7 +1055,8 @@ transcribe_status run(transcribe_session *          session,
             st != TRANSCRIBE_OK) {
             return st;
         }
-        prompt_ids = build_prompt_canary2(*cm, cm->hparams, src_id, tgt_id, task, pnc, prefix_ids);
+        prompt_ids = build_prompt_canary2(*cm, cm->hparams, src_id, tgt_id, task, pnc, prefix_ids,
+                                          canary_output_mode(cm->hparams, params));
     } else if (cm->hparams.prompt_format == "canary") {
         prompt_ids = build_prompt_canary(cm->hparams, src_id, tgt_id, task, pnc);
     }
@@ -1636,7 +1698,7 @@ transcribe_status run_batch(transcribe_session *          session,
     }
     std::vector<int32_t> prompt_ids;
     if (hp.prompt_format == "canary2") {
-        prompt_ids = build_prompt_canary2(*cm, hp, src_id, tgt_id, task, pnc, {});
+        prompt_ids = build_prompt_canary2(*cm, hp, src_id, tgt_id, task, pnc, {}, canary_output_mode(hp, params));
     } else if (hp.prompt_format == "canary") {
         prompt_ids = build_prompt_canary(hp, src_id, tgt_id, task, pnc);
     }

src/arch/canary/weights.cpp

@@ -192,6 +192,21 @@ transcribe_status read_canary_hparams(const gguf_context * gguf, CanaryHParams &
     if (auto st = opt_special("stt.canary.special.noitn_id", hp.noitn_id); st != TRANSCRIBE_OK) {
         return st;
     }
+    if (auto st = opt_special("stt.canary.special.romanized_id", hp.romanized_id); st != TRANSCRIBE_OK) {
+        return st;
+    }
+    if (auto st = opt_special("stt.canary.special.noromanized_id", hp.noromanized_id); st != TRANSCRIBE_OK) {
+        return st;
+    }
+    if (auto st = read_optional_string_kv(gguf, "stt.canary.default_mode", kFamilyTag, "native", hp.default_mode);
+        st != TRANSCRIBE_OK) {
+        return st;
+    }
+    if (hp.default_mode != "native" && hp.default_mode != "mixed" && hp.default_mode != "romanized") {
+        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "%s: stt.canary.default_mode '%s' is not native, mixed or romanized",
+                kFamilyTag, hp.default_mode.c_str());
+        return TRANSCRIBE_ERR_GGUF;
+    }
     if (auto st = opt_special("stt.canary.special.timestamp_id", hp.timestamp_id); st != TRANSCRIBE_OK) {
         return st;
     }

src/arch/canary/weights.h

@@ -75,6 +75,14 @@ struct CanaryHParams {
     int32_t     nopnc_id             = -1;
     int32_t     itn_id               = -1;
     int32_t     noitn_id             = -1;
+    // Indic-Transcribe (Bodhan AI) canary2 variant: a tenth prompt slot
+    // after <|itn|> selects native vs romanized (Latin) output. Absent on
+    // stock canary GGUFs, which keep the nine-slot prompt.
+    int32_t     romanized_id         = -1;
+    int32_t     noromanized_id       = -1;
+    // Output mode used when the caller leaves itn at DEFAULT:
+    // "native" (noitn, noromanized), "mixed" (itn), "romanized".
+    std::string default_mode = "native";
     int32_t     timestamp_id         = -1;
     int32_t     notimestamp_id       = -1;
     int32_t     diarize_id           = -1;