main.py 2.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990
  1. import os
  2. import logging
  3. from fastapi import (
  4. FastAPI,
  5. Request,
  6. Depends,
  7. HTTPException,
  8. status,
  9. UploadFile,
  10. File,
  11. Form,
  12. )
  13. from fastapi.middleware.cors import CORSMiddleware
  14. from faster_whisper import WhisperModel
  15. from constants import ERROR_MESSAGES
  16. from utils.utils import (
  17. decode_token,
  18. get_current_user,
  19. get_verified_user,
  20. get_admin_user,
  21. )
  22. from utils.misc import calculate_sha256
  23. from config import (
  24. SRC_LOG_LEVELS,
  25. CACHE_DIR,
  26. UPLOAD_DIR,
  27. WHISPER_MODEL,
  28. WHISPER_MODEL_DIR,
  29. )
  30. log = logging.getLogger(__name__)
  31. log.setLevel(SRC_LOG_LEVELS["AUDIO"])
  32. app = FastAPI()
  33. app.add_middleware(
  34. CORSMiddleware,
  35. allow_origins=["*"],
  36. allow_credentials=True,
  37. allow_methods=["*"],
  38. allow_headers=["*"],
  39. )
  40. @app.post("/transcribe")
  41. def transcribe(
  42. file: UploadFile = File(...),
  43. user=Depends(get_current_user),
  44. ):
  45. log.info(f"file.content_type: {file.content_type}")
  46. if file.content_type not in ["audio/mpeg", "audio/wav"]:
  47. raise HTTPException(
  48. status_code=status.HTTP_400_BAD_REQUEST,
  49. detail=ERROR_MESSAGES.FILE_NOT_SUPPORTED,
  50. )
  51. try:
  52. filename = file.filename
  53. file_path = f"{UPLOAD_DIR}/{filename}"
  54. contents = file.file.read()
  55. with open(file_path, "wb") as f:
  56. f.write(contents)
  57. f.close()
  58. model = WhisperModel(
  59. WHISPER_MODEL,
  60. device="auto",
  61. compute_type="int8",
  62. download_root=WHISPER_MODEL_DIR,
  63. )
  64. segments, info = model.transcribe(file_path, beam_size=5)
  65. log.info(
  66. "Detected language '%s' with probability %f"
  67. % (info.language, info.language_probability)
  68. )
  69. transcript = "".join([segment.text for segment in list(segments)])
  70. return {"text": transcript.strip()}
  71. except Exception as e:
  72. log.exception(e)
  73. raise HTTPException(
  74. status_code=status.HTTP_400_BAD_REQUEST,
  75. detail=ERROR_MESSAGES.DEFAULT(e),
  76. )