taters 0.1.0__py3-none-any.whl

taters/Taters.py

@@ -0,0 +1,115 @@
+# taters/Taters.py
+from __future__ import annotations
+from pathlib import Path
+from typing import Any
+import inspect
+
+def _abs(p: str | Path) -> str:
+    return str(Path(p).resolve())
+
+def _forward(func, kwargs: dict[str, Any]):
+    sig = inspect.signature(func)
+    try:
+        # validate names & requireds; don't execute defaults here
+        sig.bind_partial(**kwargs)
+    except TypeError as e:
+        allowed = ", ".join([str(p) for p in sig.parameters.values()])
+        raise TypeError(f"{func.__module__}.{func.__name__}: {e}\nAllowed params: {allowed}")
+    return func(**kwargs)
+
+
+class Taters:
+    def __init__(self):
+        self.audio = _AudioAPI(self)
+        self.text = _TextAPI(self)
+        self.helpers = _HelpersAPI(self)
+
+    # back-compat pass-throughs
+    
+    # audio
+    def convert_to_wav(self, **kwargs):                         return self.audio.convert_to_wav(**kwargs)
+    def extract_wavs_from_video(self, **kwargs):                return self.audio.extract_wavs_from_video(**kwargs)
+    def split_wav_by_speaker(self, **kwargs):                   return self.audio.split_wav_by_speaker(**kwargs)
+    def extract_whisper_embeddings(self, **kwargs):             return self.audio.extract_whisper_embeddings(**kwargs)
+    def diarize_with_thirdparty(self, **kwargs):                return self.audio.diarize_with_thirdparty(**kwargs)
+    
+    #text
+    def analyze_with_dictionaries(self, **kwargs):              return self.text.analyze_with_dictionaries(**kwargs)
+    def analyze_with_archetypes(self, **kwargs):                return self.text.analyze_with_archetypes(**kwargs)
+    def extract_sentence_embeddings(self, **kwargs):            return self.text.extract_sentence_embeddings(**kwargs)
+    
+    # helpers
+    def txt_folder_to_analysis_ready_csv(self, **kwargs):       return self.helpers.txt_folder_to_analysis_ready_csv(**kwargs)
+    def csv_to_analysis_ready_csv(self, **kwargs):              return self.helpers.csv_to_analysis_ready_csv(**kwargs)
+    def find_files(self, **kwargs):                             return self.helpers.find_files(**kwargs)
+    def feature_gather(self, **kwargs):                         return self.helpers.feature_gather(**kwargs)
+    
+
+    def txt_folder_to_analysis_ready_csv(self, **kwargs):
+        from .helpers.text_gather import txt_folder_to_analysis_ready_csv
+        return _forward(txt_folder_to_analysis_ready_csv, kwargs)
+    
+    def csv_to_analysis_ready_csv(self, **kwargs):
+        from .helpers.text_gather import csv_to_analysis_ready_csv
+        return _forward(csv_to_analysis_ready_csv, kwargs)
+
+
+class _AudioAPI:
+    def __init__(self, parent: Taters): self._cs = parent
+
+    def convert_to_wav(self, **kwargs):
+        from .audio.convert_to_wav import convert_audio_to_wav
+        return _forward(convert_audio_to_wav, kwargs)
+
+    def extract_wavs_from_video(self, **kwargs):
+        from .audio.extract_wav_from_video import split_audio_streams_to_wav
+        return _forward(split_audio_streams_to_wav, kwargs)
+
+    def split_wav_by_speaker(self, **kwargs):
+        from .audio.split_wav_by_speaker import make_speaker_wavs_from_csv
+        return _forward(make_speaker_wavs_from_csv, kwargs)
+
+    def extract_whisper_embeddings(self, **kwargs):
+        from .audio.extract_whisper_embeddings import extract_whisper_embeddings
+        return _forward(extract_whisper_embeddings, kwargs)
+
+    def diarize_with_thirdparty(self, **kwargs):
+        from .audio.diarizer.whisper_diar_wrapper import run_whisper_diarization_repo
+        return _forward(run_whisper_diarization_repo, kwargs)
+
+
+class _TextAPI:
+    def __init__(self, parent: Taters): self._cs = parent
+
+    def analyze_with_dictionaries(self, **kwargs):
+        from .text.analyze_with_dictionaries import analyze_with_dictionaries
+        return _forward(analyze_with_dictionaries, kwargs)
+    
+    def analyze_with_archetypes(self, **kwargs):
+        from .text.analyze_with_archetypes import analyze_with_archetypes
+        return _forward(analyze_with_archetypes, kwargs)
+    
+    def extract_sentence_embeddings(self, **kwargs):
+        from .text.extract_sentence_embeddings import analyze_with_sentence_embeddings
+        return _forward(analyze_with_sentence_embeddings, kwargs)
+
+
+class _HelpersAPI:
+    def __init__(self, parent: Taters): self._cs = parent
+    
+    def txt_folder_to_analysis_ready_csv(self, **kwargs):
+        from .helpers.text_gather import txt_folder_to_analysis_ready_csv
+        return _forward(txt_folder_to_analysis_ready_csv, kwargs)
+    
+    def csv_to_analysis_ready_csv(self, **kwargs):
+        from .helpers.text_gather import csv_to_analysis_ready_csv
+        return _forward(csv_to_analysis_ready_csv, kwargs)
+    
+    def find_files(self, **kwargs):
+        from .helpers.find_files import find_files
+        return _forward(find_files, kwargs)
+    
+    def feature_gather(self, **kwargs):
+        from .helpers import feature_gather
+        return _forward(feature_gather, kwargs)
+    

