56 lines
1.7 KiB
Python
56 lines
1.7 KiB
Python
import logging
|
|
|
|
import librosa
|
|
import torch
|
|
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
|
|
|
|
# MODEL = "m3hrdadfi/wav2vec2-large-xlsr-persian"
|
|
MODEL = "/home/reza/data/huggingface-models/04.wav2vec2-large-xlsr-persian"
|
|
|
|
|
|
def initLogger(name=__name__, level=logging.DEBUG):
|
|
if name[:2] == '__' and name[-2:] == '__':
|
|
name = name[2:-2]
|
|
logger = logging.getLogger(name)
|
|
|
|
fmt = '%(asctime)s | %(levelname)-8s | %(name)s | %(message)s'
|
|
datefmt = '%Y-%m-%d %H:%M:%S'
|
|
ch = logging.StreamHandler()
|
|
ch.setLevel(logging.DEBUG)
|
|
formatter = logging.Formatter(fmt, datefmt)
|
|
ch.setFormatter(formatter)
|
|
logger.addHandler(ch)
|
|
|
|
logger.setLevel(level)
|
|
return logger
|
|
|
|
|
|
def mp3_to_text(mp3_file_path):
|
|
# Load the MP3 file and resample to 16kHz
|
|
audio, sample_rate = librosa.load(mp3_file_path, sr=16000)
|
|
logger.info("Resampling is Done!")
|
|
|
|
# Load tokenizer and model from Hugging Face
|
|
tokenizer = Wav2Vec2Processor.from_pretrained(MODEL)
|
|
model = Wav2Vec2ForCTC.from_pretrained(MODEL)
|
|
logger.info("Loading model is Done!")
|
|
|
|
# Preprocess the audio
|
|
input_values = tokenizer(audio, sampling_rate=16000, return_tensors="pt", padding="longest").input_values
|
|
logits = model(input_values).logits
|
|
logger.info("Processing the audio is Done!")
|
|
|
|
# Decode the predicted IDs
|
|
predicted_ids = torch.argmax(logits, dim=-1)
|
|
transcription = tokenizer.batch_decode(predicted_ids)
|
|
logger.info("Decoding the prodicted IDs is Done!")
|
|
|
|
return transcription[0]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
logger = initLogger('speech2text_fa', level=logging.INFO)
|
|
text = mp3_to_text("samples/sample1.wav")
|
|
print()
|
|
print(text)
|