Update app.py

app.py

@@ -3,6 +3,8 @@ import math
 import os
 import tempfile
 import time
+from typing import Dict, Any
+
 import yt_dlp as youtube_dl
 from fastapi import FastAPI, UploadFile, Form, HTTPException
 from fastapi.responses import HTMLResponse
@@ -48,8 +50,22 @@ random_timestamps = pipeline.forward(random_inputs, batch_size=BATCH_SIZE, retur
 compile_time = time.time() - start
 logger.debug(f"Compiled in {compile_time}s")
 
+def timeit(func):
+    async def wrapper(*args, **kwargs):
+        start_time = time.time()
+        result = await func(*args, **kwargs)
+        end_time = time.time()
+        execution_time = end_time - start_time
+        if isinstance(result, dict):
+            result['total_execution_time'] = execution_time
+        else:
+            result = {'result': result, 'total_execution_time': execution_time}
+        return result
+    return wrapper
+
 @app.post("/transcribe_audio")
-async def transcribe_chunked_audio(audio_file: UploadFile, task: str = "transcribe", return_timestamps: bool = False):
+@timeit
+async def transcribe_chunked_audio(audio_file: UploadFile, task: str = "transcribe", return_timestamps: bool = False) -> Dict[str, Any]:
     logger.debug("Starting transcribe_chunked_audio function")
     logger.debug(f"Received parameters - task: {task}, return_timestamps: {return_timestamps}")
     
@@ -85,17 +101,18 @@ async def transcribe_chunked_audio(audio_file: UploadFile, task: str = "transcri
 
     logger.debug("Calling tqdm_generate to transcribe audio")
     try:
-        text, runtime = tqdm_generate(inputs, task=task, return_timestamps=return_timestamps)
+        text, runtime, timing_info = tqdm_generate(inputs, task=task, return_timestamps=return_timestamps)
         logger.debug(f"Transcription completed. Runtime: {runtime:.2f}s")
     except Exception as e:
         logger.error(f"Error in tqdm_generate: {str(e)}", exc_info=True)
         raise HTTPException(status_code=500, detail=f"Error transcribing audio: {str(e)}")
 
     logger.debug("Transcribe_chunked_audio function completed successfully")
-    return {"text": text, "runtime": runtime}
+    return {"text": text, "runtime": runtime, "timing_info": timing_info}
 
 @app.post("/transcribe_youtube")
-async def transcribe_youtube(yt_url: str = Form(...), task: str = "transcribe", return_timestamps: bool = False):
+@timeit
+async def transcribe_youtube(yt_url: str = Form(...), task: str = "transcribe", return_timestamps: bool = False) -> Dict[str, Any]:
     logger.debug("Loading YouTube file...")
     try:
         html_embed_str = _return_yt_html_embed(yt_url)
@@ -126,14 +143,15 @@ async def transcribe_youtube(yt_url: str = Form(...), task: str = "transcribe",
 
     try:
         logger.debug("Calling tqdm_generate to transcribe YouTube audio")
-        text, runtime = tqdm_generate(inputs, task=task, return_timestamps=return_timestamps)
+        text, runtime, timing_info = 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}
+    return {"html_embed": html_embed_str, "text": text, "runtime": runtime, "timing_info": timing_info}
 
 def tqdm_generate(inputs: dict, task: str, return_timestamps: bool):
+    start_time = time.time()
     logger.debug(f"Starting tqdm_generate - task: {task}, return_timestamps: {return_timestamps}")
     
     inputs_len = inputs["array"].shape[0]
@@ -153,7 +171,7 @@ def tqdm_generate(inputs: dict, task: str, return_timestamps: bool):
         raise
 
     model_outputs = []
-    start_time = time.time()
+    transcription_start_time = time.time()
     logger.debug("Starting transcription...")
     
     try:
@@ -166,8 +184,8 @@ def tqdm_generate(inputs: dict, task: str, return_timestamps: bool):
         logger.error(f"Error during batch processing: {str(e)}", exc_info=True)
         raise
 
-    runtime = time.time() - start_time
-    logger.debug(f"Transcription completed in {runtime:.2f}s")
+    transcription_runtime = time.time() - transcription_start_time
+    logger.debug(f"Transcription completed in {transcription_runtime:.2f}s")
 
     logger.debug("Post-processing transcription results")
     try:
@@ -186,8 +204,12 @@ def tqdm_generate(inputs: dict, task: str, return_timestamps: bool):
         ]
         text = "\n".join(str(feature) for feature in timestamps)
     
+    total_processing_time = time.time() - start_time
     logger.debug("tqdm_generate function completed successfully")
-    return text, runtime
+    return text, transcription_runtime, {
+        "transcription_time": transcription_runtime,
+        "total_processing_time": total_processing_time
+    }
 
 def _return_yt_html_embed(yt_url):
     video_id = yt_url.split("?v=")[-1]
@@ -246,4 +268,4 @@ def format_timestamp(seconds: float, always_include_hours: bool = False, decimal
         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 
+        return seconds