Initial commit

gradio-app.py

@@ -0,0 +1,193 @@
+import subprocess
+import time
+
+import gradio as gr
+import librosa
+import pytube as pt
+from models import asr, processor
+from utils import format_timestamp
+from vad import SpeechTimestampsMap, collect_chunks, get_speech_timestamps
+
+## details: https://huggingface.co/docs/diffusers/optimization/fp16#automatic-mixed-precision-amp
+# from torch import autocast
+
+apply_vad = True
+vad_parameters = {}
+
+# task = "transcribe"  # transcribe or translate
+# language = "bn"
+# asr.model.config.forced_decoder_ids = processor.get_decoder_prompt_ids(language=language, task=task)
+# asr.model.config.max_new_tokens = 448 #default is 448
+
+
+def _preprocess(filename):
+    audio_name = "audio.wav"
+    subprocess.call(
+        [
+            "ffmpeg",
+            "-y",
+            "-i",
+            filename,
+            "-acodec",
+            "pcm_s16le",
+            "-ar",
+            "16000",
+            "-ac",
+            "1",
+            "-loglevel",
+            "quiet",
+            audio_name,
+        ]
+    )
+    return audio_name
+
+
+def transcribe(microphone, file_upload):
+    warn_output = ""
+    if (microphone is not None) and (file_upload is not None):
+        warn_output = (
+            "WARNING: You've uploaded an audio file and used the microphone. "
+            "The recorded file from the microphone will be used and the uploaded audio will be discarded.\n"
+        )
+
+    elif (microphone is None) and (file_upload is None):
+        return "ERROR: You have to either use the microphone or upload an audio file"
+
+    file = microphone if microphone is not None else file_upload
+    print(f"\n\nFile is: {file}\n\n")
+
+    # for _preprocess(). No need if name of file provided in string format to asr pipeline as automatically uses ffmeg.
+    # Only required if ndarray given by using librosa.load() to load a file
+    start_time = time.time()
+    print("Starting Preprocessing")
+    # speech_array = _preprocess(filename=file)
+    filename = _preprocess(filename=file)
+    speech_array, sample_rate = librosa.load(f"{filename}", sr=16_000)
+    if apply_vad:
+        duration = speech_array.shape[0] / sample_rate
+        print(f"Processing audio with duration: {format_timestamp(duration)}")
+        speech_chunks = get_speech_timestamps(speech_array, **vad_parameters)
+        speech_array = collect_chunks(speech_array, speech_chunks)
+        print(f"VAD filter removed {format_timestamp(duration - (speech_array.shape[0] / sample_rate))}")
+        remaining_segments = ", ".join(
+            f'[{format_timestamp(chunk["start"] / sample_rate)} -> {format_timestamp(chunk["end"] / sample_rate)}]'
+            for chunk in speech_chunks
+        )
+        print(f"VAD filter kept the following audio segments: {remaining_segments}")
+        if not remaining_segments:
+            return "ERROR: No speech detected in the audio file"
+
+
+
+    print(f"\n Preprocessing COMPLETED in {round(time.time()-start_time, 2)}s \n")
+
+    start_time = time.time()
+    print("Starting Inference")
+    text = asr(speech_array)["text"]
+    # text = asr(file)["text"]
+    # with autocast("cuda"):
+    #     text = asr(speech_array)["text"]
+    print(f"\n Inference COMPLETED in {round(time.time()-start_time, 2)}s \n")
+
+    return warn_output + text
+
+
+def _return_yt_html_embed(yt_url):
+    if "?v=" in yt_url:
+        video_id = yt_url.split("?v=")[-1].split("&")[0]
+    else:
+        video_id = yt_url.split("/")[-1].split("?feature=")[0]
+
+    print(f"\n\nYT ID is: {video_id}\n\n")
+    return f'<center><iframe width="500" height="320" src="https://www.youtube.com/embed/{video_id}"> </iframe> </center>'
+
+
+def yt_transcribe(yt_url):
+    start_time = time.time()
+    yt = pt.YouTube(yt_url)
+    html_embed_str = _return_yt_html_embed(yt_url)
+    stream = yt.streams.filter(only_audio=True)[0]
+    filename = "audio.mp3"
+    stream.download(filename=filename)
+    print(f"\n YT Audio Downloaded in {round(time.time()-start_time, 2)}s \n")
+
+    # for _preprocess(). No need if name of file provided in string format to asr pipeline as automatically uses ffmeg.
+    # Only required if ndarray given by using librosa.load() to load a file
+    start_time = time.time()
+    # print("Starting Preprocessing")
+    # speech_array = _preprocess(filename=filename)
+    # filename = _preprocess(filename=filename)
+    # speech_array, sample_rate = librosa.load(f"{filename}", sr=16_000)
+    # print(f"\n Preprocessing COMPLETED in {round(time.time()-start_time, 2)}s \n")
+
+    start_time = time.time()
+    print("Starting Inference")
+    text = asr(filename)["text"]
+    # with autocast("cuda"):
+    #     text = asr(speech_array)["text"]
+    print(f"\n Inference COMPLETED in {round(time.time()-start_time, 2)}s \n")
+
+    return html_embed_str, text
+
+
+mf_transcribe = gr.Interface(
+    fn=transcribe,
+    inputs=[
+        gr.Audio(source="microphone", type="filepath", label="Microphone"),
+        gr.Audio(source="upload", type="filepath", label="Upload File"),
+    ],
+    outputs="text",
+    title="Bangla Demo: Transcribe Audio",
+    description=(
+        "Transcribe long-form microphone or audio inputs in BANGLA with the click of a button!"
+    ),
+    allow_flagging="never",
+)
+
+yt_transcribe = gr.Interface(
+    fn=yt_transcribe,
+    inputs=[
+        gr.Textbox(
+            lines=1,
+            placeholder="Paste the URL to a Bangla language YouTube video here",
+            label="YouTube URL",
+        )
+    ],
+    outputs=["html", "text"],
+    title="Bangla Demo: Transcribe YouTube",
+    description=(
+        "Transcribe long-form YouTube videos in BANGLA with the click of a button!"
+    ),
+    allow_flagging="never",
+)
+# def transcribe2(audio, state=""):
+#     text = "text"
+#     state += text + " "
+#     return state, state
+
+# Set the starting state to an empty string
+
+# real_transcribe = gr.Interface(
+#     fn=transcribe2,
+#     inputs=[
+#         gr.Audio(source="microphone", type="filepath", streaming=True),
+#         "state"
+#     ],
+#     outputs=[
+#         "textbox",
+#         "state"
+#     ],
+#     live=True)
+
+
+# demo = gr.TabbedInterface([mf_transcribe, yt_transcribe,real_transcribe], ["Transcribe Bangla Audio", "Transcribe Bangla YouTube Video","real time"])
+demo = gr.TabbedInterface(
+    [mf_transcribe, yt_transcribe],
+    ["Transcribe Bangla Audio", "Transcribe Bangla YouTube Video"],
+)
+
+
+if __name__ == "__main__":
+    demo.queue()
+    demo.launch(share="True")
+    # demo.launch(share='True', server_name="0.0.0.0", server_port=8080)

