File size: 2,723 Bytes
0db027e 0524ab0 0db027e 0524ab0 6db367d 0db027e 6db367d 0db027e 6db367d 0db027e 0524ab0 6db367d 0524ab0 6db367d 0524ab0 6db367d 0524ab0 6db367d 0524ab0 0db027e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | import json, os, sys, tempfile, wave
WAV_HEADER_SIZE = 44
WHISPER_CACHE = os.path.expanduser("~/.cache/whisper")
def write_wav(path: str, raw_pcm: bytes, sample_rate: int = 16000):
with wave.open(path, 'wb') as w:
w.setnchannels(1)
w.setsampwidth(2)
w.setframerate(sample_rate)
w.writeframes(raw_pcm)
def get_model_path(name):
lp = os.path.expanduser(f"~/.cache/huggingface/hub/models--Systran--faster-whisper-{name}")
if os.path.isdir(lp):
return lp
return name
def transcribe(wav_path: str, model_size: str = "small", language: str | None = None) -> dict:
try:
from faster_whisper import WhisperModel
model_path = get_model_path(model_size)
model = WhisperModel(model_path, device="cuda", compute_type="int8_float16")
opts = {"beam_size": 5}
if language:
opts["language"] = language
segments, info = model.transcribe(wav_path, **opts)
text = " ".join(seg.text for seg in segments)
return {"success": True, "text": text.strip(), "language": info.language}
except ImportError:
pass
try:
import whisper
model = whisper.load_model(model_size, device="cuda", download_root=WHISPER_CACHE)
opts = {}
if language:
opts["language"] = language
result = model.transcribe(wav_path, **opts)
return {"success": True, "text": result["text"].strip(), "language": result.get("language", "")}
except ImportError:
pass
return {"success": False, "error": "Neither faster-whisper nor openai-whisper is installed. Run: pip install faster-whisper"}
def main():
if len(sys.argv) < 2:
print(json.dumps({"success": False, "error": "Usage: transcribe.py <wav_path> [--model <size>] [--language <code>]"}))
sys.exit(1)
wav_path = sys.argv[1]
model_size = "small"
language = None
i = 2
while i < len(sys.argv):
if sys.argv[i] == "--model" and i + 1 < len(sys.argv):
model_size = sys.argv[i + 1]
i += 2
elif sys.argv[i] == "--language" and i + 1 < len(sys.argv):
language = sys.argv[i + 1]
i += 2
else:
i += 1
if sys.argv[1] == "--stdin-pcm":
sample_rate = int(sys.argv[2]) if len(sys.argv) > 2 else 16000
raw = sys.stdin.buffer.read()
tmp = tempfile.mktemp(suffix=".wav")
write_wav(tmp, raw, sample_rate)
wav_path = tmp
result = transcribe(wav_path, model_size, language)
os.unlink(tmp)
else:
result = transcribe(wav_path, model_size, language)
print(json.dumps(result))
if __name__ == "__main__":
main()
|