taters/audio/convert_to_wav.py

@@ -0,0 +1,112 @@
+# taters/audio/convert_to_wav.py
+from __future__ import annotations
+import shutil
+import subprocess
+from pathlib import Path
+from typing import Optional, Union
+
+class FFmpegNotFoundError(RuntimeError):
+    pass
+
+def _check_ffmpeg():
+    if shutil.which("ffmpeg") is None or shutil.which("ffprobe") is None:
+        raise FFmpegNotFoundError("ffmpeg and/or ffprobe not found on PATH.")
+
+def convert_audio_to_wav(
+    input_path: Union[str, Path],
+    *,
+    output_path: Optional[Union[str, Path]] = None,
+    output_dir: Optional[Union[str, Path]] = None,
+    sample_rate: int = 16000,          # common for ASR
+    bit_depth: int = 16,               # 16/24/32 signed PCM
+    channels: int = 1,                 # 1=mono, 2=stereo
+    overwrite_existing: bool = False,  # if the file already exists, let's not overwrite by default
+) -> Path:
+    """
+    Convert any audio (or A/V container) to a PCM WAV file using ffmpeg.
+
+    If output_path is None and output_dir is None, writes <input_stem>.wav next to input.
+    If output_dir is given (and output_path is None), writes <output_dir>/<input_stem>.wav.
+    If output_path is given, it takes precedence.
+
+    Returns the Path to the created WAV.
+    """
+    _check_ffmpeg()
+
+    in_path = Path(input_path).resolve()
+    if not in_path.exists():
+        raise FileNotFoundError(f"Input file not found: {in_path}")
+
+    if output_path and output_dir:
+        raise ValueError("Provide at most one of output_path or output_dir.")
+
+    if output_path:
+        out_path = Path(output_path).resolve()
+    else:
+        base = in_path.stem + ".wav"
+        out_dir = Path(output_dir).resolve() if output_dir else Path.cwd() / "audio"
+        out_dir.mkdir(parents=True, exist_ok=True)
+        out_path = out_dir / base
+
+    if not overwrite_existing and Path(out_path).is_file():
+        print("WAV file already exists; returning existing file.")
+        return out_path
+
+    pcm_map = {16: "pcm_s16le", 24: "pcm_s24le", 32: "pcm_s32le"}
+    if bit_depth not in pcm_map:
+        raise ValueError("bit_depth must be one of {16, 24, 32}.")
+    if channels not in (1, 2):
+        raise ValueError("channels must be 1 (mono) or 2 (stereo).")
+
+    cmd = [
+        "ffmpeg",
+        "-nostdin",
+        "-hide_banner", "-loglevel", "error",
+        "-y" if overwrite_existing else "-n",
+        "-i", str(in_path),
+        "-vn",                        # ignore video
+        "-acodec", pcm_map[bit_depth],
+        "-ar", str(sample_rate),
+        "-ac", str(channels),
+        str(out_path),
+    ]
+
+    result = subprocess.run(cmd, capture_output=True, text=True, stdin=subprocess.DEVNULL)
+    if result.returncode != 0:
+        if not overwrite_existing and out_path.exists():
+            raise FileExistsError(f"Target exists (use overwrite=True): {out_path}")
+        raise RuntimeError(f"ffmpeg failed: {result.stderr.strip()}")
+
+    return out_path
+
+
+# --- CLI --------------------------------------------------------------------
+def _build_arg_parser():
+    import argparse
+    p = argparse.ArgumentParser(description="Convert any audio (or A/V) file to PCM WAV via ffmpeg.")
+    p.add_argument("input", help="Input file (audio or video container)")
+    p.add_argument("--out", dest="output_path", default=None,
+                   help="Exact output .wav path (overrides --out-dir)")
+    p.add_argument("--out-dir", dest="output_dir", default=None,
+                   help="Directory for output (filename will be <input_stem>.wav)")
+    p.add_argument("--sr", dest="sample_rate", type=int, default=16000, help="Sample rate (Hz)")
+    p.add_argument("--bit-depth", type=int, choices=[16, 24, 32], default=16, help="PCM bit depth")
+    p.add_argument("--channels", type=int, choices=[1, 2], default=1, help="1=mono, 2=stereo")
+    p.add_argument("--overwrite_existing", type=bool, default=False, help="Overwrite existing output")
+    return p
+
+def main():
+    args = _build_arg_parser().parse_args()
+    out = convert_audio_to_wav(
+        args.input,
+        output_path=args.output_path,
+        output_dir=args.output_dir,
+        sample_rate=args.sample_rate,
+        bit_depth=args.bit_depth,
+        channels=args.channels,
+        overwrite_existing=args.overwrite_existing,
+    )
+    print(str(out))
+
+if __name__ == "__main__":
+    main()

