learning_rate: 0.001
train_batch_size: 8
eval_batch_size: 8
seed: 42
gradient_accumulation_steps: 2
total_train_batch_size: 16
optimizer: Use OptimizerNames.ADAMW_TORCH_FUSED with betas=(0.9,0.999) and epsilon=1e-08 and optimizer_args=No additional optimizer arguments
lr_scheduler_type: linear
lr_scheduler_warmup_steps: 50
num_epochs: 5
mixed_precision_training: Native AMP

Framework versions

Transformers 4.56.2
Pytorch 2.8.0+cu126
Datasets 4.1.1
Tokenizers 0.22.0

infrance

python

# ==========================================================
# STEP 0: INSTALL LIBRARIES
# # ==========================================================
# print("Installing necessary libraries...")
# !pip install -q --upgrade transformers datasets accelerate torch torchaudio soundfile librosa

import torch
import librosa
import numpy as np
import os
from transformers import WhisperProcessor, WhisperForConditionalGeneration

# # --- (Optional) Mount Google Drive ---
# try:
#     from google.colab import drive
#     drive.mount('/content/drive')
#     print("Google Drive mounted successfully.")
# except ImportError:
#     print("Not in a Google Colab environment. Skipping Google Drive mount.")


# ==========================================================
# STEP 1: DEFINE THE FINAL, ROBUST, PRODUCTION-READY MODEL CLASS
# ==========================================================
class WhisperForSpeakerASR(WhisperForConditionalGeneration):
    def __init__(self, config):
        super().__init__(config)
        self.diarization_head = torch.nn.Linear(config.d_model, 2)
        # ### THE FINAL, DEFINITIVE FIX - PART 1 ###
        # Create a temporary storage to capture the logits during generation
        self.latest_diarization_logits = None

    def forward(
        self,
        input_features=None,
        labels=None,
        diarization_labels=None,
        **kwargs,
    ):
        # We need to get the encoder_hidden_states.
        if "encoder_outputs" in kwargs and kwargs["encoder_outputs"] is not None:
            encoder_hidden_states = kwargs["encoder_outputs"][0]
        else:
            if input_features is None:
                raise ValueError("input_features must be provided when encoder_outputs is not.")
            encoder_outputs = self.model.encoder(
                input_features=input_features, output_hidden_states=True
            )
            encoder_hidden_states = encoder_outputs.last_hidden_state

        # Run the diarization head on the encoder states
        diarization_logits = self.diarization_head(encoder_hidden_states)
        
        # ### THE FINAL, DEFINITIVE FIX - PART 2 ###
        # Store the logits so we can access them after calling .generate()
        self.latest_diarization_logits = diarization_logits

        # We now call the ORIGINAL forward pass of the base Whisper model.
        kwargs["encoder_outputs"] = (encoder_hidden_states,)
        outputs = super().forward(
            input_features=input_features,
            labels=labels,
            **kwargs,
        )

        # Attach our custom outputs to the result for consistency (optional but good practice)
        outputs.diarization_logits = diarization_logits

        # Handle custom loss during training
        if diarization_labels is not None:
            loss_fct = torch.nn.CrossEntropyLoss()
            reshaped_logits = diarization_logits.view(-1, 2)
            reshaped_labels = diarization_labels.view(-1)
            diarization_loss = loss_fct(reshaped_logits, reshaped_labels)
            outputs.loss = diarization_loss if outputs.loss is None else outputs.loss + diarization_loss
            
        return outputs

# ==========================================================
# STEP 2: SETUP MODEL AND INFERENCE FUNCTION
# ==========================================================
MODEL_HUB_PATH = "Ahmed107/whisper-small-speaker-asr-final"

