Update app.py

app.py

@@ -4,10 +4,13 @@ import os
 import tempfile
 import time
 from typing import Dict, Any
+from functools import wraps
 
 import yt_dlp as youtube_dl
-from fastapi import FastAPI, UploadFile, Form, HTTPException
+from fastapi import FastAPI, File, UploadFile, Depends, HTTPException
 from fastapi.responses import HTMLResponse
+from fastapi.encoders import jsonable_encoder
+from pydantic import BaseModel
 import jax.numpy as jnp
 import numpy as np
 from transformers.models.whisper.tokenization_whisper import TO_LANGUAGE_CODE
@@ -50,7 +53,17 @@ random_timestamps = pipeline.forward(random_inputs, batch_size=BATCH_SIZE, retur
 compile_time = time.time() - start
 logger.debug(f"Compiled in {compile_time}s")
 
+class TranscribeAudioRequest(BaseModel):
+    task: str = "transcribe"
+    return_timestamps: bool = False
+
+class TranscribeYouTubeRequest(BaseModel):
+    yt_url: str
+    task: str = "transcribe"
+    return_timestamps: bool = False
+
 def timeit(func):
+    @wraps(func)
     async def wrapper(*args, **kwargs):
         start_time = time.time()
         result = await func(*args, **kwargs)
@@ -65,9 +78,12 @@ def timeit(func):
 
 @app.post("/transcribe_audio")
 @timeit
-async def transcribe_chunked_audio(audio_file: UploadFile, task: str = "transcribe", return_timestamps: bool = False) -> Dict[str, Any]:
+async def transcribe_chunked_audio(
+    audio_file: UploadFile = File(...),
+    request: TranscribeAudioRequest = Depends()
+) -> Dict[str, Any]:
     logger.debug("Starting transcribe_chunked_audio function")
-    logger.debug(f"Received parameters - task: {task}, return_timestamps: {return_timestamps}")
+    logger.debug(f"Received parameters - task: {request.task}, return_timestamps: {request.return_timestamps}")
     
     logger.debug("Checking for audio file...")
     if not audio_file:
@@ -101,21 +117,25 @@ async def transcribe_chunked_audio(audio_file: UploadFile, task: str = "transcri
 
     logger.debug("Calling tqdm_generate to transcribe audio")
     try:
-        text, runtime, timing_info = tqdm_generate(inputs, task=task, return_timestamps=return_timestamps)
+        text, runtime, timing_info = tqdm_generate(inputs, task=request.task, return_timestamps=request.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, "timing_info": timing_info}
+    return jsonable_encoder({
+        "text": text,
+        "runtime": runtime,
+        "timing_info": timing_info
+    })
 
 @app.post("/transcribe_youtube")
 @timeit
-async def transcribe_youtube(yt_url: str = Form(...), task: str = "transcribe", return_timestamps: bool = False) -> Dict[str, Any]:
+async def transcribe_youtube(request: TranscribeYouTubeRequest) -> Dict[str, Any]:
     logger.debug("Loading YouTube file...")
     try:
-        html_embed_str = _return_yt_html_embed(yt_url)
+        html_embed_str = _return_yt_html_embed(request.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")
@@ -124,7 +144,7 @@ async def transcribe_youtube(yt_url: str = Form(...), task: str = "transcribe",
         filepath = os.path.join(tmpdirname, "video.mp4")
         try:
             logger.debug("Downloading YouTube audio...")
-            download_yt_audio(yt_url, filepath)
+            download_yt_audio(request.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")
@@ -143,12 +163,17 @@ async def transcribe_youtube(yt_url: str = Form(...), task: str = "transcribe",
 
     try:
         logger.debug("Calling tqdm_generate to transcribe YouTube audio")
-        text, runtime, timing_info = tqdm_generate(inputs, task=task, return_timestamps=return_timestamps)
+        text, runtime, timing_info = tqdm_generate(inputs, task=request.task, return_timestamps=request.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, "timing_info": timing_info}
+    return jsonable_encoder({
+        "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()