switch to ruff and precommit (#360)

* switch to ruff and precommit

* ignore the notebook

Author: MahmoudAshraf97

.github/workflows/test_run.yml

@@ -9,8 +9,20 @@ on:
     - cron: '0 0 * * TUE'
 
 jobs:
-  build:
+  check-python-style:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python 3.10
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.10'
+
+      - uses: pre-commit/action@v3.0.1
 
+  test-run:
+    needs: [check-python-style]
     strategy:
       fail-fast: false
       matrix:
@@ -48,6 +60,6 @@ jobs:
         python -m pip install torch torchaudio --index-url https://download.pytorch.org/whl/cpu
         uv pip install --system -c constraints.txt -r requirements.txt
 
-    - name: Test running a file
+    - name: Test with msdd diarizer
       run: |
-        python diarize.py -a "./tests/assets/test.opus" --whisper-model tiny.en
+        python diarize.py -a "./tests/assets/test.opus" --whisper-model tiny.en --diarizer msdd

.pre-commit-config.yaml

@@ -0,0 +1,7 @@
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.9.9
+    hooks:
+      - id: ruff
+        args: [--fix]
+      - id: ruff-format

diarization/msdd/msdd.py

@@ -1,11 +1,11 @@
 import json
 import os
 import tempfile
+import wave
 
 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
@@ -18,12 +18,12 @@ def __init__(self, device: Union[str, torch.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,
-            )
+            pcm = (audio.cpu().numpy() * 32768).clip(-32768, 32767).astype("int16")
+            with wave.open(os.path.join(temp_path, "mono_file.wav"), "wb") as wf:
+                wf.setnchannels(1)
+                wf.setsampwidth(2)
+                wf.setframerate(16000)
+                wf.writeframes(pcm.tobytes())
 
             manifest_path = os.path.join(temp_path, "manifest.json")
             meta = {
@@ -48,9 +48,7 @@ def diarize(self, audio: torch.Tensor):
                 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.out_dir = temp_path
             self.model.clustering_embedding.clus_diar_model._diarizer_params.manifest_filepath = (
                 manifest_path
             )
@@ -74,18 +72,14 @@ def diarize(self, audio: torch.Tensor):
 
 
 def create_config():
-    config = OmegaConf.load(
-        os.path.join(os.path.dirname(__file__), "diar_infer_telephonic.yaml")
-    )
+    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.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

diarize.py

@@ -32,23 +32,18 @@
 
 mtypes = {"cpu": "int8", "cuda": "float16"}
 
-pid = os.getpid()
-temp_outputs_dir = f"temp_outputs_{pid}"
-temp_path = os.path.join(os.getcwd(), "temp_outputs")
+temp_path = os.path.join(os.getcwd(), f"temp_outputs_{os.getpid()}")
 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("-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.",
+    help="Disables source separation.This helps with long files that don't contain a lot of music.",
 )
 
 parser.add_argument(
@@ -105,7 +100,8 @@
     # 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}"'
+        f"python -m demucs.separate -n htdemucs --two-stems=vocals "
+        f'"{args.audio}" -o "{temp_path}" --device "{args.device}"'
     )
 
     if return_code != 0:
@@ -116,7 +112,7 @@
         vocal_target = args.audio
     else:
         vocal_target = os.path.join(
-            temp_outputs_dir,
+            temp_path,
             "htdemucs",
             os.path.splitext(os.path.basename(args.audio))[0],
             "vocals.wav",
@@ -133,9 +129,7 @@
 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]
+    find_numeral_symbol_tokens(whisper_model.hf_tokenizer) if args.suppress_numerals else [-1]
 )
 
 if args.batch_size > 0:
@@ -167,9 +161,7 @@
 
 emissions, stride = generate_emissions(
     alignment_model,
-    torch.from_numpy(audio_waveform)
-    .to(alignment_model.dtype)
-    .to(alignment_model.device),
+    torch.from_numpy(audio_waveform).to(alignment_model.dtype).to(alignment_model.device),
     batch_size=args.batch_size,
 )
 

diarize_parallel.py

@@ -50,9 +50,7 @@ def diarize_parallel(audio: torch.Tensor, device, queue: mp.Queue):
 
     # Initialize parser
     parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "-a", "--audio", help="name of the target audio file", required=True
-    )
+    parser.add_argument("-a", "--audio", help="name of the target audio file", required=True)
     parser.add_argument(
         "--no-stem",
         action="store_false",
@@ -116,7 +114,8 @@ def diarize_parallel(audio: torch.Tensor, device, queue: mp.Queue):
         # 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}"'
+            f"python -m demucs.separate -n htdemucs --two-stems=vocals "
+            f'"{args.audio}" -o "{temp_path}" --device "{args.device}"'
         )
 
         if return_code != 0:
@@ -156,9 +155,7 @@ def diarize_parallel(audio: torch.Tensor, device, queue: mp.Queue):
     whisper_pipeline = faster_whisper.BatchedInferencePipeline(whisper_model)
 
     suppress_tokens = (
-        find_numeral_symbol_tokens(whisper_model.hf_tokenizer)
-        if args.suppress_numerals
-        else [-1]
+        find_numeral_symbol_tokens(whisper_model.hf_tokenizer) if args.suppress_numerals else [-1]
     )
 
     if args.batch_size > 0:
@@ -190,9 +187,7 @@ def diarize_parallel(audio: torch.Tensor, device, queue: mp.Queue):
 
     emissions, stride = generate_emissions(
         alignment_model,
-        torch.from_numpy(audio_waveform)
-        .to(alignment_model.dtype)
-        .to(alignment_model.device),
+        torch.from_numpy(audio_waveform).to(alignment_model.dtype).to(alignment_model.device),
         batch_size=args.batch_size,
     )
 
@@ -261,9 +256,7 @@ def diarize_parallel(audio: torch.Tensor, device, queue: mp.Queue):
     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:
+    with open(f"{os.path.splitext(args.audio)[0]}.srt", "w", encoding="utf-8-sig") as srt:
         write_srt(ssm, srt)
 
     cleanup(temp_path)

helpers.py

@@ -137,9 +137,7 @@
     "castilian": "es",
 }
 
