PR_#1839 from nithinraok/add-ml-prompt, multilingual verbatim transcription prompt for Qwen3-Omni

Adds language-agnostic single-turn ASR pseudo-labeling prompt for non-English
audio. Unlike the English two-turn flow (transcription + disfluency followup),
this prompt combines transcription and verbatim fidelity into one instruction,
making the followup turn unnecessary for ML languages.

- examples/audio/qwen_omni_inprocess/prompts/ml_qwen3_omni_disfluency_asr.md
  (uses {language} placeholder)
- nemo_curator/models/qwen_omni.py: _resolve_prompt() helper + thread
  language through _build_messages, _build_turn2_messages, _prepare_single,
  _prepare_batch, _prepare_turn2_single, _prepare_turn2_batch, generate()
- nemo_curator/stages/audio/inference/qwen_omni.py: source_lang_key field
  pulls per-sample language from manifest and passes to model.generate()
- examples/audio/qwen_omni_inprocess/run_pipeline.py: --source_lang_key CLI

Surgical squash cherry-pick of #1839 (additive bits only).
Skipped FastTextLIDStage source_lang_key (would conflict with PR #1's
source-tracking refactor) and initialize_fields drop (already handled). #NO_PR

Signed-off-by: George Zelenfroynd <gzelenfroind@nvidia.com>

Author: Jorjeous

examples/audio/qwen_omni_inprocess/prompts/ml_qwen3_omni_disfluency_asr.md

@@ -0,0 +1 @@
+Transcribe the {language} audio into text exactly as the speaker says it. Write numbers as spoken words.

examples/audio/qwen_omni_inprocess/run_pipeline.py

@@ -113,6 +113,9 @@ def _build_arg_parser() -> argparse.ArgumentParser:
                     help="Path to regex substitution rules YAML.")
     tf.add_argument("--target_lang", type=str, default="en",
                     help="Expected language code for LID filtering.")
