from transformers import AutoFeatureExtractor, AutoTokenizer

pretrained_name = 'openai/whisper-base'
feature_extractor = AutoFeatureExtractor.from_pretrained(pretrained_name)
tokenizer = AutoTokenizer.from_pretrained(pretrained_name)

def prepare_dataset(
    sample: dict,
    labels_max_len: int = None,
):
    sample = sample['audio']
    inputs = feature_extractor(sample["array"], sampling_rate=sample["sampling_rate"])

    sample['input_features'] = inputs.get('input_features')[0]
    sample["input_length"] = len(sample["array"])

    input_str = sample['sentence']
    
    sample['labels'] = tokenizer(input_str).input_ids
    sample['labels_length'] = len(sample['labels'])  # include special characters

    sample['labels_truncated'] = 0
    # need to truncate validation and test labels that are longer that model.config.max_length.
    # can't drop such examples because this will affect validation and test scores.
    # thus need to truncate.
    if labels_max_len is not None:
        if len(sample['labels']) > labels_max_len:
            sample['labels'] = sample['labels'][:labels_max_len]
            sample['labels_truncated'] = 1

    return sample