api.py 2.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980
  1. # Set the device with environment, default is cuda:0
  2. # export SENSEVOICE_DEVICE=cuda:1
  3. import os, re
  4. from fastapi import FastAPI, File, Form
  5. from fastapi.responses import HTMLResponse
  6. from typing_extensions import Annotated
  7. from typing import List
  8. from enum import Enum
  9. import torchaudio
  10. from model import SenseVoiceSmall
  11. from funasr.utils.postprocess_utils import rich_transcription_postprocess
  12. from io import BytesIO
  13. class Language(str, Enum):
  14. auto = "auto"
  15. zh = "zh"
  16. en = "en"
  17. yue = "yue"
  18. ja = "ja"
  19. ko = "ko"
  20. nospeech = "nospeech"
  21. model_dir = "iic/SenseVoiceSmall"
  22. m, kwargs = SenseVoiceSmall.from_pretrained(model=model_dir, device=os.getenv("SENSEVOICE_DEVICE", "cuda:0"))
  23. m.eval()
  24. regex = r"<\|.*\|>"
  25. app = FastAPI()
  26. @app.get("/", response_class=HTMLResponse)
  27. async def root():
  28. return """
  29. <!DOCTYPE html>
  30. <html>
  31. <head>
  32. <meta charset=utf-8>
  33. <title>Api information</title>
  34. </head>
  35. <body>
  36. <a href='./docs'>Documents of API</a>
  37. </body>
  38. </html>
  39. """
  40. @app.post("/api/v1/asr")
  41. async def turn_audio_to_text(files: Annotated[List[bytes], File(description="wav or mp3 audios in 16KHz")], keys: Annotated[str, Form(description="name of each audio joined with comma")], lang: Annotated[Language, Form(description="language of audio content")] = "auto"):
  42. audios = []
  43. audio_fs = 0
  44. for file in files:
  45. file_io = BytesIO(file)
  46. data_or_path_or_list, audio_fs = torchaudio.load(file_io)
  47. data_or_path_or_list = data_or_path_or_list.mean(0)
  48. audios.append(data_or_path_or_list)
  49. file_io.close()
  50. if lang == "":
  51. lang = "auto"
  52. if keys == "":
  53. key = ["wav_file_tmp_name"]
  54. else:
  55. key = keys.split(",")
  56. res = m.inference(
  57. data_in=audios,
  58. language=lang, # "zh", "en", "yue", "ja", "ko", "nospeech"
  59. use_itn=False,
  60. ban_emo_unk=False,
  61. key=key,
  62. fs=audio_fs,
  63. **kwargs,
  64. )
  65. if len(res) == 0:
  66. return {"result": []}
  67. for it in res[0]:
  68. it["raw_text"] = it["text"]
  69. it["clean_text"] = re.sub(regex, "", it["text"], 0, re.MULTILINE)
  70. it["text"] = rich_transcription_postprocess(it["text"])
  71. return {"result": res[0]}