fix: dual-mode imports + complete ChatMLDatasetSample + transcribe output parsing

transcribe.py

@@ -0,0 +1,147 @@
+"""Standalone transcription helper for Higgs Audio v3 models.
+
+Bundles all preprocessing needed to run inference without boson_multimodal.
+
+Usage:
+    from transcribe import transcribe
+    
+    model = AutoModel.from_pretrained("bosonai/higgs-audio-v3-8b-stt", ...)
+    tokenizer = AutoTokenizer.from_pretrained("bosonai/higgs-audio-v3-8b-stt")
+    
+    text = transcribe(model, tokenizer, audio_np, sample_rate=16000)
+"""
+
+import torch
+import numpy as np
+from dataclasses import dataclass, asdict, field
+from typing import List, Optional
+from transformers import AutoConfig, WhisperProcessor
+
+try:
+    from .higgs_audio_collator import HiggsAudioSampleCollator, ChatMLDatasetSample
+except ImportError:
+    from higgs_audio_collator import HiggsAudioSampleCollator, ChatMLDatasetSample
+
+
+def _build_input_tokens(tokenizer, user_prompt, enable_thinking=True):
+    """Build input token sequence for transcription."""
+    # Construct ChatML format manually
+    audio_token = "<|audio_in_bos|><|AUDIO|><|audio_in_eos|>"
+    
+    messages = [
+        {"role": "user", "content": f"{user_prompt}\n{audio_token}"}
+    ]
+    
+    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    if enable_thinking:
+        text += "<think>\n"
+    
+    input_ids = tokenizer.encode(text, add_special_tokens=False)
+    return input_ids
+
+
+def _create_collator(config):
+    """Create audio collator from model config."""
+    whisper_proc = WhisperProcessor.from_pretrained("openai/whisper-large-v3")
+
+    return HiggsAudioSampleCollator(
+        whisper_processor=whisper_proc,
+        audio_in_token_id=config.audio_in_token_idx,
+        audio_out_token_id=config.audio_out_token_idx,
+        audio_stream_bos_id=config.audio_stream_bos_id,
+        audio_stream_eos_id=config.audio_stream_eos_id,
+        encode_whisper_embed=config.encode_whisper_embed,
+        pad_token_id=config.pad_token_id,
+        return_audio_in_tokens=config.encode_audio_in_tokens,
+        use_delay_pattern=config.use_delay_pattern,
+        round_to=1,
+        audio_num_codebooks=config.audio_num_codebooks,
+        chunk_size_seconds=getattr(config, "chunk_size_seconds", 30),
+        encoder_padding_method=getattr(config, "encoder_padding_method", "max_length"),
+    )
+
+
+_cached_collator = None
+
+
+def transcribe(model, tokenizer, audio_np, sample_rate=16000,
+               user_prompt="Transcribe the speech. Output only the spoken words in lowercase with no punctuation.",
+               enable_thinking=True, max_new_tokens=1024):
+    """Transcribe audio using a Higgs Audio model.
+    
+    Args:
+        model: loaded HiggsAudio3Model
+        tokenizer: loaded tokenizer
+        audio_np: numpy array of audio samples (float32, mono)
+        sample_rate: audio sample rate (default 16000)
+        user_prompt: transcription prompt
+        enable_thinking: use chain-of-thought (recommended)
+        max_new_tokens: max generation length
+    
+    Returns:
+        str: transcribed text
+    """
+    global _cached_collator
+    
+    device = next(model.parameters()).device
+    config = model.config
+    
+    # Resample if needed
+    if sample_rate != 16000:
+        import torchaudio
+        audio_tensor = torch.tensor(audio_np).unsqueeze(0)
+        audio_tensor = torchaudio.functional.resample(audio_tensor, sample_rate, 16000)
+        audio_np = audio_tensor.squeeze(0).numpy()
+    
+    audio_np = np.array(audio_np, dtype=np.float32)
+    
+    # Build input tokens
+    input_ids = _build_input_tokens(tokenizer, user_prompt, enable_thinking)
+    
+    # Create sample
+    sample = ChatMLDatasetSample(
+        input_ids=torch.LongTensor(input_ids),
+        label_ids=torch.LongTensor([-100] * len(input_ids)),
+        audio_waveforms_concat=torch.tensor(audio_np, dtype=torch.float32),
+        audio_waveforms_start=torch.tensor([0]),
+        audio_sample_rate=torch.tensor([sample_rate]),
+        audio_speaker_indices=torch.tensor([0]),
+    )
+    
+    # Collate
+    if _cached_collator is None:
+        _cached_collator = _create_collator(config)
+    
+    batch = asdict(_cached_collator([sample]))
+    batch = {k: v.to(device).contiguous() if isinstance(v, torch.Tensor) else v
+             for k, v in batch.items()}
+    
+    # Generate
+    import re
+    with torch.inference_mode():
+        outputs = model.generate(
+            **batch,
+            max_new_tokens=max_new_tokens,
+            use_cache=True,
+            do_sample=False,
+            stop_strings=["<|im_end|>", "<|endoftext|>"],
+            tokenizer=tokenizer,
+        )
+
+    output_ids = outputs[0] if isinstance(outputs, tuple) else outputs
+    full_text = tokenizer.decode(output_ids[0], skip_special_tokens=False)
+
+    # Extract assistant response
+    parts = full_text.split("assistant\n")
+    hyp = parts[-1] if len(parts) > 1 else full_text
+
+    # Remove complete <think>...</think> blocks (keep text outside)
+    hyp = re.sub(r"<think>.*?</think>", "", hyp, flags=re.DOTALL)
+
+    # Handle open <think> without </think> — the transcription IS the think content
+    if "<think>" in hyp:
+        hyp = hyp[hyp.index("<think>") + len("<think>"):]
+
+    # Remove special tokens
+    hyp = re.sub(r"<\|.*?\|>", "", hyp)
+    return hyp.strip()