speech_recognition.py 1.1 KB

123456789101112131415161718192021222324252627282930313233343536
  1. import ffmpeg
  2. import subprocess
  3. import tempfile
  4. SAMPLE_RATE = 16000
  5. def convert_audio(data: bytes) -> bytes:
  6. try:
  7. # This launches a subprocess to decode audio while down-mixing and resampling as necessary.
  8. # Requires the ffmpeg CLI and `ffmpeg-python` package to be installed.
  9. out, _ = (
  10. ffmpeg.input("pipe:", threads=0)
  11. .output("-", format="wav", acodec="pcm_s16le", ac=1, ar=SAMPLE_RATE)
  12. .run(cmd="ffmpeg", capture_stdout=True, capture_stderr=True, input=data)
  13. )
  14. except ffmpeg.Error as e:
  15. raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e
  16. return out
  17. class ASR():
  18. def __init__(self, model = "tiny"):
  19. self.model = model
  20. def transcribe(self, audio: bytes) -> str:
  21. audio = convert_audio(audio)
  22. with tempfile.NamedTemporaryFile("w+b") as file:
  23. file.write(audio)
  24. file.flush()
  25. stdout, stderr = subprocess.Popen(
  26. ["./main", "-m", f"models/ggml-{self.model}.bin", "-f", file.name],
  27. stdout=subprocess.PIPE
  28. ).communicate()
  29. if stderr:
  30. print(stderr.decode())
  31. return stdout.decode()