models.py

@@ -0,0 +1,132 @@
+import os
+
+abs_path = os.path.abspath('.')
+base_dir = os.path.dirname(os.path.dirname(abs_path))
+os.environ['TRANSFORMERS_CACHE'] = os.path.join(base_dir, 'models_cache')
+
+import torch
+# Details: https://huggingface.co/docs/diffusers/optimization/fp16#enable-cudnn-autotuner
+torch.backends.cudnn.benchmark = True
+torch.backends.cuda.matmul.allow_tf32 = True
+from transformers import pipeline, AutoTokenizer, AutoFeatureExtractor, AutoConfig, WhisperProcessor, WhisperForConditionalGeneration, WhisperTokenizer, WhisperFeatureExtractor
+from typing import Union, BinaryIO
+# from optimum.bettertransformer import BetterTransformer
+
+language = '<|bn|>'
+# language = '<|en|>'
+task = "transcribe"  # transcribe or translate
+
+# model_name = 'openai/whisper-tiny.en'
+# model_name = 'openai/whisper-base.en'
+# model_name = 'openai/whisper-small.en'
+# model_name = 'openai/whisper-medium' 
+## v2: trained on more epochs with regularization
+# model_name = 'openai/whisper-large-v2' 
+
+## bangla
+# model_name = 'Rakib/whisper-tiny-bn' 
+#model_name = 'anuragshas/whisper-small-bn' 
+# model_name = 'anuragshas/whisper-large-v2-bn'
+# model_name = "Rakib/whisper-small-bn"
+# model_name = "Rakib/whisper-small-bn-all"
+# model_name = "Rakib/whisper-small-bn-all-600"
+# model_name = "Rakib/whisper-small-bn-all-600-v2"
+model_name = "Rakib/whisper-small-bn-crblp"
+
+## lets you know the device count: cuda:0 or cuda:1
+# print(torch.cuda.device_count())
+
+device = 0 if torch.cuda.is_available() else -1
+# device = -1 #Exclusively CPU
+
+print(f"Using device: {'GPU' if device==0 else 'CPU'}")
+
+if device !=0:
+    print("[Warning!] Using CPU could hamper performance")
+
+print("Loading Tokenizer for ASR Speech-to-Text Model...\n" + "*" * 100)
+# tokenizer = AutoTokenizer.from_pretrained(model_name, language=language, task=task)
+# tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+tokenizer = WhisperTokenizer.from_pretrained(model_name)
+# tokenizer(['�', '�্র'],add_prefix_space=True, add_special_tokens=False).input_ids
+
+print("Loading Feature Extractor for ASR Speech-to-Text Model...\n" + "*" * 100)
+# feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
+feature_extractor = WhisperFeatureExtractor.from_pretrained(model_name)
+
+print("Loading Config for ASR Speech-to-Text Model...\n" + "*" * 100)
+config = AutoConfig.from_pretrained(model_name)
+
+print("Loading Processor for ASR Speech-to-Text Model...\n" + "*" * 100)
+processor = WhisperProcessor(feature_extractor=feature_extractor, tokenizer=tokenizer)
+
+print("Loading WHISPER ASR Speech-to-Text Model...\n" + "*" * 100)
+model = WhisperForConditionalGeneration.from_pretrained(model_name)
+
+## BetterTransformer (No Need if PyTorch 2.0 works!!) 
+## (currently 2secs faster inference than PyTorch 2.0 )
+# model = WhisperForConditionalGeneration.from_pretrained(model_name)
+# model = BetterTransformer.transform(model)
+
+## bitsandbytes (only Linux & GPU) (requires conda env with conda-based pytorch!!!)
+## currently only reduces size. slower inference than native models!!!
+## from_pretrained doc: https://huggingface.co/docs/transformers/v4.25.1/en/main_classes/model#transformers.PreTrainedModel.from_pretrained
+# model = WhisperForConditionalGeneration.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
+
+## For PyTorch 2.0 (Only Linux)
+# model = WhisperForConditionalGeneration.from_pretrained(model_name).to(device="cuda:0")
+##mode options are "default", "reduce-overhead" and "max-autotune". See: https://pytorch.org/get-started/pytorch-2.0/#modes
+# model = torch.compile(model, mode="default") 
+
+
+asr = pipeline(
+    task="automatic-speech-recognition",
+    model=model,
+    tokenizer=tokenizer,
+    feature_extractor=feature_extractor,
+    # processor=processor, #no effect see: https://github.com/huggingface/transformers/blob/main/src/transformers/pipelines/automatic_speech_recognition.py
+    # config=config, #no effect see: https://github.com/huggingface/transformers/blob/main/src/transformers/pipelines/automatic_speech_recognition.py
+    device=device,  # for gpu 1 for cpu -1
+    ## chunk files longer than 30s into shorted samples
+    chunk_length_s=30, 
+    ## the amount of overlap (in secs) to be discarded while stitching the inferenced chunks
+    ## stride_length_s is a tuple of the left and right stride(overlap) length.
+    ## With only 1 number, both sides get the same stride, by default
+    ## The stride_length on one side is 1/6th of the chunk_length_s if stride_length no provided
+    # stride_length_s=[8, 8],
+    stride_length_s=[5, 5],
+    # stride_length_s=[6,0],
+    batch_size=16,
+    ignore_warning=True,
+    ## force whisper to generate timestamps so that the chunking and stitching can be accurate
+    # return_timestamps=True, 
+    generate_kwargs = {
+                       'language':language, 
+                       'task':task, 
+                       'repetition_penalty':1.8,
+                       'num_beams':2,
+                       'max_new_tokens':448,
+                       'early_stopping':True,
+                    #    'renormalize_logits':True,
+                       # [16867]: �, [16867, 156, 100, 235, 156, 12811]: �্র
+                       'bad_words_ids':[[16867], [16867, 156, 100, 235, 156, 12811]],
+                    #    'supress_tokens': [16867, 156, 100, 235, 156, 12811],
+                       }
+)
+
+
+def transcribe(speech_array: Union[str, BinaryIO], language: str = "en") -> str:
+    """
+    Transcribes an audio array to text
+    Args:
+        speech_array (np.ndarray): audio in numpy array format
+        language (str): "sv" or "en"
+    Returns:
+        a string containing transcription
+    """
+    asr.model.config.forced_decoder_ids = processor.get_decoder_prompt_ids(language=language, task=task)
+    # asr.model.config.max_new_tokens = 448 #default is 448
+    
+    result = asr(speech_array)
+
+    return str(result["text"])

