Files
annescribe/transcription/transcription_helper.py
T
2024-12-29 23:43:26 +01:00

156 lines
6.4 KiB
Python

import concurrent.futures
import os
from multiprocessing import Lock
from concurrent.futures import ProcessPoolExecutor
from utils import get_logger
from transcription import job_utils
from transcription.job_utils import is_job, remove_job, create_job, get_processing, set_processing
from transcription.audio_processing import split_audio, convert_to_wav
from transcription.transcription import transcribe_audio
job_update_lock = Lock()
default_input_audio_file_name = 'input_audio'
logger = get_logger(__name__)
# define a global executor for non-blocking transcription jobs
executor = ProcessPoolExecutor()
def transcribe(audio_file_path: str,
config: dict,
blocking: bool = True,
overwrite: bool = False,
executor: concurrent.futures.Executor = executor) -> str:
"""
Transcribe an audio file.
:param audio_file_path: path to the audio file
:param config: app configuration
:param blocking: block until the job is finished
:param overwrite: overwrite existing job
:param executor: executor for non-blocking transcription
:return: name of the created job
"""
job_name = os.path.basename(audio_file_path).split('.')[0]
job_dir = os.path.join(config['jobs']['root_directory'], job_name)
logger.info(f"Transcribing {audio_file_path} to {job_dir}, blocking={blocking}, overwrite={overwrite}")
# check, if a job with the same audio file already exists
job_already_exists = is_job(job_name, config)
if job_already_exists and not overwrite:
logger.error(f'Job for {audio_file_path} already exists and overwrite is not set')
raise ValueError(f'Job for {audio_file_path} already exists')
# remove, if overwrite is set
if job_already_exists and overwrite:
logger.info(f'Removing existing job {job_name} since overwrite is set')
remove_job(job_name, config)
# create a new base job
logger.info(f'Creating job {job_name}')
create_job(job_name, config)
# copy the audio file to the job directory
logger.info(f'Copying audio file {audio_file_path} to {job_dir}')
input_audio_file_name = os.path.basename(audio_file_path).split('.')[0]
input_audio_file_type = os.path.splitext(audio_file_path)[1]
audio_file_destination = os.path.join(job_dir, f"{input_audio_file_name}{input_audio_file_type}")
os.system(f'cp {audio_file_path} {audio_file_destination}')
# convert the audio file to wav, if it is not already
if not audio_file_destination.endswith('.wav'):
logger.info(f'Converting audio file {audio_file_destination} to wav')
convert_to_wav(audio_file_path,
os.path.join(job_dir, f"{default_input_audio_file_name}.wav"))
# run the transcription job
logger.info(f'Running transcription job {job_name}')
if blocking:
logger.info(f'Starting blocking transcription job {job_name}')
run_transcription_job(job_name, config)
else:
logger.info(f'Starting async transcription job {job_name}')
executor.submit(run_transcription_job, job_name, config, 100)
logger.info(f'Started async transcription job {job_name}')
return job_name
def run_transcription_job(job_name: str,
config: dict,
num_chunks=100) -> None:
"""
Start a transcription job.
:param job_name: id of the job
:param config: app configuration
:param num_chunks: number of chunks to split the audio file into
"""
# check if the job is already running
if get_processing(job_name, config):
raise ValueError(f'Job {job_name} is already running')
# set the job to processing
set_processing(job_name, config, processing=True)
try:
# split the audio file into chunks
logger.info(f'Splitting audio file into chunks for job {job_name}')
audio_file_path = os.path.join(config['jobs']['root_directory'], job_name,
f"{default_input_audio_file_name}.wav")
audio_chunk_folder = os.path.join(config['jobs']['root_directory'], job_name, job_utils.chunk_input_subdir)
audio_chunks = split_audio(audio_file_path,
audio_chunk_folder,
num_chunks,
config["audio"]["chunking"]["min_silence_level"],
config["audio"]["chunking"]["min_silence_length"],
config["audio"]["chunking"]["ms_silence_to_keep"])
# transcribe the audio chunks
logger.info(f'Transcribing {len(audio_chunks)} audio chunks for job {job_name}')
# create chunk output folder, if not done yet
chunk_output_folder = os.path.join(config['jobs']['root_directory'], job_name, job_utils.chunk_output_subdir)
if not os.path.exists(chunk_output_folder):
os.makedirs(chunk_output_folder)
for i, audio_chunk in enumerate(audio_chunks):
# transcribe the audio
logger.info(f'Transcribing audio chunk {i + 1} of {len(audio_chunks)} for job {job_name}')
output_file_path = os.path.join(chunk_output_folder, f'output_{i}.txt')
transcribe_audio(audio_chunk,
output_file_path,
config)
# combine the output text chunks
logger.info(f'Combining output text chunks for job {job_name}')
output_chunk_folder = os.path.join(config['jobs']['root_directory'], job_name, job_utils.chunk_output_subdir)
output_text = ''
for text_chunk_file in sorted(os.listdir(output_chunk_folder), key=lambda x: int(x.split('_')[1].split('.')[0])):
with open(os.path.join(output_chunk_folder, text_chunk_file), 'r') as text_chunk:
output_text += text_chunk.read() + ' '
# remove trailing space
output_text = output_text.strip()
# write the output to a file
logger.info(f'Writing output text to file for job {job_name}')
output_file = os.path.join(config['jobs']['root_directory'], job_name, 'output.txt')
with open(output_file, 'w') as f:
f.write(output_text)
# set the job to not processing
set_processing(job_name, config, processing=False)
logger.info(f'Job {job_name} completed')
except Exception as e:
# set the job to not processing
set_processing(job_name, config, processing=False)
raise e