| import os | |
| import sys | |
| import torch | |
| import argparse | |
| from halo import Halo | |
| import soundfile as sf | |
| from transformers import AutoModelForCTC, AutoProcessor, Wav2Vec2Processor | |
| path_this_file = os.path.dirname(os.path.abspath(__file__)) | |
| pat_project_root = os.path.join(path_this_file, "..") | |
| sys.path.append(pat_project_root) | |
| class Wav2Vec2Inference: | |
| def __init__(self,model_name, hotwords=[], use_lm_if_possible=True, use_gpu=True): | |
| """ | |
| Initializes the class with the provided parameters. | |
| Args: | |
| model_name (str): The name of the model to be used. | |
| hotwords (list, optional): A list of hotwords. Defaults to an empty list. | |
| use_lm_if_possible (bool, optional): Specifies whether to use a language model if possible. | |
| Defaults to True. | |
| use_gpu (bool, optional): Specifies whether to use the GPU. Defaults to True. | |
| Returns: | |
| None | |
| """ | |
| self.device = "cuda" if use_gpu and torch.cuda.is_available() else "cpu" | |
| if use_lm_if_possible: | |
| self.processor = AutoProcessor.from_pretrained(model_name) | |
| else: | |
| self.processor = Wav2Vec2Processor.from_pretrained(model_name) | |
| self.model = AutoModelForCTC.from_pretrained(model_name) | |
| self.model.to(self.device) | |
| self.hotwords = hotwords | |
| self.use_lm_if_possible = use_lm_if_possible | |
| def buffer_to_text(self, audio_buffer): | |
| """ | |
| Transcribes the given audio buffer into text. | |
| Args: | |
| audio_buffer (list): A list representing the audio buffer. | |
| Returns: | |
| tuple: A tuple containing the transcribed text (str) and the confidence score (float). | |
| """ | |
| spinner = Halo(text="Transcribing audio...", spinner="dots") | |
| spinner.start() | |
| try: | |
| if len(audio_buffer) == 0: | |
| return "" | |
| inputs = self.processor(torch.tensor(audio_buffer), sampling_rate=16_000, return_tensors="pt", padding=True) | |
| with torch.no_grad(): | |
| logits = self.model(inputs.input_values.to(self.device), | |
| attention_mask=inputs.attention_mask.to(self.device)).logits | |
| if hasattr(self.processor, 'decoder') and self.use_lm_if_possible: | |
| transcription = \ | |
| self.processor.decode(logits[0].cpu().numpy(), | |
| hotwords=self.hotwords, | |
| output_word_offsets=True, | |
| ) | |
| confidence = transcription.lm_score / len(transcription.text.split(" ")) | |
| transcription = transcription.text | |
| else: | |
| predicted_ids = torch.argmax(logits, dim=-1) | |
| transcription = self.processor.batch_decode(predicted_ids)[0] | |
| confidence = self.confidence_score(logits,predicted_ids) | |
| spinner.succeed("Audio transcribed successfully!") | |
| return transcription, confidence.item() | |
| except Exception as e: | |
| spinner.fail(f"Error during transcription: {str(e)}") | |
| return "", 0.0 | |
| def confidence_score(self, logits, predicted_ids): | |
| """ | |
| Calculate the confidence score for the predicted IDs based on the logits. | |
| Parameters: | |
| logits (torch.Tensor): The logits tensor. | |
| predicted_ids (torch.Tensor): The predicted IDs tensor. | |
| Returns: | |
| float: The average confidence score for the predicted IDs. | |
| """ | |
| scores = torch.nn.functional.softmax(logits, dim=-1) | |
| pred_scores = scores.gather(-1, predicted_ids.unsqueeze(-1))[:, :, 0] | |
| mask = torch.logical_and( | |
| predicted_ids.not_equal(self.processor.tokenizer.word_delimiter_token_id), | |
| predicted_ids.not_equal(self.processor.tokenizer.pad_token_id)) | |
| character_scores = pred_scores.masked_select(mask) | |
| total_average = torch.sum(character_scores) / len(character_scores) | |
| return total_average | |
| def file_to_text(self, filename): | |
| """ | |
| Reads an audio file and converts it to text using the buffer_to_text method. | |
| Parameters: | |
| filename (str): The path to the audio file. | |
| Returns: | |
| tuple: A tuple containing the transcription (str) and the confidence (float) of the transcription. If there is an error reading the audio file, an empty string and a confidence of 0.0 will be returned. | |
| """ | |
| spinner = Halo(text="Reading audio file...", spinner="dots") | |
| spinner.start() | |
| try: | |
| audio_input, samplerate = sf.read(filename) | |
| assert samplerate == 16000 | |
| transcription, confidence = self.buffer_to_text(audio_input) | |
| spinner.succeed("File read successfully!") | |
| return transcription, confidence | |
| except Exception as e: | |
| spinner.fail(f"Error reading audio file: {str(e)}") | |
| return "", 0.0 | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--model_name", type=str, default="arifagustyawan/wav2vec2-large-xlsr-common_voice_13_0-id") | |
| parser.add_argument("--filename", type=str, default="assets/halo.wav") | |
| args = parser.parse_args() | |
| with Halo(text="Initializing Wav2Vec2 Inference...", spinner="dots") as init_spinner: | |
| try: | |
| asr = Wav2Vec2Inference(args.model_name) | |
| init_spinner.succeed("Wav2Vec2 Inference initialized successfully!") | |
| except Exception as e: | |
| init_spinner.fail(f"Error initializing Wav2Vec2 Inference: {str(e)}") | |
| sys.exit(1) | |
| with Halo(text="Performing audio transcription...", spinner="dots") as transcribe_spinner: | |
| transcription, confidence = asr.file_to_text(args.filename) | |
| print("\033[94mTranscription:\033[0m", transcription) | |
| print("\033[94mConfidence:\033[0m", confidence) | |