aligner.py

import dataclasses
import logging
import os
import re
import tempfile
import warnings
from typing import List, Iterable, Optional, Union, Tuple

import numpy as np
from json_database import JsonStorage
from ovos_plugin_manager.templates.stt import STT
from ovos_stt_plugin_whisper import MyNorthAISTT
from ovos_utils.parse import fuzzy_match, MatchStrategy
from quebra_frases import sentence_tokenize
from scipy.io import wavfile
from scipy.io.wavfile import write
from speech_recognition import Recognizer, AudioFile, AudioData
from tqdm import tqdm  # Progress bar

# Suppress specific FutureWarning
warnings.filterwarnings("ignore", category=FutureWarning)
logging.getLogger("transformers").setLevel("ERROR")


def file2audiodata(file_name: str) -> AudioData:
    with AudioFile(file_name) as source:
        rec = Recognizer()
        return rec.record(source)


def buffer2audiodata(buffer: np.ndarray, sample_rate) -> AudioData:
    # Use a temporary file for the chunk
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_wav:
        write(temp_wav.name, sample_rate, buffer)
        return file2audiodata(temp_wav.name)


@dataclasses.dataclass
class AudioSplit:
    start_time: float = 0.0
    end_time: float = -1.0
    target_text: str = ""
    best_transcript: str = ""
    sample_rate: int = 16000
    confidence: float = 0.0
    full_audio: Optional[np.ndarray] = None

    def save(self, output_path):
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        wavfile.write(output_path, self.sample_rate, self.buffer)

    @property
    def buffer(self) -> np.ndarray:
        return self._get_split(self.start_time, self.end_time)

    def _get_split(self, s, e):
        # Convert times to sample indices
        start_sample = int(s * self.sample_rate)
        end_sample = int(e * self.sample_rate)
        # chunk
        return self.full_audio[start_sample:end_sample]

    @classmethod
    def calc_score(cls, n1: str, n2: str) -> float:
        i = int(len(n1) / 3)  # last 1/3 of the text
        n1, n2 = cls.norm_text(n1), cls.norm_text(n2)
        s = fuzzy_match(n1, n2, strategy=MatchStrategy.DAMERAU_LEVENSHTEIN_SIMILARITY)
        # add a weight based on word ratio  len(txt.split()) / len(stx.split())
        if len(n2) > len(n1) + i:
            s = 0
        return min(s, 1.0)

    @staticmethod
    def norm_text(text: str) -> str:
        # Convert to lowercase
        text = text.lower().replace("-", " ")

        # Remove common punctuation and symbols
        text = re.sub(r'[,.!?;:\-\(\)\[\]\"\'“”‘’]', '', text)

        # Replace multiple spaces with a single space
        # ignore stopwords (speech particles)
        stopwords = ["hum", "ham"]
        text = " ".join([w for w in text.split() if w and w not in stopwords])

        return text.strip()

    @property
    def ending_matches(self) -> bool:
        if not self.norm_text(self.best_transcript) or not self.norm_text(self.target_text):
            return False
        lastw = self.norm_text(self.target_text).split()[-1]
        lw = self.norm_text(self.best_transcript).split()[-1]
        return fuzzy_match(lastw, lw) >= 0.9

    @property
    def start_matches(self) -> bool:
        if not self.norm_text(self.best_transcript) or not self.norm_text(self.target_text):
            return False
        firstw = self.norm_text(self.target_text).split()[0]
        fw = self.norm_text(self.best_transcript).split()[0]
        return fuzzy_match(firstw, fw) >= 0.9

    def _crop_left(self, step, stt, lang):
        new_start = self.start_time
        new_start += step
        if new_start < self.end_time:
            buffer = self._get_split(new_start, self.end_time)
            audio = buffer2audiodata(buffer, self.sample_rate)
            transcript = stt.execute(audio, language=lang)
            score = self.calc_score(transcript, self.target_text)
        else:
            return self.confidence, self.best_transcript, self.start_time
        return score, transcript, new_start

    def _extend_left(self, step, stt, lang):
        new_start = self.start_time
        new_start -= step
        buffer = self._get_split(new_start, self.end_time)
        audio = buffer2audiodata(buffer, self.sample_rate)
        transcript = stt.execute(audio, language=lang)
        score = self.calc_score(transcript, self.target_text)
        return score, transcript, new_start

    def adjust_start_word(self, stt: STT, lang: str, step: float = 0.1):
        # - remove N miliseconds from start of buffer
        # - perform ASR
        # - calc new score
        # - if score decreases -> stop
        # - if score increases -> continue
        # - if last word matches -> stop
        start = self.start_time
        old_score = self.confidence
        twords = self.norm_text(self.target_text).split()
        firstw = twords[0]

        tqdm.write(f"⭐️ Target text | \"{self.target_text}\"")

        for i in range(20):  # TODO tqdm
            words = self.norm_text(self.best_transcript).split()
            # decide which direction of the buffer to check
            if (firstw in words[:5] or
                    (self.ending_matches and len(words) > len(twords))):
                tqdm.write("⚙️ ---> cropping left side of buffer")
                score, transcript, start = self._crop_left(step, stt, lang)
            else:
                tqdm.write("⚙️ <--- growing left side of buffer")
                score, transcript, start = self._extend_left(step, stt, lang)

            tqdm.write(f"📝 Transcription | \"{transcript}\"")

            if self.norm_text(transcript):

                fw = self.norm_text(transcript).split()[0]

                if not self.start_matches and fw == firstw:
                    # finished
                    tqdm.write("✔️ first word now matches!")
                    self.start_time = start
                    self.best_transcript = transcript
                    self.confidence = score
                    break

            if score < old_score:
                # making it worse
                tqdm.write(f"⬇️ New score: {score} (-{old_score - score})")
                break
            if score > old_score:
                # keep adjusting
                tqdm.write(f"⬆️ New score: {score} (+{score - old_score})")
            self.start_time = start
            self.best_transcript = transcript
            self.confidence = old_score = score
            tqdm.write(f"⏱ Timespan [{round(self.start_time, 2)}, {round(self.end_time, 2)}]")

    def _extend_right(self, step, stt, lang):
        new_end = self.end_time
        new_end += step
        buffer = self._get_split(self.start_time, new_end)
        audio = buffer2audiodata(buffer, self.sample_rate)
        transcript = stt.execute(audio, language=lang)
        score = self.calc_score(transcript, self.target_text)
        return score, transcript, new_end

    def _crop_right(self, step, stt, lang):
        new_end = self.end_time
        new_end -= step
        if new_end > self.start_time:
            buffer = self._get_split(self.start_time, new_end)
            audio = buffer2audiodata(buffer, self.sample_rate)
            transcript = stt.execute(audio, language=lang)
            score = self.calc_score(transcript, self.target_text)
        else:  # can't crop anymore
            return self.confidence, self.best_transcript, self.end_time
        return score, transcript, new_end

    def adjust_trailing_word(self, stt: STT, lang: str, step: float = 0.1):
        # - remove N miliseconds from end of buffer
        # - perform ASR
        # - calc new score
        # - if score decreases -> stop
        # - if score increases -> continue
        # - if last word matches -> stop
        old_score = self.confidence
        twords = self.norm_text(self.target_text).split()

        tqdm.write(f"⭐️ Target text | \"{self.target_text}\"")

        for i in range(20):  # TODO tqdm
            words = self.norm_text(self.best_transcript).split()
            # decide which direction of the buffer to check
            if (twords[-1] in words or
                    len(words) > len(self.norm_text(self.target_text).split())):
                tqdm.write("⚙️ cropping right side of buffer <---")
                score, transcript, end = self._crop_right(step, stt, lang)
            else:
                tqdm.write("⚙️ growing right side of buffer --->")
                score, transcript, end = self._extend_right(step, stt, lang)

            tqdm.write(f"📝 Transcription | \"{transcript}\"")

            if self.norm_text(transcript):
                lw = self.norm_text(transcript).split()[-1]

                if not self.ending_matches and lw == twords[-1]:
                    # finished
                    tqdm.write("✔️ last word now matches!")
                    self.end_time = end
                    self.best_transcript = transcript
                    self.confidence = score
                    break

            if score < old_score:
                # making it worse
                tqdm.write(f"⬇️ New score: {score} (-{old_score - score})")
                break
            if score > old_score:
                # keep adjusting
                tqdm.write(f"⬆️ New score: {score} (+{score - old_score})")
            self.end_time = end
            self.best_transcript = transcript
            self.confidence = old_score = score
            tqdm.write(f"⏱ Timespan [{round(self.start_time, 2)}, {round(self.end_time, 2)}]")