-whisper_langs = sorted(LANGUAGES.keys()) + sorted(
-    [k.title() for k in TO_LANGUAGE_CODE.keys()]
-)
+whisper_langs = sorted(LANGUAGES.keys()) + sorted([k.title() for k in TO_LANGUAGE_CODE.keys()])
 
 langs_to_iso = {
     "af": "afr",
@@ -270,19 +268,15 @@ def get_words_speaker_mapping(wrd_ts, spk_ts, word_anchor_option="start"):
             s, e, sp = spk_ts[turn_idx]
             if turn_idx == len(spk_ts) - 1:
                 e = get_word_ts_anchor(ws, we, option="end")
-        wrd_spk_mapping.append(
-            {"word": wrd, "start_time": ws, "end_time": we, "speaker": sp}
-        )
+        wrd_spk_mapping.append({"word": wrd, "start_time": ws, "end_time": we, "speaker": sp})
     return wrd_spk_mapping
 
 
 sentence_ending_punctuations = ".?!"
 
 
 def get_first_word_idx_of_sentence(word_idx, word_list, speaker_list, max_words):
-    is_word_sentence_end = (
-        lambda x: x >= 0 and word_list[x][-1] in sentence_ending_punctuations
-    )
+    is_word_sentence_end = lambda x: x >= 0 and word_list[x][-1] in sentence_ending_punctuations
     left_idx = word_idx
     while (
         left_idx > 0
@@ -296,9 +290,7 @@ def get_first_word_idx_of_sentence(word_idx, word_list, speaker_list, max_words)
 
 
 def get_last_word_idx_of_sentence(word_idx, word_list, max_words):
-    is_word_sentence_end = (
-        lambda x: x >= 0 and word_list[x][-1] in sentence_ending_punctuations
-    )
+    is_word_sentence_end = lambda x: x >= 0 and word_list[x][-1] in sentence_ending_punctuations
     right_idx = word_idx
     while (
         right_idx < len(word_list) - 1
@@ -307,19 +299,12 @@ def get_last_word_idx_of_sentence(word_idx, word_list, max_words):
     ):
         right_idx += 1
 
-    return (
-        right_idx
-        if right_idx == len(word_list) - 1 or is_word_sentence_end(right_idx)
-        else -1
-    )
+    return right_idx if right_idx == len(word_list) - 1 or is_word_sentence_end(right_idx) else -1
 
 
-def get_realigned_ws_mapping_with_punctuation(
-    word_speaker_mapping, max_words_in_sentence=50
-):
+def get_realigned_ws_mapping_with_punctuation(word_speaker_mapping, max_words_in_sentence=50):
     is_word_sentence_end = (
-        lambda x: x >= 0
-        and word_speaker_mapping[x]["word"][-1] in sentence_ending_punctuations
+        lambda x: x >= 0 and word_speaker_mapping[x]["word"][-1] in sentence_ending_punctuations
     )
     wsp_len = len(word_speaker_mapping)
 
@@ -357,9 +342,7 @@ def get_realigned_ws_mapping_with_punctuation(
                 k += 1
                 continue
 
-            speaker_list[left_idx : right_idx + 1] = [mod_speaker] * (
-                right_idx - left_idx + 1
-            )
+            speaker_list[left_idx : right_idx + 1] = [mod_speaker] * (right_idx - left_idx + 1)
             k = right_idx
 
         k += 1
@@ -434,9 +417,7 @@ def format_timestamp(
     milliseconds -= seconds * 1_000
 
     hours_marker = f"{hours:02d}:" if always_include_hours or hours > 0 else ""
-    return (
-        f"{hours_marker}{minutes:02d}:{seconds:02d}{decimal_marker}{milliseconds:03d}"
-    )
+    return f"{hours_marker}{minutes:02d}:{seconds:02d}{decimal_marker}{milliseconds:03d}"
 
 
 def write_srt(transcript, file):
@@ -490,17 +471,11 @@ def _get_next_start_timestamp(word_timestamps, current_word_index, final_timesta
             return word_timestamps[next_word_index]["start"]
 
 
-def filter_missing_timestamps(
-    word_timestamps, initial_timestamp=0, final_timestamp=None
-):
+def filter_missing_timestamps(word_timestamps, initial_timestamp=0, final_timestamp=None):
     # handle the first and last word
     if word_timestamps[0].get("start") is None:
-        word_timestamps[0]["start"] = (
-            initial_timestamp if initial_timestamp is not None else 0
-        )
-        word_timestamps[0]["end"] = _get_next_start_timestamp(
-            word_timestamps, 0, final_timestamp
-        )
+        word_timestamps[0]["start"] = initial_timestamp if initial_timestamp is not None else 0
+        word_timestamps[0]["end"] = _get_next_start_timestamp(word_timestamps, 0, final_timestamp)
 
     result = [
         word_timestamps[0],

pyproject.toml

@@ -0,0 +1,10 @@
+[tool.ruff]
+line-length = 100
+extend-exclude = ["*.ipynb"]
+
+[tool.ruff.lint]
+select = ["E", "F", "W", "I"]
+ignore = ["E203", "E731"]
+
+[tool.ruff.lint.isort]
+lines-between-types = 1