aleger's picture
add bentoml files
d8d26b1
import os
import tempfile
import bentoml
from bentoml.io import JSON, File
from runners.audio_transcriber import AudioTranscriber
from runners.keyword_extractor import KeywordExtractor
from runners.transcription_zipper import TranscriptionZipper
from runners.audio_amplitude import AudioAmplitude
runner_audio_transcriber = bentoml.Runner(
AudioTranscriber,
name="audio_transcriber",
)
runner_audio_amplitude = bentoml.Runner(
AudioAmplitude,
name="audio_amplitude",
)
runner_keyword_extractor = bentoml.Runner(
KeywordExtractor,
name="keyword_extractor",
)
runner_transcription_zipper = bentoml.Runner(
TranscriptionZipper,
name="transcription_zipper"
)
svc = bentoml.Service(
"speech_to_text_pipeline",
runners=[
runner_audio_transcriber,
runner_audio_amplitude,
runner_keyword_extractor,
runner_transcription_zipper,
],
)
@svc.api(input=File(), output=JSON())
async def process_uploaded_file(file):
with tempfile.NamedTemporaryFile(delete=False) as temp_file:
temp_file.write(file.read())
temp_file_path = temp_file.name
transcript, info, segment_info = await runner_audio_transcriber.transcribe_audio_faster.async_run(temp_file_path)
# amplitudes = await runner_audio_amplitude.get_audio_amplitude.async_run(temp_file_path)
output = {"file_name": file.name, "transcript": transcript, "language": info.language,
"file_size": os.stat(temp_file_path).st_size, "segments": segment_info}
return output
@svc.api(input=JSON(), output=File())
async def zip_transcription(transcription):
zip_file = await runner_transcription_zipper.zip_transcription.async_run(transcription)
return zip_file