Files
room_glitch/labs/glitch_odin/prepare_teddy_dataset.py
T

87 lines
4.6 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""Create conservative draft Piper clips using local ASR word timestamps.
Run with .training/venv/Scripts/python.exe. No audio is uploaded.
Only sentences matching the supplied transcript exactly after normalization
are exported. Review timing and excluded sentences before training.
"""
import difflib
import json
import re
from pathlib import Path
import numpy as np
from scipy.io import wavfile
from scipy.signal import resample_poly
from faster_whisper import WhisperModel
ROOT = Path(__file__).resolve().parent / ".training" / "teddy-v1"
def tokens(text):
return re.findall(r"[a-z0-9]+(?:'[a-z]+)?", text.lower().replace("’", "'"))
def main():
text = (ROOT / "teddy-transcript.txt").read_text(encoding="utf-8-sig")
cache = ROOT / "asr-words.json"
if cache.exists():
words = json.loads(cache.read_text(encoding="utf-8"))
else:
model = WhisperModel("small.en", device="cpu", compute_type="int8", cpu_threads=8,
download_root=str(ROOT.parent / "models"))
rate, source = wavfile.read(ROOT / "TeddyFull.wav")
assert source.ndim == 1 and source.dtype == np.int16
recognition_audio = resample_poly(source.astype(np.float32) / 32768, 16000, rate)
segments, _ = model.transcribe(recognition_audio, language="en",
word_timestamps=True, vad_filter=True, beam_size=5)
words = []
for segment in segments:
for word in segment.words:
words.append({"text":word.word, "start":word.start, "end":word.end})
print(f"Aligned through {segment.end:.1f} seconds", flush=True)
cache.write_text(json.dumps(words, indent=2), encoding="utf-8")
recognized = []
for word in words:
for token in tokens(word["text"]):
recognized.append(dict(word, token=token))
reference = tokens(text)
matcher = difflib.SequenceMatcher(None, reference, [w["token"] for w in recognized], autojunk=False)
mapping = {}
for block in matcher.get_matching_blocks():
for offset in range(block.size):
mapping[block.a+offset] = block.b+offset
rate, pcm = wavfile.read(ROOT / "TeddyFull.wav")
assert pcm.ndim == 1 and pcm.dtype == np.int16
audio = pcm.astype(np.float64) / 32768
clips = ROOT / "audio"
clips.mkdir(exist_ok=True)
accepted, excluded, rows = [], [], []
cursor = 0
for index, sentence in enumerate(re.split(r"(?<=[.!?])\s+", text.strip()), 1):
count = len(tokens(sentence))
positions = [mapping.get(i) for i in range(cursor,cursor+count)]
cursor += count
if not count: continue
if any(i is None for i in positions) or positions != list(range(positions[0],positions[0]+count)):
excluded.append({"sentence":index,"text":sentence,"reason":"Transcript / recognition mismatch; review required"})
continue
first,last = positions[0],positions[-1]
start = max(0,recognized[first]["start"]-.10)
end = min(len(audio)/rate,recognized[last]["end"]+.10)
# Do not include neighboring spoken words in the padding.
if first: start=max(start,(recognized[first-1]["end"]+recognized[first]["start"])/2)
if last+1<len(recognized): end=min(end,(recognized[last]["end"]+recognized[last+1]["start"])/2)
if not .7 <= end-start <= 15:
excluded.append({"sentence":index,"text":sentence,"reason":"Duration outside 0.7-15 seconds"})
continue
clip = resample_poly(audio[int(start*rate):int(end*rate)],22050,rate)
filename=f"teddy_{index:04}.wav"
wavfile.write(clips/filename,22050,np.clip(clip*32768,-32768,32767).astype(np.int16))
rows.append(filename+"|"+sentence.replace("\n"," "))
accepted.append({"file":filename,"start":start,"end":end,"text":sentence})
(ROOT/"metadata.draft.csv").write_text("\n".join(rows)+"\n",encoding="utf-8")
report={"status":"DRAFT - inspect timing and transcript before training", "matched_word_ratio":len(mapping)/len(reference),
"source_seconds":len(audio)/rate,"source_peak":float(np.max(np.abs(audio))),
"source_clipped_samples":int(np.sum(np.abs(audio)>=32767/32768)),
"clip_count":len(accepted),"clip_seconds":sum(c["end"]-c["start"] for c in accepted),
"clips":accepted,"excluded":excluded}
(ROOT/"alignment-report.json").write_text(json.dumps(report,indent=2),encoding="utf-8")
print(json.dumps({k:v for k,v in report.items() if k not in ("clips","excluded")},indent=2))
print(f"Excluded {len(excluded)} sentences for review")
if __name__ == "__main__": main()