Update app.py

app.py

@@ -12,15 +12,6 @@ from transformers.models.whisper.tokenization_whisper import TO_LANGUAGE_CODE
 from transformers.pipelines.audio_utils import ffmpeg_read
 from whisper_jax import FlaxWhisperPipline
 
-cc.initialize_cache("./jax_cache")
-checkpoint = "openai/whisper-large-v3"
-
-BATCH_SIZE = 32
-CHUNK_LENGTH_S = 30
-NUM_PROC = 32
-FILE_LIMIT_MB = 10000
-YT_LENGTH_LIMIT_S = 15000  # limit to 2 hour YouTube files
-
 app = FastAPI(title="Whisper JAX: The Fastest Whisper API ⚡️")
 
 logger = logging.getLogger("whisper-jax-app")
@@ -31,6 +22,14 @@ formatter = logging.Formatter("%(asctime)s;%(levelname)s;%(message)s", "%Y-%m-%d
 ch.setFormatter(formatter)
 logger.addHandler(ch)
 
+checkpoint = "openai/whisper-large-v3"
+
+BATCH_SIZE = 32
+CHUNK_LENGTH_S = 30
+NUM_PROC = 32
+FILE_LIMIT_MB = 10000
+YT_LENGTH_LIMIT_S = 15000  # limit to 2 hour YouTube files
+
 pipeline = FlaxWhisperPipline(checkpoint, dtype=jnp.bfloat16, batch_size=BATCH_SIZE)
 stride_length_s = CHUNK_LENGTH_S / 6
 chunk_len = round(CHUNK_LENGTH_S * pipeline.feature_extractor.sampling_rate)
@@ -149,4 +148,23 @@ def download_yt_audio(yt_url, filename):
         try:
             ydl.download([yt_url])
         except youtube_dl.utils.ExtractorError as err:
-            raise HTTPException(status_code=400, detail=str(err))
+            raise HTTPException(status_code=400, detail=str(err))
+
+def format_timestamp(seconds: float, always_include_hours: bool = False, decimal_marker: str = "."):
+    if seconds is not None:
+        milliseconds = round(seconds * 1000.0)
+
+        hours = milliseconds // 3_600_000
+        milliseconds -= hours * 3_600_000
+
+        minutes = milliseconds // 60_000
+        milliseconds -= minutes * 60_000
+
+        seconds = milliseconds // 1_000
+        milliseconds -= seconds * 1_000
+
+        hours_marker = f"{hours:02d}:" if always_include_hours or hours > 0 else ""
+        return f"{hours_marker}{minutes:02d}:{seconds:02d}{decimal_marker}{milliseconds:03d}"
+    else:
+        # we have a malformed timestamp so just return it as is
+        return seconds