dora-distil-whisper 0.3.10rc0__py3-none-any.whl → 0.3.11__py3-none-any.whl

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/main.py

@@ -1,5 +1,9 @@
+"""TODO: Add docstring."""
+
 import os
+import re
 import sys
+import time
 from pathlib import Path
 
 import pyarrow as pa
@@ -11,30 +15,102 @@ 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)
+    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)
+        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,
+        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(
+    processor = AutoProcessor.from_pretrained(model_name_or_path)
+    return pipeline(
         "automatic-speech-recognition",
         model=model,
         tokenizer=processor.tokenizer,
@@ -43,7 +119,6 @@ def load_model():
         torch_dtype=torch_dtype,
         device=device,
     )
-    return pipe
 
 
 BAD_SENTENCES = [
@@ -66,10 +141,12 @@ BAD_SENTENCES = [
     " 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)
@@ -103,37 +180,61 @@ def cut_repetition(text, min_repeat_length=4, max_repeat_length=50):
 
 
 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":
-            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=".",
+            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:
-                result = pipe(
-                    audio,
-                    generate_kwargs=confg,
+                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},
                 )
-            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.dist-info → dora_distil_whisper-0.3.11.dist-info}/METADATA

@@ -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.11.dist-info/RECORD

@@ -0,0 +1,7 @@
+dora_distil_whisper/__init__.py,sha256=tF7WHhHiDweUUzyHsbmFe_ktphE08aA5j33E4ja1udA,381
+dora_distil_whisper/main.py,sha256=wOOJ9TeigWm-byZ0wGVJbgS15fLRSMH2cGFRHXcj9eg,7643
+dora_distil_whisper-0.3.11.dist-info/METADATA,sha256=FwdwRV5UY0Wd4dDswj1tdiO5nXQf9asKyLHFhgcemFQ,1254
+dora_distil_whisper-0.3.11.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+dora_distil_whisper-0.3.11.dist-info/entry_points.txt,sha256=c6QYCZs0YVR2uasYHES67JhOOvTm5QbcwGk-9IrG9oM,70
+dora_distil_whisper-0.3.11.dist-info/top_level.txt,sha256=h5QH64SWnqZA83bx740-NTxfQKdeiKTLAdGqhnwKhuQ,20
+dora_distil_whisper-0.3.11.dist-info/RECORD,,

{dora_distil_whisper-0.3.10rc0.dist-info → dora_distil_whisper-0.3.11.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.8.0)
+Generator: setuptools (78.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

dora_distil_whisper-0.3.10rc0.dist-info/RECORD

@@ -1,7 +0,0 @@
-dora_distil_whisper/__init__.py,sha256=HuSK3dnyI9Pb5QAuaKFwQQ3J5SIZnLcKHPJO0norGzc,353
-dora_distil_whisper/main.py,sha256=MbT9nsEHxpyzcFzkDe4FjITzUeemh8LzCMmQSRL4xqo,4083
-dora_distil_whisper-0.3.10rc0.dist-info/METADATA,sha256=0elfHye3_gyyVg5dWnWcaDaPddXLZUUO2SRJxoDIwGA,1257
-dora_distil_whisper-0.3.10rc0.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-dora_distil_whisper-0.3.10rc0.dist-info/entry_points.txt,sha256=c6QYCZs0YVR2uasYHES67JhOOvTm5QbcwGk-9IrG9oM,70
-dora_distil_whisper-0.3.10rc0.dist-info/top_level.txt,sha256=h5QH64SWnqZA83bx740-NTxfQKdeiKTLAdGqhnwKhuQ,20
-dora_distil_whisper-0.3.10rc0.dist-info/RECORD,,