class OVOSAlignedAudioChunker:
    def __init__(self, stt: STT, lang="pt",
                 stt_medium: Optional[STT] = None,
                 stt_big: Optional[STT] = None,
                 step=500,
                 look_ahead_threshold=0.15,
                 max_look_ahead=5,  # seconds
                 stop_thresh=0.88,
                 max_extra_words=3,
                 window_size=200  # milliseconds
                 ):
        self.stt = stt  # best performance
        self.stt_med = stt_medium  # middle ground
        self.stt_big = stt_big  # best accuracy
        self.lang = lang
        self.STEP = step
        self.WINDOW = window_size
        self.MAX_EXTRA_WORDS = max_extra_words
        self.MAX_LOOK_AHEAD = max_look_ahead
        self.STOP_THRESH = stop_thresh  # stop looking ahead if score reaches this value
        self.THRESH = look_ahead_threshold  # stop looking ahead when score drops by this factor (begin of next sentence)

    def transcribe(self, file_name: Union[str, AudioData], preferred_model: str = "small") -> str:
        """Transcribe audio file using the specified STT."""
        if preferred_model == "large" and self.stt_big is None:
            preferred_model = "medium"
        if preferred_model == "medium" and self.stt_med is None:
            preferred_model = "small"

        if isinstance(file_name, str):
            audio = file2audiodata(file_name)
        else:
            audio = file_name

        transcript = ""
        try:
            if preferred_model == "large":
                transcript = self.stt_big.execute(audio, language=self.lang)
            elif preferred_model == "medium":
                transcript = self.stt_med.execute(audio, language=self.lang)
            else:
                transcript = self.stt.execute(audio, language=self.lang)
        except Exception as e:
            pass
        return transcript

    @staticmethod
    def _get_audio_chunk(start_t, end_t, wav, sample_rate):
        # Convert times to sample indices
        start_sample = int(start_t * sample_rate)
        end_sample = int(end_t * sample_rate)
        # save the chunk
        return wav[start_sample:end_sample]

    def _tx(self, target_txt, wav, start_t, end_t, sample_rate, high_score) -> Tuple[str, float, str]:

        # Use a temporary file for the chunk
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_wav:
            # get the chunk and save it
            chunk = OVOSAlignedAudioChunker._get_audio_chunk(start_t, end_t, wav, sample_rate)
            write(temp_wav.name, sample_rate, chunk)

            # Transcribe chunk
            model = "small"
            transcription = self.transcribe(temp_wav.name)
            if not transcription:
                current_score = 0
            else:
                current_score = AudioSplit.calc_score(target_txt, transcription)

            if current_score < high_score and self.stt_med is not None:
                model = "medium"
                transcription = self.transcribe(temp_wav.name, preferred_model=model)
                current_score = AudioSplit.calc_score(target_txt, transcription)

            if current_score < high_score and self.stt_big is not None:
                model = "large"
                transcription = self.transcribe(temp_wav.name, preferred_model=model)
                current_score = AudioSplit.calc_score(target_txt, transcription)

        return transcription, current_score, model

    def _process_sentence(self, target_txt: str,
                          audio_path: str,
                          start_t: float = 0) -> AudioSplit:
        target_nw = len(target_txt.split(" "))
        target_txt = target_txt.replace('"', "'")
        best_tx = ""
        high_score = 0
        end_t = 0

        sample_rate, wav = wavfile.read(audio_path)
        total_samples = len(wav)  # Get the total number of samples in the wav data
        duration_seconds = total_samples / sample_rate  # Calculate total duration in seconds
        tqdm.write(f"🎵 Audio duration: {duration_seconds} seconds")

        fake_bar = tqdm(total=100, desc=f"Similarity score", ncols=80, leave=False,
                        bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt}\n')

        start_ms = int(start_t * 1000)
        end_ms = int(duration_seconds * 1000)
        tqdm.write(f"⚙️ Parsing audio from {start_ms} to {end_ms}ms")
        for e in range(start_ms, end_ms, self.STEP):
            if e == start_ms:
                continue
            e = e / 1000  # back to seconds
            stx, s, m = self._tx(target_txt, wav, start_t, e, sample_rate, high_score)
            if stx == best_tx:
                tqdm.write(f"🤷‍♂️️ Transcription didn't change with extra {self.STEP} milliseconds of audio")
            else:
                tqdm.write(f"📝 Transcription | {stx}")

            if s < high_score and end_t:
                tqdm.write(
                    f"⚠️ Score: {round(s, 3)} (-{round(high_score - s, 3)}%)  | ⏱ Timespan [{round(start_t, 2)}, {round(e, 2)}] | Model: {m}")
                if (s < high_score - self.THRESH and
                        m == "large" and
                        len(stx) >= len(target_txt)):
                    # tqdm.write(f"⚠️ moving to next sentence | ... {stx[-20:]}")
                    fake_bar.close()
                    break

            nw = len(stx.split())
            if best_tx and nw > target_nw + self.MAX_EXTRA_WORDS:
                tqdm.write(f"‼️ transcription is adding too many extra words (+{nw - target_nw})")
                fake_bar.close()
                break

            if s > high_score or not best_tx:
                best_tx = stx
                high_score = s
                end_t = e
                tqdm.write(
                    f"🔍 Score: {round(s, 4)} | ⏱ Timespan [{round(start_t, 2)}, {round(end_t, 2)}] | Model: {m}")

                # Update fake progress bar
                fake_bar.n = int(s * 100)
                fake_bar.refresh()

                if high_score >= self.STOP_THRESH:
                    tqdm.write(f"⚠️ Score reached threshold: {high_score}")
                    fake_bar.close()
                    break

            elif end_t and e - end_t >= self.MAX_LOOK_AHEAD:
                tqdm.write(f"‼️️ Score didn't improve after {e - end_t} seconds, moving to next sentence")
                fake_bar.close()
                break

        audio_chunk = AudioSplit(start_time=start_t,
                                 end_time=end_t,
                                 target_text=target_txt,
                                 best_transcript=best_tx,
                                 confidence=high_score,
                                 full_audio=wav,
                                 sample_rate=sample_rate)

        if not audio_chunk.start_matches:
            tqdm.write(
                f"❌️ Split failed validation 1/3, start of transcript doesnt match start of text | "
                f"🔍 Score: {round(audio_chunk.confidence, 4)}")

            tqdm.write("🔧 trying to adjust start word")
            audio_chunk.adjust_start_word(stt=self.stt_big or self.stt_med or self.stt,
                                          lang=self.lang)
        else:
            tqdm.write("✔️ Split passed validation 1/3")

        if not audio_chunk.ending_matches:
            tqdm.write(
                f"❌️ Split failed validation 2/3, end of transcript doesnt match end of text | "
                f"🔍 Score: {round(audio_chunk.confidence, 4)}")

            tqdm.write("🔧 trying to adjust trailing word")
            audio_chunk.adjust_trailing_word(stt=self.stt_big or self.stt_med or self.stt,
                                             lang=self.lang)
        else:
            tqdm.write("✔️ Split passed validation 2/3")

        # reached end of file?
        fake_bar.close()
        return audio_chunk

    def _get_input(self, audio_chunk: AudioSplit, output_path: str):
        tqdm.write(f"⭐️ Target text | \"{audio_chunk.target_text}\"")
        tqdm.write(f"📝 Transcription | \"{audio_chunk.best_transcript}\"")

        # whats wrong:
        # 1 - missing audio at start
        # 2 - missing audio at end
        # 3 - extra audio at start
        # 4 - extra audio at end
        # 5 - low quality/uninteligible sample (skip)
        write(output_path, audio_chunk.sample_rate, audio_chunk.buffer)
        opt = -1
        while opt > 6 or opt < 0:
            tqdm.write(f"⏱ Timespan [{round(audio_chunk.start_time, 2)}, {round(audio_chunk.end_time, 2)}]")
            tqdm.write(f"🎶 Listen to audio sample: \"{output_path}\"")
            tqdm.write("0 - Looks good!")
            tqdm.write("1 - Search for missing audio before start (auto)")  # TODO
            tqdm.write("2 - Search for missing audio after end (auto)")  # TODO
            tqdm.write("3 - Remove extra audio at beginning (auto)")  # TODO
            tqdm.write("4 - Remove extra audio at end (auto)")  # TODO
            tqdm.write(f"5 - Skip text and keep audio (next start_time: {audio_chunk.start_time})")
            tqdm.write(f"6 - Skip text and discard audio (next start_time: {audio_chunk.end_time})")
            # tqdm.write(f"7 - Retry text and discard audio (next start_time: {audio_chunk.end_time})")  # TODO
            try:
                opt = int(input("✍️ Enter your choice (0-6): "))
            except:
                pass
            if opt > 6 or opt < 0:
                tqdm.write("❌️ Invalid option, try again")
        return opt

    def _validate(self, audio_chunk: AudioSplit,
                  prev_audio_chunk: AudioSplit,
                  output_path: str):

        if audio_chunk.confidence < 0.9:
            tqdm.write(
                f"❌️ Split failed validation 3/3, confidence too low | "
                f"🔍 Score: {round(audio_chunk.confidence, 4)}")

            selected_opt = self._get_input(audio_chunk, output_path)

            def do_action(opt, audio: AudioSplit):
                stt = self.stt_big or self.stt_med or self.stt
                if opt == 1:
                    tqdm.write("🔧 searching missing audio before start")
                    audio.confidence, audio.best_transcript, audio.start_time = audio._extend_left(1, stt, self.lang)
                    tqdm.write(f"⏱ Timespan [{round(audio.start_time, 2)}, {round(audio.end_time, 2)}]")
                    audio.adjust_start_word(stt, self.lang)
                elif opt == 2:
                    tqdm.write("🔧 searching missing audio after chunk end")
                    audio.confidence, audio.best_transcript, audio.end_time = audio._extend_right(1, stt, self.lang)
                    tqdm.write(f"⏱ Timespan [{round(audio.start_time, 2)}, {round(audio.end_time, 2)}]")
                    audio.adjust_trailing_word(stt, self.lang)
                elif opt == 3:
                    tqdm.write("🔧 removing extra audio at beginning of buffer")
                    audio.confidence, audio.best_transcript, audio.start_time = audio._crop_left(0.5, stt, self.lang)
                    tqdm.write(f"⏱ Timespan [{round(audio.start_time, 2)}, {round(audio.end_time, 2)}]")
                    audio.adjust_start_word(stt, self.lang)
                elif opt == 4:
                    tqdm.write("🔧 removing extra audio at end of buffer")
                    audio.confidence, audio.best_transcript, audio.end_time = audio._crop_right(0.5, stt, self.lang)
                    tqdm.write(f"⏱ Timespan [{round(audio.start_time, 2)}, {round(audio.end_time, 2)}]")
                    audio.adjust_trailing_word(stt, self.lang)

                tqdm.write(f"⏱ Timespan [{round(audio.start_time, 2)}, {round(audio.end_time, 2)}]")
                return audio

            def do_skip_action(opt, audio: AudioSplit):
                if opt == 5:
                    tqdm.write(f"❌️ Skipping text, rolling back audio to: {audio_chunk.start_time}s")
                    audio.end_time = audio.start_time
                elif opt == 6:
                    tqdm.write(f"❌️ Skipping text and discarding audio, advancing to: {audio_chunk.end_time}s")
                    # audio.end_time = audio.end_time
                elif opt == 7:
                    tqdm.write(f"❌️ Retrying text and discarding audio, advancing to: {audio_chunk.end_time}s")
                    # audio.end_time = audio.end_time
                return audio

            if selected_opt == 0:
                tqdm.write("✔️ Split passed validation 3/3")
                return audio_chunk, prev_audio_chunk
            elif selected_opt < 5:
                audio_chunk = prev_audio_chunk = do_action(selected_opt, audio_chunk)
                # validate until OK or Skip is selected
                return self._validate(audio_chunk, prev_audio_chunk, output_path)
            else:
                prev_audio_chunk = do_skip_action(selected_opt, prev_audio_chunk)
        else:
            tqdm.write("✔️ Split passed validation 3/3")
        return audio_chunk, prev_audio_chunk

    def align(self, audio_path: str,
              output_folder: str,
              sentences: List[str],
              cache: dict) -> Iterable[AudioSplit]:

        os.makedirs(output_folder, exist_ok=True)

        prev_audio_chunk = AudioSplit(end_time=0)
        for idx, target_txt in enumerate(sentences):
            tqdm.write(f"⭐️ New chunk | Target text: \"{target_txt}\"")

            if f"{target_txt}_{idx}" in cache:
                start_t, end_t, target_txt, best_tx, high_score = cache[f"{target_txt}_{idx}"]
                tqdm.write(f"✔️ Found audio split in cache | {target_txt}_{idx}")

                sample_rate, wav = wavfile.read(audio_path)
                prev_audio_chunk = AudioSplit(start_time=start_t,
                                              end_time=end_t,
                                              full_audio=wav,
                                              sample_rate=sample_rate,
                                              target_text=target_txt,
                                              best_transcript=best_tx,
                                              confidence=high_score)
                yield prev_audio_chunk
                continue

            output_path = f'{output_folder}/chunk_{idx + 1}.wav'
            audio_chunk = prev_audio_chunk = self._process_sentence(target_txt, audio_path,
                                                                    start_t=prev_audio_chunk.end_time)
            if audio_chunk.buffer is None:
                continue

            audio_chunk, prev_audio_chunk = self._validate(audio_chunk, prev_audio_chunk, output_path)
            tqdm.write(f"✅️ Found audio split | {target_txt}")
            yield audio_chunk


