From 02cd2e67a33a5e748ee1e17d2a88f6d7e1b77e0a Mon Sep 17 00:00:00 2001
From: Matjaz Debelak <matjaz.debelak@bcc.no>
Date: Fri, 30 Aug 2024 15:34:15 +0200
Subject: [PATCH] Respect the choice of model via API

Fallback to large-v3 if the privided model is not valid
---
 transcription-handler.go |  2 +-
 transcription-runner.go  | 19 ++++++++++++++++++-
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/transcription-handler.go b/transcription-handler.go
index 2eed0f6..3392721 100644
--- a/transcription-handler.go
+++ b/transcription-handler.go
@@ -66,7 +66,7 @@ func (h *handlers) SubmitJob(c *gin.Context) {
 	}
 
 	if job.Model == "" {
-		job.Model = "large-v2"
+		job.Model = "openai/whisper-large-v3"
 	}
 
 	if job.Priority >= 500 {
diff --git a/transcription-runner.go b/transcription-runner.go
index f1feb6a..eceaff1 100644
--- a/transcription-runner.go
+++ b/transcription-runner.go
@@ -11,6 +11,7 @@ import (
 	"time"
 
 	"github.com/bcc-code/mediabank-bridge/log"
+	"github.com/samber/lo"
 )
 
 func doCallback(job *Job) {
@@ -47,7 +48,23 @@ func runJob(job *Job) {
 		return
 	}
 
-	cmd := exec.Command("python3", "bcc-whisper/main.py", "-l", job.Language, "-m", "openai/whisper-large-v3", job.Path, job.OutputPath)
+	model := "openai/whisper-large-v3"
+	if lo.Contains([]string{
+		"openai/whisper-large-v2",
+		"openai/whisper-large-v3",
+		"openai/whisper-large",
+		"openai/whisper-medium",
+		"openai/whisper-small",
+		"openai/whisper-tiny",
+		"NbAiLab/nb-whisper-large",
+		"NbAiLab/nb-whisper-medium",
+		"NbAiLab/nb-whisper-small",
+		"NbAiLab/nb-whisper-tiny",
+	}, job.Model) {
+		model = job.Model
+	}
+
+	cmd := exec.Command("python3", "bcc-whisper/main.py", "-l", job.Language, "-m", model, job.Path, job.OutputPath)
 	cmd.Env = append(os.Environ(), "PYTHONUNBUFFERED=1")
 
 	stderr, _ := cmd.StderrPipe()