Setting the file. One moment. def _format_output
— line 134
This file
- Number
- 36.2
- Position
- 2 of 5
- Type
- Python
- Size
- 8 KB
- Lines
- 276
scripts/transcribe_diarize.py
Python·276 lines·8 KB
from
typing
import
Any, Dict, List, Optional, Tuple
14
15DEFAULT_MODEL = "gpt-4o-mini-transcribe"
16DEFAULT_RESPONSE_FORMAT = "text"
17DEFAULT_CHUNKING_STRATEGY = "auto"
18MAX_AUDIO_BYTES = 25 * 1024 * 1024
19MAX_KNOWN_SPEAKERS = 4
20
21ALLOWED_RESPONSE_FORMATS = {"text", "json", "diarized_json"}
22
23
24def _die(message: str, code: int = 1) -> None:
25 print(f"Error: {message}", file=sys.stderr)
26 raise SystemExit(code)
27
28
29def _warn(message: str) -> None:
30 print(f"Warning: {message}", file=sys.stderr)
31
32
33def _ensure_api_key(dry_run: bool) -> None:
34 if os.getenv("OPENAI_API_KEY"):
35 print("OPENAI_API_KEY is set.", file=sys.stderr)
36 return
37 if dry_run:
38 _warn("OPENAI_API_KEY is not set; dry-run only.")
39 return
40 _die("OPENAI_API_KEY is not set. Export it before running.")
41
42
43def _normalize_response_format(value: Optional[str]) -> str:
44 if not value:
45 return DEFAULT_RESPONSE_FORMAT
46 fmt = value.strip().lower()
47 if fmt not in ALLOWED_RESPONSE_FORMATS:
48 _die(
49 "response-format must be one of: "
50 + ", ".join(sorted(ALLOWED_RESPONSE_FORMATS))
51 )
52 return fmt
53
54
55def _normalize_chunking_strategy(value: Optional[str]) -> Any:
56 if not value:
57 return DEFAULT_CHUNKING_STRATEGY
58 raw = str(value).strip()
59 if raw.startswith("{"):
60 try:
61 return json.loads(raw)
62 except json.JSONDecodeError:
63 _die("chunking-strategy JSON is invalid")
64 return raw
65
66
67def _guess_mime_type(path: Path) -> str:
68 mime, _ = mimetypes.guess_type(str(path))
69 if mime:
70 return mime
71 return "audio/wav"
72
73
74def _encode_data_url(path: Path) -> str:
75 data = path.read_bytes()
76 mime = _guess_mime_type(path)
77 encoded = base64.b64encode(data).decode("ascii")
78 return f"data:{mime};base64,{encoded}"
79
80
81def _parse_known_speakers(raw_items: List[str]) -> Tuple[List[str], List[str]]:
82 names: List[str] = []
83 refs: List[str] = []
84 for raw in raw_items:
85 if "=" not in raw:
86 _die("known-speaker must be NAME=PATH")
87 name, path_str = raw.split("=", 1)
88 name = name.strip()
89 path = Path(path_str.strip())
90 if not name or not path_str.strip():
91 _die("known-speaker must be NAME=PATH")
92 if not path.exists():
93 _die(f"Known speaker file not found: {path}")
94 names.append(name)
95 refs.append(_encode_data_url(path))
96 if len(names) > MAX_KNOWN_SPEAKERS:
97 _die(f"known speakers must be <= {MAX_KNOWN_SPEAKERS}")
98 return names, refs
99
100
101def _output_extension(response_format: str) -> str:
102 return "txt" if response_format == "text" else "json"
103
104
105def _build_output_path(
106 audio_path: Path,
107 response_format: str,
108 out: Optional[str],
109 out_dir: Optional[str],
110) -> Path:
111 ext = "." + _output_extension(response_format)
112 if out:
113 path = Path(out)
114 if path.exists() and path.is_dir():
115 return path / f"{audio_path.stem}.transcript{ext}"
116 if path.suffix == "":
117 return path.with_suffix(ext)
118 return path
119 if out_dir:
120 base = Path(out_dir)
121 base.mkdir(parents=True, exist_ok=True)
122 return base / f"{audio_path.stem}.transcript{ext}"
123 return Path(f"{audio_path.stem}.transcript{ext}")
124
125
126def _create_client():
127 try:
128 from openai import OpenAI
129 except ImportError:
130 _die("openai SDK not installed. Install with `uv pip install openai`.")
131 return OpenAI()
132
133
134def _format_output(result: Any, response_format: str) -> str:
135 if response_format == "text":
136 text = getattr(result, "text", None)
137 return text if isinstance(text, str) else str(result)
138 if hasattr(result, "model_dump"):
139 return json.dumps(result.model_dump(), indent=2)
140 if isinstance(result, (dict, list)):
141 return json.dumps(result, indent=2)
142 return json.dumps({"text": getattr(result, "text", str(result))}, indent=2)
143
144
145def _validate_audio(path: Path) -> None:
146 if not path.exists():
147 _die(f"Audio file not found: {path}")
148 size = path.stat().st_size
149 if size > MAX_AUDIO_BYTES:
150 _warn(
151 f"Audio file exceeds 25MB limit ({size} bytes): {path}"
152 )
153
154
155def _build_payload(
156 args: argparse.Namespace,
157 known_speaker_names: List[str],
158 known_speaker_refs: List[str],
159) -> Dict[str, Any]:
160 payload: Dict[str, Any] = {
161 "model": args.model,
162 "response_format": args.response_format,
163 "chunking_strategy": args.chunking_strategy,
164 }
165 if args.language:
166 payload["language"] = args.language
167 if args.prompt:
168 payload["prompt"] = args.prompt
169 if known_speaker_names:
170 payload["extra_body"] = {
171 "known_speaker_names": known_speaker_names,
172 "known_speaker_references": known_speaker_refs,
173 }
174 return payload
175
176
177def _run_one(
178 client: Any,
179 audio_path: Path,
180 payload: Dict[str, Any],
181) -> Any:
182 with audio_path.open("rb") as audio_file:
183 return client.audio.transcriptions.create(
184 file=audio_file,
185 **payload,
186 )
187
188
189def main() -> None:
190 parser = argparse.ArgumentParser(
191 description="Transcribe audio (optionally with speaker diarization) using OpenAI."
192 )
193 parser.add_argument("audio", nargs="+", help="Audio file(s) to transcribe")
194 parser.add_argument(
195 "--model",
196 default=DEFAULT_MODEL,
197 help=f"Model to use (default: {DEFAULT_MODEL})",
198 )
199 parser.add_argument(
200 "--response-format",
201 default=DEFAULT_RESPONSE_FORMAT,
202 help="Response format: text, json, or diarized_json",
203 )
204 parser.add_argument(
205 "--chunking-strategy",
206 default=DEFAULT_CHUNKING_STRATEGY,
207 help="Chunking strategy (use 'auto' for long audio)",
208 )
209 parser.add_argument("--language", help="Optional language hint (e.g. 'en')")
210 parser.add_argument("--prompt", help="Optional prompt to guide transcription")
211 parser.add_argument(
212 "--known-speaker",
213 action="append",
214 default=[],
215 help="Known speaker reference as NAME=PATH (repeatable, max 4)",
216 )
217 parser.add_argument("--out", help="Output file path (single audio only)")
218 parser.add_argument("--out-dir", help="Output directory for transcripts")
219 parser.add_argument(
220 "--stdout",
221 action="store_true",
222 help="Write transcript to stdout instead of a file",
223 )
224 parser.add_argument(
225 "--dry-run",
226 action="store_true",
227 help="Validate inputs and print payload without calling the API",
228 )
229
230 args = parser.parse_args()
231 args.response_format = _normalize_response_format(args.response_format)
232 args.chunking_strategy = _normalize_chunking_strategy(args.chunking_strategy)
233
234 if args.out and len(args.audio) > 1:
235 _die("--out only supports a single audio file")
236 if args.stdout and (args.out or args.out_dir):
237 _die("--stdout cannot be combined with --out or --out-dir")
238 if args.stdout and len(args.audio) > 1:
239 _die("--stdout only supports a single audio file")
240
241 if args.prompt and "transcribe-diarize" in args.model:
242 _die("prompt is not supported with gpt-4o-transcribe-diarize")
243 if args.response_format == "diarized_json" and "transcribe-diarize" not in args.model:
244 _die("diarized_json requires gpt-4o-transcribe-diarize")
245
246 _ensure_api_key(args.dry_run)
247
248 audio_paths = [Path(p) for p in args.audio]
249 for path in audio_paths:
250 _validate_audio(path)
251
252 known_names, known_refs = _parse_known_speakers(args.known_speaker)
253 if known_names and "transcribe-diarize" not in args.model:
254 _warn("known-speaker references are only supported for gpt-4o-transcribe-diarize")
255 payload = _build_payload(args, known_names, known_refs)
256
257 if args.dry_run:
258 print(json.dumps(payload, indent=2))
259 return
260
261 client = _create_client()
262
263 for path in audio_paths:
264 result = _run_one(client, path, payload)
265 output = _format_output(result, args.response_format)
266 if args.stdout:
267 print(output)
268 continue
269 out_path = _build_output_path(path, args.response_format, args.out, args.out_dir)
270 out_path.parent.mkdir(parents=True, exist_ok=True)
271 out_path.write_text(output, encoding="utf-8")
272 print(f"Wrote {out_path}")
273
274
275if __name__ == "__main__":
276 main()
