Switch to Docker

Dockerfile

@@ -0,0 +1,30 @@
+FROM ubuntu
+
+# Based on https://huggingface.co/spaces/radames/nginx-gradio-reverse-proxy/blob/main/Dockerfile
+USER root
+
+RUN apt-get -y update && apt-get -y install nginx
+RUN mkdir -p /var/cache/nginx \
+             /var/log/nginx \
+             /var/lib/nginx
+RUN touch /var/run/nginx.pid
+
+RUN chown -R 1000:1000 /var/cache/nginx \
+                       /var/log/nginx \
+                       /var/lib/nginx \
+                       /var/run/nginx.pid
+
+RUN useradd -m -u 1000 user
+
+USER user
+ENV HOME=/home/user
+
+RUN mkdir $HOME/app
+WORKDIR $HOME/app
+
+# Copy nginx configuration
+COPY --chown=user nginx.conf /etc/nginx/sites-available/default
+COPY --chown=user . .
+
+CMD ["bash", "run.sh"]
+

README.md

@@ -3,9 +3,7 @@ title: Whisper JAX
 emoji: ⚡️
 colorFrom: yellow
 colorTo: indigo
-sdk: gradio
-sdk_version: 3.24.1
-app_file: app.py
+sdk: docker
 pinned: false
 ---
 

app.py

