From c10bc56d8272a8c68944e6f7d91b0c74b2b50533 Mon Sep 17 00:00:00 2001 From: Aidan Crowther <aidan.crowther@outlook.com> Date: Fri, 13 Dec 2024 11:58:57 -0500 Subject: [PATCH] Added language detection confidence to returned JSON --- CHANGELOG.md | 4 ++++ app/faster_whisper/core.py | 3 ++- app/openai_whisper/core.py | 2 +- app/webservice.py | 4 ++-- 4 files changed, 9 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9c8c9c4..d2b4ec2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,10 @@ Changelog Unreleased ---------- +### Added + + - Added detection confidence to langauge detection endpoint + [1.6.0] (2024-10-06) -------------------- diff --git a/app/faster_whisper/core.py b/app/faster_whisper/core.py index 1d48f2d..c99e049 100644 --- a/app/faster_whisper/core.py +++ b/app/faster_whisper/core.py @@ -70,8 +70,9 @@ def language_detection(audio): with model_lock: segments, info = model.transcribe(audio, beam_size=5) detected_lang_code = info.language + detected_language_confidence = info.language_probability - return detected_lang_code + return detected_lang_code, detected_language_confidence def write_result(result: dict, file: BinaryIO, output: Union[str, None]): diff --git a/app/openai_whisper/core.py b/app/openai_whisper/core.py index 88bde4f..6ccee13 100644 --- a/app/openai_whisper/core.py +++ b/app/openai_whisper/core.py @@ -55,7 +55,7 @@ def language_detection(audio): _, probs = model.detect_language(mel) detected_lang_code = max(probs, key=probs.get) - return detected_lang_code + return detected_lang_code, probs[max(probs)] def write_result(result: dict, file: BinaryIO, output: Union[str, None]): diff --git a/app/webservice.py b/app/webservice.py index 5391afe..84efa8c 100644 --- a/app/webservice.py +++ b/app/webservice.py @@ -89,8 +89,8 @@ async def detect_language( audio_file: UploadFile = File(...), # noqa: B008 encode: bool = Query(default=True, description="Encode audio first through FFmpeg"), ): - detected_lang_code = language_detection(load_audio(audio_file.file, encode)) - return {"detected_language": tokenizer.LANGUAGES[detected_lang_code], "language_code": detected_lang_code} + detected_lang_code, confidence = language_detection(load_audio(audio_file.file, encode)) + return {"detected_language": tokenizer.LANGUAGES[detected_lang_code], "language_code": detected_lang_code, "confidence": confidence} def load_audio(file: BinaryIO, encode=True, sr: int = SAMPLE_RATE): -- GitLab