taters/audio/diarize_with_thirdparty.py

@@ -0,0 +1,5 @@
+# thin alias (args still pass through)
+from .diarizer.whisper_diar_wrapper import main as main
+
+if __name__ == "__main__":
+    main()

taters/audio/diarizer/whisper-diarization/diarization/__init__.py

@@ -0,0 +1,3 @@
+from .msdd.msdd import MSDDDiarizer
+
+__all__ = ["MSDDDiarizer"]

taters/audio/diarizer/whisper-diarization/diarization/msdd/msdd.py

@@ -0,0 +1,100 @@
+import json
+import os
+import tempfile
+
+from typing import Union
+
+import torch
+import torchaudio
+
+from nemo.collections.asr.models.msdd_models import NeuralDiarizer
+from nemo.collections.asr.parts.utils.speaker_utils import rttm_to_labels
+from omegaconf import OmegaConf
+
+
+class MSDDDiarizer:
+    def __init__(self, device: Union[str, torch.device]):
+        self.model: NeuralDiarizer = NeuralDiarizer(cfg=create_config()).to(device)
+
+    def diarize(self, audio: torch.Tensor):
+        with tempfile.TemporaryDirectory() as temp_path:
+            torchaudio.save(
+                os.path.join(temp_path, "mono_file.wav"),
+                audio,
+                16000,
+                channels_first=True,
+            )
+
+            manifest_path = os.path.join(temp_path, "manifest.json")
+            meta = {
+                "audio_filepath": os.path.join(temp_path, "mono_file.wav"),
+                "offset": 0,
+                "duration": None,
+                "label": "infer",
+                "text": "-",
+                "rttm_filepath": None,
+                "uem_filepath": None,
+            }
+
+            with open(manifest_path, "w") as f:
+                json.dump(meta, f)
+
+            self.model._initialize_configs(
+                manifest_path=manifest_path,
+                max_speakers=8,
+                num_speakers=None,
+                tmpdir=temp_path,
+                batch_size=24,
+                num_workers=0,
+                verbose=True,
+            )
+            self.model.clustering_embedding.clus_diar_model._diarizer_params.out_dir = (
+                temp_path
+            )
+            self.model.clustering_embedding.clus_diar_model._diarizer_params.manifest_filepath = (
+                manifest_path
+            )
+            self.model.msdd_model.cfg.test_ds.manifest_filepath = manifest_path
+            self.model.diarize()
+
+            pred_labels_clus = rttm_to_labels(
+                os.path.join(temp_path, "pred_rttms", "mono_file.rttm")
+            )
+
+            labels = []
+            for label in pred_labels_clus:
+                start, end, speaker = label.split()
+                start, end = float(start), float(end)
+                start, end = int(start * 1000), int(end * 1000)
+                labels.append((start, end, int(speaker.split("_")[1])))
+
+            labels = sorted(labels, key=lambda x: x[0])
+
+        return labels
+
+
+def create_config():
+    config = OmegaConf.load(
+        os.path.join(os.path.dirname(__file__), "diar_infer_telephonic.yaml")
+    )
+    pretrained_vad = "vad_multilingual_marblenet"
+    pretrained_speaker_model = "titanet_large"
+
+    config.diarizer.out_dir = None
+    config.diarizer.manifest_filepath = None
+    config.diarizer.speaker_embeddings.model_path = pretrained_speaker_model
+    config.diarizer.oracle_vad = (
+        False  # compute VAD provided with model_path to vad config
+    )
+    config.diarizer.clustering.parameters.oracle_num_speakers = False
+
+    # Here, we use our in-house pretrained NeMo VAD model
+    config.diarizer.vad.model_path = pretrained_vad
+    config.diarizer.vad.parameters.onset = 0.8
+    config.diarizer.vad.parameters.offset = 0.6
+    config.diarizer.vad.parameters.pad_offset = -0.05
+    config.diarizer.msdd_model.model_path = (
+        "diar_msdd_telephonic"  # Telephonic speaker diarization model
+    )
+
+    return config