@@ -1,177 +0,0 @@
-import base64
-import os
-from functools import partial
-from multiprocessing import Pool
-
-import gradio as gr
-import numpy as np
-import requests
-from processing_whisper import WhisperPrePostProcessor
-from transformers.models.whisper.tokenization_whisper import TO_LANGUAGE_CODE
-from transformers.pipelines.audio_utils import ffmpeg_read
-
-
-title = "Whisper JAX: The Fastest Whisper API ⚡️"
-
-description = """Whisper JAX is an optimised implementation of the [Whisper model](https://huggingface.co/openai/whisper-large-v2) by OpenAI. It runs on JAX with a TPU v4-8 in the backend. Compared to PyTorch on an A100 GPU, it is over [**70x faster**](https://github.com/sanchit-gandhi/whisper-jax#benchmarks), making it the fastest Whisper API available.
-
-Note that using microphone or audio file requires the audio input to be transferred from the Gradio demo to the TPU, which for large audio files can be slow. We recommend using YouTube where possible, since this directly downloads the audio file to the TPU, skipping the file transfer step.
-"""
-
-API_URL = os.getenv("API_URL")
-API_URL_FROM_FEATURES = os.getenv("API_URL_FROM_FEATURES")
-
-article = "Whisper large-v2 model by OpenAI. Backend running JAX on a TPU v4-8 through the generous support of the [TRC](https://sites.research.google/trc/about/) programme. Whisper JAX [code](https://github.com/sanchit-gandhi/whisper-jax) and Gradio demo by 🤗 Hugging Face."
-
-language_names = sorted(TO_LANGUAGE_CODE.keys())
-CHUNK_LENGTH_S = 30
-BATCH_SIZE = 16
-NUM_PROC = 16
-FILE_LIMIT_MB = 1000
-
-
-def query(payload):
-    response = requests.post(API_URL, json=payload)
-    return response.json(), response.status_code
-
-
-def inference(inputs, language=None, task=None, return_timestamps=False):
-    payload = {"inputs": inputs, "task": task, "return_timestamps": return_timestamps}
-
-    # langauge can come as an empty string from the Gradio `None` default, so we handle it separately
-    if language:
-        payload["language"] = language
-
-    data, status_code = query(payload)
-
-    if status_code == 200:
-        text = data["text"]
-    else:
-        text = data["detail"]
-
-    if return_timestamps:
-        timestamps = data["chunks"]
-    else:
-        timestamps = None
-
-    return text, timestamps
-
-
-def chunked_query(payload):
-    response = requests.post(API_URL_FROM_FEATURES, json=payload)
-    return response.json()
-
-
-def forward(batch, task=None, return_timestamps=False):
-    feature_shape = batch["input_features"].shape
-    batch["input_features"] = base64.b64encode(batch["input_features"].tobytes()).decode()
-    outputs = chunked_query(
-        {"batch": batch, "task": task, "return_timestamps": return_timestamps, "feature_shape": feature_shape}
-    )
-    outputs["tokens"] = np.asarray(outputs["tokens"])
-    return outputs
-
-
-if __name__ == "__main__":
-    processor = WhisperPrePostProcessor.from_pretrained("openai/whisper-large-v2")
-    pool = Pool(NUM_PROC)
-
-    def transcribe_chunked_audio(inputs, task, return_timestamps):
-        file_size_mb = os.stat(inputs).st_size / (1024 * 1024)
-        if file_size_mb > FILE_LIMIT_MB:
-            return f"ERROR: File size exceeds file size limit. Got file of size {file_size_mb:.2f}MB for a limit of {FILE_LIMIT_MB}MB.", None
-
-        with open(inputs, "rb") as f:
-            inputs = f.read()
-
-        inputs = ffmpeg_read(inputs, processor.feature_extractor.sampling_rate)
-        inputs = {"array": inputs, "sampling_rate": processor.feature_extractor.sampling_rate}
-
-        dataloader = processor.preprocess_batch(inputs, chunk_length_s=CHUNK_LENGTH_S, batch_size=BATCH_SIZE)
-
-        try:
-            model_outputs = pool.map(partial(forward, task=task, return_timestamps=return_timestamps), dataloader)
-        except ValueError as err:
-            # pre-processor does all the necessary compatibility checks for our audio inputs
-            return err, None
-
-        post_processed = processor.postprocess(model_outputs, return_timestamps=return_timestamps)
-        timestamps = post_processed.get("chunks")
-        return post_processed["text"], timestamps
-
-    def _return_yt_html_embed(yt_url):
-        video_id = yt_url.split("?v=")[-1]
-        HTML_str = (
-            f'<center> <iframe width="500" height="320" src="https://www.youtube.com/embed/{video_id}"> </iframe>'
-            " </center>"
-        )
-        return HTML_str
-
-    def transcribe_youtube(yt_url, task, return_timestamps):
-        html_embed_str = _return_yt_html_embed(yt_url)
-
-        text, timestamps = inference(inputs=yt_url, task=task, return_timestamps=return_timestamps)
-
-        return html_embed_str, text, timestamps
-
-    microphone_chunked = gr.Interface(
-        fn=transcribe_chunked_audio,
-        inputs=[
-            gr.inputs.Audio(source="microphone", optional=True, type="filepath"),
-            gr.inputs.Radio(["transcribe", "translate"], label="Task", default="transcribe"),
-            gr.inputs.Checkbox(default=False, label="Return timestamps"),
-        ],
-        outputs=[
-            gr.outputs.Textbox(label="Transcription"),
-            gr.outputs.Textbox(label="Timestamps"),
-        ],
-        allow_flagging="never",
-        title=title,
-        description=description,
-        article=article,
-    )
-
-    audio_chunked = gr.Interface(
-        fn=transcribe_chunked_audio,
-        inputs=[
-            gr.inputs.Audio(source="upload", optional=True, label="Audio file", type="filepath"),
-            gr.inputs.Radio(["transcribe", "translate"], label="Task", default="transcribe"),
-            gr.inputs.Checkbox(default=False, label="Return timestamps"),
-        ],
-        outputs=[
-            gr.outputs.Textbox(label="Transcription"),
-            gr.outputs.Textbox(label="Timestamps"),
-        ],
-        allow_flagging="never",
-        title=title,
-        description=description,
-        article=article,
-    )
-
-    youtube = gr.Interface(
-        fn=transcribe_youtube,
-        inputs=[
-            gr.inputs.Textbox(lines=1, placeholder="Paste the URL to a YouTube video here", label="YouTube URL"),
-            gr.inputs.Radio(["transcribe", "translate"], label="Task", default="transcribe"),
-            gr.inputs.Checkbox(default=False, label="Return timestamps"),
-        ],
-        outputs=[
-            gr.outputs.HTML(label="Video"),
-            gr.outputs.Textbox(label="Transcription"),
-            gr.outputs.Textbox(label="Timestamps"),
-        ],
-        allow_flagging="never",
-        title=title,
-        examples=[["https://www.youtube.com/watch?v=m8u-18Q0s7I", "transcribe", False]],
-        cache_examples=False,
-        description=description,
-        article=article,
-    )
-
-    demo = gr.Blocks()
-
-    with demo:
-        gr.TabbedInterface([microphone_chunked, audio_chunked, youtube], ["Microphone", "Audio File", "YouTube"])
-
-    demo.queue(max_size=3)
-    demo.launch(show_api=False)