def test_and_format_from_path(audio_path):
    print("="*60)
    print("🚀 Starting Final Inference Test...")
    print("="*60)

    if not os.path.exists(audio_path):
        print(f"❌ ERROR: Audio file not found: {audio_path}")
        return

    print(f"Loading model from: {MODEL_HUB_PATH}")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    processor = WhisperProcessor.from_pretrained(MODEL_HUB_PATH)
    model = WhisperForSpeakerASR.from_pretrained(MODEL_HUB_PATH).to(device).eval()

    print(f"Loading audio file: {audio_path}")
    audio_array, _ = librosa.load(audio_path, sr=16000)
    print("Audio loaded successfully.")

    chunk_duration_s = 30
    samples_per_chunk = chunk_duration_s * 16000
    time_per_token = chunk_duration_s / 1500
    speaker_map = {0: "SPEAKER_A", 1: "SPEAKER_B"}

    print("Running model for transcription and diarization...")
    full_transcription = ""
    diarization_timeline = []
    
    with torch.no_grad():
        for i in range(0, len(audio_array), samples_per_chunk):
            chunk_start_time = i / 16000
            chunk_audio = audio_array[i : i + samples_per_chunk]
            if len(chunk_audio) < samples_per_chunk:
                chunk_audio = np.pad(chunk_audio, (0, samples_per_chunk - len(chunk_audio)))

            inputs = processor(chunk_audio, return_tensors="pt", sampling_rate=16000)
            input_features = inputs.input_features.to(device)
            
            # --- Transcription will now work correctly ---
            generated_ids = model.generate(input_features, language="en")
            transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
            full_transcription += transcription + " "
            
            # ### THE FINAL, DEFINITIVE FIX - PART 3 ###
            # Instead of calling the model again, retrieve the stored logits
            predicted_ids = torch.argmax(model.latest_diarization_logits, dim=-1).squeeze(0).cpu()

            # Process timeline logic (remains the same)
            current_speaker_id, segment_start_token = -1, 0
            for token_idx, speaker_id in enumerate(predicted_ids):
                speaker_id = speaker_id.item()
                if speaker_id != current_speaker_id:
                    if current_speaker_id != -1:
                        end_time = chunk_start_time + (token_idx * time_per_token)
                        start_time = chunk_start_time + (segment_start_token * time_per_token)
                        diarization_timeline.append({"speaker": speaker_map.get(current_speaker_id, "UNKNOWN"), "start": start_time, "end": end_time})
                    current_speaker_id, segment_start_token = speaker_id, token_idx
            if current_speaker_id != -1:
                end_time = chunk_start_time + (len(predicted_ids) * time_per_token)
                start_time = chunk_start_time + (segment_start_token * time_per_token)
                diarization_timeline.append({"speaker": speaker_map.get(current_speaker_id, "UNKNOWN"), "start": start_time, "end": end_time})
    
    # Merge adjacent segments logic (remains the same)
    merged_timeline = []
    if diarization_timeline:
        merged_timeline.append(diarization_timeline[0])
        for segment in diarization_timeline[1:]:
            if segment["speaker"] == merged_timeline[-1]["speaker"] and segment["start"] - merged_timeline[-1]["end"] < 0.1:
                merged_timeline[-1]["end"] = segment["end"]
            else:
                merged_timeline.append(segment)

    # Final output for README
    print("\n\n" + "="*25 + " COPY BELOW THIS LINE " + "="*25)
    print("\n### Example Output\n")
    print("```text")
    print("Full Transcription:")
    print(full_transcription.strip())
    print("\nDiarization Timeline:")
    for segment in merged_timeline:
        if segment['start'] < segment['end']:
            print(f"[{segment['start']:0>6.2f}s - {segment['end']:0>6.2f}s] {segment['speaker']}")
    print("```")
    print("\n" + "="*26 + " COPY ABOVE THIS LINE " + "="*26 + "\n")

# ==========================================================
# STEP 3: SPECIFY YOUR FILE PATH AND RUN THE TEST
# ==========================================================
# --- ⬇️⬇️⬇️ EDIT THIS LINE ⬇️⬇️⬇️ ---
AUDIO_FILE_PATH = "/content/audio.wav"
# --- ⬆️⬆️⬆️ EDIT THIS LINE ⬆️⬆️⬆️ ---

# Run the function
test_and_format_from_path(AUDIO_FILE_PATH)