Configurer le beam size via WHISPER_BEAM_SIZE

app.py

@@ -18,6 +18,9 @@ LOG_LEVEL = os.getenv("WHISPER_LOG_LEVEL", "INFO").upper()
 VAD_ENABLED = env_bool("WHISPER_VAD", True)
 MODEL_SIZE = str(os.getenv("MODEL_SIZE", "small").lower())
 COMPUTE_TYPE = os.getenv("WHISPER_COMPUTE_TYPE", "float32").strip().lower()
+BEAM_SIZE = int(os.getenv("WHISPER_BEAM_SIZE", "1"))
+if BEAM_SIZE < 1:
+    raise ValueError("WHISPER_BEAM_SIZE doit être supérieur ou égal à 1")
 MIN_AUDIO_SEC = float(os.getenv("MIN_AUDIO_SEC", "2.0"))
 MAX_AUDIO_SEC = float(os.getenv("MAX_AUDIO_SEC", "10.0"))
 
@@ -34,6 +37,7 @@ logger.setLevel(getattr(logging, LOG_LEVEL, logging.INFO))
 logger.info(f"Log level = {LOG_LEVEL}")
 logger.info(f"VAD enabled = {VAD_ENABLED}")
 logger.info("Model = %s, compute type = %s", MODEL_SIZE, COMPUTE_TYPE)
+logger.info("Beam size = %d", BEAM_SIZE)
 
 app = FastAPI()
 
@@ -42,7 +46,7 @@ app = FastAPI()
 # =========================
 model = WhisperModel(MODEL_SIZE, device="cuda", compute_type=COMPUTE_TYPE)
 # Vérifier les kernels CUDA avant de déclarer le service prêt.
-list(model.transcribe(np.zeros(16000, dtype=np.float32), language="fr", beam_size=1)[0])
+list(model.transcribe(np.zeros(16000, dtype=np.float32), language="fr", beam_size=BEAM_SIZE)[0])
 model_lock = Lock()
 
 def transcribe_audio(audio):
@@ -50,13 +54,13 @@ def transcribe_audio(audio):
     with model_lock:
         segments, info = model.transcribe(
             audio, language="fr", vad_filter=VAD_ENABLED,
-            beam_size=1, condition_on_previous_text=False,
+            beam_size=BEAM_SIZE, condition_on_previous_text=False,
         )
         return " ".join(seg.text for seg in segments).strip(), info.language
 
 @app.get("/healthz")
 def healthz():
-    return {"status": "ready", "model": MODEL_SIZE, "compute_type": COMPUTE_TYPE}
+    return {"status": "ready", "model": MODEL_SIZE, "compute_type": COMPUTE_TYPE, "beam_size": BEAM_SIZE}
 
 # =========================
 # REST API (Swagger / tests)