#!/bin/bash
# register omniASR_CTC_1B_v2 (arch 1b_v2 = 1b with vocab 10288, tokenizer omniASR_tokenizer_written_v2) in venv_omni 0.1.0
P=/home/dev/venv_omni/lib/python3.12/site-packages/omnilingual_asr
grep -q '"1b_v2"' $P/models/wav2vec2_asr/config.py || python3 - <<'PY'
p="/home/dev/venv_omni/lib/python3.12/site-packages/omnilingual_asr/models/wav2vec2_asr/config.py"
s=open(p).read()
add='''
    @arch("1b_v2", advanced=True)
    def _1b_asr_v2(resolver: DependencyResolver) -> Wav2Vec2AsrConfig:
        config = get_config(resolver, Wav2Vec2AsrConfig, "base_10h")
        config.encoder_config = get_config(resolver, Wav2Vec2Config, "1b").encoder_config
        config.encoder_config.dropout_p = 0.0
        config.encoder_config.attn_dropout_p = 0.0
        config.encoder_config.ffn_inner_dropout_p = 0.1
        config.encoder_config.layer_drop_p = 0.1
        config.use_masking = False
        config.max_temporal_mask_prob = 0.0
        config.max_spatial_mask_prob = 0.0
        config.target_vocab_size = 10288
        return config
'''
i=s.index('    @arch("1b", advanced=True)')
s=s[:i]+add.lstrip("\n")+"\n"+s[i:]
open(p,"w").write(s); print("arch patched")
PY
cat > $P/cards/models/rc_models_v2_local.yaml <<'Y'
name: omniASR_tokenizer_written_v2
tokenizer_family: char_tokenizer
tokenizer: https://dl.fbaipublicfiles.com/mms/omniASR_tokenizer_written_v2.model

---

name: omniASR_CTC_1B_v2
model_family: wav2vec2_asr
model_arch: 1b_v2
checkpoint: https://dl.fbaipublicfiles.com/mms/omniASR-CTC-1B-v2.pt
tokenizer_ref: omniASR_tokenizer_written_v2
Y
echo card written