hubert-base-los-2k

from datasets import load_dataset, Audio
import torch
from transformers import AutoFeatureExtractor, AutoModel

#Load the dataset
dataset = load_dataset("projecte-aina/ib3_ca_asr", split='train[:1%]', trust_remote_code=True)

#Downsample to 16kHz
dataset = dataset.cast_column("audio", Audio(sampling_rate=16_000))

# Hugginface pre-trained model path
MODEL_NAME = "BSC-LT/hubert-base-los-2k"

# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using {device} device.")

# Load feature extractor
feature_extractor = AutoFeatureExtractor.from_pretrained(MODEL_NAME)

# Load model
model = AutoModel.from_pretrained(MODEL_NAME)
model = model.to(device)

def map_to_speech_representations(batch):

  #Process the dataset
  audio = batch["audio"]
  input_features = feature_extractor(audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt").input_values
  input_features = input_features.to(device)

  # Extract HuBERT's Speech Representations
  with torch.no_grad():
    outputs = model(
        input_features,
        output_hidden_states = True,
        )
    speech_representations = outputs.last_hidden_state
    hidden_states = outputs.hidden_states

  batch["speech_representations"] = speech_representations
  batch["hidden_states"] = hidden_states

  return batch

dataset = dataset.map(map_to_speech_representations)

print(dataset)