File size: 10,934 Bytes
c173d9b
 
8da2b37
c173d9b
 
8da2b37
 
 
c173d9b
 
 
 
 
 
8da2b37
 
 
d63d47a
8da2b37
d63d47a
8da2b37
 
 
 
2ecbad4
 
 
 
 
 
 
 
8da2b37
 
 
 
 
 
 
d63d47a
8da2b37
 
 
 
 
 
 
 
d63d47a
8da2b37
 
 
7826b7e
 
 
 
8da2b37
 
 
7826b7e
 
 
 
920d6ef
 
 
 
 
7826b7e
920d6ef
 
7826b7e
8da2b37
7826b7e
8da2b37
 
7826b7e
 
920d6ef
7826b7e
 
 
 
 
d63d47a
7826b7e
d63d47a
 
7826b7e
d63d47a
7826b7e
 
d63d47a
7826b7e
8da2b37
c173d9b
 
8da2b37
d63d47a
 
 
 
 
 
 
8da2b37
 
d63d47a
 
 
 
 
 
8da2b37
d63d47a
118252d
d63d47a
 
 
 
 
8da2b37
 
 
d63d47a
 
 
118252d
d63d47a
 
 
 
 
8da2b37
 
 
7826b7e
 
8da2b37
7826b7e
 
8da2b37
 
 
7826b7e
8da2b37
118252d
7826b7e
 
 
 
 
 
 
8da2b37
 
7826b7e
 
 
 
 
 
 
 
 
 
 
 
8da2b37
7826b7e
8da2b37
118252d
d63d47a
 
7826b7e
d63d47a
7826b7e
 
 
8da2b37
 
 
 
 
 
 
 
7826b7e
 
8da2b37
 
 
 
 
 
 
 
 
 
 
 
c173d9b
d63d47a
8da2b37
 
d63d47a
8da2b37
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c173d9b
d63d47a
8da2b37
 
d63d47a
2ecbad4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
import logging
import math
import os
import tempfile
import time
import yt_dlp as youtube_dl
from fastapi import FastAPI, UploadFile, Form, HTTPException
from fastapi.responses import HTMLResponse
import jax.numpy as jnp
import numpy as np
from transformers.models.whisper.tokenization_whisper import TO_LANGUAGE_CODE
from transformers.pipelines.audio_utils import ffmpeg_read
from whisper_jax import FlaxWhisperPipline

app = FastAPI(title="Whisper JAX: The Fastest Whisper API ⚡️")

logger = logging.getLogger("whisper-jax-app")
logger.setLevel(logging.DEBUG)
ch = logging.StreamHandler()
ch.setLevel(logging.DEBUG)
formatter = logging.Formatter("%(asctime)s;%(levelname)s;%(message)s", "%Y-%m-%d %H:%M:%S")
ch.setFormatter(formatter)
logger.addHandler(ch)

checkpoint = "openai/whisper-large-v3"

BATCH_SIZE = 32
CHUNK_LENGTH_S = 30
NUM_PROC = 32
FILE_LIMIT_MB = 10000
YT_LENGTH_LIMIT_S = 15000  # limit to 2 hour YouTube files

pipeline = FlaxWhisperPipline(checkpoint, dtype=jnp.bfloat16, batch_size=BATCH_SIZE)
stride_length_s = CHUNK_LENGTH_S / 6
chunk_len = round(CHUNK_LENGTH_S * pipeline.feature_extractor.sampling_rate)
stride_left = stride_right = round(stride_length_s * pipeline.feature_extractor.sampling_rate)
step = chunk_len - stride_left - stride_right

# do a pre-compile step so that the first user to use the demo isn't hit with a long transcription time
logger.debug("Compiling forward call...")
start = time.time()
random_inputs = {
    "input_features": np.ones(
        (BATCH_SIZE, pipeline.model.config.num_mel_bins, 2 * pipeline.model.config.max_source_positions)
    )
}
random_timestamps = pipeline.forward(random_inputs, batch_size=BATCH_SIZE, return_timestamps=True)
compile_time = time.time() - start
logger.debug(f"Compiled in {compile_time}s")

