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
Unreleased
----------
### Added
- Added detection confidence to langauge detection endpoint
[1.6.0] (2024-10-06)
--------------------
......
......@@ -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]):
......
......@@ -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]):
......
......@@ -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):
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment