Spaces:
Sleeping
Sleeping
| import pandas as pd | |
| import numpy as np | |
| import os | |
| import sys | |
| from models.parakeet import parakeet_ctc_model, parakeet_ctc_process | |
| from models.es_fastconformer import stt_es_model, stt_es_process | |
| from pydub import AudioSegment | |
| import json | |
| class Data_Pipeline(): | |
| """ | |
| A class to handle data processing and transcription using various models. | |
| Args: | |
| wav_dir (str): Directory containing the .wav files. | |
| Example: "path/to/wav_files/" | |
| lang (str): Language of the audio files, either 'en' for English or 'es' for Spanish. | |
| Example: "en" | |
| """ | |
| def __init__(self, wav_dir, lang): | |
| self.wav_dir = wav_dir | |
| self.lang = lang | |
| def get_wav_files(self): | |
| """ | |
| Retrieve all .wav files from the specified directory. | |
| Returns: | |
| list: A list of paths to the .wav files. | |
| Example: ["path/to/file1.wav", "path/to/file2.wav"] | |
| """ | |
| self.wav_files = [] | |
| for root, dirs, files in os.walk(self.wav_dir): | |
| for file in files: | |
| if file.endswith('.wav'): | |
| # make the path into the full path to the file | |
| file = os.path.join(self.wav_dir, file) | |
| self.wav_files.append(file) | |
| return self.wav_files | |
| def get_combined_wav_lengths(self): | |
| """ | |
| Calculate the total length of all .wav files in the directory. | |
| Returns: | |
| float: The total length of all .wav files in seconds. | |
| Example: 123.45 | |
| """ | |
| # returns a float number of the total length of all the wav files in the directory | |
| total_length = 0 | |
| for file in self.wav_files: | |
| if file.endswith('.wav'): | |
| audio = AudioSegment.from_file(file) | |
| total_length += len(audio) | |
| return total_length/1000 | |
| def estimated_transcription_time(self): | |
| """ | |
| Estimate the transcription time based on the real-time factor. | |
| Returns: | |
| float: The estimated transcription time in seconds. | |
| Example: 185.175 | |
| str: Error message if the language is not supported. | |
| Example: "Error: Language not supported." | |
| """ | |
| # returns a float number of the estimated time based on the real time factor | |
| # real time factor is 1.5 | |
| if self.lang == "en": | |
| return self.get_combined_wav_lengths() * 1.5 | |
| elif self.lang == "es": | |
| return self.get_combined_wav_lengths() * 1.5 | |
| return "Error: Language not supported." | |
| def load_models(self): | |
| """ | |
| Load the appropriate ASR model based on the language. | |
| Returns: | |
| object: The loaded ASR model. | |
| Example: <nemo.collections.asr.models.EncDecCTCModelBPE object> | |
| str: Error message if the language is not supported. | |
| Example: "Error: Language not supported." | |
| """ | |
| if self.lang == "en": | |
| self.parakeet_model = parakeet_ctc_model() | |
| return self.parakeet_model | |
| elif self.lang == "es": | |
| self.es_model = stt_es_model() | |
| return self.es_model | |
| return "Error: Language not supported." | |
| def en_transcribe(self, audio_file): | |
| """ | |
| Transcribe an English audio file using the Parakeet CTC model. | |
| Args: | |
| audio_file (str): Path to the audio file. | |
| Example: "path/to/audio_file.wav" | |
| Returns: | |
| list: A list containing the transcribed text. | |
| Example: ["transcribed text"] | |
| """ | |
| text = parakeet_ctc_process(self.parakeet_model, audio_file) | |
| return text | |
| def es_transcribe(self, audio_file): | |
| """ | |
| Transcribe a Spanish audio file using the FastConformer model. | |
| Args: | |
| audio_file (str): Path to the audio file. | |
| Example: "path/to/audio_file.wav" | |
| Returns: | |
| list: A list containing the transcribed text. | |
| Example: ["transcribed text"] | |
| """ | |
| text = stt_es_process(self.es_model, audio_file) | |
| return text | |
| def read_transcriptions(self, json_path): | |
| """ | |
| Read transcriptions from a JSON file. | |
| Args: | |
| json_path (str): Path to the JSON file. | |
| Example: "path/to/data.json" | |
| Returns: | |
| dict: The data read from the JSON file. | |
| Example: {"text": ["text1", "text2"], "original_path": ["path1", "path2"]} | |
| """ | |
| # read the json file | |
| with open(json_path) as f: | |
| self.data = json.load(f) | |
| return self.data | |
| def get_transcription(self, file_path): | |
| """ | |
| Get the transcription for a specific file from the JSON data. | |
| Args: | |
| file_path (str): Path to the original audio file. | |
| Example: "path/to/audio_file.wav" | |
| Returns: | |
| str: The transcription for the specified file. | |
| Example: "This is the transcription." | |
| str: Error message if no transcription is found. | |
| Example: "Error: No transcription found." | |
| """ | |
| # the json file has the following keys: text, original_path, path_to_save, language, order, original_text | |
| # get the "original_text" of the element that has the "original_path" equal to the file_path | |
| for i in range(len(self.data['original_path'])): | |
| if self.data['original_path'][i] == file_path: | |
| return self.data['original_text'][i] | |
| return "Error: No transcription found." | |
| def data_formatter_with_models(self): | |
| """ | |
| Format data by transcribing audio files using the appropriate models. | |
| Returns: | |
| pd.DataFrame: A DataFrame containing the transcriptions. | |
| Example: pd.DataFrame({'wav_file': ["file1", "file2"], 'transcription': ["text1", "text2"], 'transcription2': ["text1", "text2"], 'speaker_name': ["user0", "user0"]}) | |
| str: Error message if the language is not supported. | |
| Example: "Error: Language not supported." | |
| """ | |
| self.transcriptions_df = pd.DataFrame(columns = ['wav_file', 'transcription','transcription2' ]) | |
| if self.lang == "en": | |
| self.load_models() | |
| self.get_wav_files() | |
| for file in self.wav_files: | |
| if file.endswith('.wav'): | |
| transcription = parakeet_ctc_process(self.parakeet_model, file) | |
| # append transcriptions_df with the wav_file and transcription | |
| self.transcriptions_df = self.transcriptions_df.append({'wav_file': file, 'transcription': transcription[0], 'transcription2': transcription[0],'speaker_name': "user0"}, ignore_index=True) | |
| return self.transcriptions_df | |
| elif self.lang == "es": | |
| self.load_models() | |
| self.get_wav_files() | |
| for file in self.wav_files: | |
| if file.endswith('.wav'): | |
| self.transcriptions_df = stt_es_process(self.es_model, file) | |
| # make the path into the full path to the file | |
| file = os.path.join(self.wav_dir, file) | |
| # append transcriptions_df with the wav_file and transcription | |
| self.transcriptions_df = self.transcriptions_df.append({'wav_file': file, 'transcription': transcription[0], 'transcription2': transcription[0], 'speaker_name':"user0"}, ignore_index=True) | |
| return self.transcriptions_df | |
| return "Error: Language not supported." | |
| def data_formatter_without_models(self): | |
| """ | |
| Format data by retrieving transcriptions from a JSON file and transcribing any missing data. | |
| Returns: | |
| pd.DataFrame: A DataFrame containing the transcriptions. | |
| Example: pd.DataFrame({'wav_file': ["file1", "file2"], 'transcription': ["text1", "text2"], 'transcription2': ["text1", "text2"], 'speaker_name': ["user0", "user0"]}) | |
| """ | |
| self.transcriptions_df = pd.DataFrame(columns = ['wav_file', 'transcription','transcription2' ]) | |
| self.get_wav_files() | |
| for file in self.wav_files: | |
| if file.endswith('.wav'): | |
| transcription = self.get_transcription(file) | |
| if transcription == "Error: No transcription found." and self.lang == "en": | |
| transcription = parakeet_ctc_process(self.parakeet_model, file) | |
| elif transcription == "Error: No transcription found." and self.lang == "es": | |
| transcription = stt_es_process(self.es_model, file) | |
| # make the path into the full path to the file | |
| #file = os.path.join(self.wav_dir, file) | |
| # append transcriptions_df with the wav_file and transcription | |
| self.transcriptions_df = self.transcriptions_df.append({'wav_file': file, 'transcription': transcription[0], 'transcription2': transcription[0], 'speaker_name': "user0"}, ignore_index=True) | |
| return self.transcriptions_df | |
| def save_transcriptions(self, output_file): | |
| """ | |
| Save the transcriptions to CSV files, splitting into training and evaluation datasets. | |
| Args: | |
| output_file (str): Base path for the output CSV files. | |
| Example: "path/to/output" | |
| Returns: | |
| tuple: A tuple containing a success message and paths to the training and evaluation CSV files. | |
| Example: ("Data saved successfully.", "path/to/output_train.csv", "path/to/output_eval.csv") | |
| """ | |
| # split the data into two data, train and eval data | |
| from sklearn.model_selection import train_test_split | |
| train_data, eval_data = train_test_split(self.transcriptions_df, test_size=0.2, random_state=42) | |
| # save the data into csv files | |
| self.path_to_train_data = output_file + "_train.csv" | |
| self.path_to_eval_data = output_file + "_eval.csv" | |
| train_data.to_csv(output_file + "_train.csv" , index=False, sep='|') | |
| eval_data.to_csv(output_file + "_eval.csv", index=False, sep='|') | |
| return "Data saved successfully.", self.path_to_train_data, self.path_to_eval_data | |
| def get_paths(self): | |
| """ | |
| Retrieve the paths to the training and evaluation CSV files and the .wav files directory. | |
| Returns: | |
| tuple: A tuple containing paths to the training data, evaluation data, and .wav files directory. | |
| Example: ("path/to/train.csv", "path/to/eval.csv", "path/to/wav_files/") | |
| """ | |
| # csv files, wav files directory | |
| return self.path_to_train_data, self.path_to_eval_data, self.wav_dir | |