Przeglądaj źródła

Merge branch 'codex' of penta/whisper-stt-chatbot into master

penta 18 godzin temu
rodzic
commit
d4a2235130
4 zmienionych plików z 68 dodań i 1 usunięć
  1. 21 1
      app.py
  2. 29 0
      test_transcription.py
  3. 11 0
      whisper-vocabulary-patch.yaml
  4. 7 0
      whisper-vocabulary.yaml

+ 21 - 1
app.py

@@ -1,4 +1,6 @@
 import tempfile
+import re
+import unicodedata
 import os
 import numpy as np
 import time
@@ -19,6 +21,7 @@ VAD_ENABLED = env_bool("WHISPER_VAD", True)
 MODEL_SIZE = str(os.getenv("MODEL_SIZE", "small").lower())
 COMPUTE_TYPE = os.getenv("WHISPER_COMPUTE_TYPE", "float32").strip().lower()
 BEAM_SIZE = int(os.getenv("WHISPER_BEAM_SIZE", "1"))
+HOTWORDS = os.getenv("WHISPER_HOTWORDS", "Satoko").strip()
 if BEAM_SIZE < 1:
     raise ValueError("WHISPER_BEAM_SIZE doit être supérieur ou égal à 1")
 MIN_AUDIO_SEC = float(os.getenv("MIN_AUDIO_SEC", "2.0"))
@@ -49,14 +52,31 @@ model = WhisperModel(MODEL_SIZE, device="cuda", compute_type=COMPUTE_TYPE)
 list(model.transcribe(np.zeros(16000, dtype=np.float32), language="fr", beam_size=BEAM_SIZE)[0])
 model_lock = Lock()
 
+def is_subtitle_artifact(text):
+    normalized = unicodedata.normalize("NFKD", text.lower())
+    normalized = "".join(c for c in normalized if not unicodedata.combining(c))
+    normalized = re.sub(r"[^a-z0-9]+", " ", normalized).strip()
+    return bool(re.fullmatch(
+        r"(?:sous titrage|sous titres?) (?:st ?501|mfp|www amara org|amara org)"
+        r"|sous titres? (?:realises? par la communaute )?d? ?amara org",
+        normalized,
+    ))
+
 def transcribe_audio(audio):
     # Consommer le générateur sous le verrou, hors de la boucle asyncio.
     with model_lock:
         segments, info = model.transcribe(
             audio, language="fr", vad_filter=VAD_ENABLED,
             beam_size=BEAM_SIZE, condition_on_previous_text=False,
+            hotwords=HOTWORDS or None,
         )
-        return " ".join(seg.text for seg in segments).strip(), info.language
+        texts = []
+        for segment in segments:
+            if is_subtitle_artifact(segment.text):
+                logger.info("Segment parasite de sous-titrage ignoré")
+            else:
+                texts.append(segment.text)
+        return " ".join(texts).strip(), info.language
 
 @app.get("/healthz")
 def healthz():

+ 29 - 0
test_transcription.py

@@ -0,0 +1,29 @@
+import ast
+import re
+import unicodedata
+import unittest
+from pathlib import Path
+from types import SimpleNamespace
+from threading import Lock
+from unittest.mock import Mock
+
+
+class TranscriptionTests(unittest.TestCase):
+    def test_artifacts_are_removed_without_losing_real_speech(self):
+        tree = ast.parse(Path('app.py').read_text(encoding='utf-8'))
+        functions = [node for node in tree.body if isinstance(node, ast.FunctionDef)
+                     and node.name in ('is_subtitle_artifact', 'transcribe_audio')]
+        model = Mock()
+        texts = ["Sous-titrage ST' 501", 'Bonjour Satoko', 'Sous-titres MFP',
+                 'Peux-tu afficher les sous-titres ?', 'Sous-titres réalisés par la communauté d’Amara.org']
+        model.transcribe.return_value = ([SimpleNamespace(text=t) for t in texts], SimpleNamespace(language='fr'))
+        ns = dict(re=re, unicodedata=unicodedata, model=model, model_lock=Lock(),
+                  VAD_ENABLED=True, BEAM_SIZE=3, HOTWORDS='Satoko', logger=Mock())
+        exec(compile(ast.Module(body=functions, type_ignores=[]), '<stt>', 'exec'), ns)
+        self.assertEqual(ns['transcribe_audio'](None),
+                         ('Bonjour Satoko Peux-tu afficher les sous-titres ?', 'fr'))
+        self.assertFalse(ns['is_subtitle_artifact']('Explique le sous-titrage ST 501'))
+
+
+if __name__ == '__main__':
+    unittest.main()

+ 11 - 0
whisper-vocabulary-patch.yaml

@@ -0,0 +1,11 @@
+spec:
+  template:
+    spec:
+      containers:
+        - name: whisper-stt-chatbot
+          env:
+            - name: WHISPER_HOTWORDS
+              valueFrom:
+                configMapKeyRef:
+                  name: whisper-vocabulary
+                  key: WHISPER_HOTWORDS

+ 7 - 0
whisper-vocabulary.yaml

@@ -0,0 +1,7 @@
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: whisper-vocabulary
+  namespace: localai
+data:
+  WHISPER_HOTWORDS: "Satoko, Rika, Penta, Johnounet, Jupi, lulu"