   100→            if not any(AUDIO_PLACEHOLDER in message["content"] for message in messages):
   101→                messages[0]["content"] = AUDIO_PLACEHOLDER * len(audios) + messages[0]["content"]
   102→
   103→        messages = template.mm_plugin.process_messages(
   104→            messages, mm_input_dict["images"], mm_input_dict["videos"], mm_input_dict["audios"], processor
   105→        )
   106→        paired_messages = messages + [{"role": "assistant", "content": ""}]
   107→        prompt_ids, _ = template.encode_oneturn(tokenizer, paired_messages, system, tools)
   108→        prompt_ids, _ = template.mm_plugin.process_token_ids(
   109→            prompt_ids,
   110→            None,
   111→            mm_input_dict["images"],
   112→            mm_input_dict["videos"],
   113→            mm_input_dict["audios"],
   114→            tokenizer,
   115→            processor,
   116→        )
   117→        prompt_length = len(prompt_ids)
   118→        inputs = torch.tensor([prompt_ids], device=model.device)
   119→        attention_mask = torch.ones_like(inputs, dtype=torch.long)
   120→
   121→        do_sample: Optional[bool] = input_kwargs.pop("do_sample", None)
   122→        temperature: Optional[float] = input_kwargs.pop("temperature", None)
   123→        top_p: Optional[float] = input_kwargs.pop("top_p", None)
   124→        top_k: Optional[float] = input_kwargs.pop("top_k", None)
   125→        num_return_sequences: int = input_kwargs.pop("num_return_sequences", 1)
   126→        repetition_penalty: Optional[float] = input_kwargs.pop("repetition_penalty", None)
   127→        length_penalty: Optional[float] = input_kwargs.pop("length_penalty", None)
   128→        skip_special_tokens: Optional[bool] = input_kwargs.pop("skip_special_tokens", None)
   129→        max_length: Optional[int] = input_kwargs.pop("max_length", None)
   130→        max_new_tokens: Optional[int] = input_kwargs.pop("max_new_tokens", None)
   131→        stop: Optional[Union[str, list[str]]] = input_kwargs.pop("stop", None)
   132→
   133→        if stop is not None:
   134→            logger.warning_rank0("Stop parameter is not supported by the huggingface engine yet.")
   135→
   136→        generating_args = generating_args.copy()
   137→        generating_args.update(
   138→            dict(
   139→                do_sample=do_sample if do_sample is not None else generating_args["do_sample"],
   140→                temperature=temperature if temperature is not None else generating_args["temperature"],
   141→                top_p=top_p if top_p is not None else generating_args["top_p"],
   142→                top_k=top_k if top_k is not None else generating_args["top_k"],
   143→                num_return_sequences=num_return_sequences,
   144→                repetition_penalty=repetition_penalty
   145→                if repetition_penalty is not None
   146→                else generating_args["repetition_penalty"],
   147→                length_penalty=length_penalty if length_penalty is not None else generating_args["length_penalty"],
   148→                skip_special_tokens=skip_special_tokens
   149→                if skip_special_tokens is not None
   150→                else generating_args["skip_special_tokens"],
   151→                eos_token_id=template.get_stop_token_ids(tokenizer),
   152→                pad_token_id=tokenizer.pad_token_id,
   153→            )
   154→        )
   155→
   156→        if isinstance(num_return_sequences, int) and num_return_sequences > 1:  # do_sample needs temperature > 0
   157→            generating_args["do_sample"] = True
   158→            generating_args["temperature"] = generating_args["temperature"] or 1.0
   159→
   160→        if not generating_args["temperature"]:
   161→            generating_args["do_sample"] = False
   162→
   163→        if not generating_args["do_sample"]:
   164→            generating_args.pop("temperature", None)
   165→            generating_args.pop("top_p", None)
   166→
   167→        if max_length:
   168→            generating_args.pop("max_new_tokens", None)
   169→            generating_args["max_length"] = max_length
   170→
   171→        if max_new_tokens:
   172→            generating_args.pop("max_length", None)
   173→            generating_args["max_new_tokens"] = max_new_tokens
   174→
   175→        gen_kwargs = dict(
   176→            inputs=inputs,
   177→            attention_mask=attention_mask,
   178→            generation_config=GenerationConfig(**generating_args),
   179→        )
   180→
   181→        mm_inputs = template.mm_plugin.get_mm_inputs(**mm_input_dict, batch_ids=[prompt_ids], processor=processor)
   182→        for key, value in mm_inputs.items():
   183→            if isinstance(value, list) and isinstance(value[0], torch.Tensor):  # for pixtral inputs
   184→                value = torch.stack(value)  # assume they have same sizes
   185→            elif (
   186→                isinstance(value, list) and isinstance(value[0], list) and isinstance(value[0][0], torch.Tensor)
   187→            ):  # for minicpmv inputs
   188→                value = torch.stack([torch.stack(v) for v in value])
   189→            elif not isinstance(value, torch.Tensor):
   190→                value = torch.tensor(value)
   191→
   192→            if torch.is_floating_point(value):  # cast data dtype for paligemma
   193→                value = value.to(model.dtype)
   194→
   195→            if key == "second_per_grid_ts":  # qwen2.5vl special case
   196→                gen_kwargs[key] = value.tolist()
   197→            else:
   198→                gen_kwargs[key] = value.to(model.device)
   199→

            if not any(AUDIO_PLACEHOLDER in message["content"] for message in messages):
                messages[0]["content"] = AUDIO_PLACEHOLDER * len(audios) + messages[0]["content"]

        messages = template.mm_plugin.process_messages(
            messages, mm_input_dict["images"], mm_input_dict["videos"], mm_input_dict["audios"], processor
        )
        paired_messages = messages + [{"role": "assistant", "content": ""}]
        prompt_ids, _ = template.encode_oneturn(tokenizer, paired_messages, system, tools)
        prompt_ids, _ = template.mm_plugin.process_token_ids(
            prompt_ids,
            None,
            mm_input_dict["images"],
            mm_input_dict["videos"],
            mm_input_dict["audios"],
            tokenizer,
            processor,
        )
        prompt_length = len(prompt_ids)
        inputs = torch.tensor([prompt_ids], device=model.device)
        attention_mask = torch.ones_like(inputs, dtype=torch.long)

        do_sample: Optional[bool] = input_kwargs.pop("do_sample", None)
        temperature: Optional[float] = input_kwargs.pop("temperature", None)
        top_p: Optional[float] = input_kwargs.pop("top_p", None)
        top_k: Optional[float] = input_kwargs.pop("top_k", None)
        num_return_sequences: int = input_kwargs.pop("num_return_sequences", 1)
        repetition_penalty: Optional[float] = input_kwargs.pop("repetition_penalty", None)
        length_penalty: Optional[float] = input_kwargs.pop("length_penalty", None)
        skip_special_tokens: Optional[bool] = input_kwargs.pop("skip_special_tokens", None)
        max_length: Optional[int] = input_kwargs.pop("max_length", None)
        max_new_tokens: Optional[int] = input_kwargs.pop("max_new_tokens", None)
        stop: Optional[Union[str, list[str]]] = input_kwargs.pop("stop", None)

        if stop is not None:
            logger.warning_rank0("Stop parameter is not supported by the huggingface engine yet.")

        generating_args = generating_args.copy()
        generating_args.update(
            dict(
                do_sample=do_sample if do_sample is not None else generating_args["do_sample"],
                temperature=temperature if temperature is not None else generating_args["temperature"],
                top_p=top_p if top_p is not None else generating_args["top_p"],
                top_k=top_k if top_k is not None else generating_args["top_k"],
                num_return_sequences=num_return_sequences,
                repetition_penalty=repetition_penalty
                if repetition_penalty is not None
                else generating_args["repetition_penalty"],
                length_penalty=length_penalty if length_penalty is not None else generating_args["length_penalty"],
                skip_special_tokens=skip_special_tokens
                if skip_special_tokens is not None
                else generating_args["skip_special_tokens"],
                eos_token_id=template.get_stop_token_ids(tokenizer),
                pad_token_id=tokenizer.pad_token_id,
            )
        )

        if isinstance(num_return_sequences, int) and num_return_sequences > 1:  # do_sample needs temperature > 0
            generating_args["do_sample"] = True
            generating_args["temperature"] = generating_args["temperature"] or 1.0

        if not generating_args["temperature"]:
            generating_args["do_sample"] = False

        if not generating_args["do_sample"]:
            generating_args.pop("temperature", None)
            generating_args.pop("top_p", None)

        if max_length:
            generating_args.pop("max_new_tokens", None)
            generating_args["max_length"] = max_length

        if max_new_tokens:
            generating_args.pop("max_length", None)
            generating_args["max_new_tokens"] = max_new_tokens

        gen_kwargs = dict(
            inputs=inputs,
            attention_mask=attention_mask,
            generation_config=GenerationConfig(**generating_args),
        )

        mm_inputs = template.mm_plugin.get_mm_inputs(**mm_input_dict, batch_ids=[prompt_ids], processor=processor)
        for key, value in mm_inputs.items():
            if isinstance(value, list) and isinstance(value[0], torch.Tensor):  # for pixtral inputs
                value = torch.stack(value)  # assume they have same sizes
            elif (
                isinstance(value, list) and isinstance(value[0], list) and isinstance(value[0][0], torch.Tensor)
            ):  # for minicpmv inputs
                value = torch.stack([torch.stack(v) for v in value])
            elif not isinstance(value, torch.Tensor):
                value = torch.tensor(value)

            if torch.is_floating_point(value):  # cast data dtype for paligemma
                value = value.to(model.dtype)

            if key == "second_per_grid_ts":  # qwen2.5vl special case
                gen_kwargs[key] = value.tolist()
            else:
                gen_kwargs[key] = value.to(model.device)


