diff --git a/scripts/convert-indic-canary.py b/scripts/convert-indic-canary.py new file mode 100644 index 00000000..dac3bc1b --- /dev/null +++ b/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 == "" + } + self.bos_id = self.special_tokens["<|startoftranscript|>"] + self.eos_id = self.special_tokens["<|endoftext|>"] + self.pad_id = self.special_tokens[""] + + 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} + # "-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()) diff --git a/src/arch/canary/model.cpp b/src/arch/canary/model.cpp index cdaf3ac6..088c14f1 100644 --- a/src/arch/canary/model.cpp +++ b/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 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 build_prompt_canary2(const CanaryModel & cm, int tgt_lang_id, const char * /*task*/, bool pnc, - const std::vector & prefix_ids) { + const std::vector & 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 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 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); } diff --git a/src/arch/canary/weights.cpp b/src/arch/canary/weights.cpp index bb3d18d9..012b7b02 100644 --- a/src/arch/canary/weights.cpp +++ b/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; } diff --git a/src/arch/canary/weights.h b/src/arch/canary/weights.h index 979fe4ab..300fcbba 100644 --- a/src/arch/canary/weights.h +++ b/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;