Update main.py

main.py

@@ -1,7 +1,11 @@
 import torch
-from fastapi import BackgroundTasks, File, UploadFile, FastAPI
 from transformers import pipeline
 
+from fastapi import BackgroundTasks, File, UploadFile, FastAPI
+from fastapi.openapi.utils import get_openapi
+from fastapi.staticfiles import StaticFiles
+from fastapi.openapi.docs import get_swagger_ui_html
+
 MODEL_NAME = "openai/whisper-large-v2"
 
 device = 0 if torch.cuda.is_available() else "cpu"
@@ -27,6 +31,16 @@ def transcribe(file):
 
 app = FastAPI()
 
+app.mount("/", StaticFiles(directory="swagger"), name="swagger")
+
+@app.get("/docs", include_in_schema=False)
+async def override_swagger():
+    return get_swagger_ui_html(openapi_url="/openapi.json", title="API Docs")
+
+@app.get("/openapi.json", include_in_schema=False)
+async def get_open_api_endpoint():
+    return get_openapi(title="API Docs", version="1.0.0", routes=app.routes)
+
 @app.get("/transcribe")
 def transcribe(background_tasks: BackgroundTasks, file: UploadFile = File(...)):
     background_tasks.add_task(transcribe, file)