taters/audio/diarizer/whisper-diarization/diarize.py

@@ -0,0 +1,247 @@
+import argparse
+import logging
+import os
+import re
+
+import faster_whisper
+import torch
+
+from ctc_forced_aligner import (
+    generate_emissions,
+    get_alignments,
+    get_spans,
+    load_alignment_model,
+    postprocess_results,
+    preprocess_text,
+)
+from deepmultilingualpunctuation import PunctuationModel
+
+from helpers import (
+    cleanup,
+    find_numeral_symbol_tokens,
+    get_realigned_ws_mapping_with_punctuation,
+    get_sentences_speaker_mapping,
+    get_speaker_aware_transcript,
+    get_words_speaker_mapping,
+    langs_to_iso,
+    process_language_arg,
+    punct_model_langs,
+    whisper_langs,
+    write_srt,
+)
+
+mtypes = {"cpu": "int8", "cuda": "float16"}
+
+pid = os.getpid()
+temp_outputs_dir = f"temp_outputs_{pid}"
+temp_path = os.path.join(os.getcwd(), "temp_outputs")
+os.makedirs(temp_path, exist_ok=True)
+
+# Initialize parser
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "-a", "--audio", help="name of the target audio file", required=True
+)
+parser.add_argument(
+    "--no-stem",
+    action="store_false",
+    dest="stemming",
+    default=True,
+    help="Disables source separation."
+    "This helps with long files that don't contain a lot of music.",
+)
+
+parser.add_argument(
+    "--suppress_numerals",
+    action="store_true",
+    dest="suppress_numerals",
+    default=False,
+    help="Suppresses Numerical Digits."
+    "This helps the diarization accuracy but converts all digits into written text.",
+)
+
+parser.add_argument(
+    "--whisper-model",
+    dest="model_name",
+    default="medium.en",
+    help="name of the Whisper model to use",
+)
+
+parser.add_argument(
+    "--batch-size",
+    type=int,
+    dest="batch_size",
+    default=8,
+    help="Batch size for batched inference, reduce if you run out of memory, "
+    "set to 0 for original whisper longform inference",
+)
+
+parser.add_argument(
+    "--language",
+    type=str,
+    default=None,
+    choices=whisper_langs,
+    help="Language spoken in the audio, specify None to perform language detection",
+)
+
+parser.add_argument(
+    "--device",
+    dest="device",
+    default="cuda" if torch.cuda.is_available() else "cpu",
+    help="if you have a GPU use 'cuda', otherwise 'cpu'",
+)
+
+parser.add_argument(
+    "--diarizer",
+    default="msdd",
+    choices=["msdd"],
+    help="Choose the diarization model to use",
+)
+
+args = parser.parse_args()
+language = process_language_arg(args.language, args.model_name)
+
+if args.stemming:
+    # Isolate vocals from the rest of the audio
+
+    return_code = os.system(
+        f'python -m demucs.separate -n htdemucs --two-stems=vocals "{args.audio}" -o "{temp_outputs_dir}" --device "{args.device}"'
+    )
+
+    if return_code != 0:
+        logging.warning(
+            "Source splitting failed, using original audio file. "
+            "Use --no-stem argument to disable it."
+        )
+        vocal_target = args.audio
+    else:
+        vocal_target = os.path.join(
+            temp_outputs_dir,
+            "htdemucs",
+            os.path.splitext(os.path.basename(args.audio))[0],
+            "vocals.wav",
+        )
+else:
+    vocal_target = args.audio
+
+
+# Transcribe the audio file
+
+whisper_model = faster_whisper.WhisperModel(
+    args.model_name, device=args.device, compute_type=mtypes[args.device]
+)
+whisper_pipeline = faster_whisper.BatchedInferencePipeline(whisper_model)
+audio_waveform = faster_whisper.decode_audio(vocal_target)
+suppress_tokens = (
+    find_numeral_symbol_tokens(whisper_model.hf_tokenizer)
+    if args.suppress_numerals
+    else [-1]
+)
+
+if args.batch_size > 0:
+    transcript_segments, info = whisper_pipeline.transcribe(
+        audio_waveform,
+        language,
+        suppress_tokens=suppress_tokens,
+        batch_size=args.batch_size,
+    )
+else:
+    transcript_segments, info = whisper_model.transcribe(
+        audio_waveform,
+        language,
+        suppress_tokens=suppress_tokens,
+        vad_filter=True,
+    )
+
+full_transcript = "".join(segment.text for segment in transcript_segments)
+
+# clear gpu vram
+del whisper_model, whisper_pipeline
+torch.cuda.empty_cache()
+
+# Forced Alignment
+alignment_model, alignment_tokenizer = load_alignment_model(
+    args.device,
+    dtype=torch.float16 if args.device == "cuda" else torch.float32,
+)
+
+emissions, stride = generate_emissions(
+    alignment_model,
+    torch.from_numpy(audio_waveform)
+    .to(alignment_model.dtype)
+    .to(alignment_model.device),
+    batch_size=args.batch_size,
+)
+
+del alignment_model
+torch.cuda.empty_cache()
+
+tokens_starred, text_starred = preprocess_text(
+    full_transcript,
+    romanize=True,
+    language=langs_to_iso[info.language],
+)
+
+segments, scores, blank_token = get_alignments(
+    emissions,
+    tokens_starred,
+    alignment_tokenizer,
+)
+
+spans = get_spans(tokens_starred, segments, blank_token)
+
+word_timestamps = postprocess_results(text_starred, spans, stride, scores)
+
+if args.diarizer == "msdd":
+    from diarization import MSDDDiarizer
+
+    diarizer_model = MSDDDiarizer(device=args.device)
+
+speaker_ts = diarizer_model.diarize(torch.from_numpy(audio_waveform).unsqueeze(0))
+del diarizer_model
+torch.cuda.empty_cache()
+
+wsm = get_words_speaker_mapping(word_timestamps, speaker_ts, "start")
+
+if info.language in punct_model_langs:
+    # restoring punctuation in the transcript to help realign the sentences
+    punct_model = PunctuationModel(model="kredor/punctuate-all")
+
+    words_list = list(map(lambda x: x["word"], wsm))
+
+    labled_words = punct_model.predict(words_list, chunk_size=230)
+
+    ending_puncts = ".?!"
+    model_puncts = ".,;:!?"
+
+    # We don't want to punctuate U.S.A. with a period. Right?
+    is_acronym = lambda x: re.fullmatch(r"\b(?:[a-zA-Z]\.){2,}", x)
+
+    for word_dict, labeled_tuple in zip(wsm, labled_words):
+        word = word_dict["word"]
+        if (
+            word
+            and labeled_tuple[1] in ending_puncts
+            and (word[-1] not in model_puncts or is_acronym(word))
+        ):
+            word += labeled_tuple[1]
+            if word.endswith(".."):
+                word = word.rstrip(".")
+            word_dict["word"] = word
+
+else:
+    logging.warning(
+        f"Punctuation restoration is not available for {info.language} language."
+        " Using the original punctuation."
+    )
+
+wsm = get_realigned_ws_mapping_with_punctuation(wsm)
+ssm = get_sentences_speaker_mapping(wsm, speaker_ts)
+
+with open(f"{os.path.splitext(args.audio)[0]}.txt", "w", encoding="utf-8-sig") as f:
+    get_speaker_aware_transcript(ssm, f)
+
+with open(f"{os.path.splitext(args.audio)[0]}.srt", "w", encoding="utf-8-sig") as srt:
+    write_srt(ssm, srt)
+
+cleanup(temp_path)