22-
23-
24:AUDIO_PLACEHOLDER = os.getenv("AUDIO_PLACEHOLDER", "<audio>")
25-
26-CHECKPOINT_NAMES = {


975-    ),
976-    format_prefix=EmptyFormatter(slots=[{"bos_token"}]),
977-    stop_words=["<end_of_turn>"],
978-    replace_eos=True,
979-    mm_plugin=get_mm_plugin("gemma3", image_token="<image_soft_token>"),
980-    template_class=Llama2Template,
981-)
982-
983-
984-register_template(
985:    name="gemma3n",
986-    format_user=StringFormatter(slots=["<start_of_turn>user\n{{content}}<end_of_turn>\n<start_of_turn>model\n"]),
987-    format_assistant=StringFormatter(slots=["{{content}}<end_of_turn>\n"]),
988-    format_system=StringFormatter(slots=["{{content}}\n\n"]),
989-    format_observation=StringFormatter(
990-        slots=["<start_of_turn>tool\n{{content}}<end_of_turn>\n<start_of_turn>model\n"]
991-    ),
992-    format_prefix=EmptyFormatter(slots=[{"bos_token"}]),
993-    stop_words=["<end_of_turn>"],
994-    replace_eos=True,
995:    mm_plugin=get_mm_plugin("gemma3n", image_token="<image_soft_token>", audio_token="<audio_soft_token>"),
996-    template_class=Llama2Template,
997-)
998-
999-
1000-register_template(
1001-    name="glm4",
1002-    format_user=StringFormatter(slots=["<|user|>\n{{content}}<|assistant|>"]),
1003-    format_assistant=StringFormatter(slots=["\n{{content}}"]),
1004-    format_system=StringFormatter(slots=["<|system|>\n{{content}}"]),
1005-    format_function=FunctionFormatter(slots=["{{content}}"], tool_format="glm4"),


Now let me understand the gemma3n mm_plugin for audio processing.


721-        return mm_inputs
722-
723-
724:class Gemma3nPlugin(Gemma3Plugin):
725-    @override
726-    def process_messages(
727-        self,
--
2364-    "ernie_vl": ErnieVLPlugin,
2365-    "gemma3": Gemma3Plugin,
2366-    "glm4v": GLM4VPlugin,
2367:    "gemma3n": Gemma3nPlugin,
2368-    "intern_vl": InternVLPlugin,
2369-    "kimi_vl": KimiVLPlugin,
2370-    "llama4": Llama4Plugin,


   724→class Gemma3nPlugin(Gemma3Plugin):
   725→    @override
   726→    def process_messages(
   727→        self,
   728→        messages: list[dict[str, str]],
   729→        images: list["ImageInput"],
   730→        videos: list["VideoInput"],
   731→        audios: list["AudioInput"],
   732→        processor: Optional["MMProcessor"],
   733→    ) -> list[dict[str, str]]:
   734→        self._validate_input(processor, images, videos, audios)
   735→        self._validate_messages(messages, images, videos, audios)
   736→        messages = deepcopy(messages)
   737→        boi_token: str = getattr(processor, "boi_token")
   738→        boa_token: str = getattr(processor, "boa_token")
   739→        full_image_sequence: str = getattr(processor, "full_image_sequence")
   740→        full_audio_sequence: str = getattr(processor, "full_audio_sequence")
   741→        image_str = full_image_sequence if self.expand_mm_tokens else boi_token
   742→        audio_str = full_audio_sequence if self.expand_mm_tokens else boa_token
   743→
   744→        for message in messages:
   745→            content = message["content"]
   746→            while IMAGE_PLACEHOLDER in content:
   747→                content = content.replace(IMAGE_PLACEHOLDER, image_str, 1)
   748→
   749→            while AUDIO_PLACEHOLDER in content:
   750→                content = content.replace(AUDIO_PLACEHOLDER, audio_str, 1)
   751→
   752→            message["content"] = content
   753→
   754→        return messages
   755→
   756→
   757→@dataclass
   758→class InternVLPlugin(BasePlugin):
   759→    @override
   760→    def _get_mm_inputs(
   761→        self,
   762→        images: list["ImageInput"],
   763→        videos: list["VideoInput"],
   764→        audios: list["AudioInput"],
   765→        processor: "ProcessorMixin",
   766→        **kwargs,
   767→    ) -> dict[str, "torch.Tensor"]:
   768→        image_processor: BaseImageProcessor = getattr(processor, "image_processor")
   769→        image_processor_kwargs = {}
   770→        if getattr(processor, "crop_to_patches", False):
   771→            image_processor_kwargs.update(
   772→                {
   773→                    "crop_to_patches": True,
   774→                    "max_patches": 12,
   775→                    "min_patches": 1,
   776→                }
   777→            )
   778→
   779→        mm_inputs = {}
   780→        image_video_patches = []
   781→
   782→        if len(images) != 0:
   783→            images = self._regularize_images(
   784→                images,
   785→                image_max_pixels=getattr(processor, "image_max_pixels", 1024 * 1024),
   786→                image_min_pixels=getattr(processor, "image_min_pixels", 32 * 32),
   787→            )["images"]
   788→
   789→        if len(videos) != 0:
   790→            videos = self._regularize_videos(
   791→                videos,
   792→                image_max_pixels=getattr(processor, "video_max_pixels", 256 * 256),
   793→                image_min_pixels=getattr(processor, "video_min_pixels", 16 * 16),
   794→                video_fps=getattr(processor, "video_fps", 2.0),
   795→                video_maxlen=getattr(processor, "video_maxlen", 128),
   796→            )["videos"]
   797→
   798→        if len(images) != 0:
   799→            images = make_flat_list_of_images(images)
   800→            image_inputs = image_processor(images=images, return_tensors="pt", **image_processor_kwargs)
   801→            image_num_patches = image_inputs.pop("num_patches")
   802→            image_pixel_values = image_inputs.pop("pixel_values")
   803→            image_num_patches_indices = np.cumsum(image_num_patches)
   804→
   805→        if len(videos) != 0:
   806→            videos = make_batched_videos(videos)
   807→            num_frames_per_video = [len(video) for video in videos]
   808→            patch_indices = np.cumsum(num_frames_per_video)
   809→            image_processor_kwargs["crop_to_patches"] = False
   810→            video_inputs = image_processor(images=videos, return_tensors="pt", **image_processor_kwargs)
   811→            video_num_patches = video_inputs.pop("num_patches")
   812→            video_pixel_values = video_inputs.pop("pixel_values")
   813→            video_num_patches_indices = np.cumsum(video_num_patches)
   814→
   815→        # NOT SUPPORT IMAGE VIDEO INTERLEAVED
   816→        if len(images) != 0 and image_pixel_values is not None:
   817→            for i in range(len(images)):
   818→                start_index = image_num_patches_indices[i - 1] if i > 0 else 0
   819→                end_index = image_num_patches_indices[i]
   820→                image_video_patches.append(image_pixel_values[start_index:end_index])
   821→
   822→        if len(videos) != 0 and video_pixel_values is not None:
   823→            patch_indices_with_prefix = [0] + list(patch_indices)
   824→            for i in range(len(videos)):
   825→                current_patch_index = patch_indices_with_prefix[i]
   826→                end_patch_index = patch_indices_with_prefix[i + 1]
   827→                start_index = video_num_patches_indices[current_patch_index - 1] if i > 0 else 0
   828→                end_index = video_num_patches_indices[end_patch_index - 1]
   829→                image_video_patches.append(video_pixel_values[start_index:end_index])
   830→
   831→        if len(images) != 0 or len(videos) != 0:
   832→            mm_inputs["pixel_values"] = torch.cat(image_video_patches, dim=0)
   833→
   834→        if len(images) != 0:
   835→            mm_inputs.update({"image_num_patches": image_num_patches})
   836→
   837→        if len(videos) != 0:
   838→            mm_inputs.update({"video_patch_indices": patch_indices})
   839→            mm_inputs.update({"video_num_patches": video_num_patches})
   840→
   841→        # Audio processing via feature extractor (produces input_audio_features + audio_mel_mask)
   842→        if len(audios) != 0:
   843→            feature_extractor = getattr(processor, "feature_extractor", None)
   844→            if feature_extractor is not None:
   845→                audio_sr = getattr(processor, "audio_sampling_rate", None) or getattr(
   846→                    feature_extractor, "sampling_rate", 16000
   847→                )
   848→                regularized = self._regularize_audios(audios, sampling_rate=audio_sr)
   849→                audio_inputs = feature_extractor(
   850→                    regularized["audios"],
   851→                    sampling_rate=audio_sr,
   852→                    return_attention_mask=True,
   853→                    return_tensors="pt",
   854→                )
   855→                mm_inputs.update(audio_inputs)
   856→
   857→        return mm_inputs
   858→
   859→    @override
   860→    def process_messages(
   861→        self,
   862→        messages: list[dict[str, str]],
   863→        images: list["ImageInput"],
   864→        videos: list["VideoInput"],
   865→        audios: list["AudioInput"],
   866→        processor: Optional["ProcessorMixin"],
   867→    ) -> list[dict[str, str]]:
   868→        self._validate_input(processor, images, videos, audios)
   869→        self._validate_messages(messages, images, videos, audios)
   870→        num_image_tokens, num_video_tokens = 0, 0
   871→        image_seqlen = getattr(processor, "image_seq_length") if self.expand_mm_tokens else 1
   872→        messages = deepcopy(messages)
   873→        mm_inputs = self._get_mm_inputs(images, videos, audios, processor)

class Gemma3nPlugin(Gemma3Plugin):
    @override
    def process_messages(
        self,
        messages: list[dict[str, str]],
        images: list["ImageInput"],
        videos: list["VideoInput"],
        audios: list["AudioInput"],
        processor: Optional["MMProcessor"],
    ) -> list[dict[str, str]]:
        self._validate_input(processor, images, videos, audios)
        self._validate_messages(messages, images, videos, audios)
        messages = deepcopy(messages)
        boi_token: str = getattr(processor, "boi_token")
        boa_token: str = getattr(processor, "boa_token")
        full_image_sequence: str = getattr(processor, "full_image_sequence")
        full_audio_sequence: str = getattr(processor, "full_audio_sequence")
        image_str = full_image_sequence if self.expand_mm_tokens else boi_token
        audio_str = full_audio_sequence if self.expand_mm_tokens else boa_token

        for message in messages:
            content = message["content"]
            while IMAGE_PLACEHOLDER in content:
                content = content.replace(IMAGE_PLACEHOLDER, image_str, 1)

            while AUDIO_PLACEHOLDER in content:
                content = content.replace(AUDIO_PLACEHOLDER, audio_str, 1)

            message["content"] = content

        return messages

@dataclass
class InternVLPlugin(BasePlugin):
    @override
    def _get_mm_inputs(
        self,
        images: list["ImageInput"],
        videos: list["VideoInput"],
        audios: list["AudioInput"],
        processor: "ProcessorMixin",
        **kwargs,
    ) -> dict[str, "torch.Tensor"]:
        image_processor: BaseImageProcessor = getattr(processor, "image_processor")
        image_processor_kwargs = {}
        if getattr(processor, "crop_to_patches", False):
            image_processor_kwargs.update(
                {
                    "crop_to_patches": True,
                    "max_patches": 12,
                    "min_patches": 1,
                }
            )

        mm_inputs = {}
        image_video_patches = []

        if len(images) != 0:
            images = self._regularize_images(
                images,
                image_max_pixels=getattr(processor, "image_max_pixels", 1024 * 1024),
                image_min_pixels=getattr(processor, "image_min_pixels", 32 * 32),
            )["images"]

        if len(videos) != 0:
            videos = self._regularize_videos(
                videos,
                image_max_pixels=getattr(processor, "video_max_pixels", 256 * 256),
                image_min_pixels=getattr(processor, "video_min_pixels", 16 * 16),
                video_fps=getattr(processor, "video_fps", 2.0),
                video_maxlen=getattr(processor, "video_maxlen", 128),
            )["videos"]

        if len(images) != 0:
            images = make_flat_list_of_images(images)
            image_inputs = image_processor(images=images, return_tensors="pt", **image_processor_kwargs)
            image_num_patches = image_inputs.pop("num_patches")
            image_pixel_values = image_inputs.pop("pixel_values")
            image_num_patches_indices = np.cumsum(image_num_patches)

        if len(videos) != 0:
            videos = make_batched_videos(videos)
            num_frames_per_video = [len(video) for video in videos]
            patch_indices = np.cumsum(num_frames_per_video)
            image_processor_kwargs["crop_to_patches"] = False
            video_inputs = image_processor(images=videos, return_tensors="pt", **image_processor_kwargs)
            video_num_patches = video_inputs.pop("num_patches")
            video_pixel_values = video_inputs.pop("pixel_values")
            video_num_patches_indices = np.cumsum(video_num_patches)

        # NOT SUPPORT IMAGE VIDEO INTERLEAVED
        if len(images) != 0 and image_pixel_values is not None:
            for i in range(len(images)):
                start_index = image_num_patches_indices[i - 1] if i > 0 else 0
                end_index = image_num_patches_indices[i]
                image_video_patches.append(image_pixel_values[start_index:end_index])

        if len(videos) != 0 and video_pixel_values is not None:
            patch_indices_with_prefix = [0] + list(patch_indices)
            for i in range(len(videos)):
                current_patch_index = patch_indices_with_prefix[i]
                end_patch_index = patch_indices_with_prefix[i + 1]
                start_index = video_num_patches_indices[current_patch_index - 1] if i > 0 else 0
                end_index = video_num_patches_indices[end_patch_index - 1]
                image_video_patches.append(video_pixel_values[start_index:end_index])

        if len(images) != 0 or len(videos) != 0:
            mm_inputs["pixel_values"] = torch.cat(image_video_patches, dim=0)

        if len(images) != 0:
            mm_inputs.update({"image_num_patches": image_num_patches})

        if len(videos) != 0:
            mm_inputs.update({"video_patch_indices": patch_indices})
            mm_inputs.update({"video_num_patches": video_num_patches})

        # Audio processing via feature extractor (produces input_audio_features + audio_mel_mask)
        if len(audios) != 0:
            feature_extractor = getattr(processor, "feature_extractor", None)
            if feature_extractor is not None:
                audio_sr = getattr(processor, "audio_sampling_rate", None) or getattr(
                    feature_extractor, "sampling_rate", 16000
                )
                regularized = self._regularize_audios(audios, sampling_rate=audio_sr)
                audio_inputs = feature_extractor(
                    regularized["audios"],
                    sampling_rate=audio_sr,
                    return_attention_mask=True,
                    return_tensors="pt",
                )
                mm_inputs.update(audio_inputs)

        return mm_inputs

    @override
    def process_messages(
        self,
        messages: list[dict[str, str]],
        images: list["ImageInput"],
        videos: list["VideoInput"],
        audios: list["AudioInput"],
        processor: Optional["ProcessorMixin"],
    ) -> list[dict[str, str]]:
        self._validate_input(processor, images, videos, audios)
        self._validate_messages(messages, images, videos, audios)
        num_image_tokens, num_video_tokens = 0, 0
        image_seqlen = getattr(processor, "image_seq_length") if self.expand_mm_tokens else 1
        messages = deepcopy(messages)
        mm_inputs = self._get_mm_inputs(images, videos, audios, processor)


663-
664-@dataclass
665:class Gemma3Plugin(BasePlugin):
666-    @override
667-    def process_messages(


   665→class Gemma3Plugin(BasePlugin):
   666→    @override
   667→    def process_messages(
   668→        self,
   669→        messages: list[dict[str, str]],
   670→        images: list["ImageInput"],
   671→        videos: list["VideoInput"],
   672→        audios: list["AudioInput"],
   673→        processor: Optional["MMProcessor"],
   674→    ) -> list[dict[str, str]]:
   675→        self._validate_input(processor, images, videos, audios)
   676→        self._validate_messages(messages, images, videos, audios)
   677→        num_image_tokens = 0
   678→        messages = deepcopy(messages)
   679→        boi_token: str = getattr(processor, "boi_token")
   680→        full_image_sequence: str = getattr(processor, "full_image_sequence")
   681→        image_str = full_image_sequence if self.expand_mm_tokens else boi_token
   682→
   683→        do_pan_and_scan: bool = getattr(processor, "image_do_pan_and_scan", False)
   684→        if do_pan_and_scan:
   685→            mm_inputs = self._get_mm_inputs(images, videos, audios, processor)
   686→
   687→        for message in messages:
   688→            content = message["content"]
   689→            while IMAGE_PLACEHOLDER in content:
   690→                if do_pan_and_scan:
   691→                    image_placeholder_str = (
   692→                        "Here is the original image {{image}} and here are some crops to help you see better "
   693→                        + " ".join(["{{image}}"] * mm_inputs["num_crops"][0][num_image_tokens])
   694→                    )
   695→                else:
   696→                    image_placeholder_str = "{{image}}"
   697→
   698→                content = content.replace(IMAGE_PLACEHOLDER, image_placeholder_str, 1)
   699→                num_image_tokens += 1
   700→
   701→            message["content"] = content.replace("{{image}}", image_str)
   702→
   703→        return messages
   704→
   705→    @override
   706→    def get_mm_inputs(
   707→        self,
   708→        images: list["ImageInput"],
   709→        videos: list["VideoInput"],
   710→        audios: list["AudioInput"],
   711→        imglens: list[int],
   712→        vidlens: list[int],
   713→        audlens: list[int],
   714→        batch_ids: list[list[int]],
   715→        processor: Optional["MMProcessor"],
   716→    ) -> dict[str, Union[list[int], "torch.Tensor"]]:
   717→        self._validate_input(processor, images, videos, audios)
   718→        mm_inputs = self._get_mm_inputs(images, videos, audios, processor)
   719→        mm_inputs.pop("num_crops", None)
   720→        mm_inputs["token_type_ids"] = _get_gemma3_token_type_ids(batch_ids, processor)
   721→        return mm_inputs
   722→
   723→
   724→class Gemma3nPlugin(Gemma3Plugin):

class Gemma3Plugin(BasePlugin):
    @override
    def process_messages(
        self,
        messages: list[dict[str, str]],
        images: list["ImageInput"],
        videos: list["VideoInput"],
        audios: list["AudioInput"],
        processor: Optional["MMProcessor"],
    ) -> list[dict[str, str]]:
        self._validate_input(processor, images, videos, audios)
        self._validate_messages(messages, images, videos, audios)
        num_image_tokens = 0
        messages = deepcopy(messages)
        boi_token: str = getattr(processor, "boi_token")
        full_image_sequence: str = getattr(processor, "full_image_sequence")
        image_str = full_image_sequence if self.expand_mm_tokens else boi_token

        do_pan_and_scan: bool = getattr(processor, "image_do_pan_and_scan", False)
        if do_pan_and_scan:
            mm_inputs = self._get_mm_inputs(images, videos, audios, processor)

        for message in messages:
            content = message["content"]
            while IMAGE_PLACEHOLDER in content:
                if do_pan_and_scan:
                    image_placeholder_str = (
                        "Here is the original image {{image}} and here are some crops to help you see better "
                        + " ".join(["{{image}}"] * mm_inputs["num_crops"][0][num_image_tokens])
                    )
                else:
                    image_placeholder_str = "{{image}}"

                content = content.replace(IMAGE_PLACEHOLDER, image_placeholder_str, 1)
                num_image_tokens += 1

            message["content"] = content.replace("{{image}}", image_str)

        return messages

    @override
    def get_mm_inputs(
        self,
        images: list["ImageInput"],
        videos: list["VideoInput"],
        audios: list["AudioInput"],
        imglens: list[int],
        vidlens: list[int],
        audlens: list[int],
        batch_ids: list[list[int]],
        processor: Optional["MMProcessor"],
    ) -> dict[str, Union[list[int], "torch.Tensor"]]:
        self._validate_input(processor, images, videos, audios)
        mm_inputs = self._get_mm_inputs(images, videos, audios, processor)
        mm_inputs.pop("num_crops", None)
        mm_inputs["token_type_ids"] = _get_gemma3_token_type_ids(batch_ids, processor)
        return mm_inputs

class Gemma3nPlugin(Gemma3Plugin):


466-        return {"audios": results, "sampling_rates": sampling_rates}
467-
468:    def _get_mm_inputs(
469-        self,
470-        images: list["ImageInput"],
--
758-class InternVLPlugin(BasePlugin):
759-    @override
760:    def _get_mm_inputs(
761-        self,
762-        images: list["ImageInput"],
--
1154-class MiniCPMVPlugin(BasePlugin):
1155-    @override
1156:    def _get_mm_inputs(
1157-        self,
1158-        images: list["ImageInput"],
--
1689-
1690-    @override
1691:    def _get_mm_inputs(
1692-        self,
1693-        images: list["ImageInput"],
--
1774-class Qwen3VLPlugin(Qwen2VLPlugin):
1775-    @override
1776:    def _get_mm_inputs(
1777-        self,
1778-        images: list["ImageInput"],
--
1897-class GLM4VPlugin(Qwen2VLPlugin):
1898-    @override
1899:    def _get_mm_inputs(
1900-        self,
1901-        images: list["ImageInput"],
--
2031-
2032-    @override
2033:    def _get_mm_inputs(
2034-        self,
2035-        images: list["ImageInput"],
--
2269-
2270-    @override
2271:    def _get_mm_inputs(
2272-        self,
2273-        images: list["ImageInput"],


   468→    def _get_mm_inputs(
   469→        self,
   470→        images: list["ImageInput"],
   471→        videos: list["VideoInput"],
   472→        audios: list["AudioInput"],
   473→        processor: "MMProcessor",
   474→        imglens: list[int] | None = None,
   475→    ) -> dict[str, "torch.Tensor"]:
   476→        r"""Process visual inputs.
   477→
   478→        Returns: (llava and paligemma)
   479→            pixel_values: tensor with shape (B, C, H, W)
   480→
   481→        Returns: (qwen2-vl)
   482→            pixel_values: tensor with shape (num_patches, patch_dim)
   483→            image_grid_thw: tensor with shape (num_images, 3), where the three numbers are time, width, height
   484→                            where num_patches == torch.prod(image_grid_thw)
   485→
   486→        Returns: (mllama)
   487→            pixel_values: tensor with shape
   488→                          (batch_size, max_num_images, max_image_tiles, channels, tile_height, tile_width)
   489→                          For example, (2, 1, 4, 3, 560, 560).
   490→            aspect_ratio_ids: tensor with shape (batch_size, max_num_images). For example, (2, 1).
   491→            aspect_ratio_mask: tensor with shape (batch_size, max_num_images, max_image_tiles). For example, (2, 1, 4).
   492→            num_tiles: List[List[int]] with shape (batch_size, num_images_in_batch). For example, (2, 1).
   493→
   494→        """
   495→        mm_inputs = {}
   496→        if len(images) != 0:
   497→            image_processor: BaseImageProcessor = getattr(processor, "image_processor", None)
   498→            images = self._regularize_images(
   499→                images,
   500→                image_max_pixels=getattr(processor, "image_max_pixels", 768 * 768),
   501→                image_min_pixels=getattr(processor, "image_min_pixels", 32 * 32),
   502→            )["images"]
   503→            if imglens is not None:  # if imglens are provided, make batched images
   504→                images = _make_batched_images(images, imglens)
   505→
   506→            image_processor_kwargs = {}
   507→            if getattr(processor, "image_do_pan_and_scan", False):  # gemma3 image processor
   508→                image_processor_kwargs.update(
   509→                    {
   510→                        "do_pan_and_scan": True,
   511→                        "pan_and_scan_min_crop_size": 256,
   512→                        "pan_and_scan_max_num_crops": 4,
   513→                        "pan_and_scan_min_ratio_to_activate": 1.2,
   514→                    }
   515→                )
   516→
   517→            mm_inputs.update(image_processor(images, return_tensors="pt", **image_processor_kwargs))
   518→
   519→        if len(videos) != 0:
   520→            video_processor: BaseImageProcessor = getattr(
   521→                processor, "video_processor", getattr(processor, "image_processor", None)
   522→            )
   523→            videos = self._regularize_videos(
   524→                videos,
   525→                image_max_pixels=getattr(processor, "video_max_pixels", 256 * 256),
   526→                image_min_pixels=getattr(processor, "video_min_pixels", 16 * 16),
   527→                video_fps=getattr(processor, "video_fps", 2.0),
   528→                video_maxlen=getattr(processor, "video_maxlen", 128),
   529→            )["videos"]
   530→            if "videos" in inspect.signature(video_processor.preprocess).parameters:  # for qwen2_vl and video_llava
   531→                mm_inputs.update(video_processor(images=None, videos=videos, return_tensors="pt"))
   532→            else:  # for llava_next_video
   533→                mm_inputs.update(video_processor(videos, return_tensors="pt"))
   534→
   535→        if len(audios) != 0:
   536→            feature_extractor: SequenceFeatureExtractor = getattr(processor, "feature_extractor", None) or getattr(
   537→                processor, "audio_processor", None
   538→            )
   539→            audios = self._regularize_audios(
   540→                audios,
   541→                sampling_rate=getattr(processor, "audio_sampling_rate", 16000),
   542→            )["audios"]
   543→            mm_inputs.update(
   544→                feature_extractor(
   545→                    audios,
   546→                    sampling_rate=getattr(processor, "audio_sampling_rate", 16000),
   547→                    return_attention_mask=True,

    def _get_mm_inputs(
        self,
        images: list["ImageInput"],
        videos: list["VideoInput"],
        audios: list["AudioInput"],
        processor: "MMProcessor",
        imglens: list[int] | None = None,
    ) -> dict[str, "torch.Tensor"]:
        r"""Process visual inputs.

        Returns: (llava and paligemma)
            pixel_values: tensor with shape (B, C, H, W)

        Returns: (qwen2-vl)
            pixel_values: tensor with shape (num_patches, patch_dim)
            image_grid_thw: tensor with shape (num_images, 3), where the three numbers are time, width, height
                            where num_patches == torch.prod(image_grid_thw)

        Returns: (mllama)
            pixel_values: tensor with shape
                          (batch_size, max_num_images, max_image_tiles, channels, tile_height, tile_width)
                          For example, (2, 1, 4, 3, 560, 560).
            aspect_ratio_ids: tensor with shape (batch_size, max_num_images). For example, (2, 1).
            aspect_ratio_mask: tensor with shape (batch_size, max_num_images, max_image_tiles). For example, (2, 1, 4).
            num_tiles: List[List[int]] with shape (batch_size, num_images_in_batch). For example, (2, 1).

        """
        mm_inputs = {}
        if len(images) != 0:
            image_processor: BaseImageProcessor = getattr(processor, "image_processor", None)
            images = self._regularize_images(
                images,
                image_max_pixels=getattr(processor, "image_max_pixels", 768 * 768),
                image_min_pixels=getattr(processor, "image_min_pixels", 32 * 32),
            )["images"]
            if imglens is not None:  # if imglens are provided, make batched images
                images = _make_batched_images(images, imglens)

            image_processor_kwargs = {}
            if getattr(processor, "image_do_pan_and_scan", False):  # gemma3 image processor
                image_processor_kwargs.update(
                    {
                        "do_pan_and_scan": True,
                        "pan_and_scan_min_crop_size": 256,
                        "pan_and_scan_max_num_crops": 4,
                        "pan_and_scan_min_ratio_to_activate": 1.2,
                    }
                )

            mm_inputs.update(image_processor(images, return_tensors="pt", **image_processor_kwargs))

        if len(videos) != 0:
            video_processor: BaseImageProcessor = getattr(
                processor, "video_processor", getattr(processor, "image_processor", None)
            )
            videos = self._regularize_videos(
                videos,
                image_max_pixels=getattr(processor, "video_max_pixels", 256 * 256),
                image_min_pixels=getattr(processor, "video_min_pixels", 16 * 16),
                video_fps=getattr(processor, "video_fps", 2.0),
                video_maxlen=getattr(processor, "video_maxlen", 128),
            )["videos"]
            if "videos" in inspect.signature(video_processor.preprocess).parameters:  # for qwen2_vl and video_llava
                mm_inputs.update(video_processor(images=None, videos=videos, return_tensors="pt"))
            else:  # for llava_next_video
                mm_inputs.update(video_processor(videos, return_tensors="pt"))

        if len(audios) != 0:
            feature_extractor: SequenceFeatureExtractor = getattr(processor, "feature_extractor", None) or getattr(
                processor, "audio_processor", None
            )
            audios = self._regularize_audios(
                audios,
                sampling_rate=getattr(processor, "audio_sampling_rate", 16000),
            )["audios"]
            mm_inputs.update(
                feature_extractor(
                    audios,
                    sampling_rate=getattr(processor, "audio_sampling_rate", 16000),
                    return_attention_mask=True,


Now I have a clear picture. The model uses `Gemma3nProcessor` with a `Gemma3nAudioFeatureExtractor`. Audio is passed as numpy arrays at 16kHz. Rather than using the LlamaFactory `ChatModel` API (which requires complex args parsing), I can write a simpler direct inference script using HuggingFace transformers directly.

Let me also check how `_regularize_audios` works and what `AudioInput` type is.


121-
122-_audio_decode_error_count = 0
123-_audio_decode_error_lock = threading.Lock()
124-
125-
126:def _log_decode_fallback(audio: "AudioInput", err: Exception, fallback_len: int, sampling_rate: float) -> None:
127-    global _audio_decode_error_count
128-
129-    with _audio_decode_error_lock:
130-        _audio_decode_error_count += 1
131-        err_idx = _audio_decode_error_count
--
173-        path: str | None
174-        bytes: bytes | None
175-
176-    ImageInput = Union[str, bytes, EncodedImage, BinaryIO, ImageObject]
177-    VideoInput = Union[str, BinaryIO, list[list[ImageInput]]]
178:    AudioInput = Union[str, BinaryIO, NDArray]
179-
180-    class RegularizedImageOutput(TypedDict):
181-        images: list[ImageObject]
182-
183-    class RegularizedVideoOutput(TypedDict):
--
259-    def _validate_input(
260-        self,
261-        processor: Optional["MMProcessor"],
262-        images: list["ImageInput"],
263-        videos: list["VideoInput"],
264:        audios: list["AudioInput"],

[Showing results with pagination = limit: 30, offset: 0]


   430→
   431→            # Load audio from file path or tar:// URI.
   432→            # If decode fails and fallback is enabled, return zero-audio so one bad
   433→            # sample does not crash distributed training.
   434→            try:
   435→                if isinstance(audio, str) and audio.startswith(_TAR_URI_PREFIX):
   436→                    audio_tensor, sr = _load_audio_from_tar_uri(audio)
   437→                else:
   438→                    audio_tensor, sr = torchaudio.load(audio)
   439→            except Exception as err:
   440→                if not _should_fallback_on_decode_error():
   441→                    raise
   442→
   443→                fallback_len = _get_decode_fallback_num_samples(sampling_rate)
   444→                _log_decode_fallback(audio, err, fallback_len, sampling_rate)
   445→                results.append(np.zeros(fallback_len, dtype=np.float32))
   446→                sampling_rates.append(sampling_rate)
   447→                continue
   448→
   449→            # Convert to mono
   450→            if audio_tensor.shape[0] > 1:
   451→                audio_tensor = audio_tensor.mean(dim=0, keepdim=True)
   452→
   453→            # SR check — no silent runtime resample unless opted in
   454→            if sr != sampling_rate:
   455→                if os.environ.get("LLAMAFACTORY_ALLOW_RESAMPLE", "0") == "1":
   456→                    audio_tensor = torchaudio.functional.resample(audio_tensor, sr, sampling_rate)
   457→                else:
   458→                    raise ValueError(
   459→                        f"Audio sample rate {sr} != expected {sampling_rate}. "
   460→                        f"Pre-convert to {sampling_rate}Hz or set LLAMAFACTORY_ALLOW_RESAMPLE=1."
   461→                    )
   462→
   463→            results.append(audio_tensor.squeeze(0).numpy())
   464→            sampling_rates.append(sampling_rate)
   465→
   466→        return {"audios": results, "sampling_rates": sampling_rates}
   467→
   468→    def _get_mm_inputs(
   469→        self,

            # Load audio from file path or tar:// URI.
            # If decode fails and fallback is enabled, return zero-audio so one bad
            # sample does not crash distributed training.
            try:
                if isinstance(audio, str) and audio.startswith(_TAR_URI_PREFIX):
                    audio_tensor, sr = _load_audio_from_tar_uri(audio)
                else:
                    audio_tensor, sr = torchaudio.load(audio)
            except Exception as err:
                if not _should_fallback_on_decode_error():
                    raise

                fallback_len = _get_decode_fallback_num_samples(sampling_rate)
                _log_decode_fallback(audio, err, fallback_len, sampling_rate)
                results.append(np.zeros(fallback_len, dtype=np.float32))
                sampling_rates.append(sampling_rate)
                continue

            # Convert to mono
            if audio_tensor.shape[0] > 1:
                audio_tensor = audio_tensor.mean(dim=0, keepdim=True)

            # SR check — no silent runtime resample unless opted in
            if sr != sampling_rate:
                if os.environ.get("LLAMAFACTORY_ALLOW_RESAMPLE", "0") == "1":
                    audio_tensor = torchaudio.functional.resample(audio_tensor, sr, sampling_rate)
                else:
                    raise ValueError(
                        f"Audio sample rate {sr} != expected {sampling_rate}. "
                        f"Pre-convert to {sampling_rate}Hz or set LLAMAFACTORY_ALLOW_RESAMPLE=1."
                    )

            results.append(audio_tensor.squeeze(0).numpy())
            sampling_rates.append(sampling_rate)

        return {"audios": results, "sampling_rates": sampling_rates}

    def _get_mm_inputs(
        self,


     1→{{ '<bos>' }}{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0]['content'] %}{% else %}{% set loop_messages = messages %}{% endif %}{% for message in loop_messages %}{% if loop.index0 == 0 and system_message is defined %}{% set content = system_message + '
     2→
     3→' + message['content'] %}{% else %}{% set content = message['content'] %}{% endif %}{% if message['role'] == 'user' %}{{ '<start_of_turn>user
     4→' + content + '<end_of_turn>
     5→<start_of_turn>model
     6→' }}{% elif message['role'] == 'assistant' %}{{ content + '<end_of_turn>
     7→' }}{% endif %}{% endfor %}

{{ '<bos>' }}{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0]['content'] %}{% else %}{% set loop_messages = messages %}{% endif %}{% for message in loop_messages %}{% if loop.index0 == 0 and system_message is defined %}{% set content = system_message + '

' + message['content'] %}{% else %}{% set content = message['content'] %}{% endif %}{% if message['role'] == 'user' %}{{ '<start_of_turn>user
' + content + '<end_of_turn>
<start_of_turn>model
' }}{% elif message['role'] == 'assistant' %}{{ content + '<end_of_turn>
' }}{% endif %}{% endfor %}


     1→{
     2→  "bos_token_id": 2,
     3→  "cache_implementation": "hybrid",
     4→  "do_sample": true,
     5→  "eos_token_id": [
     6→    106,
     7→    1
     8→  ],
     9→  "pad_token_id": 0,
    10→  "top_k": 64,
    11→  "top_p": 0.95,
    12→  "transformers_version": "5.2.0"
    13→}
    14→


Bash

python3 -c "
from datasets import load_dataset
ds = load_dataset('BayAreaBoys/indic-asr-benchmark-6k', cache_dir='/home/ubuntu/training/datasets/indic-asr-benchmark-6k', split='train')
print('Columns:', ds.column_names)
print('N:', len(ds))
print('Sample:', ds[0])
print('Languages:', sorted(set(ds['language'])))
" 2>&1 | head -30

Check dataset structure


Using the latest cached version of the dataset since BayAreaBoys/indic-asr-benchmark-6k couldn't be found on the Hugging Face Hub
Found the latest cached dataset configuration 'default' at /home/ubuntu/training/datasets/indic-asr-benchmark-6k/BayAreaBoys___indic-asr-benchmark-6k/default/0.0.0/4d4bc89b7c915a6d80a1efbee3a708d58e688c81 (last modified on Thu Mar 26 00:49:07 2026).
Columns: ['id', 'filename', 'reference', 'language', 'primary_language', 'lang_code', 'source', 'duration', 'speaker_id', 'scenario', 'm4a_path', 'audio']
N: 6000
Sample: {'id': 'as_0000', 'filename': 'as_0000.wav', 'reference': 'অঁ দুদিনীয়া মই দুদিনীয়া বুলিয়ে শুনিছিলোঁ এনেকৈ ভালে লাগিব', 'language': 'assamese', 'primary_language': 'assamese', 'lang_code': 'as', 'source': 'IndicVoices (valid)', 'duration': 3.292, 'speaker_id': 'S4257954800309252', 'scenario': 'Conversation', 'm4a_path': '', 'audio': <datasets.features._torchcodec.AudioDecoder object at 0x7e1e74b53160>}
Languages: ['assamese', 'bengali', 'english', 'gujarati', 'hindi', 'kannada', 'malayalam', 'marathi', 'odia', 'punjabi', 'tamil', 'telugu']


Now I have all the info needed. Let me check download progress and write the benchmark script.


du -sh /home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000/ 2>&1 && ls /home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000/ 2>&1

Check download progress


19G	/home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000/
chat_template.jinja
config.json
generation_config.json
model.safetensors.5C8B6FA1
optimizer.pt.AE1ed226
processor_config.json
rng_state_0.pth
rng_state_1.pth
rng_state_2.pth
rng_state_3.pth
rng_state_4.pth
rng_state_5.pth
rng_state_6.pth
rng_state_7.pth
scheduler.pt
tokenizer_config.json
tokenizer.json
trainer_state.json
training_args.bin


ls -lh /home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000/model.safetensors* 2>&1

Check model file status


-rw-rw-r-- 1 ubuntu ubuntu 9.4G Mar 26 00:53 /home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000/model.safetensors.5C8B6FA1


The model safetensors is still downloading (~9.4G of ~20GB). Let me write the benchmark script while it downloads.


/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-28-    r"""Dataset attributes."""
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-29-
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-30-    # basic configs
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-31-    load_from: Literal["hf_hub", "ms_hub", "om_hub", "script", "file"]
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-32-    dataset_name: str
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py:33:    formatting: Literal["alpaca", "sharegpt", "openai"] = "alpaca"
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-34-    ranking: bool = False
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-35-    # extra configs
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-36-    subset: str | None = None
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-37-    split: str = "train"
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-38-    folder: str | None = None
--
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-45-    audios: str | None = None
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-46-    # dpo columns
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-47-    chosen: str | None = None
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-48-    rejected: str | None = None
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-49-    kto_tag: str | None = None
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py:50:    # alpaca columns
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-51-    prompt: str | None = "instruction"
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-52-    query: str | None = "input"
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-53-    response: str | None = "output"
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-54-    history: str | None = None
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-55-    # sharegpt columns
--
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-68-
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-69-    def set_attr(self, key: str, obj: dict[str, Any], default: Any | None = None) -> None:
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-70-        setattr(self, key, obj.get(key, default))
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-71-
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py-72-    def join(self, attr: dict[str, Any]) -> None:
/home/ubuntu/training/gemma3n-E2B-asr-ft/vendor/LlamaFactory/src/llamafactory/data/parser.py:73:        self.set_attr("formatting", attr, default="alpaca")

[Showing results with pagination = limit: 30, offset: 0]


No matches found

[Showing results with pagination = limit: 40, offset: 0]


No files found


107-        if self.dataset_attr.query and example[self.dataset_attr.query]:
108-            query.append(example[self.dataset_attr.query])
109-
110:        prompt.append({"role": Role.USER.value, "content": "\n".join(query)})  # "prompt\nquery"
111-
112-        if self.dataset_attr.kto_tag and isinstance(example[self.dataset_attr.kto_tag], bool):  # kto example
113-            response = [{"role": Role.ASSISTANT.value, "content": example[self.dataset_attr.response]}]

[Showing results with pagination = limit: 40, offset: 0]


Line 110 confirms: the instruction and input are joined with `\n`. So the user content during training was:

```
<audio>
Transcribe the audio.
Language: Hindi
```

Now I have everything. Let me write the benchmark script.


Write

#!/usr/bin/env python3
"""
Benchmark Gemma3n-E2B-ASR checkpoints on indic-asr-benchmark-6k dataset.

Usage:
    python benchmark_gemma3n_asr.py \
        --checkpoint /home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000 \
        --checkpoint-name ckpt-10000 \
        --output-dir /home/ubuntu/training/benchmark_results_gemma3n \
        --batch-size 1

Metrics computed per language and overall:
    - WER  (Word Error Rate)
    - CER  (Character Error Rate)
    - WER-norm (WER after unicode/punctuation/space normalization)
    - CER-norm (CER after normalization)
"""

import argparse
import json
import os
import re
import sys
import time
import unicodedata
from collections import defaultdict
from pathlib import Path

import numpy as np
import torch
import torchaudio
from datasets import load_dataset as hf_load_dataset
from jiwer import cer, wer
from transformers import AutoModelForImageTextToText, AutoProcessor

# Language code → display name mapping
LANG_CODE_TO_NAME = {
    "as": "Assamese", "bn": "Bengali", "en": "English", "gu": "Gujarati",
    "hi": "Hindi", "kn": "Kannada", "ml": "Malayalam", "mr": "Marathi",
    "or": "Odia", "pa": "Punjabi", "ta": "Tamil", "te": "Telugu",
}

# ── Text normalization (standard ASR eval practice) ──────────────────────

def normalize_text(text: str) -> str:
    """Normalize text for fair WER/CER comparison.
    Unicode NFC → lowercase → remove punctuation → collapse whitespace → strip.
    """
    text = unicodedata.normalize("NFC", text)
    text = text.lower()
    text = re.sub(r"[^\w\s]", "", text, flags=re.UNICODE)
    text = re.sub(r"\s+", " ", text).strip()
    return text

def load_audio_array(sample, target_sr=16000):
    """Extract audio as numpy float32 array at target sample rate."""
    audio = sample["audio"]
    wav = np.array(audio["array"], dtype=np.float32)
    sr = audio["sampling_rate"]

    if sr != target_sr:
        wav_tensor = torch.from_numpy(wav).unsqueeze(0)
        wav_tensor = torchaudio.functional.resample(wav_tensor, sr, target_sr)
        wav = wav_tensor.squeeze(0).numpy()

    return wav

def run_inference(model, processor, dataset, batch_size: int = 1, max_new_tokens: int = 256):
    """Run inference on HF dataset and return list of result dicts."""
    results = []
    total = len(dataset)

    for start in range(0, total, batch_size):
        end = min(start + batch_size, total)
        batch = dataset.select(range(start, end))

        for i, sample in enumerate(batch):
            lang_name = LANG_CODE_TO_NAME.get(sample["lang_code"], sample["language"].capitalize())

            # Construct prompt matching training format
            user_content = f"<audio>\nTranscribe the audio.\nLanguage: {lang_name}"
            messages = [{"role": "user", "content": user_content}]

            # Get audio array
            wav = load_audio_array(sample)

            try:
                # Apply chat template
                prompt_text = processor.apply_chat_template(
                    messages, add_generation_prompt=True, tokenize=False
                )

                # Process inputs (text + audio)
                inputs = processor(
                    text=prompt_text,
                    audios=[wav],
                    sampling_rate=16000,
                    return_tensors="pt",
                )

                # Move to model device
                inputs = {k: v.to(model.device) if isinstance(v, torch.Tensor) else v
                          for k, v in inputs.items()}

                # Generate
                with torch.no_grad():
                    output_ids = model.generate(
                        **inputs,
                        max_new_tokens=max_new_tokens,
                        do_sample=False,
                    )

                # Decode only the new tokens
                input_len = inputs["input_ids"].shape[1]
                generated_ids = output_ids[0][input_len:]
                hypothesis = processor.tokenizer.decode(generated_ids, skip_special_tokens=True).strip()

            except Exception as e:
                print(f"  Sample {start + i} ({sample['id']}) FAILED: {e}")
                hypothesis = ""

            results.append({
                "id": sample["id"],
                "language": sample["language"],
                "lang_code": sample["lang_code"],
                "source": sample["source"],
                "duration": sample["duration"],
                "reference": sample["reference"],
                "hypothesis": hypothesis,
            })

        elapsed_pct = end / total * 100
        if end % max(batch_size * 10, 10) == 0 or end == total:
            print(f"  [{end}/{total}] ({elapsed_pct:.0f}%) processed", flush=True)

    return results

def compute_metrics(results: list) -> dict:
    """Compute WER, CER, normalized-WER/CER per language and overall."""
    by_lang = defaultdict(lambda: {"refs": [], "hyps": [], "refs_norm": [], "hyps_norm": []})

    for r in results:
        lang = r["language"]
        ref = r["reference"]
        hyp = r.get("hypothesis", "")

        ref_norm = normalize_text(ref)
        hyp_norm = normalize_text(hyp)

        if not ref_norm:
            continue

        by_lang[lang]["refs"].append(ref)
        by_lang[lang]["hyps"].append(hyp)
        by_lang[lang]["refs_norm"].append(ref_norm)
        by_lang[lang]["hyps_norm"].append(hyp_norm)

    metrics = {}
    all_refs, all_hyps = [], []
    all_refs_norm, all_hyps_norm = [], []

    for lang in sorted(by_lang.keys()):
        d = by_lang[lang]
        n = len(d["refs"])
        lang_wer = wer(d["refs"], d["hyps"])
        lang_cer = cer(d["refs"], d["hyps"])
        lang_wer_norm = wer(d["refs_norm"], d["hyps_norm"])
        lang_cer_norm = cer(d["refs_norm"], d["hyps_norm"])

        empty_count = sum(1 for h in d["hyps"] if not h.strip())

        metrics[lang] = {
            "n_samples": n,
            "wer": round(lang_wer * 100, 2),
            "cer": round(lang_cer * 100, 2),
            "wer_normalized": round(lang_wer_norm * 100, 2),
            "cer_normalized": round(lang_cer_norm * 100, 2),
            "empty_hypotheses": empty_count,
        }

        all_refs.extend(d["refs"])
        all_hyps.extend(d["hyps"])
        all_refs_norm.extend(d["refs_norm"])
        all_hyps_norm.extend(d["hyps_norm"])

    # Overall (corpus-level)
    metrics["__overall__"] = {
        "n_samples": len(all_refs),
        "wer": round(wer(all_refs, all_hyps) * 100, 2),
        "cer": round(cer(all_refs, all_hyps) * 100, 2),
        "wer_normalized": round(wer(all_refs_norm, all_hyps_norm) * 100, 2),
        "cer_normalized": round(cer(all_refs_norm, all_hyps_norm) * 100, 2),
    }

    # Macro average across languages
    lang_keys = [k for k in metrics if not k.startswith("__")]
    metrics["__macro_avg__"] = {
        "n_languages": len(lang_keys),
        "wer": round(np.mean([metrics[k]["wer"] for k in lang_keys]), 2),
        "cer": round(np.mean([metrics[k]["cer"] for k in lang_keys]), 2),
        "wer_normalized": round(np.mean([metrics[k]["wer_normalized"] for k in lang_keys]), 2),
        "cer_normalized": round(np.mean([metrics[k]["cer_normalized"] for k in lang_keys]), 2),
    }

    return metrics

def print_metrics_table(metrics: dict, checkpoint_name: str):
    """Pretty print metrics as a table."""
    print(f"\n{'='*85}")
    print(f"  BENCHMARK RESULTS: {checkpoint_name}")
    print(f"{'='*85}")
    print(f"{'Language':<14} {'N':>5} {'WER%':>8} {'CER%':>8} {'WER-n%':>8} {'CER-n%':>8}")
    print(f"{'-'*14} {'-'*5} {'-'*8} {'-'*8} {'-'*8} {'-'*8}")

    for lang in sorted(k for k in metrics if not k.startswith("__")):
        m = metrics[lang]
        print(f"{lang:<14} {m['n_samples']:>5} {m['wer']:>8.2f} {m['cer']:>8.2f} "
              f"{m['wer_normalized']:>8.2f} {m['cer_normalized']:>8.2f}")

    print(f"{'-'*14} {'-'*5} {'-'*8} {'-'*8} {'-'*8} {'-'*8}")

    ov = metrics["__overall__"]
    print(f"{'OVERALL':<14} {ov['n_samples']:>5} {ov['wer']:>8.2f} {ov['cer']:>8.2f} "
          f"{ov['wer_normalized']:>8.2f} {ov['cer_normalized']:>8.2f}")

    ma = metrics["__macro_avg__"]
    print(f"{'MACRO-AVG':<14} {ma['n_languages']:>5} {ma['wer']:>8.2f} {ma['cer']:>8.2f} "
          f"{ma['wer_normalized']:>8.2f} {ma['cer_normalized']:>8.2f}")
    print(f"{'='*85}\n")

def save_sample_analysis(results: list, output_dir: Path, n_per_lang: int = 5):
    """Save a few sample predictions per language for manual inspection."""
    by_lang = defaultdict(list)
    for r in results:
        by_lang[r["language"]].append(r)

    samples = []
    for lang in sorted(by_lang.keys()):
        lang_results = by_lang[lang]
        for r in lang_results[:n_per_lang]:
            samples.append({
                "id": r["id"],
                "language": r["language"],
                "reference": r["reference"],
                "hypothesis": r["hypothesis"],
                "ref_normalized": normalize_text(r["reference"]),
                "hyp_normalized": normalize_text(r["hypothesis"]),
            })

    path = output_dir / "sample_analysis.json"
    with open(path, "w", encoding="utf-8") as f:
        json.dump(samples, f, ensure_ascii=False, indent=2)
    print(f"Sample analysis saved: {path}")

def main():
    parser = argparse.ArgumentParser(description="Benchmark Gemma3n-E2B-ASR on Indic ASR dataset")
    parser.add_argument("--checkpoint", required=True, help="Path to model checkpoint directory")
    parser.add_argument("--checkpoint-name", default=None,
                        help="Name for this checkpoint in output files (default: dirname)")
    parser.add_argument("--hf-dataset", default="BayAreaBoys/indic-asr-benchmark-6k",
                        help="HuggingFace dataset name")
    parser.add_argument("--hf-token", default=None, help="HuggingFace token (or set HF_TOKEN env)")
    parser.add_argument("--cache-dir", default="/home/ubuntu/training/datasets/indic-asr-benchmark-6k")
    parser.add_argument("--output-dir", default="/home/ubuntu/training/benchmark_results_gemma3n")
    parser.add_argument("--batch-size", type=int, default=1,
                        help="Samples per batch (1 for sequential inference)")
    parser.add_argument("--max-new-tokens", type=int, default=256)
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--languages", nargs="+", default=None,
                        help="Filter to specific language codes (e.g. hi en ta)")
    args = parser.parse_args()

    ckpt_name = args.checkpoint_name or Path(args.checkpoint).name
    output_dir = Path(args.output_dir) / ckpt_name
    output_dir.mkdir(parents=True, exist_ok=True)

    # Load model and processor
    print(f"Loading checkpoint: {args.checkpoint}")
    t0 = time.time()

    processor = AutoProcessor.from_pretrained(args.checkpoint, trust_remote_code=True)
    model = AutoModelForImageTextToText.from_pretrained(
        args.checkpoint,
        torch_dtype=torch.bfloat16,
        device_map=args.device,
        trust_remote_code=True,
    )
    model.eval()

    load_time = time.time() - t0
    print(f"Model loaded in {load_time:.1f}s")

    # Load dataset
    hf_token = args.hf_token or os.environ.get("HF_TOKEN")
    print(f"Loading dataset: {args.hf_dataset}")
    dataset = hf_load_dataset(
        args.hf_dataset, token=hf_token, cache_dir=args.cache_dir, split="train"
    )

    # Filter languages if specified
    if args.languages:
        dataset = dataset.filter(lambda x: x["lang_code"] in args.languages)
        print(f"Filtered to {len(dataset)} samples for languages: {args.languages}")

    print(f"Dataset: {len(dataset)} samples")

    # Run inference
    print(f"\nRunning inference (batch_size={args.batch_size})...")
    t0 = time.time()
    results = run_inference(model, processor, dataset, args.batch_size, args.max_new_tokens)
    inference_time = time.time() - t0

    total_audio_sec = sum(r["duration"] for r in results)
    rtf = inference_time / total_audio_sec if total_audio_sec > 0 else 0
    print(f"Inference done in {inference_time:.1f}s for {total_audio_sec:.0f}s audio (RTF={rtf:.4f})")

    # Save raw predictions
    preds_file = output_dir / "predictions.json"
    with open(preds_file, "w", encoding="utf-8") as f:
        json.dump(results, f, ensure_ascii=False, indent=2)
    print(f"Predictions saved: {preds_file}")

    # Compute metrics
    metrics = compute_metrics(results)
    metrics["__meta__"] = {
        "checkpoint": args.checkpoint,
        "checkpoint_name": ckpt_name,
        "model_type": "gemma3n-E2B-asr",
        "dataset": args.hf_dataset,
        "batch_size": args.batch_size,
        "inference_time_sec": round(inference_time, 2),
        "total_audio_sec": round(total_audio_sec, 2),
        "rtf": round(rtf, 4),
    }

    metrics_file = output_dir / "metrics.json"
    with open(metrics_file, "w") as f:
        json.dump(metrics, f, indent=2, ensure_ascii=False)
    print(f"Metrics saved: {metrics_file}")

    # Print table
    print_metrics_table(metrics, ckpt_name)

    # Save sample analysis
    save_sample_analysis(results, output_dir)

if __name__ == "__main__":
    main()


File created successfully at: /home/ubuntu/training/benchmark_gemma3n_asr.py

#!/usr/bin/env python3
"""
Benchmark Gemma3n-E2B-ASR checkpoints on indic-asr-benchmark-6k dataset.

Usage:
    python benchmark_gemma3n_asr.py \
        --checkpoint /home/ubuntu/training/checkpoints/gemma3n-e2b-ckpt-10000 \
        --checkpoint-name ckpt-10000 \
        --output-dir /home/ubuntu/training/benchmark_results_gemma3n \
        --batch-size 1

Metrics computed per language and overall:
    - WER  (Word Error Rate)
    - CER  (Character Error Rate)
    - WER-norm (WER after unicode/punctuation/space normalization)
    - CER-norm (CER after normalization)
"""

import argparse
import json
import os
import re
import sys
import time
import unicodedata
from collections import defaultdict
from pathlib import Path

import numpy as np
import torch
import torchaudio
from datasets import load_dataset as hf_load_dataset
from jiwer import cer, wer
from transformers import AutoModelForImageTextToText, AutoProcessor

# Language code → display name mapping
LANG_CODE_TO_NAME = {
    "as": "Assamese", "bn": "Bengali", "en": "English", "gu": "Gujarati",
    "hi": "Hindi", "kn": "Kannada", "ml": "Malayalam", "mr": "Marathi",
    "or": "Odia", "pa": "Punjabi", "ta": "Tamil", "te": "Telugu",
}

# ── Text normalization (standard ASR eval practice) ──────────────────────

def normalize_text(text: str) -> str:
    """Normalize text for fair WER/CER comparison.
    Unicode NFC → lowercase → remove punctuation → collapse whitespace → strip.
    """
    text = unicodedata.normalize("NFC", text)
    text = text.lower()
    text = re.sub(r"[^\w\s]", "", text, flags=re.UNICODE)
    text = re.sub(r"\s+", " ", text).strip()
    return text

def load_audio_array(sample, target_sr=16000):
    """Extract audio as numpy float32 array at target sample rate."""
    audio = sample["audio"]
    wav = np.array(audio["array"], dtype=np.float32)
    sr = audio["sampling_rate"]

    if sr != target_sr:
        wav_tensor = torch.from_numpy(wav).unsqueeze(0)
        wav_tensor = torchaudio.functional.resample(wav_tensor, sr, target_sr)
        wav = wav_tensor.squeeze(0).numpy()

    return wav

def run_inference(model, processor, dataset, batch_size: int = 1, max_new_tokens: int = 256):
    """Run inference on HF dataset and return list of result dicts."""
    results = []
    total = len(dataset)

    for start in range(0, total, batch_size):
        end = min(start + batch_size, total)
        batch = dataset.select(range(start, end))

        for i, sample in enumerate(batch):
            lang_name = LANG_CODE_TO_NAME.get(sample["lang_code"], sample["language"].capitalize())

            # Construct prompt matching training format
            user_content = f"<audio>\nTranscribe the audio.\nLanguage: {lang_name}"
            messages = [{"role": "user", "content": user_content}]

            # Get audio array
            wav = load_audio_array(sample)

            try:
                # Apply chat template
                prompt_text = processor.apply_chat_template(
                    messages, add_generation_prompt=True, tokenize=False
                )

                # Process inputs (text + audio)
                inputs = processor(
                    text=prompt_text,
                    audios=[wav],
                    sampling_rate=16000,
                    return_tensors="pt",
                )

                # Move to model device
                inputs = {k: v.to(model.device) if isinstance(v, torch.Tensor) else v
                          for k, v in inputs.items()}

                # Generate
                with torch.no_grad():
                    output_ids = model.generate(
                        **inputs,
                        max_new_tokens=max_new_tokens,
                        do_sample=False,
                    )

                # Decode only the new tokens
                input_len = inputs["input_ids"].shape[1]
                generated_ids = output_ids[0][input_len:]
                hypothesis = processor.tokenizer.decode(generated_ids, skip_special_tokens=True).strip()

            except Exception as e:
                print(f"  Sample {start + i} ({sample['id']}) FAILED: {e}")
                hypothesis = ""

            results.append({
                "id": sample["id"],
                "language": sample["language"],
                "lang_code": sample["lang_code"],
                "source": sample["source"],
                "duration": sample["duration"],
                "reference": sample["reference"],
                "hypothesis": hypothesis,
            })

        elapsed_pct = end / total * 100
        if end % max(batch_size * 10, 10) == 0 or end == total:
            print(f"  [{end}/{total}] ({elapsed_pct:.0f}%) processed", flush=True)

    return results

def compute_metrics(results: list) -> dict:
    """Compute WER, CER, normalized-WER/CER per language and overall."""
    by_lang = defaultdict(lambda: {"refs": [], "hyps": [], "refs_norm": [], "hyps_norm": []})

    for r in results:
        lang = r["language"]
        ref = r["reference"]
        hyp = r.get("hypothesis", "")

        ref_norm = normalize_text(ref)
        hyp_norm = normalize_text(hyp)

        if not ref_norm:
            continue

        by_lang[lang]["refs"].append(ref)
        by_lang[lang]["hyps"].append(hyp)
        by_lang[lang]["refs_norm"].append(ref_norm)
        by_lang[lang]["hyps_norm"].append(hyp_norm)

    metrics = {}
    all_refs, all_hyps = [], []
    all_refs_norm, all_hyps_norm = [], []

    for lang in sorted(by_lang.keys()):
        d = by_lang[lang]
        n = len(d["refs"])
        lang_wer = wer(d["refs"], d["hyps"])
        lang_cer = cer(d["refs"], d["hyps"])
        lang_wer_norm = wer(d["refs_norm"], d["hyps_norm"])
        lang_cer_norm = cer(d["refs_norm"], d["hyps_norm"])

        empty_count = sum(1 for h in d["hyps"] if not h.strip())

        metrics[lang] = {
            "n_samples": n,
            "wer": round(lang_wer * 100, 2),
            "cer": round(lang_cer * 100, 2),
            "wer_normalized": round(lang_wer_norm * 100, 2),
            "cer_normalized": round(lang_cer_norm * 100, 2),
            "empty_hypotheses": empty_count,
        }

        all_refs.extend(d["refs"])
        all_hyps.extend(d["hyps"])
        all_refs_norm.extend(d["refs_norm"])
        all_hyps_norm.extend(d["hyps_norm"])

    # Overall (corpus-level)
    metrics["__overall__"] = {
        "n_samples": len(all_refs),
        "wer": round(wer(all_refs, all_hyps) * 100, 2),
        "cer": round(cer(all_refs, all_hyps) * 100, 2),
        "wer_normalized": round(wer(all_refs_norm, all_hyps_norm) * 100, 2),
        "cer_normalized": round(cer(all_refs_norm, all_hyps_norm) * 100, 2),
    }

    # Macro average across languages
    lang_keys = [k for k in metrics if not k.startswith("__")]
    metrics["__macro_avg__"] = {
        "n_languages": len(lang_keys),
        "wer": round(np.mean([metrics[k]["wer"] for k in lang_keys]), 2),
        "cer": round(np.mean([metrics[k]["cer"] for k in lang_keys]), 2),
        "wer_normalized": round(np.mean([metrics[k]["wer_normalized"] for k in lang_keys]), 2),
        "cer_normalized": round(np.mean([metrics[k]["cer_normalized"] for k in lang_keys]), 2),
    }

    return metrics

def print_metrics_table(metrics: dict, checkpoint_name: str):
    """Pretty print metrics as a table."""
    print(f"\n{'='*85}")
    print(f"  BENCHMARK RESULTS: {checkpoint_name}")
    print(f"{'='*85}")
    print(f"{'Language':<14} {'N':>5} {'WER%':>8} {'CER%':>8} {'WER-n%':>8} {'CER-n%':>8}")
    print(f"{'-'*14} {'-'*5} {'-'*8} {'-'*8} {'-'*8} {'-'*8}")

    for lang in sorted(k for k in metrics if not k.startswith("__")):
        m = metrics[lang]
        print(f"{lang:<14} {m['n_samples']:>5} {m['wer']:>8.2f} {m['cer']:>8.2f} "
              f"{m['wer_normalized']:>8.2f} {m['cer_normalized']:>8.2f}")

    print(f"{'-'*14} {'-'*5} {'-'*8} {'-'*8} {'-'*8} {'-'*8}")

    ov = metrics["__overall__"]
    print(f"{'OVERALL':<14} {ov['n_samples']:>5} {ov['wer']:>8.2f} {ov['cer']:>8.2f} "
          f"{ov['wer_normalized']:>8.2f} {ov['cer_normalized']:>8.2f}")

    ma = metrics["__macro_avg__"]
    print(f"{'MACRO-AVG':<14} {ma['n_languages']:>5} {ma['wer']:>8.2f} {ma['cer']:>8.2f} "
          f"{ma['wer_normalized']:>8.2f} {ma['cer_normalized']:>8.2f}")
    print(f"{'='*85}\n")

def save_sample_analysis(results: list, output_dir: Path, n_per_lang: int = 5):
    """Save a few sample predictions per language for manual inspection."""
    by_lang = defaultdict(list)
    for r in results:
        by_lang[r["language"]].append(r)

    samples = []
    for lang in sorted(by_lang.keys()):
        lang_results = by_lang[lang]
        for r in lang_results[:n_per_lang]:
            samples.append({
                "id": r["id"],
                "language": r["language"],
                "reference": r["reference"],
                "hypothesis": r["hypothesis"],
                "ref_normalized": normalize_text(r["reference"]),
                "hyp_normalized": normalize_text(r["hypothesis"]),
            })

    path = output_dir / "sample_analysis.json"
    with open(path, "w", encoding="utf-8") as f:
        json.dump(samples, f, ensure_ascii=False, indent=2)
    print(f"Sample analysis saved: {path}")

def main():
    parser = argparse.ArgumentParser(description="Benchmark Gemma3n-E2B-ASR on Indic ASR dataset")
    parser.add_argument("--checkpoint", required=True, help="Path to model checkpoint directory")
    parser.add_argument("--checkpoint-name", default=None,
                        help="Name for this checkpoint in output files (default: dirname)")
    parser.add_argument("--hf-dataset", default="BayAreaBoys/indic-asr-benchmark-6k",
                        help="HuggingFace dataset name")
    parser.add_argument("--hf-token", default=None, help="HuggingFace token (or set HF_TOKEN env)")
    parser.add_argument("--cache-dir", default="/home/ubuntu/training/datasets/indic-asr-benchmark-6k")
    parser.add_argument("--output-dir", default="/home/ubuntu/training/benchmark_results_gemma3n")
    parser.add_argument("--batch-size", type=int, default=1,
                        help="Samples per batch (1 for sequential inference)")
    parser.add_argument("--max-new-tokens", type=int, default=256)
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--languages", nargs="+", default=None,
                        help="Filter to specific language codes (e.g. hi en ta)")
    args = parser.parse_args()

    ckpt_name = args.checkpoint_name or Path(args.checkpoint).name
    output_dir = Path(args.output_dir) / ckpt_name
    output_dir.mkdir(parents=True, exist_ok=True)

    # Load model and processor
    print(f"Loading checkpoint: {args.checkpoint}")
    t0 = time.time()

    processor = AutoProcessor.from_pretrained(args.checkpoint, trust_remote_code=True)
    model = AutoModelForImageTextToText.from_pretrained(
        args.checkpoint,
        torch_dtype=torch.bfloat16,
        device_map=args.device,
        trust_remote_code=True,
    )
    model.eval()

    load_time = time.time() - t0
    print(f"Model loaded in {load_time:.1f}s")

    # Load dataset
    hf_token = args.hf_token or os.environ.get("HF_TOKEN")
    print(f"Loading dataset: {args.hf_dataset}")
    dataset = hf_load_dataset(
        args.hf_dataset, token=hf_token, cache_dir=args.cache_dir, split="train"
    )

    # Filter languages if specified
    if args.languages:
        dataset = dataset.filter(lambda x: x["lang_code"] in args.languages)
        print(f"Filtered to {len(dataset)} samples for languages: {args.languages}")

    print(f"Dataset: {len(dataset)} samples")

    # Run inference
    print(f"\nRunning inference (batch_size={args.batch_size})...")
    t0 = time.time()
    results = run_inference(model, processor, dataset, args.batch_size, args.max_new_tokens)
    inference_time = time.time() - t0

    total_audio_sec = sum(r["duration"] for r in results)
    rtf = inference_time / total_audio_sec if total_audio_sec > 0 else 0
    print(f"Inference done in {inference_time:.1f}s for {total_audio_sec:.0f}s audio (RTF={rtf:.4f})")

    # Save raw predictions
    preds_file = output_dir / "predictions.json"
    with open(preds_file, "w", encoding="utf-8") as f:
        json.dump(results, f, ensure_ascii=False, indent=2)
    print(f"Predictions saved: {preds_file}")

    # Compute metrics
    metrics = compute_metrics(results)
    metrics["__meta__"] = {
        "checkpoint": args.checkpoint,
        "checkpoint_name": ckpt_name,
        "model_type": "gemma3n-E2B-asr",
        "dataset": args.hf_dataset,
        "batch_size": args.batch_size,
        "inference_time_sec": round(inference_time, 2),
        "total_audio_sec": round(total_audio_sec, 2),
        "rtf": round(rtf, 4),
    }

    metrics_file = output_dir / "metrics.json"
    with open(metrics_file, "w") as f:
        json.dump(metrics, f, indent=2, ensure_ascii=False)
    print(f"Metrics saved: {metrics_file}")

    # Print table
    print_metrics_table(metrics, ckpt_name)

    # Save sample analysis
    save_sample_analysis(results, output_dir)

if __name__ == "__main__":
    main()


Now let me check the download status and see if there are any issues with the checkpoint files — the model safetensors needs to be named correctly.


Bash

cat /tmp/claude-1001/-home-ubuntu-training/5e207b1e-98a7-43a4-9947-9579680cf6bd/tasks/b5dil7b5b.output 2>&1 | tail -5

Check R2 download status