pushing code to the hub

README.md

@@ -30,7 +30,6 @@ It is based on the [mHuBERT-147](https://huggingface.co/utter-project/mHuBERT-14
 
 ## Training Parameters
 The training parameters are available in config.yaml.
-We downsample the commonvoice dataset to 70,000 utterances.
 
 ## ASR Model class
 
@@ -41,4 +40,15 @@ The code is available in [CTC_model.py](https://huggingface.co/naver/mHuBERT-147
 ## Running inference
 
 The run_asr.py file illustrates how to load the model for inference (**load_asr_model**), and how to produce transcription for a file (**run_asr_inference**).
-Please follow the [requirements file](https://huggingface.co/naver/mHuBERT-147-ASR-fr/blob/main/requirements.txt) to avoid incorrect model loading.
+Please follow the [requirements file](https://huggingface.co/naver/mHuBERT-147-ASR-fr/blob/main/requirements.txt) to avoid incorrect model loading.
+
+Here is a simple example of the inference loop. Please notice that the sampling rate must be 16,000Hz.
+
+```
+from inference_code.run_inference import load_asr_model, run_asr_inference
+
+model, processor = load_asr_model()
+
+prediction = run_inference(model, processor, your_audio_file)
+
+```

config.yaml

@@ -0,0 +1,24 @@
+group_by_length: True
+evaluation_strategy: "steps"
+num_train_epochs: 100
+fp16: False
+gradient_checkpointing: True
+eval_steps: 10000
+save_steps: 10000
+logging_steps: 10000
+learning_rate: 1e-4
+adam_beta1: 0.9
+adam_beta2: 0.98
+adam_epsilon: 1e-08
+warmup_ratio: 0.2
+save_total_limit: 4
+load_best_model_at_end: True
+per_device_train_batch_size: 8
+per_device_eval_batch_size: 2
+metric_for_best_model: "cer"
+greater_is_better: False
+gradient_accumulation_steps: 8
+final_dropout: 0.3
+seed: 3452
+add_interface_layer: True
+num_interface_layers: 3

inference_code/run_inference.py

@@ -0,0 +1,80 @@
+"""
+Inference main class.
+
+Author: Marcely Zanon Boito, 2024
+"""
+
+from .CTC_model import mHubertForCTC
+
+import torch
+from transformers import Wav2Vec2CTCTokenizer, Wav2Vec2FeatureExtractor, Wav2Vec2Processor
+from transformers import HubertConfig
+
+from datasets import load_dataset
+
+fbk_test_id = 'FBK-MT/Speech-MASSIVE-test'
+mhubert_id = 'utter-project/mHuBERT-147'
+
+def load_asr_model():
+	def init_config():
+		config = HubertConfig.from_pretrained(mhubert_id)
+		config.pad_token_id = processor.tokenizer.pad_token_id
+		config.ctc_token_id = processor.tokenizer.convert_tokens_to_ids('[CTC]')
+		config.vocab_size = len(processor.tokenizer)
+
+		config.output_hidden_states = False
+		config.add_interface = True
+		config.num_interface_layers = 3
+		return config
+
+	# Load the ASR model
+	tokenizer = Wav2Vec2CTCTokenizer('vocab.json', unk_token="[UNK]", pad_token="[PAD]", word_delimiter_token="|")
+	feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(mhubert_id)
+	processor = Wav2Vec2Processor(feature_extractor=feature_extractor, tokenizer=tokenizer)
+
+	config = init_config()
+	model = mHubertForCTC.from_pretrained("naver/mHuBERT-147-ASR-fr", config=config)
+	model.eval()
+	return model, processor
+
+def run_asr_inference(model, processor, example):
+	audio = processor(example["array"], sampling_rate=example["sampling_rate"]).input_values[0]
+	input_values = torch.tensor(audio).unsqueeze(0) 
+
+	with torch.no_grad():
+		logits = model(input_values).logits
+	
+	pred_ids = torch.argmax(logits, dim=-1)
+
+	prediction = processor.batch_decode(pred_ids)[0].replace('[CTC]', "")
+	return prediction
+
+if __name__ == '__main__':
+
+	# Load the dataset in streaming mode
+	dataset = load_dataset(fbk_test_id, 'fr-FR', streaming=True)
+	dataset = dataset['test']
+	generator = iter(dataset)
+	
+	# load model
+	model, processor = load_asr_model()
+	print(model)
+
+	# decode 10 examples from speech-MASSIVE
+	num_examples= 10
+	while num_examples >= 0:
+		example = next(generator)
+		
+		prediction = run_inference(model, processor, example['audio'])
+		
+		gold_standard = example['utt']
+
+		print("Gold standard:", gold_standard)
+		print("Prediction:", prediction)
+		print()
+		num_examples-=1
+
+
+	
+
+