feat(tools): 2.0 phase 5 wave 5b - ai pool: ocr-pdf, transcription, background composites (5 tools) (#226)

2026-08-03 07:46:42 +02:00 · 2026-06-13 10:19:47 +08:00
parent 6e1b9865f1
commit 51666cdd5f
59 changed files with 5423 additions and 611 deletions
@@ -0,0 +1,85 @@
+"""Speech-to-text transcription using faster-whisper (CTranslate2)."""
+import sys
+import json
+import os
+
+from gpu import gpu_available
+
+
+MODELS_PATH = os.environ.get(
+    "MODELS_PATH",
+    os.path.join(os.environ.get("DATA_DIR", "/data"), "ai", "models"),
+)
+
+
+def emit_progress(percent, stage):
+    """Emit structured progress to stderr for bridge.ts to capture."""
+    print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True)
+
+
+def main():
+    input_path = sys.argv[1]
+    settings = json.loads(sys.argv[2]) if len(sys.argv) > 2 else {}
+
+    language = settings.get("language", "auto")
+    # "task" is accepted for forward-compat but only "transcribe" is used today
+    _task = settings.get("task", "transcribe")
+
+    try:
+        emit_progress(5, "Loading model")
+
+        # Lazy import -- faster_whisper is only available when the
+        # transcription bundle is installed; keeping it lazy lets
+        # py_compile succeed without the dependency.
+        from faster_whisper import WhisperModel
+
+        model_dir = os.path.join(MODELS_PATH, "faster-whisper-small")
+
+        if gpu_available():
+            device, compute_type = "cuda", "float16"
+        else:
+            device, compute_type = "cpu", "int8"
+
+        model = WhisperModel(model_dir, device=device, compute_type=compute_type)
+
+        emit_progress(20, "Transcribing")
+
+        lang_arg = None if language == "auto" else language
+        segments_iter, info = model.transcribe(
+            input_path,
+            language=lang_arg,
+            vad_filter=True,
+        )
+
+        detected_language = info.language if info else (language if language != "auto" else "en")
+
+        segments = []
+        batch_count = 0
+        for seg in segments_iter:
+            segments.append({
+                "start": round(seg.start, 3),
+                "end": round(seg.end, 3),
+                "text": seg.text.strip(),
+            })
+            batch_count += 1
+            if batch_count % 5 == 0:
+                emit_progress(min(20 + batch_count, 90), "Transcribing")
+
+        emit_progress(95, "Done")
+
+        full_text = " ".join(s["text"] for s in segments)
+
+        print(json.dumps({
+            "success": True,
+            "language": detected_language,
+            "segments": segments,
+            "text": full_text,
+        }))
+
+    except Exception as e:
+        print(json.dumps({"error": str(e)}))
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()