Update app.py

app.py

@@ -15,9 +15,9 @@ from whisper_jax import FlaxWhisperPipline
 app = FastAPI(title="Whisper JAX: The Fastest Whisper API ⚡️")
 
 logger = logging.getLogger("whisper-jax-app")
-logger.setLevel(logging.INFO)
+logger.setLevel(logging.DEBUG)
 ch = logging.StreamHandler()
-ch.setLevel(logging.INFO)
+ch.setLevel(logging.DEBUG)
 formatter = logging.Formatter("%(asctime)s;%(levelname)s;%(message)s", "%Y-%m-%d %H:%M:%S")
 ch.setFormatter(formatter)
 logger.addHandler(ch)
@@ -37,7 +37,7 @@ stride_left = stride_right = round(stride_length_s * pipeline.feature_extractor.
 step = chunk_len - stride_left - stride_right
 
 # do a pre-compile step so that the first user to use the demo isn't hit with a long transcription time
-logger.info("compiling forward call...")
+logger.debug("Compiling forward call...")
 start = time.time()
 random_inputs = {
     "input_features": np.ones(
@@ -46,11 +46,11 @@ random_inputs = {
 }
 random_timestamps = pipeline.forward(random_inputs, batch_size=BATCH_SIZE, return_timestamps=True)
 compile_time = time.time() - start
-logger.info(f"compiled in {compile_time}s")
+logger.debug(f"Compiled in {compile_time}s")
 
 @app.post("/transcribe_audio")
 async def transcribe_chunked_audio(audio_file: UploadFile, task: str = "transcribe", return_timestamps: bool = False):
-    logger.info("loading audio file...")
+    logger.debug("Loading audio file...")
     if not audio_file:
         logger.warning("No audio file")
         raise HTTPException(status_code=400, detail="No audio file submitted!")
@@ -59,30 +59,60 @@ async def transcribe_chunked_audio(audio_file: UploadFile, task: str = "transcri
         logger.warning("Max file size exceeded")
         raise HTTPException(status_code=400, detail=f"File size exceeds file size limit. Got file of size {file_size_mb:.2f}MB for a limit of {FILE_LIMIT_MB}MB.")
 
-    with open(audio_file.filename, "rb") as f:
-        inputs = f.read()
+    try:
+        with open(audio_file.filename, "rb") as f:
+            inputs = f.read()
+    except Exception as e:
+        logger.error("Error reading audio file:", exc_info=True)
+        raise HTTPException(status_code=500, detail="Error reading audio file")
 
     inputs = ffmpeg_read(inputs, pipeline.feature_extractor.sampling_rate)
     inputs = {"array": inputs, "sampling_rate": pipeline.feature_extractor.sampling_rate}
-    logger.info("done loading")
-    text, runtime = tqdm_generate(inputs, task=task, return_timestamps=return_timestamps)
+    logger.debug("Done loading audio file")
+
+    try:
+        text, runtime = tqdm_generate(inputs, task=task, return_timestamps=return_timestamps)
+    except Exception as e:
+        logger.error("Error transcribing audio:", exc_info=True)
+        raise HTTPException(status_code=500, detail="Error transcribing audio")
+
     return {"text": text, "runtime": runtime}
 
 @app.post("/transcribe_youtube")
 async def transcribe_youtube(yt_url: str = Form(...), task: str = "transcribe", return_timestamps: bool = False):
-    logger.info("loading youtube file...")
-    html_embed_str = _return_yt_html_embed(yt_url)
+    logger.debug("Loading YouTube file...")
+    try:
+        html_embed_str = _return_yt_html_embed(yt_url)
+    except Exception as e:
+        logger.error("Error generating YouTube HTML embed:", exc_info=True)
+        raise HTTPException(status_code=500, detail="Error generating YouTube HTML embed")
+
     with tempfile.TemporaryDirectory() as tmpdirname:
         filepath = os.path.join(tmpdirname, "video.mp4")
-        download_yt_audio(yt_url, filepath)
+        try:
+            logger.debug("Downloading YouTube audio...")
+            download_yt_audio(yt_url, filepath)
+        except Exception as e:
+            logger.error("Error downloading YouTube audio:", exc_info=True)
+            raise HTTPException(status_code=500, detail="Error downloading YouTube audio")
 
-        with open(filepath, "rb") as f:
-            inputs = f.read()
+        try:
+            with open(filepath, "rb") as f:
+                inputs = f.read()
+        except Exception as e:
+            logger.error("Error reading downloaded audio file:", exc_info=True)
+            raise HTTPException(status_code=500, detail="Error reading downloaded audio file")
 
     inputs = ffmpeg_read(inputs, pipeline.feature_extractor.sampling_rate)
     inputs = {"array": inputs, "sampling_rate": pipeline.feature_extractor.sampling_rate}
-    logger.info("done loading...")
-    text, runtime = tqdm_generate(inputs, task=task, return_timestamps=return_timestamps)
+    logger.debug("Done loading YouTube file")
+
+    try:
+        text, runtime = tqdm_generate(inputs, task=task, return_timestamps=return_timestamps)
+    except Exception as e:
+        logger.error("Error transcribing YouTube audio:", exc_info=True)
+        raise HTTPException(status_code=500, detail="Error transcribing YouTube audio")
+
     return {"html_embed": html_embed_str, "text": text, "runtime": runtime}
 
 def tqdm_generate(inputs: dict, task: str, return_timestamps: bool):
@@ -94,15 +124,19 @@ def tqdm_generate(inputs: dict, task: str, return_timestamps: bool):
     dataloader = pipeline.preprocess_batch(inputs, chunk_length_s=CHUNK_LENGTH_S, batch_size=BATCH_SIZE)
     model_outputs = []
     start_time = time.time()
-    logger.info("transcribing...")
+    logger.debug("Transcribing...")
     # iterate over our chunked audio samples - always predict timestamps to reduce hallucinations
     for batch in dataloader:
         model_outputs.append(pipeline.forward(batch, batch_size=BATCH_SIZE, task=task, return_timestamps=True))
     runtime = time.time() - start_time
-    logger.info("done transcription")
+    logger.debug("Done transcription")
 
-    logger.info("post-processing...")
-    post_processed = pipeline.postprocess(model_outputs, return_timestamps=True)
+    logger.debug("Post-processing...")
+    try:
+        post_processed = pipeline.postprocess(model_outputs, return_timestamps=True)
+    except Exception as e:
+        logger.error("Error post-processing transcription:", exc_info=True)
+        raise HTTPException(status_code=500, detail="Error post-processing transcription")
     text = post_processed["text"]
     if return_timestamps:
         timestamps = post_processed.get("chunks")
@@ -111,7 +145,7 @@ def tqdm_generate(inputs: dict, task: str, return_timestamps: bool):
             for chunk in timestamps
         ]
         text = "\n".join(str(feature) for feature in timestamps)
-    logger.info("done post-processing")
+    logger.debug("Done post-processing")
     return text, runtime
 
 def _return_yt_html_embed(yt_url):
@@ -125,8 +159,10 @@ def _return_yt_html_embed(yt_url):
 def download_yt_audio(yt_url, filename):
     info_loader = youtube_dl.YoutubeDL()
     try:
+        logger.debug(f"Extracting info for YouTube URL: {yt_url}")
         info = info_loader.extract_info(yt_url, download=False)
     except youtube_dl.utils.DownloadError as err:
+        logger.error("Error extracting YouTube info:", exc_info=True)
         raise HTTPException(status_code=400, detail=str(err))
 
     file_length = info["duration_string"]
@@ -146,8 +182,10 @@ def download_yt_audio(yt_url, filename):
     ydl_opts = {"outtmpl": filename, "format": "worstvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best"}
     with youtube_dl.YoutubeDL(ydl_opts) as ydl:
         try:
+            logger.debug(f"Downloading YouTube audio to {filename}")
             ydl.download([yt_url])
         except youtube_dl.utils.ExtractorError as err:
+            logger.error("Error downloading YouTube audio:", exc_info=True)
             raise HTTPException(status_code=400, detail=str(err))
 
 def format_timestamp(seconds: float, always_include_hours: bool = False, decimal_marker: str = "."):