import torch
import numpy as np
from transformers import WhisperProcessor
from dataset import DataCollatorSpeechSeq2SeqWithPadding

def debug_collator():
    print("Loading processor...")
    # using a standard model name just for tokenizer/processor loading
    model_name = "openai/whisper-tiny" 
    processor = WhisperProcessor.from_pretrained(model_name)
    
    collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)
    
    # Create dummy features simulating dataset output
    # input_features: (80, 3000) - standard whisper mel spec
    dummy_features = torch.randn(80, 3000)
    dummy_labels = torch.tensor([1, 2, 3, 4, 50257])
    
    features = [
        {
            "input_features": dummy_features,
            "labels": dummy_labels,
            "text": "dummy text",
            "language": "hindi",
            "is_code_switched": False
        },
        {
            "input_features": dummy_features,
            "labels": torch.tensor([1, 2, 3]),
            "text": "dummy text 2",
            "language": "hindi",
            "is_code_switched": False
        }
    ]
    
    print("Running collator...")
    batch = collator(features)
    
    print("Batch keys:", batch.keys())
    for k, v in batch.items():
        if torch.is_tensor(v):
            print(f"{k}: shape={v.shape}")
        else:
            print(f"{k}: {type(v)}")

if __name__ == "__main__":
    debug_collator()