if __name__ == "__main__":

    def get_sentences(f) -> Iterable[str]:
        """implementation specific, contains sentences that should correspond to final chunks"""
        SPKRS = f"{os.path.dirname(__file__)}/speakers.json"
        transcripts = JsonStorage(SPKRS)
        k = [_ for _ in transcripts if _.endswith(f)][0]
        for s, txt in transcripts[k]:
            for t in sentence_tokenize(txt.replace("...", ",")):
                if AudioSplit.norm_text(t):
                    yield t


    DATASET = "/home/miro/PycharmProjects/dataset_pt/Spoken Portuguese Geographical and Social Varieties"
    OUT = "/tmp"

    stt = MyNorthAISTT({"model": "my-north-ai/whisper-small-pt",
                        "use_cuda": True})
    stt2 = MyNorthAISTT({"model": "my-north-ai/whisper-medium-pt",
                         "use_cuda": True})
    stt3 = MyNorthAISTT({"model": "my-north-ai/whisper-large-v3-pt",
                         "use_cuda": False})  # doesnt fit in my GPU....
    CHUNKER = OVOSAlignedAudioChunker(stt,
                                      stt_medium=stt2,
                                      stt_big=stt3)

    for region in os.listdir(DATASET):
        CACHE = JsonStorage(f"{os.path.dirname(__file__)}/splits_{region}.json")
        for root, _, files in os.walk(f"{DATASET}/{region}"):
            for f in files:
                if f"{region}/{f}" not in CACHE:
                    CACHE[f"{region}/{f}"] = {}

                if f.endswith(".wav"):
                    idx = 0
                    for audio_split in CHUNKER.align(
                            audio_path=f"{root}/{f}",
                            output_folder=f"{OUT}/{region}/{f.split('.wav')[0]}",
                            sentences=list(get_sentences(f)),
                            cache=CACHE[f"{region}/{f}"]):

                        tqdm.write(f"🟢 Best Transcription | {audio_split.best_transcript}")
                        tqdm.write(
                            f"✔️ Processed segment:"
                            f" start={round(audio_split.start_time, 2)}s,"
                            f" end={round(audio_split.end_time, 2)}s |"
                            f" Score: {round(audio_split.confidence, 3)}")

                        if f"{audio_split.target_text}_{idx}" not in CACHE[f"{region}/{f}"]:
                            CACHE[f"{region}/{f}"][f"{audio_split.target_text}_{idx}"] = [audio_split.start_time,
                                                                                          audio_split.end_time,
                                                                                          audio_split.target_text,
                                                                                          audio_split.best_transcript,
                                                                                          audio_split.confidence]
                            CACHE.store()
                            tqdm.write(f"🟢 Cached new split data | {region}/{f}")
                        idx += 1