feat: word timestamps & languages

LICENSE

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2024, Doppeltilde
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

src/routes/api/asr.py

@@ -4,23 +4,36 @@
 import io
 import traceback
 import time
+from src.shared.shared import LANGUAGE_CODES
 
 router = APIRouter()
 
 
-@router.post("/api/auto-asr", dependencies=[Depends(get_api_key)])
+@router.post("/api/asr", dependencies=[Depends(get_api_key)])
 async def asr(
     file: UploadFile = File(),
     model_name: str = Query(None),
+    word_timestamps: bool = Query(False),
+    language: str = Query(
+        None,
+        enum=LANGUAGE_CODES,
+        description="If not provided, will fallback to detected language.",
+    ),
 ):
+
     start_time = time.time()
     model = asr_model(model_name)
 
     try:
         content = await file.read()
         audio_file = io.BytesIO(content)
 
-        segments, info = model.transcribe(audio_file, beam_size=5, word_timestamps=True)
+        segments, info = model.transcribe(
+            audio_file,
+            beam_size=5,
+            word_timestamps=True if word_timestamps else False,
+            language="en" if language is None else language,
+        )
 
         segment_words = []
         entire_transcription = ""
@@ -32,10 +45,11 @@ async def asr(
                 trimmed_segment_text = segment.text
             entire_transcription += trimmed_segment_text
 
-            for word in segment.words:
-                segment_words.append(
-                    {"start": word.start, "end": word.end, "word": word.word}
-                )
+            if segment.words is not None:
+                for word in segment.words:
+                    segment_words.append(
+                        {"start": word.start, "end": word.end, "word": word.word}
+                    )
 
         return {
             "res": {

src/shared/shared.py

@@ -31,3 +31,107 @@ def asr_model(model_name):
     except Exception as e:
         print(e)
         return {"error": str(e)}
+
+
+LANGUAGE_CODES = (
+    "af",
+    "am",
+    "ar",
+    "as",
+    "az",
+    "ba",
+    "be",
+    "bg",
+    "bn",
+    "bo",
+    "br",
+    "bs",
+    "ca",
+    "cs",
+    "cy",
+    "da",
+    "de",
+    "el",
+    "en",
+    "es",
+    "et",
+    "eu",
+    "fa",
+    "fi",
+    "fo",
+    "fr",
+    "gl",
+    "gu",
+    "ha",
+    "haw",
+    "he",
+    "hi",
+    "hr",
+    "ht",
+    "hu",
+    "hy",
+    "id",
+    "is",
+    "it",
+    "ja",
+    "jw",
+    "ka",
+    "kk",
+    "km",
+    "kn",
+    "ko",
+    "la",
+    "lb",
+    "ln",
+    "lo",
+    "lt",
+    "lv",
+    "mg",
+    "mi",
+    "mk",
+    "ml",
+    "mn",
+    "mr",
+    "ms",
+    "mt",
+    "my",
+    "ne",
+    "nl",
+    "nn",
+    "no",
+    "oc",
+    "pa",
+    "pl",
+    "ps",
+    "pt",
+    "ro",
+    "ru",
+    "sa",
+    "sd",
+    "si",
+    "sk",
+    "sl",
+    "sn",
+    "so",
+    "sq",
+    "sr",
+    "su",
+    "sv",
+    "sw",
+    "ta",
+    "te",
+    "tg",
+    "th",
+    "tk",
+    "tl",
+    "tr",
+    "tt",
+    "uk",
+    "ur",
+    "uz",
+    "vi",
+    "yi",
+    "yo",
+    "zh",
+    "yue",
+)