batchalign 0.7.11b3__tar.gz → 0.7.11b4__tar.gz

{batchalign-0.7.11b3/batchalign.egg-info → batchalign-0.7.11b4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: batchalign
-Version: 0.7.11b3
+Version: 0.7.11b4
 Summary: Python Speech Language Sample Analysis
 Author: Brian MacWhinney, Houjun Liu
 Author-email: macw@cmu.edu, houjun@cmu.edu

{batchalign-0.7.11b3 → batchalign-0.7.11b4}/batchalign/cli/cli.py

@@ -107,10 +107,12 @@ batchalign.add_command(train, "models")
 @common_options
 @click.option("--whisper/--rev",
               default=False, help="For utterance timing recovery, OpenAI Whisper (ASR) instead of Rev.AI (default).")
+@click.option("--wav2vec/--whisper_fa",
+              default=False, help="Use Whisper instead of Wav2Vec for English (defaults for Whisper for non-English)")
 @click.option("--pauses", type=bool, default=False, help="Should we try to bullet each word or should we try to add pauses in between words by grouping them? Default: no pauses.", is_flag=True)
 
 @click.pass_context
-def align(ctx, in_dir, out_dir, whisper, **kwargs):
+def align(ctx, in_dir, out_dir, whisper, wav2vec, **kwargs):
     """Align transcripts against corresponding media files."""
     def loader(file):
         return (
@@ -121,12 +123,22 @@ def align(ctx, in_dir, out_dir, whisper, **kwargs):
     def writer(doc, output):
         CHATFile(doc=doc).write(output)
 
-    _dispatch("align", "eng", 1,
-              ["cha"], ctx,
-              in_dir, out_dir,
-              loader, writer, C,
-              utr="whisper_utr" if whisper else "rev_utr",
-              **kwargs)
+    if not wav2vec:
+        _dispatch("align", "eng", 1,
+                  ["cha"], ctx,
+                  in_dir, out_dir,
+                  loader, writer, C,
+                  fa="whisper_fa",
+                  utr="whisper_utr" if whisper else "rev_utr",
+                  **kwargs)
+    else:
+        _dispatch("align", "eng", 1,
+                  ["cha"], ctx,
+                  in_dir, out_dir,
+                  loader, writer, C,
+                  fa="wav2vec_fa",
+                  utr="whisper_utr" if whisper else "rev_utr",
+                  **kwargs)
 
 #################### TRANSCRIBE ################################
 

{batchalign-0.7.11b3 → batchalign-0.7.11b4}/batchalign/models/__init__.py

@@ -3,4 +3,4 @@ from .whisper import WhisperASRModel, WhisperFAModel
 from .speaker import NemoSpeakerModel
 from .utils import ASRAudioFile
 from .resolve import resolve
-
+from .wave2vec import Wave2VecFAModel

batchalign-0.7.11b4/batchalign/models/wave2vec/__init__.py

@@ -0,0 +1 @@
+from .infer_fa import Wave2VecFAModel

batchalign-0.7.11b4/batchalign/models/wave2vec/infer_fa.py

@@ -0,0 +1,135 @@
+from transformers import WhisperProcessor, WhisperTokenizer, WhisperForConditionalGeneration
+
+import torch
+from torchaudio import load
+from torchaudio import transforms as T
+from batchalign.models.utils import ASRAudioFile
+
+import torchaudio
+bundle = torchaudio.pipelines.MMS_FA
+import torchaudio.functional as AF
+
+import numpy as np
+
+import logging
+L = logging.getLogger("batchalign")
+
+# DEVICE = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+DEVICE = torch.device('cuda') if torch.cuda.is_available() else torch.device("mps") if torch.backends.mps.is_available() else torch.device('cpu')
+TIME_PRECISION = 0.02
+
+# inference engine
+class Wave2VecFAModel(object):
+    """An Forced Alignment engine built out of whisper
+
+    Parameters
+    ----------
+    model : str
+        The model path to load from.
+    target_sample_rate : optional, int
+        The sample rate to cast to. Defaults 16000 by Whisper.
+
+    Example
+    -------
+    >>> engine = Wave2VecFAModel()
+    >>> file = engine.load("./data/myfile.wav")
+    >>> timestamps = engine(audio=file.chunk(0, 1500), text="this is my transcript") # FA
+    """
+
+    def __init__(self, target_sample_rate=16000):
+        L.debug("Initializing Wave2vec FA model")
+        self.model = bundle.get_model().to(DEVICE)
+        L.debug("Wave2Vec FA initialization done.")
+
+        # save the target sample rate
+        self.sample_rate = target_sample_rate
+
+    def load(self, f):
+        """Load an audio file for procesing.
+
+        Parameters
+        ----------
+        f : str
+            The audio .wav file name to process.
+        num_speakers : int
+            The number of speakers
+
+        Returns
+        -------
+        Tuple[ASRAudioFile, List[dict]]
+            Return processed audio file and speaker segments.
+        """
+
+        # function: load and resample audio
+        audio_arr, rate = load(f)
+
+        # resample if needed
+        if rate != self.sample_rate:
+            audio_arr = T.Resample(rate, self.sample_rate)(audio_arr)
+
+        # transpose and mean
+        resampled = torch.mean(audio_arr.transpose(0,1), dim=1)
+
+        # and return the audio file
+        return ASRAudioFile(f, resampled, self.sample_rate)
+
+    def __call__(self, audio, text):
+        """Run forced alignment on the audio file.
+
+        Arguments
+        ----------
+        audio : tensor
+            The audio file to process.
+        text : str
+            The transcript to align to.
+
+        Returns
+        -------
+        List[Tuple[str, Tuple[int, int]]]
+                A list of speaker segments
+        """
+
+        L.debug("Running Wav2Vec word-level forced alignment...")
+
+        # complete the call function, don't write anything else
+        L.debug("Running Wav2Vec word-level forced alignment...")
+
+        # Move audio to device and normalize
+        audio = audio.to(DEVICE)
+
+        # Get emission matrix from model
+        emission, _ = self.model(audio.unsqueeze(0))
+        emission = emission.cpu().detach()
+
+        # Get tokens and transcript
+        dictionary = bundle.get_dict()
+
+        # Convert text to tokens 
+        transcript = torch.tensor([dictionary.get(c, dictionary["*"])
+                                   for word in text
+                                   for c in word.lower()])
+
+        # Run forced alignment
+        path, scores = AF.forced_align(emission, transcript.unsqueeze(0))
+        alignments, scores = path[0], scores[0]
+        scores = scores.exp()
+
+        # Merge repeated tokens and remove blanks
+        path = AF.merge_tokens(alignments, scores)
+
+        def unflatten(list_, lengths):
+            assert len(list_) == sum(lengths)
+            i = 0
+            ret = []
+            for l in lengths:
+                ret.append(list_[i : i + l])
+                i += l
+            return ret
+
+        # Unflatten to get character-level alignments
+        word_spans = unflatten(path, [len(word) for word in text])
+        ratio = audio.size(0)/emission.size(1)
+        word_spans = [(int(((spans[0].start*ratio)/self.sample_rate)*1000),
+                       int(((spans[-1].end*ratio)/self.sample_rate)*1000)) for spans in word_spans]
+
+        return list(zip(text, word_spans))

{batchalign-0.7.11b3 → batchalign-0.7.11b4}/batchalign/pipelines/__init__.py

@@ -6,7 +6,7 @@ from .morphosyntax import StanzaEngine, CorefEngine
 from .cleanup import NgramRetraceEngine, DisfluencyReplacementEngine
 from .speaker import NemoSpeakerEngine
 
-from .fa import WhisperFAEngine
+from .fa import WhisperFAEngine, Wave2VecFAEngine
 from .utr import WhisperUTREngine, RevUTREngine
 
 from .analysis import EvaluationEngine

{batchalign-0.7.11b3 → batchalign-0.7.11b4}/batchalign/pipelines/dispatch.py

@@ -6,7 +6,7 @@ Tabulate default packages and options.
 from batchalign import (WhisperEngine, WhisperFAEngine, StanzaEngine, RevEngine,
                         NgramRetraceEngine, DisfluencyReplacementEngine, WhisperUTREngine,
                         RevUTREngine, EvaluationEngine, WhisperXEngine, NemoSpeakerEngine,
-                        StanzaUtteranceEngine, CorefEngine)
+                        StanzaUtteranceEngine, CorefEngine, Wave2VecFAEngine)
 from batchalign import BatchalignPipeline
 from batchalign.models import resolve
 
@@ -127,7 +127,8 @@ def dispatch_pipeline(pkg_str, lang, num_speakers=None, **arg_overrides):
             engines.append(StanzaUtteranceEngine())
         elif engine == "stanza_coref":
             engines.append(CorefEngine())
-
+        elif engine == "wav2vec_fa":
+            engines.append(Wave2VecFAEngine())
 
     L.debug(f"Done initalizing packages.")
     return BatchalignPipeline(*engines)

batchalign-0.7.11b4/batchalign/pipelines/fa/__init__.py

@@ -0,0 +1,2 @@
+from .whisper_fa import WhisperFAEngine
+from .wave2vec_fa import Wave2VecFAEngine

batchalign-0.7.11b4/batchalign/pipelines/fa/wave2vec_fa.py

@@ -0,0 +1,162 @@
+from batchalign.models import Wave2VecFAModel
+from batchalign.document import *
+from batchalign.pipelines.base import *
+from batchalign.utils import *
+from batchalign.utils.dp import *
+from batchalign.constants import *
+
+import logging
+L = logging.getLogger("batchalign")
+
+import re
+
+import pycountry
+import warnings
+
+class Wave2VecFAEngine(BatchalignEngine):
+    tasks = [ Task.FORCED_ALIGNMENT ]
+
+    def _hook_status(self, status_hook):
+        self.status_hook = status_hook
+
+    def __init__(self):
+        self.status_hook = None
+        self.__wav2vec = Wave2VecFAModel()
+
+    def process(self, doc:Document, **kwargs):
+        # check that the document has a media path to align to
+        assert doc.media != None and doc.media.url != None, f"We cannot forced-align something that doesn't have a media path! Provided media tier='{doc.media}'"
+
+        if doc.langs[0] != "eng":
+            warnings.warn("Looks like you are not aligning English with wav2vec; this works for a lot of Roman languages, but outside of that your milage may vary.")
+
+        # load the audio file
+        L.debug(f"Wave2Vec FA is loading url {doc.media.url}...")
+        f = self.__wav2vec.load(doc.media.url)
+        L.debug(f"Wav2Vec FA finished loading media.")
+
+        # collect utterances 30 secondish segments to be aligned for whisper
+        # we have to do this because whisper does poorly with very short segments
+        groups = []
+        group = []
+        seg_start = 0
+
+        L.debug(f"Wav2Vec FA finished loading media.")
+
+        for i in doc.content:
+            if not isinstance(i, Utterance):
+                continue
+            if i.alignment == None:
+                warnings.warn("We found at least one utterance without utterance-level alignment; this is usually not an issue, but if the entire transcript is unaligned, it means that utterance level timing recovery (which is fuzzy using ASR) failed due to the audio clarity. On this transcript, before running forced-alignment, please supply utterance-level links.")
+                continue
+
+            # pop the previous group onto the stack
+            if (i.alignment[-1] - seg_start) > 20*1000:
+                groups.append(group)
+                group = []
+                seg_start = i.alignment[0]
+
+            # append the contents to the running group
+            for word in i.content:
+                group.append((word, i.alignment))
+
+        groups.append(group)
+
+        L.debug(f"Begin Wav2Vec Inference...")
+
+        for indx, grp in enumerate(groups):
+            L.info(f"Wave2Vec FA processing segment {indx+1}/{len(groups)}...")
+            if self.status_hook != None:
+                self.status_hook(indx+1, len(groups))
+
+            # perform alignment
+            # we take a 2 second buffer in each direction
+            try:
+                transcript = [word[0].text for word in grp]
+                # replace ANY punctuation
+                for p in MOR_PUNCT + ENDING_PUNCT:
+                    transcript = [i.replace("_", " ") for i in transcript if i.strip() != p]
+                # if "noone's" in detokenized:
+                    # breakpoint()
+                res = self.__wav2vec(audio=f.chunk(grp[0][1][0], grp[-1][1][1]), text=transcript)
+            except IndexError:
+                # utterance contains nothing
+                continue
+
+            # create reference backplates, which are the word ids to set the timing for
+            ref_targets = []
+            for indx, (word, _) in enumerate(grp):
+                for char in word.text:
+                    ref_targets.append(ReferenceTarget(char, payload=indx))
+            # create target backplates for the timings
+            payload_targets = []
+            timings = []
+            for indx, (word, time) in enumerate(res):
+                timings.append(time)
+                for char in word:
+                    payload_targets.append(PayloadTarget(char, payload=indx))
+            # alignment!
+            alignments = align(payload_targets, ref_targets, tqdm=False)
+
+            # set the ids back to the text ids
+            # we do this BACKWARDS because we went to have the first timestamp
+            # we get about a word first
+            alignments.reverse()
+            for indx,elem in enumerate(alignments):
+                if isinstance(elem, Match):
+                    grp[elem.reference_payload][0].time = (int(round((timings[elem.payload][0] +
+                                                                      grp[0][1][0]))),
+                                                           int(round((timings[elem.payload][1] +
+                                                                      grp[0][1][0]))))
+
+        L.debug(f"Correcting text...")
+
+        # we now set the end alignment of each word to the start of the next
+        for doc_ut, ut in enumerate(doc.content):
+            if not isinstance(ut, Utterance):
+                continue
+
+            # correct each word by bumping it forward
+            # and if its not a word we remove the timing
+            for indx, w in enumerate(ut.content):
+                if w.type in [TokenType.PUNCT, TokenType.FEAT, TokenType.ANNOT]:
+                    w.time = None
+                elif indx == len(ut.content)-1 and w.text in ENDING_PUNCT:
+                    w.time = None
+                elif indx != len(ut.content)-1:
+                    # search forward for the next compatible time
+                    tmp = indx+1
+                    while tmp < len(ut.content)-1 and ut.content[tmp].time == None:
+                        tmp += 1
+                    if w.time == None:
+                        continue
+                    if ut.content[tmp].time == None:
+                        # seek forward one utterance to find their start time
+                        next_ut = doc_ut + 1 
+                        while next_ut < len(doc.content)-1 and (not isinstance(doc.content, Utterance) or doc.content[next_ut].alignment == None):
+                            next_ut += 1
+                        if next_ut < len(doc.content) and isinstance(doc.content, Utterance) and doc.content[next_ut].alignment:
+                            w.time = (w.time[0], doc.content[next_ut].alignment[0])
+                        else:
+                            w.time = (w.time[0], w.time[0]+500) # give half a second because we don't know
+
+                    # just in case, bound the time by the utterance derived timings
+                    if ut.alignment and ut.alignment[0] != None:
+                        w.time = (max(w.time[0], ut.alignment[0]), min(w.time[1], ut.alignment[1]))
+                    # if we ended up with timings that don't make sense, drop it
+                    if w.time and w.time[0] >= w.time[1]:
+                        w.time = None
+
+            # clear any built-in timing (i.e. we should use utterance-derived timing)
+            ut.time = None
+            # correct the text 
+            if ut.alignment and ut.text != None:
+                if '\x15' not in ut.text:
+                    ut.text = (ut.text+f" \x15{ut.alignment[0]}_{ut.alignment[1]}\x15").strip()
+                else:
+                    ut.text = re.sub("\x15\d+_\d+\x15",
+                                     f"\x15{ut.alignment[0]}_{ut.alignment[1]}\x15", ut.text).strip()
+            elif ut.text != None:
+                ut.text = re.sub("\x15\d+_\d+\x15", f"", ut.text).strip()
+
+        return doc

batchalign-0.7.11b4/batchalign/version

@@ -0,0 +1,3 @@
+0.7.11-beta.4
+Feburary 6nd, 2025
+Wav2vec support!

{batchalign-0.7.11b3 → batchalign-0.7.11b4/batchalign.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: batchalign
-Version: 0.7.11b3
+Version: 0.7.11b4
 Summary: Python Speech Language Sample Analysis
 Author: Brian MacWhinney, Houjun Liu
 Author-email: macw@cmu.edu, houjun@cmu.edu

{batchalign-0.7.11b3 → batchalign-0.7.11b4}/batchalign.egg-info/SOURCES.txt

@@ -45,6 +45,8 @@ batchalign/models/utterance/execute.py
 batchalign/models/utterance/infer.py
 batchalign/models/utterance/prep.py
 batchalign/models/utterance/train.py
+batchalign/models/wave2vec/__init__.py
+batchalign/models/wave2vec/infer_fa.py
 batchalign/models/whisper/__init__.py
 batchalign/models/whisper/infer_asr.py
 batchalign/models/whisper/infer_fa.py
@@ -68,6 +70,7 @@ batchalign/pipelines/cleanup/support/filled_pauses.eng
 batchalign/pipelines/cleanup/support/replacements.eng
 batchalign/pipelines/cleanup/support/test.test
 batchalign/pipelines/fa/__init__.py
+batchalign/pipelines/fa/wave2vec_fa.py
 batchalign/pipelines/fa/whisper_fa.py
 batchalign/pipelines/morphosyntax/__init__.py
 batchalign/pipelines/morphosyntax/coref.py

batchalign-0.7.11b3/batchalign/pipelines/fa/__init__.py

@@ -1 +0,0 @@
-from .whisper_fa import WhisperFAEngine

batchalign-0.7.11b3/batchalign/version

@@ -1,3 +0,0 @@
-0.7.11-beta.3
-Feburary 2nd, 2025
-Incorporate additional pauses