+    tf.add_argument("--source_lang_key", type=str, default="",
+                    help="Per-sample language key in manifest for {language} prompt substitution. "
+                         "Empty (default) disables per-sample language threading.")
     tf.add_argument("--min_lang_prob", type=float, default=0.8,
                     help="Minimum FastText language probability to keep an entry.")
     tf.add_argument("--unique_words_threshold", type=float, default=0.4,
@@ -224,6 +227,7 @@ def main() -> None:
             max_num_seqs=args.max_num_seqs,
             gpu_memory_utilization=args.gpu_memory_utilization,
             prep_workers=args.prep_workers,
+            source_lang_key=args.source_lang_key,
             pred_text_key="qwen3_prediction_s1",
             disfluency_text_key="qwen3_prediction_s2",
             keep_waveform=bool(args.asr_model_id),

nemo_curator/models/qwen_omni.py

@@ -161,49 +161,66 @@ def _resample(waveform: np.ndarray, orig_sr: int, target_sr: int = _QWEN_SAMPLE_
 
         return librosa.resample(waveform, orig_sr=orig_sr, target_sr=target_sr)
 
-    def _build_messages(self, waveform: np.ndarray) -> list[dict[str, Any]]:
-        """Build Turn 1 chat messages with an in-memory waveform (numpy array at 16 kHz)."""
+    def _resolve_prompt(self, template: str, language: str | None) -> str:
+        """Replace ``{language}`` placeholder if *language* is provided."""
+        if language and template and "{language}" in template:
+            return template.replace("{language}", language)
+        return template
+
+    def _build_messages(self, waveform: np.ndarray, language: str | None = None) -> list[dict[str, Any]]:
+        """Build Turn 1 chat messages with an in-memory waveform (numpy array at 16 kHz).
+
+        Prompts may contain a ``{language}`` placeholder which is replaced
+        with *language* (e.g., ``"French"``) when provided.
+        """
+        prompt = self._resolve_prompt(self.prompt_text, language)
         messages: list[dict[str, Any]] = []
         if self.system_prompt:
-            messages.append({"role": "system", "content": [{"type": "text", "text": self.system_prompt}]})
+            sys_prompt = self._resolve_prompt(self.system_prompt, language)
+            messages.append({"role": "system", "content": [{"type": "text", "text": sys_prompt}]})
         messages.append({
             "role": "user",
             "content": [
-                {"type": "text", "text": self.prompt_text},
+                {"type": "text", "text": prompt},
                 {"type": "audio", "audio": waveform},
             ],
         })
         return messages
 
-    def _build_turn2_messages(self, waveform: np.ndarray, pred_text: str) -> list[dict[str, Any]]:
-        """Build Turn 2 messages: full Turn 1 conversation history + follow-up promt."""
+    def _build_turn2_messages(
+        self, waveform: np.ndarray, pred_text: str, language: str | None = None,
+    ) -> list[dict[str, Any]]:
+        """Build Turn 2 messages: full Turn 1 conversation history + follow-up prompt."""
+        prompt = self._resolve_prompt(self.prompt_text, language)
+        followup = self._resolve_prompt(self.followup_prompt, language)
         messages: list[dict[str, Any]] = []
         if self.system_prompt:
-            messages.append({"role": "system", "content": [{"type": "text", "text": self.system_prompt}]})
+            sys_prompt = self._resolve_prompt(self.system_prompt, language)
+            messages.append({"role": "system", "content": [{"type": "text", "text": sys_prompt}]})
         messages.append({
             "role": "user",
             "content": [
-                {"type": "text", "text": self.prompt_text},
+                {"type": "text", "text": prompt},
                 {"type": "audio", "audio": waveform},
             ],
         })
         messages.append({"role": "assistant", "content": [{"type": "text", "text": pred_text}]})
         messages.append({
             "role": "user",
             "content": [
-                {"type": "text", "text": self.followup_prompt},
+                {"type": "text", "text": followup},
             ],
         })
         return messages
 
     def _prepare_single(
-        self, waveform: np.ndarray, sample_rate: int,
+        self, waveform: np.ndarray, sample_rate: int, language: str | None = None,
     ) -> tuple[dict[str, Any], np.ndarray] | None:
         from qwen_omni_utils import process_mm_info
 
         try:
             waveform_16k = self._resample(waveform, sample_rate)
-            messages = self._build_messages(waveform_16k)
+            messages = self._build_messages(waveform_16k, language)
             text = self._processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
             audios, images, videos = process_mm_info(messages, use_audio_in_video=False)
         except Exception:  # noqa: BLE001
@@ -227,18 +244,23 @@ def _prepare_batch(
         self,
         waveforms: list[np.ndarray],
         sample_rates: list[int],
+        languages: list[str | None] | None = None,
     ) -> list[tuple[dict[str, Any], np.ndarray] | None]:
+        langs = languages if languages is not None else [None] * len(waveforms)
         if self._prep_pool is None:
-            return [self._prepare_single(w, sr) for w, sr in zip(waveforms, sample_rates, strict=False)]
-        return list(self._prep_pool.map(self._prepare_single, waveforms, sample_rates))
+            return [
+                self._prepare_single(w, sr, lang)
+                for w, sr, lang in zip(waveforms, sample_rates, langs, strict=False)
+            ]
+        return list(self._prep_pool.map(self._prepare_single, waveforms, sample_rates, langs))
 
     def _prepare_turn2_single(
-        self, waveform_16k: np.ndarray, pred_text: str,
+        self, waveform_16k: np.ndarray, pred_text: str, language: str | None = None,
     ) -> dict[str, Any] | None:
         from qwen_omni_utils import process_mm_info
 
         try:
-            messages = self._build_turn2_messages(waveform_16k, pred_text)
+            messages = self._build_turn2_messages(waveform_16k, pred_text, language)
             text = self._processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
             audios, images, videos = process_mm_info(messages, use_audio_in_video=False)
         except Exception:  # noqa: BLE001
@@ -262,13 +284,15 @@ def _prepare_turn2_batch(
         self,
         waveforms_16k: list[np.ndarray],
         pred_texts: list[str],
+        languages: list[str | None] | None = None,
     ) -> list[dict[str, Any] | None]:
+        langs = languages if languages is not None else [None] * len(waveforms_16k)
         if self._prep_pool is None:
             return [
-                self._prepare_turn2_single(w, pt)
-                for w, pt in zip(waveforms_16k, pred_texts, strict=False)
+                self._prepare_turn2_single(w, pt, lang)
+                for w, pt, lang in zip(waveforms_16k, pred_texts, langs, strict=False)
             ]
-        return list(self._prep_pool.map(self._prepare_turn2_single, waveforms_16k, pred_texts))
+        return list(self._prep_pool.map(self._prepare_turn2_single, waveforms_16k, pred_texts, langs))
 
     # ------------------------------------------------------------------
     # Generation
@@ -278,6 +302,7 @@ def generate(
         self,
         waveforms: list[np.ndarray],
         sample_rates: list[int],
+        languages: list[str | None] | None = None,
     ) -> tuple[list[str], list[str]]:
         """Run batched two-turn inference on in-memory audio waveforms.
 
@@ -288,6 +313,9 @@ def generate(
         Args:
             waveforms: List of 1-D mono numpy float32 arrays.
             sample_rates: Corresponding sample rates for each waveform.
+            languages: Optional per-sample language strings for ``{language}``
+                placeholder substitution in prompts. Length must match
+                ``waveforms``. Pass ``None`` (default) to skip substitution.
 
         Returns:
             ``(pred_texts, disfluency_texts)`` — one string per input for
@@ -301,7 +329,7 @@ def generate(
         n = len(waveforms)
 
         # -- Turn 1 ----------------------------------------------------------
-        prepared = self._prepare_batch(waveforms, sample_rates)
+        prepared = self._prepare_batch(waveforms, sample_rates, languages)
         valid_indices = [i for i, p in enumerate(prepared) if p is not None]
         valid_inputs = [prepared[i][0] for i in valid_indices]
         waveforms_16k: dict[int, np.ndarray] = {i: prepared[i][1] for i in valid_indices}
@@ -327,9 +355,13 @@ def generate(
         if not t2_indices:
             return pred_texts, [""] * n
 
+        t2_languages = (
+            [languages[i] for i in t2_indices] if languages is not None else None
+        )
         t2_prepared = self._prepare_turn2_batch(
             [waveforms_16k[i] for i in t2_indices],
             [pred_texts[i] for i in t2_indices],
+            t2_languages,
         )
 
         t2_valid = [(i, p) for i, p in zip(t2_indices, t2_prepared, strict=False) if p is not None]

nemo_curator/stages/audio/inference/qwen_omni.py

@@ -71,6 +71,7 @@ class InferenceQwenOmniStage(ProcessingStage[AudioTask, AudioTask]):
     system_prompt: str | None = None
     waveform_key: str = "waveform"
     sample_rate_key: str = "sample_rate"
+    source_lang_key: str = ""
     pred_text_key: str = "qwen3_prediction_s1"
     disfluency_text_key: str = "qwen3_prediction_s2"
     max_model_len: int = 32768
@@ -166,8 +167,11 @@ def process_batch(self, tasks: list[AudioTask]) -> list[AudioTask]:
 
         waveforms = [t.data[self.waveform_key] for t in tasks]
         sample_rates = [t.data[self.sample_rate_key] for t in tasks]
+        languages: list[str | None] | None = None
+        if self.source_lang_key:
+            languages = [t.data.get(self.source_lang_key) or None for t in tasks]
 
-        pred_texts, disfluency_texts = self._model.generate(waveforms, sample_rates)
+        pred_texts, disfluency_texts = self._model.generate(waveforms, sample_rates, languages)
 
         for task, pred, disfl in zip(tasks, pred_texts, disfluency_texts, strict=True):
             task.data[self.pred_text_key] = pred