model parameter added

app/api/endpoints/transcribe.py

@@ -6,10 +6,14 @@ from pydantic import BaseModel
 
 from app.core.database import SessionLocal
 
-from app.utils.utils import save_audio_file, transcribe_file, get_audio_duration
+from app.utils.utils import (
+    save_audio_file,
+    transcribe_file,
+    get_audio_duration,
+    get_model_name,
+)
 from app.core.models import AuthTokenController, TranscribeController
 
-
 router = APIRouter()
 database = SessionLocal()
 
@@ -24,12 +28,13 @@ async def post_audio(
     background_tasks: BackgroundTasks,
     request: Request,
     file: UploadFile = File(...),
+    model: str = "tiny.en.q5",
     Authentication: Annotated[Union[str, None], Header()] = None,
 ):
     try:
         userId = AuthTokenController(database).get_userid_from_token(Authentication)
         file_path = save_audio_file(file)
-        [data, output_audio_path] = transcribe_file(file_path)
+        [data, output_audio_path] = transcribe_file(file_path, get_model_name(model))
         background_tasks.add_task(
             create_transcribe_record, database, userId, data, output_audio_path
         )

app/utils/constant.py

@@ -0,0 +1,5 @@
+model_names = {
+    "tiny.en": "ggml-tiny.en.bin",
+    "tiny.en.q5": "ggml-model-whisper-tiny.en-q5_1.bin",
+    "base.en.q5": "ggml-model-whisper-base.en-q5_1.bin",
+}

app/utils/utils.py

@@ -4,6 +4,8 @@ import uuid
 import logging
 import wave
 
+from .constant import model_names
+
 
 def get_all_routes(app):
     routes = []
@@ -88,3 +90,13 @@ def get_audio_duration(audio_file):
         rounded_duration = int(round(duration, 0))
 
     return rounded_duration
+
+
+def get_model_name(model: str = None):
+    if model is None:
+        model_names["tiny.en.q5"]
+
+    if model in model_names.keys():
+        return model_names[model]
+
+    return model_names["tiny.en.q5"]