третий шаг - транскрипция аудио в формат датасетаHEADmain

1 files changed, 98 insertions, 0 deletions

diff --git a/scripts/02_transcribe_audio.py b/scripts/02_transcribe_audio.py
new file mode 100755
index 0000000..3e513e4
--- /dev/null
+++ b/scripts/02_transcribe_audio.py
@@ -0,0 +1,98 @@
+#!/usr/local/bin/python3.12
+
+from vosk import KaldiRecognizer, Model
+import os.path
+import json
+import os
+
+
+FRAME = 4_000
+
+SAMPLE_RATE = 16_000
+
+MODELS = list(
+  map(
+    lambda model: os.path.join(os.getcwd(), 'third-party/vosk-models', model),
+    [
+      'vosk-model-en-us-0.22-lgraph', # English (United States)
+      'vosk-model-small-ru-0.22', # Russian
+      'vosk-model-small-uk-v3-small', # Ukrainian
+    ]
+  )
+)
+
+DATASET = os.path.join(os.getcwd(), 'files/dataset.json')
+
+AUDIOS = os.path.join(os.getcwd(), 'files/audios')
+
+
+def score_transcription(result: dict) -> float:
+  words = result.get("result", [])
+  values = list(
+    map(lambda word: word.get("conf", 0.0), words)
+  )
+  try:
+    return sum(values) / len(values)
+  except ZeroDivisionError:
+    return 0.0
+
+
+def main() -> None:
+  recognizers = []
+  transcriptions = []
+
+  for model in MODELS:
+    model = Model(model)
+    recognizer = KaldiRecognizer(model, SAMPLE_RATE)
+    recognizer.SetWords(True)
+    recognizers.append(recognizer)
+
+  for child in os.listdir(AUDIOS):
+    if child == ".gitkeep": continue
+
+    file = os.path.join(AUDIOS, child)
+    if not os.path.isfile(file): continue
+
+    with open(file, 'rb') as audio:
+      while True:
+        frame = audio.read(FRAME)
+        if len(frame) == 0:
+          break
+        for recognizer in recognizers:
+          recognizer.AcceptWaveform(frame)
+
+      results = []
+      for i, recognizer in enumerate(recognizers):
+        result = json.loads(recognizer.FinalResult())
+        if result.get("text", "") == "": continue
+
+        score = score_transcription(result)
+        if score == 0.0: continue
+
+        results.append({'model': MODELS[i], 'score': score, 'result': result})
+
+      if len(results) == 0:
+        continue
+      results = list(
+        sorted(results, key=lambda result: result['score'], reverse=True)
+      )
+
+      transcription = results[0]
+      transcription.update({'video': child})
+
+      transcriptions.append(transcription)
+
+  dataset = []
+  for transcription in transcriptions:
+    words = transcription['result']['result']
+    message = []
+    for i, word in enumerate(words):
+      message.append(f'[{i + 1}]{word['word']}')
+    dataset.append({'user': ' '.join(message), 'assistant': None})
+
+  with open(DATASET, 'w') as file:
+    json.dump(dataset, file, ensure_ascii=False, indent=2)
+
+
+if __name__ == "__main__":
+  main()