@app.post("/transcribe_audio")
async def transcribe_chunked_audio(audio_file: UploadFile, task: str = "transcribe", return_timestamps: bool = False):
    logger.debug("Starting transcribe_chunked_audio function")
    logger.debug(f"Received parameters - task: {task}, return_timestamps: {return_timestamps}")
    
    logger.debug("Checking for audio file...")
    if not audio_file:
        logger.warning("No audio file")
        raise HTTPException(status_code=400, detail="No audio file submitted!")
    
    logger.debug(f"Audio file received: {audio_file.filename}")
    
    try:
        # Read the file content
        file_content = await audio_file.read()
        file_size = len(file_content)
        file_size_mb = file_size / (1024 * 1024)
        logger.debug(f"File size: {file_size} bytes ({file_size_mb:.2f}MB)")
    except Exception as e:
        logger.error(f"Error reading file: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"Error reading file: {str(e)}")

    if file_size_mb > FILE_LIMIT_MB:
        logger.warning(f"Max file size exceeded: {file_size_mb:.2f}MB > {FILE_LIMIT_MB}MB")
        raise HTTPException(status_code=400, detail=f"File size exceeds file size limit. Got file of size {file_size_mb:.2f}MB for a limit of {FILE_LIMIT_MB}MB.")

    try:
        logger.debug("Performing ffmpeg read on audio file")
        inputs = ffmpeg_read(file_content, pipeline.feature_extractor.sampling_rate)
        inputs = {"array": inputs, "sampling_rate": pipeline.feature_extractor.sampling_rate}
        logger.debug("ffmpeg read completed successfully")
    except Exception as e:
        logger.error(f"Error in ffmpeg read: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"Error processing audio file: {str(e)}")

    logger.debug("Calling tqdm_generate to transcribe audio")
    try:
        text, runtime = tqdm_generate(inputs, task=task, return_timestamps=return_timestamps)
        logger.debug(f"Transcription completed. Runtime: {runtime:.2f}s")
    except Exception as e:
        logger.error(f"Error in tqdm_generate: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"Error transcribing audio: {str(e)}")

    logger.debug("Transcribe_chunked_audio function completed successfully")
    return {"text": text, "runtime": runtime}

@app.post("/transcribe_youtube")
async def transcribe_youtube(yt_url: str = Form(...), task: str = "transcribe", return_timestamps: bool = False):
    logger.debug("Loading YouTube file...")
    try:
        html_embed_str = _return_yt_html_embed(yt_url)
    except Exception as e:
        logger.error("Error generating YouTube HTML embed:", exc_info=True)
        raise HTTPException(status_code=500, detail="Error generating YouTube HTML embed")

    with tempfile.TemporaryDirectory() as tmpdirname:
        filepath = os.path.join(tmpdirname, "video.mp4")
        try:
            logger.debug("Downloading YouTube audio...")
            download_yt_audio(yt_url, filepath)
        except Exception as e:
            logger.error("Error downloading YouTube audio:", exc_info=True)
            raise HTTPException(status_code=500, detail="Error downloading YouTube audio")

        try:
            logger.debug(f"Opening downloaded audio file: {filepath}")
            with open(filepath, "rb") as f:
                inputs = f.read()
        except Exception as e:
            logger.error("Error reading downloaded audio file:", exc_info=True)
            raise HTTPException(status_code=500, detail="Error reading downloaded audio file")

    inputs = ffmpeg_read(inputs, pipeline.feature_extractor.sampling_rate)
    inputs = {"array": inputs, "sampling_rate": pipeline.feature_extractor.sampling_rate}
    logger.debug("Done loading YouTube file")

    try:
        logger.debug("Calling tqdm_generate to transcribe YouTube audio")
        text, runtime = tqdm_generate(inputs, task=task, return_timestamps=return_timestamps)
    except Exception as e:
        logger.error("Error transcribing YouTube audio:", exc_info=True)
        raise HTTPException(status_code=500, detail="Error transcribing YouTube audio")

    return {"html_embed": html_embed_str, "text": text, "runtime": runtime}

def tqdm_generate(inputs: dict, task: str, return_timestamps: bool):
    logger.debug(f"Starting tqdm_generate - task: {task}, return_timestamps: {return_timestamps}")
    
    inputs_len = inputs["array"].shape[0]
    logger.debug(f"Input array length: {inputs_len}")
    
    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)
    logger.debug(f"Number of samples: {num_samples}, Number of batches: {num_batches}")

    logger.debug("Preprocessing audio for inference")
    try:
        dataloader = pipeline.preprocess_batch(inputs, chunk_length_s=CHUNK_LENGTH_S, batch_size=BATCH_SIZE)
        logger.debug("Preprocessing completed successfully")
    except Exception as e:
        logger.error(f"Error in preprocessing: {str(e)}", exc_info=True)
        raise

    model_outputs = []
    start_time = time.time()
    logger.debug("Starting transcription...")
    
    try:
        for i, batch in enumerate(dataloader):
            logger.debug(f"Processing batch {i+1}/{num_batches} with {len(batch)} samples")
            batch_output = pipeline.forward(batch, batch_size=BATCH_SIZE, task=task, return_timestamps=True)
            model_outputs.append(batch_output)
            logger.debug(f"Batch {i+1} processed successfully")
    except Exception as e:
        logger.error(f"Error during batch processing: {str(e)}", exc_info=True)
        raise

    runtime = time.time() - start_time
    logger.debug(f"Transcription completed in {runtime:.2f}s")

    logger.debug("Post-processing transcription results")
    try:
        post_processed = pipeline.postprocess(model_outputs, return_timestamps=True)
        logger.debug("Post-processing completed successfully")
    except Exception as e:
        logger.error(f"Error in post-processing: {str(e)}", exc_info=True)
        raise

    text = post_processed["text"]
    if return_timestamps:
        timestamps = post_processed.get("chunks")
        timestamps = [
            f"[{format_timestamp(chunk['timestamp'][0])} -> {format_timestamp(chunk['timestamp'][1])}] {chunk['text']}"
            for chunk in timestamps
        ]
        text = "\n".join(str(feature) for feature in timestamps)
    
    logger.debug("tqdm_generate function completed successfully")
    return text, runtime

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 download_yt_audio(yt_url, filename):
    info_loader = youtube_dl.YoutubeDL()
    try:
        logger.debug(f"Extracting info for YouTube URL: {yt_url}")
        info = info_loader.extract_info(yt_url, download=False)
    except youtube_dl.utils.DownloadError as err:
        logger.error("Error extracting YouTube info:", exc_info=True)
        raise HTTPException(status_code=400, detail=str(err))

    file_length = info["duration_string"]
    file_h_m_s = file_length.split(":")
    file_h_m_s = [int(sub_length) for sub_length in file_h_m_s]
    if len(file_h_m_s) == 1:
        file_h_m_s.insert(0, 0)
    if len(file_h_m_s) == 2:
        file_h_m_s.insert(0, 0)

    file_length_s = file_h_m_s[0] * 3600 + file_h_m_s[1] * 60 + file_h_m_s[2]
    if file_length_s > YT_LENGTH_LIMIT_S:
        yt_length_limit_hms = time.strftime("%HH:%MM:%SS", time.gmtime(YT_LENGTH_LIMIT_S))
        file_length_hms = time.strftime("%HH:%MM:%SS", time.gmtime(file_length_s))
        raise HTTPException(status_code=400, detail=f"Maximum YouTube length is {yt_length_limit_hms}, got {file_length_hms} YouTube video.")

    ydl_opts = {"outtmpl": filename, "format": "worstvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best"}
    with youtube_dl.YoutubeDL(ydl_opts) as ydl:
        try:
            logger.debug(f"Downloading YouTube audio to {filename}")
            ydl.download([yt_url])
        except youtube_dl.utils.ExtractorError as err:
            logger.error("Error downloading YouTube audio:", exc_info=True)
            raise HTTPException(status_code=400, detail=str(err))

def format_timestamp(seconds: float, always_include_hours: bool = False, decimal_marker: str = "."):
    if seconds is not None:
        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}"
    else:
        # we have a malformed timestamp so just return it as is
        return seconds