Skip to content
Snippets Groups Projects
Commit c10bc56d authored by Aidan Crowther's avatar Aidan Crowther
Browse files

Added language detection confidence to returned JSON

parent 539f5852
No related branches found
No related tags found
No related merge requests found
...@@ -4,6 +4,10 @@ Changelog ...@@ -4,6 +4,10 @@ Changelog
Unreleased Unreleased
---------- ----------
### Added
- Added detection confidence to langauge detection endpoint
[1.6.0] (2024-10-06) [1.6.0] (2024-10-06)
-------------------- --------------------
......
...@@ -70,8 +70,9 @@ def language_detection(audio): ...@@ -70,8 +70,9 @@ def language_detection(audio):
with model_lock: with model_lock:
segments, info = model.transcribe(audio, beam_size=5) segments, info = model.transcribe(audio, beam_size=5)
detected_lang_code = info.language 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]): def write_result(result: dict, file: BinaryIO, output: Union[str, None]):
......
...@@ -55,7 +55,7 @@ def language_detection(audio): ...@@ -55,7 +55,7 @@ def language_detection(audio):
_, probs = model.detect_language(mel) _, probs = model.detect_language(mel)
detected_lang_code = max(probs, key=probs.get) 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]): def write_result(result: dict, file: BinaryIO, output: Union[str, None]):
......
...@@ -89,8 +89,8 @@ async def detect_language( ...@@ -89,8 +89,8 @@ async def detect_language(
audio_file: UploadFile = File(...), # noqa: B008 audio_file: UploadFile = File(...), # noqa: B008
encode: bool = Query(default=True, description="Encode audio first through FFmpeg"), encode: bool = Query(default=True, description="Encode audio first through FFmpeg"),
): ):
detected_lang_code = language_detection(load_audio(audio_file.file, encode)) 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} 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): def load_audio(file: BinaryIO, encode=True, sr: int = SAMPLE_RATE):
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment