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