dora-distil-whisper 0.3.10rc0__tar.gz → 0.3.11__tar.gz

{dora_distil_whisper-0.3.10rc0 → dora_distil_whisper-0.3.11}/PKG-INFO

@@ -1,12 +1,12 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: dora-distil-whisper
-Version: 0.3.10rc0
+Version: 0.3.11
 Summary: Dora dora-distil-whisper
 Author-email: Haixuan Xavier Tao <tao.xavier@outlook.com>, Enzo Le Van <dev@enzo-le-van.fr>
 License: MIT
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
-Requires-Dist: dora-rs>=0.3.6
+Requires-Dist: dora-rs>=0.3.9
 Requires-Dist: numpy<2.0.0
 Requires-Dist: pyarrow>=5.0.0
 Requires-Dist: transformers>=4.0.0

{dora_distil_whisper-0.3.10rc0 → dora_distil_whisper-0.3.11}/dora_distil_whisper/__init__.py

@@ -1,3 +1,5 @@
+"""TODO: Add docstring."""
+
 import os
 
 # Define the path to the README file relative to the package directory

dora_distil_whisper-0.3.11/dora_distil_whisper/main.py

@@ -0,0 +1,240 @@
+"""TODO: Add docstring."""
+
+import os
+import re
+import sys
+import time
+from pathlib import Path
+
+import pyarrow as pa
+import torch
+from dora import Node
+
+DEFAULT_PATH = "openai/whisper-large-v3-turbo"
+TARGET_LANGUAGE = os.getenv("TARGET_LANGUAGE", "english")
+TRANSLATE = bool(os.getenv("TRANSLATE", "False") in ["True", "true"])
+
+
+def remove_text_noise(text: str, text_noise="") -> str:
+    """Remove noise from text.
+
+    Args:
+        text (str): Original text
+        text_noise (str): text to remove from the original text
+
+    Returns:
+        str: Cleaned text
+
+    """
+    # Handle the case where text_noise is empty
+    if not text_noise.strip():
+        return (
+            text  # Return the original text if text_noise is empty or just whitespace
+        )
+
+    # Helper function to normalize text (remove punctuation, make lowercase, and handle hyphens)
+    def normalize(s):
+        # Replace hyphens with spaces to treat "Notre-Dame" and "notre dame" as equivalent
+        s = re.sub(r"-", " ", s)
+        # Remove other punctuation and convert to lowercase
+        return re.sub(r"[^\w\s]", "", s).lower()
+
+    # Normalize both text and text_noise
+    normalized_text = normalize(text)
+    normalized_noise = normalize(text_noise)
+
+    # Split into words
+    text_words = normalized_text.split()
+    noise_words = normalized_noise.split()
+
+    # Function to find and remove noise sequence flexibly
+    def remove_flexible(text_list, noise_list):
+        i = 0
+        while i <= len(text_list) - len(noise_list):
+            match = True
+            extra_words = 0
+            for j, noise_word in enumerate(noise_list):
+                if i + j + extra_words >= len(text_list):
+                    match = False
+                    break
+                # Allow skipping extra words in text_list
+                while (
+                    i + j + extra_words < len(text_list)
+                    and text_list[i + j + extra_words] != noise_word
+                ):
+                    extra_words += 1
+                    if i + j + extra_words >= len(text_list):
+                        match = False
+                        break
+                if not match:
+                    break
+            if match:
+                # Remove matched part
+                del text_list[i : i + len(noise_list) + extra_words]
+                i = max(0, i - len(noise_list))  # Adjust index after removal
+            else:
+                i += 1
+        return text_list
+
+    # Only remove parts of text_noise that are found in text
+    cleaned_words = text_words[:]
+    for noise_word in noise_words:
+        if noise_word in cleaned_words:
+            cleaned_words.remove(noise_word)
+
+    # Reconstruct the cleaned text
+    return " ".join(cleaned_words)
+
+
+def load_model():
+    """TODO: Add docstring."""
+    from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
+
+    model_name_or_path = os.getenv("MODEL_NAME_OR_PATH", DEFAULT_PATH)
+
+    if bool(os.getenv("USE_MODELSCOPE_HUB") in ["True", "true"]):
+        from modelscope import snapshot_download
+
+        if not Path(model_name_or_path).exists():
+            model_name_or_path = snapshot_download(model_name_or_path)
+
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+
+    model = AutoModelForSpeechSeq2Seq.from_pretrained(
+        model_name_or_path,
+        torch_dtype=torch_dtype,
+        low_cpu_mem_usage=True,
+        use_safetensors=True,
+    )
+    model.to(device)
+
+    processor = AutoProcessor.from_pretrained(model_name_or_path)
+    return pipeline(
+        "automatic-speech-recognition",
+        model=model,
+        tokenizer=processor.tokenizer,
+        feature_extractor=processor.feature_extractor,
+        max_new_tokens=400,
+        torch_dtype=torch_dtype,
+        device=device,
+    )
+
+
+BAD_SENTENCES = [
+    "",
+    " so",
+    " so so",
+    "字幕",
+    "字幕志愿",
+    "中文字幕",
+    "我",
+    "你",
+    " you",
+    "!",
+    "THANK YOU",
+    " Thank you.",
+    " www.microsoft.com",
+    " The",
+    " BANG",
+    " Silence.",
+    " Sous-titrage Société Radio-Canada",
+    " Sous",
+    " Sous-",
+    " i'm going to go to the next one.",
+]
+
+
+def cut_repetition(text, min_repeat_length=4, max_repeat_length=50):
+    """TODO: Add docstring."""
+    if len(text) == 0:
+        return text
+    # Check if the text is primarily Chinese (you may need to adjust this threshold)
+    if sum(1 for char in text if "\u4e00" <= char <= "\u9fff") / len(text) > 0.5:
+        # Chinese text processing
+        for repeat_length in range(
+            min_repeat_length,
+            min(max_repeat_length, len(text) // 2),
+        ):
+            for i in range(len(text) - repeat_length * 2 + 1):
+                chunk1 = text[i : i + repeat_length]
+                chunk2 = text[i + repeat_length : i + repeat_length * 2]
+
+                if chunk1 == chunk2:
+                    return text[: i + repeat_length]
+    else:
+        # Non-Chinese (space-separated) text processing
+        words = text.split()
+        for repeat_length in range(
+            min_repeat_length,
+            min(max_repeat_length, len(words) // 2),
+        ):
+            for i in range(len(words) - repeat_length * 2 + 1):
+                chunk1 = " ".join(words[i : i + repeat_length])
+                chunk2 = " ".join(words[i + repeat_length : i + repeat_length * 2])
+
+                if chunk1 == chunk2:
+                    return " ".join(words[: i + repeat_length])
+
+    return text
+
+
+def main():
+    """TODO: Add docstring."""
+    node = Node()
+    text_noise = ""
+    noise_timestamp = time.time()
+    # For macos use mlx:
+    if sys.platform != "darwin":
+        pipe = load_model()
+
+    for event in node:
+        if event["type"] == "INPUT":
+            if "text_noise" in event["id"]:
+                text_noise = event["value"][0].as_py()
+                text_noise = (
+                    text_noise.replace("(", "")
+                    .replace(")", "")
+                    .replace("[", "")
+                    .replace("]", "")
+                )
+                noise_timestamp = time.time()
+            else:
+                audio = event["value"].to_numpy()
+                confg = (
+                    {"language": TARGET_LANGUAGE, "task": "translate"}
+                    if TRANSLATE
+                    else {
+                        "language": TARGET_LANGUAGE,
+                    }
+                )
+                if sys.platform == "darwin":
+                    import mlx_whisper
+
+                    result = mlx_whisper.transcribe(
+                        audio,
+                        path_or_hf_repo="mlx-community/whisper-large-v3-turbo",
+                        append_punctuations=".",
+                    )
+
+                else:
+                    result = pipe(
+                        audio,
+                        generate_kwargs=confg,
+                    )
+                if result["text"] in BAD_SENTENCES:
+                    continue
+                text = cut_repetition(result["text"])
+
+                # Remove noise filter after some time
+                if time.time() - noise_timestamp > (len(text_noise.split()) / 2):  # WPS
+                    text_noise = ""
+
+                ## Remove text noise independently of casing
+                text = remove_text_noise(text, text_noise)
+
+                if text.strip() == "" or text.strip() == ".":
+                    continue
+                node.send_output(
+                    "text", pa.array([text]), {"language": TARGET_LANGUAGE},
+                )

{dora_distil_whisper-0.3.10rc0 → dora_distil_whisper-0.3.11}/dora_distil_whisper.egg-info/PKG-INFO

@@ -1,12 +1,12 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: dora-distil-whisper
-Version: 0.3.10rc0
+Version: 0.3.11
 Summary: Dora dora-distil-whisper
 Author-email: Haixuan Xavier Tao <tao.xavier@outlook.com>, Enzo Le Van <dev@enzo-le-van.fr>
 License: MIT
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
-Requires-Dist: dora-rs>=0.3.6
+Requires-Dist: dora-rs>=0.3.9
 Requires-Dist: numpy<2.0.0
 Requires-Dist: pyarrow>=5.0.0
 Requires-Dist: transformers>=4.0.0

{dora_distil_whisper-0.3.10rc0 → dora_distil_whisper-0.3.11}/dora_distil_whisper.egg-info/requires.txt

@@ -1,4 +1,4 @@
-dora-rs>=0.3.6
+dora-rs>=0.3.9
 numpy<2.0.0
 pyarrow>=5.0.0
 transformers>=4.0.0

dora_distil_whisper-0.3.11/pyproject.toml

@@ -0,0 +1,40 @@
+[project]
+name = "dora-distil-whisper"
+version = "0.3.11"
+authors = [
+  { name = "Haixuan Xavier Tao", email = "tao.xavier@outlook.com" },
+  { name = "Enzo Le Van", email = "dev@enzo-le-van.fr" },
+]
+description = "Dora dora-distil-whisper"
+license = { text = "MIT" }
+readme = "README.md"
+requires-python = ">=3.8"
+
+dependencies = [
+  "dora-rs >= 0.3.9",
+  "numpy < 2.0.0",
+  "pyarrow >= 5.0.0",
+  "transformers >= 4.0.0",
+  "accelerate >= 0.29.2",
+  "torch >= 2.2.0",
+  "modelscope >= 1.18.1",
+  "mlx-whisper >= 0.4.1; sys_platform == 'darwin'",
+]
+
+[dependency-groups]
+dev = ["pytest >=8.1.1", "ruff >=0.9.1"]
+
+[project.scripts]
+dora-distil-whisper = "dora_distil_whisper.main:main"
+
+[tool.ruff.lint]
+extend-select = [
+  "D",    # pydocstyle
+  "UP",   # Ruff's UP rule
+  "PERF", # Ruff's PERF rule
+  "RET",  # Ruff's RET rule
+  "RSE",  # Ruff's RSE rule
+  "NPY",  # Ruff's NPY rule
+  "N",    # Ruff's N rule
+  "I",    # Ruff's I rule
+]

{dora_distil_whisper-0.3.10rc0 → dora_distil_whisper-0.3.11}/tests/test_distil_whisper.py

@@ -1,7 +1,10 @@
+"""TODO: Add docstring."""
+
 import pytest
 
 
 def test_import_main():
+    """TODO: Add docstring."""
     from dora_distil_whisper.main import main
 
     # Check that everything is working, and catch dora Runtime Exception as we're not running in a dora dataflow.

dora_distil_whisper-0.3.10rc0/dora_distil_whisper/main.py

@@ -1,139 +0,0 @@
-import os
-import sys
-from pathlib import Path
-
-import pyarrow as pa
-import torch
-from dora import Node
-
-DEFAULT_PATH = "openai/whisper-large-v3-turbo"
-TARGET_LANGUAGE = os.getenv("TARGET_LANGUAGE", "english")
-TRANSLATE = bool(os.getenv("TRANSLATE", "False") in ["True", "true"])
-
-
-def load_model():
-    from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
-
-    MODEL_NAME_OR_PATH = os.getenv("MODEL_NAME_OR_PATH", DEFAULT_PATH)
-
-    if bool(os.getenv("USE_MODELSCOPE_HUB") in ["True", "true"]):
-        from modelscope import snapshot_download
-
-        if not Path(MODEL_NAME_OR_PATH).exists():
-            MODEL_NAME_OR_PATH = snapshot_download(MODEL_NAME_OR_PATH)
-
-    device = "cuda:0" if torch.cuda.is_available() else "cpu"
-    torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-
-    model = AutoModelForSpeechSeq2Seq.from_pretrained(
-        MODEL_NAME_OR_PATH,
-        torch_dtype=torch_dtype,
-        low_cpu_mem_usage=True,
-        use_safetensors=True,
-    )
-    model.to(device)
-
-    processor = AutoProcessor.from_pretrained(MODEL_NAME_OR_PATH)
-    pipe = pipeline(
-        "automatic-speech-recognition",
-        model=model,
-        tokenizer=processor.tokenizer,
-        feature_extractor=processor.feature_extractor,
-        max_new_tokens=400,
-        torch_dtype=torch_dtype,
-        device=device,
-    )
-    return pipe
-
-
-BAD_SENTENCES = [
-    "",
-    " so",
-    " so so",
-    "字幕",
-    "字幕志愿",
-    "中文字幕",
-    "我",
-    "你",
-    " you",
-    "!",
-    "THANK YOU",
-    " Thank you.",
-    " www.microsoft.com",
-    " The",
-    " BANG",
-    " Silence.",
-    " Sous-titrage Société Radio-Canada",
-    " Sous",
-    " Sous-",
-]
-
-
-def cut_repetition(text, min_repeat_length=4, max_repeat_length=50):
-    if len(text) == 0:
-        return text
-    # Check if the text is primarily Chinese (you may need to adjust this threshold)
-    if sum(1 for char in text if "\u4e00" <= char <= "\u9fff") / len(text) > 0.5:
-        # Chinese text processing
-        for repeat_length in range(
-            min_repeat_length,
-            min(max_repeat_length, len(text) // 2),
-        ):
-            for i in range(len(text) - repeat_length * 2 + 1):
-                chunk1 = text[i : i + repeat_length]
-                chunk2 = text[i + repeat_length : i + repeat_length * 2]
-
-                if chunk1 == chunk2:
-                    return text[: i + repeat_length]
-    else:
-        # Non-Chinese (space-separated) text processing
-        words = text.split()
-        for repeat_length in range(
-            min_repeat_length,
-            min(max_repeat_length, len(words) // 2),
-        ):
-            for i in range(len(words) - repeat_length * 2 + 1):
-                chunk1 = " ".join(words[i : i + repeat_length])
-                chunk2 = " ".join(words[i + repeat_length : i + repeat_length * 2])
-
-                if chunk1 == chunk2:
-                    return " ".join(words[: i + repeat_length])
-
-    return text
-
-
-def main():
-    node = Node()
-
-    # For macos use mlx:
-    if sys.platform != "darwin":
-        pipe = load_model()
-
-    for event in node:
-        if event["type"] == "INPUT":
-            audio = event["value"].to_numpy()
-            confg = (
-                {"language": TARGET_LANGUAGE, "task": "translate"}
-                if TRANSLATE
-                else {
-                    "language": TARGET_LANGUAGE,
-                }
-            )
-            if sys.platform == "darwin":
-                import mlx_whisper
-
-                result = mlx_whisper.transcribe(
-                    audio,
-                    path_or_hf_repo="mlx-community/whisper-large-v3-turbo",
-                    append_punctuations=".",
-                )
-
-            else:
-                result = pipe(
-                    audio,
-                    generate_kwargs=confg,
-                )
-            if result["text"] in BAD_SENTENCES:
-                continue
-            text = cut_repetition(result["text"])
-            node.send_output("text", pa.array([text]), {"language": TARGET_LANGUAGE})

dora_distil_whisper-0.3.10rc0/pyproject.toml

@@ -1,29 +0,0 @@
-[project]
-name = "dora-distil-whisper"
-version = "0.3.10-rc0"
-authors = [
-    { name = "Haixuan Xavier Tao", email = "tao.xavier@outlook.com" },
-    { name = "Enzo Le Van", email = "dev@enzo-le-van.fr" },
-]
-description = "Dora dora-distil-whisper"
-license = { text = "MIT" }
-readme = "README.md"
-requires-python = ">=3.8"
-
-dependencies = [
-    "dora-rs >= 0.3.6",
-    "numpy < 2.0.0",
-    "pyarrow >= 5.0.0",
-    "transformers >= 4.0.0",
-    "accelerate >= 0.29.2",
-    "torch >= 2.2.0",
-    "modelscope >= 1.18.1",
-    "mlx-whisper >= 0.4.1; sys_platform == 'darwin'",
-]
-
-
-[dependency-groups]
-dev = ["pytest >=8.1.1", "ruff >=0.9.1"]
-
-[project.scripts]
-dora-distil-whisper = "dora_distil_whisper.main:main"