nginx.conf

@@ -0,0 +1,23 @@
+server {
+    listen 7860 default_server;
+    listen [::]:7860 default_server;
+    
+    root /usr/share/nginx/html;
+    index index.html index.htm;
+
+    server_name _;
+    location / {
+        proxy_pass https://whisper-jax.ngrok.io;
+        proxy_set_header    Host    whisper-jax.ngrok.io;
+        proxy_set_header    X-Real-IP   $remote_addr;
+        proxy_set_header    X-Forwarded-For $proxy_add_x_forwarded_for;
+        #proxy_set_header   X-Forwarded-Proto   $scheme;
+        proxy_set_header    X-Forwarded-Proto   http;
+        proxy_set_header    X-Forwarded-Ssl     off;
+        proxy_set_header    X-Url-Scheme        http;
+        proxy_buffering     off;
+        proxy_http_version  1.1;
+        proxy_set_header    Upgrade $http_upgrade;
+        proxy_set_header    Connection  "upgrade";
+    }
+}

packages.txt

@@ -1 +0,0 @@
-ffmpeg

processing_whisper.py

@@ -1,145 +0,0 @@
-import math
-
-import numpy as np
-from transformers import WhisperProcessor
-
-
-class WhisperPrePostProcessor(WhisperProcessor):
-    def chunk_iter_with_batch(self, inputs, chunk_len, stride_left, stride_right, batch_size):
-        inputs_len = inputs.shape[0]
-        step = chunk_len - stride_left - stride_right
-
-        all_chunk_start_idx = np.arange(0, inputs_len, step)
-        num_samples = len(all_chunk_start_idx)
-
-        num_batches = math.ceil(num_samples / batch_size)
-        batch_idx = np.array_split(np.arange(num_samples), num_batches)
-
-        for i, idx in enumerate(batch_idx):
-            chunk_start_idx = all_chunk_start_idx[idx]
-
-            chunk_end_idx = chunk_start_idx + chunk_len
-
-            chunks = [inputs[chunk_start:chunk_end] for chunk_start, chunk_end in zip(chunk_start_idx, chunk_end_idx)]
-            processed = self.feature_extractor(
-                chunks, sampling_rate=self.feature_extractor.sampling_rate, return_tensors="np"
-            )
-
-            _stride_left = np.where(chunk_start_idx == 0, 0, stride_left)
-            is_last = np.where(stride_right > 0, chunk_end_idx > inputs_len, chunk_end_idx >= inputs_len)
-            _stride_right = np.where(is_last, 0, stride_right)
-
-            chunk_lens = [chunk.shape[0] for chunk in chunks]
-            strides = [
-                (int(chunk_l), int(_stride_l), int(_stride_r))
-                for chunk_l, _stride_l, _stride_r in zip(chunk_lens, _stride_left, _stride_right)
-            ]
-
-            yield {"stride": strides, **processed}
-
-    def preprocess_batch(self, inputs, chunk_length_s=0, stride_length_s=None, batch_size=None):
-        stride = None
-        if isinstance(inputs, dict):
-            stride = inputs.pop("stride", None)
-            # Accepting `"array"` which is the key defined in `datasets` for
-            # better integration
-            if not ("sampling_rate" in inputs and ("raw" in inputs or "array" in inputs)):
-                raise ValueError(
-                    "When passing a dictionary to FlaxWhisperPipline, the dict needs to contain a "
-                    '"raw" or "array" key containing the numpy array representing the audio, and a "sampling_rate" key '
-                    "containing the sampling rate associated with the audio array."
-                )
-
-            _inputs = inputs.pop("raw", None)
-            if _inputs is None:
-                # Remove path which will not be used from `datasets`.
-                inputs.pop("path", None)
-                _inputs = inputs.pop("array", None)
-            in_sampling_rate = inputs.pop("sampling_rate")
-            inputs = _inputs
-
-            if in_sampling_rate != self.feature_extractor.sampling_rate:
-                try:
-                    import librosa
-                except ImportError as err:
-                    raise ImportError(
-                        "To support resampling audio files, please install 'librosa' and 'soundfile'."
-                    ) from err
-
-                inputs = librosa.resample(
-                    inputs, orig_sr=in_sampling_rate, target_sr=self.feature_extractor.sampling_rate
-                )
-                ratio = self.feature_extractor.sampling_rate / in_sampling_rate
-            else:
-                ratio = 1
-
-        if not isinstance(inputs, np.ndarray):
-            raise ValueError(f"We expect a numpy ndarray as input, got `{type(inputs)}`.")
-        if len(inputs.shape) != 1:
-            raise ValueError(
-                f"We expect a single channel audio input for the Flax Whisper API, got {len(inputs.shape)} channels."
-            )
-
-        if stride is not None:
-            if stride[0] + stride[1] > inputs.shape[0]:
-                raise ValueError("Stride is too large for input.")
-
-            # Stride needs to get the chunk length here, it's going to get
-            # swallowed by the `feature_extractor` later, and then batching
-            # can add extra data in the inputs, so we need to keep track
-            # of the original length in the stride so we can cut properly.
-            stride = (inputs.shape[0], int(round(stride[0] * ratio)), int(round(stride[1] * ratio)))
-
-        if chunk_length_s:
-            if stride_length_s is None:
-                stride_length_s = chunk_length_s / 6
-
-            if isinstance(stride_length_s, (int, float)):
-                stride_length_s = [stride_length_s, stride_length_s]
-
-            chunk_len = round(chunk_length_s * self.feature_extractor.sampling_rate)
-            stride_left = round(stride_length_s[0] * self.feature_extractor.sampling_rate)
-            stride_right = round(stride_length_s[1] * self.feature_extractor.sampling_rate)
-
-            if chunk_len < stride_left + stride_right:
-                raise ValueError("Chunk length must be superior to stride length.")
-
-            for item in self.chunk_iter_with_batch(
-                inputs,
-                chunk_len,
-                stride_left,
-                stride_right,
-                batch_size,
-            ):
-                yield item
-        else:
-            processed = self.feature_extractor(
-                inputs, sampling_rate=self.feature_extractor.sampling_rate, return_tensors="np"
-            )
-            if stride is not None:
-                processed["stride"] = stride
-            yield processed
-
-    def postprocess(self, model_outputs, return_timestamps=None, return_language=None):
-        # unpack the outputs from list(dict(list)) to list(dict)
-        model_outputs = [dict(zip(output, t)) for output in model_outputs for t in zip(*output.values())]
-
-        time_precision = self.feature_extractor.chunk_length / 1500  # max source positions = 1500
-        # Send the chunking back to seconds, it's easier to handle in whisper
-        sampling_rate = self.feature_extractor.sampling_rate
-        for output in model_outputs:
-            if "stride" in output:
-                chunk_len, stride_left, stride_right = output["stride"]
-                # Go back in seconds
-                chunk_len /= sampling_rate
-                stride_left /= sampling_rate
-                stride_right /= sampling_rate
-                output["stride"] = chunk_len, stride_left, stride_right
-
-        text, optional = self.tokenizer._decode_asr(
-            model_outputs,
-            return_timestamps=return_timestamps,
-            return_language=return_language,
-            time_precision=time_precision,
-        )
-        return {"text": text, **optional}

requirements.txt

@@ -1,3 +0,0 @@
-transformers
-pytube
-requests>=2.28.2

run.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+
+service nginx start
+sleep infinity