utils.py

@@ -0,0 +1,23 @@
+
+
+def format_timestamp(
+    seconds: float,
+    always_include_hours: bool = False,
+    decimal_marker: str = ".",
+) -> str:
+    assert seconds >= 0, "non-negative timestamp expected"
+    milliseconds = round(seconds * 1000.0)
+
+    hours = milliseconds // 3_600_000
+    milliseconds -= hours * 3_600_000
+
+    minutes = milliseconds // 60_000
+    milliseconds -= minutes * 60_000
+
+    seconds = milliseconds // 1_000
+    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}"
+    )

vad.py

@@ -0,0 +1,273 @@
+import bisect
+import functools
+import os
+import warnings
+
+from typing import List, Optional
+
+import numpy as np
+
+# The code below is adapted from https://github.com/snakers4/silero-vad.
+
+def get_assets_path():
+    """Returns the path to the assets directory."""
+    return os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets")
+
+
+def get_speech_timestamps(
+    audio: np.ndarray,
+    *,
+    threshold: float = 0.5,
+    # min_speech_duration_ms: int = 250,
+    min_speech_duration_ms: int = 800,
+    max_speech_duration_s: float = float("inf"),
+    # min_silence_duration_ms: int = 2000,
+    min_silence_duration_ms: int = 1000,
+    window_size_samples: int = 1024,
+    speech_pad_ms: int = 200,
+) -> List[dict]:
+    """This method is used for splitting long audios into speech chunks using silero VAD.
+    Args:
+      audio: One dimensional float array.
+      threshold: Speech threshold. Silero VAD outputs speech probabilities for each audio chunk,
+        probabilities ABOVE this value are considered as SPEECH. It is better to tune this
+        parameter for each dataset separately, but "lazy" 0.5 is pretty good for most datasets.
+      min_speech_duration_ms: Final speech chunks shorter min_speech_duration_ms are thrown out.
+      max_speech_duration_s: Maximum duration of speech chunks in seconds. Chunks longer
+        than max_speech_duration_s will be split at the timestamp of the last silence that
+        lasts more than 100s (if any), to prevent agressive cutting. Otherwise, they will be
+        split aggressively just before max_speech_duration_s.
+      min_silence_duration_ms: In the end of each speech chunk wait for min_silence_duration_ms
+        before separating it
+      window_size_samples: Audio chunks of window_size_samples size are fed to the silero VAD model.
+        WARNING! Silero VAD models were trained using 512, 1024, 1536 samples for 16000 sample rate.
+        Values other than these may affect model perfomance!!
+      speech_pad_ms: Final speech chunks are padded by speech_pad_ms each side
+    Returns:
+      List of dicts containing begin and end samples of each speech chunk.
+    """
+    if window_size_samples not in [512, 1024, 1536]:
+        warnings.warn(
+            "Unusual window_size_samples! Supported window_size_samples:\n"
+            " - [512, 1024, 1536] for 16000 sampling_rate"
+        )
+
+    sampling_rate = 16000
+    min_speech_samples = sampling_rate * min_speech_duration_ms / 1000
+    speech_pad_samples = sampling_rate * speech_pad_ms / 1000
+    max_speech_samples = (
+        sampling_rate * max_speech_duration_s
+        - window_size_samples
+        - 2 * speech_pad_samples
+    )
+    min_silence_samples = sampling_rate * min_silence_duration_ms / 1000
+    min_silence_samples_at_max_speech = sampling_rate * 98 / 1000
+
+    audio_length_samples = len(audio)
+
+    model = get_vad_model()
+    state = model.get_initial_state(batch_size=1)
+
+    speech_probs = []
+    for current_start_sample in range(0, audio_length_samples, window_size_samples):
+        chunk = audio[current_start_sample : current_start_sample + window_size_samples]
+        if len(chunk) < window_size_samples:
+            chunk = np.pad(chunk, (0, int(window_size_samples - len(chunk))))
+        speech_prob, state = model(chunk, state, sampling_rate)
+        speech_probs.append(speech_prob)
+
+    triggered = False
+    speeches = []
+    current_speech = {}
+    neg_threshold = threshold - 0.15
+
+    # to save potential segment end (and tolerate some silence)
+    temp_end = 0
+    # to save potential segment limits in case of maximum segment size reached
+    prev_end = next_start = 0
+
+    for i, speech_prob in enumerate(speech_probs):
+        if (speech_prob >= threshold) and temp_end:
+            temp_end = 0
+            if next_start < prev_end:
+                next_start = window_size_samples * i
+
+        if (speech_prob >= threshold) and not triggered:
+            triggered = True
+            current_speech["start"] = window_size_samples * i
+            continue
+
+        if (
+            triggered
+            and (window_size_samples * i) - current_speech["start"] > max_speech_samples
+        ):
+            if prev_end:
+                current_speech["end"] = prev_end
+                speeches.append(current_speech)
+                current_speech = {}
+                # previously reached silence (< neg_thres) and is still not speech (< thres)
+                if next_start < prev_end:
+                    triggered = False
+                else:
+                    current_speech["start"] = next_start
+                prev_end = next_start = temp_end = 0
+            else:
+                current_speech["end"] = window_size_samples * i
+                speeches.append(current_speech)
+                current_speech = {}
+                prev_end = next_start = temp_end = 0
+                triggered = False
+                continue
+
+        if (speech_prob < neg_threshold) and triggered:
+            if not temp_end:
+                temp_end = window_size_samples * i
+            # condition to avoid cutting in very short silence
+            if (window_size_samples * i) - temp_end > min_silence_samples_at_max_speech:
+                prev_end = temp_end
+            if (window_size_samples * i) - temp_end < min_silence_samples:
+                continue
+            else:
+                current_speech["end"] = temp_end
+                if (
+                    current_speech["end"] - current_speech["start"]
+                ) > min_speech_samples:
+                    speeches.append(current_speech)
+                current_speech = {}
+                prev_end = next_start = temp_end = 0
+                triggered = False
+                continue
+
+    if (
+        current_speech
+        and (audio_length_samples - current_speech["start"]) > min_speech_samples
+    ):
+        current_speech["end"] = audio_length_samples
+        speeches.append(current_speech)
+
+    for i, speech in enumerate(speeches):
+        if i == 0:
+            speech["start"] = int(max(0, speech["start"] - speech_pad_samples))
+        if i != len(speeches) - 1:
+            silence_duration = speeches[i + 1]["start"] - speech["end"]
+            if silence_duration < 2 * speech_pad_samples:
+                speech["end"] += int(silence_duration // 2)
+                speeches[i + 1]["start"] = int(
+                    max(0, speeches[i + 1]["start"] - silence_duration // 2)
+                )
+            else:
+                speech["end"] = int(
+                    min(audio_length_samples, speech["end"] + speech_pad_samples)
+                )
+                speeches[i + 1]["start"] = int(
+                    max(0, speeches[i + 1]["start"] - speech_pad_samples)
+                )
+        else:
+            speech["end"] = int(
+                min(audio_length_samples, speech["end"] + speech_pad_samples)
+            )
+
+    return speeches
+
+
+def collect_chunks(audio: np.ndarray, chunks: List[dict]) -> np.ndarray:
+    """Collects and concatenates audio chunks."""
+    if not chunks:
+        return np.array([], dtype=np.float32)
+
+    return np.concatenate([audio[chunk["start"] : chunk["end"]] for chunk in chunks])
+
+
+class SpeechTimestampsMap:
+    """Helper class to restore original speech timestamps."""
+
+    def __init__(self, chunks: List[dict], sampling_rate: int, time_precision: int = 2):
+        self.sampling_rate = sampling_rate
+        self.time_precision = time_precision
+        self.chunk_end_sample = []
+        self.total_silence_before = []
+
+        previous_end = 0
+        silent_samples = 0
+
+        for chunk in chunks:
+            silent_samples += chunk["start"] - previous_end
+            previous_end = chunk["end"]
+
+            self.chunk_end_sample.append(chunk["end"] - silent_samples)
+            self.total_silence_before.append(silent_samples / sampling_rate)
+
+    def get_original_time(
+        self,
+        time: float,
+        chunk_index: Optional[int] = None,
+    ) -> float:
+        if chunk_index is None:
+            chunk_index = self.get_chunk_index(time)
+
+        total_silence_before = self.total_silence_before[chunk_index]
+        return round(total_silence_before + time, self.time_precision)
+
+    def get_chunk_index(self, time: float) -> int:
+        sample = int(time * self.sampling_rate)
+        return min(
+            bisect.bisect(self.chunk_end_sample, sample),
+            len(self.chunk_end_sample) - 1,
+        )
+
+
+@functools.lru_cache
+def get_vad_model():
+    """Returns the VAD model instance."""
+    path = os.path.join(get_assets_path(), "silero_vad.onnx")
+    return SileroVADModel(path)
+
+
+class SileroVADModel:
+    def __init__(self, path):
+        try:
+            import onnxruntime
+        except ImportError as e:
+            raise RuntimeError(
+                "Applying the VAD filter requires the onnxruntime package"
+            ) from e
+
+        opts = onnxruntime.SessionOptions()
+        opts.inter_op_num_threads = 1
+        opts.intra_op_num_threads = 1
+        opts.log_severity_level = 4
+
+        self.session = onnxruntime.InferenceSession(
+            path,
+            providers=["CPUExecutionProvider"],
+            sess_options=opts,
+        )
+
+    def get_initial_state(self, batch_size: int):
+        h = np.zeros((2, batch_size, 64), dtype=np.float32)
+        c = np.zeros((2, batch_size, 64), dtype=np.float32)
+        return h, c
+
+    def __call__(self, x, state, sr: int):
+        if len(x.shape) == 1:
+            x = np.expand_dims(x, 0)
+        if len(x.shape) > 2:
+            raise ValueError(
+                f"Too many dimensions for input audio chunk {len(x.shape)}"
+            )
+        if sr / x.shape[1] > 31.25:
+            raise ValueError("Input audio chunk is too short")
+
+        h, c = state
+
+        ort_inputs = {
+            "input": x,
+            "h": h,
+            "c": c,
+            "sr": np.array(sr, dtype="int64"),
+        }
+
+        out, h, c = self.session.run(None, ort_inputs)
+        state = (h, c)
+
+        return out, state