This repository has no description
7.4 kB
237 lines
1#!/usr/bin/env python3
2"""
3Real-time streaming speech-to-text using Vosk.
4Reads audio from stdin and outputs words as they're recognized.
5"""
6
7import argparse
8import json
9import os
10import struct
11import sys
12import zipfile
13from pathlib import Path
14from urllib.request import urlretrieve
15
16from vosk import KaldiRecognizer, Model, SetLogLevel
17
18MODEL_DIR = Path(os.environ.get("XDG_DATA_HOME", Path.home() / ".local/share")) / "vosk"
19
20# Models organized by size: small (~40MB), medium (~128MB), large (~1-2GB)
21MODELS = {
22 "en": {
23 "small": "vosk-model-small-en-us-0.15",
24 "medium": "vosk-model-en-us-0.22-lgraph",
25 "large": "vosk-model-en-us-0.22",
26 },
27 "de": {
28 "small": "vosk-model-small-de-0.15",
29 "medium": "vosk-model-de-0.21",
30 "large": "vosk-model-de-0.21",
31 },
32 "fr": {
33 "small": "vosk-model-small-fr-0.22",
34 "medium": "vosk-model-fr-0.22",
35 "large": "vosk-model-fr-0.22",
36 },
37 "es": {
38 "small": "vosk-model-small-es-0.42",
39 "medium": "vosk-model-es-0.42",
40 "large": "vosk-model-es-0.42",
41 },
42 "ru": {
43 "small": "vosk-model-small-ru-0.22",
44 "medium": "vosk-model-ru-0.42",
45 "large": "vosk-model-ru-0.42",
46 },
47 "cn": {
48 "small": "vosk-model-small-cn-0.22",
49 "medium": "vosk-model-cn-0.22",
50 "large": "vosk-model-cn-0.22",
51 },
52 "ja": {
53 "small": "vosk-model-small-ja-0.22",
54 "medium": "vosk-model-ja-0.22",
55 "large": "vosk-model-ja-0.22",
56 },
57}
58
59
60def download_model(lang: str, size: str = "small") -> Path:
61 """Download vosk model if not present."""
62 lang_models = MODELS.get(lang)
63 if not lang_models:
64 # Fall back to English if language not available
65 print(f"No model for '{lang}', using English", file=sys.stderr)
66 lang_models = MODELS["en"]
67 lang = "en"
68
69 model_name = lang_models.get(size, lang_models["small"])
70
71 model_path = MODEL_DIR / model_name
72 if model_path.exists():
73 return model_path
74
75 print(f"Downloading model {model_name}...", file=sys.stderr)
76 MODEL_DIR.mkdir(parents=True, exist_ok=True)
77
78 url = f"https://alphacephei.com/vosk/models/{model_name}.zip"
79 zip_path = MODEL_DIR / f"{model_name}.zip"
80
81 urlretrieve(url, zip_path)
82
83 print("Extracting...", file=sys.stderr)
84 with zipfile.ZipFile(zip_path, "r") as z:
85 z.extractall(MODEL_DIR)
86 zip_path.unlink()
87
88 print("Model ready", file=sys.stderr)
89 return model_path
90
91
92def parse_wav_header(data: bytes) -> tuple[int, int, int, int]:
93 """Parse WAV header, return (sample_rate, channels, bits_per_sample, data_offset)."""
94 if len(data) < 44 or data[:4] != b"RIFF" or data[8:12] != b"WAVE":
95 raise ValueError("Invalid WAV header")
96
97 pos = 12
98 sample_rate = channels = bits_per_sample = 0
99 while pos < len(data) - 8:
100 chunk_id = data[pos : pos + 4]
101 chunk_size = struct.unpack("<I", data[pos + 4 : pos + 8])[0]
102 if chunk_id == b"fmt ":
103 channels = struct.unpack("<H", data[pos + 10 : pos + 12])[0]
104 sample_rate = struct.unpack("<I", data[pos + 12 : pos + 16])[0]
105 bits_per_sample = struct.unpack("<H", data[pos + 22 : pos + 24])[0]
106 elif chunk_id == b"data":
107 return sample_rate, channels, bits_per_sample, pos + 8
108 pos += 8 + chunk_size
109
110 raise ValueError("No data chunk found")
111
112
113def convert_audio(data: bytes, in_rate: int, in_channels: int, in_bits: int) -> bytes:
114 """Convert audio to 16kHz mono 16-bit for Vosk."""
115 import array
116
117 # Align data to sample boundary
118 bytes_per_sample = in_bits // 8
119 frame_size = in_channels * bytes_per_sample
120 aligned_len = (len(data) // frame_size) * frame_size
121 if aligned_len == 0:
122 return b""
123 data = data[:aligned_len]
124
125 # Convert to samples
126 if in_bits == 16:
127 samples = array.array("h", data)
128 elif in_bits == 8:
129 samples = array.array("h", [(s - 128) * 256 for s in data])
130 else:
131 raise ValueError(f"Unsupported bit depth: {in_bits}")
132
133 # Convert to mono
134 if in_channels == 2:
135 mono = array.array("h")
136 for i in range(0, len(samples) - 1, 2):
137 mono.append((samples[i] + samples[i + 1]) // 2)
138 samples = mono
139
140 # Resample to 16kHz if needed
141 if in_rate != 16000:
142 ratio = 16000 / in_rate
143 new_len = int(len(samples) * ratio)
144 resampled = array.array("h")
145 for i in range(new_len):
146 src_idx = min(int(i / ratio), len(samples) - 1)
147 resampled.append(samples[src_idx])
148 samples = resampled
149
150 return samples.tobytes()
151
152
153def main():
154 parser = argparse.ArgumentParser(description="Real-time STT using Vosk")
155 parser.add_argument(
156 "-l",
157 "--language",
158 default=os.environ.get("VOSK_LANG", "en"),
159 help="Language: en, de, fr, es, ru, cn, ja (default: en)",
160 )
161 parser.add_argument(
162 "-s",
163 "--size",
164 default=os.environ.get("VOSK_SIZE", "small"),
165 choices=["small", "medium", "large"],
166 help="Model size: small (~40MB), medium (~128MB), large (~1-2GB) (default: small)",
167 )
168 parser.add_argument(
169 "-w",
170 "--words",
171 action="store_true",
172 help="Output individual words as they're recognized",
173 )
174 args = parser.parse_args()
175
176 # Suppress Vosk logging
177 SetLogLevel(-1)
178
179 model_path = download_model(args.language, args.size)
180 model = Model(str(model_path))
181
182 # Read WAV header
183 stdin = sys.stdin.buffer
184 header_buf = stdin.read(4096)
185 sample_rate, channels, bits, data_offset = parse_wav_header(header_buf)
186
187 print(f"Audio: {sample_rate}Hz, {channels}ch, {bits}bit", file=sys.stderr)
188 print("Listening...", file=sys.stderr)
189
190 # Create recognizer for 16kHz audio
191 rec = KaldiRecognizer(model, 16000)
192 rec.SetWords(True)
193
194 # Process initial audio from header buffer
195 initial_audio = header_buf[data_offset:]
196 if initial_audio:
197 converted = convert_audio(initial_audio, sample_rate, channels, bits)
198 rec.AcceptWaveform(converted)
199
200 # Process audio stream
201 chunk_size = int(sample_rate * channels * (bits // 8) * 0.1) # 100ms chunks
202
203 while True:
204 data = stdin.read(chunk_size)
205 if not data:
206 break
207
208 converted = convert_audio(data, sample_rate, channels, bits)
209
210 if rec.AcceptWaveform(converted):
211 result = json.loads(rec.Result())
212 text = result.get("text", "")
213 if text:
214 if args.words:
215 # Show final sentence on stderr with period and newline
216 print(f"\r\033[K{text}.", file=sys.stderr, flush=True)
217 print(text, flush=True)
218 elif args.words:
219 partial = json.loads(rec.PartialResult())
220 text = partial.get("partial", "")
221 if text:
222 # Show partial with corrections on stderr
223 print(f"\r\033[K{text}", end="", file=sys.stderr, flush=True)
224
225 # Final result
226 result = json.loads(rec.FinalResult())
227 text = result.get("text", "")
228 if text:
229 if args.words:
230 print(f"\r\033[K{text}.", file=sys.stderr, flush=True)
231 print(text, flush=True)
232
233 print("Stream ended", file=sys.stderr)
234
235
236if __name__ == "__main__":
237 main()