File size: 4,188 Bytes
6c226f9
 
33ac9e4
6c226f9
33ac9e4
6c226f9
7097513
9d6fa91
66efbc3
 
6c226f9
 
 
 
 
 
 
 
 
 
 
0a7fcda
 
 
 
 
7097513
6c226f9
 
 
 
 
 
 
 
66efbc3
 
33ac9e4
66efbc3
 
 
 
6c226f9
 
 
7097513
0a7fcda
9d6fa91
6c226f9
 
 
 
 
 
 
 
 
 
 
 
 
33ac9e4
 
66efbc3
33ac9e4
 
 
 
 
 
 
7097513
33ac9e4
 
6c226f9
 
 
 
 
 
 
 
 
 
609dcbe
6c226f9
 
 
 
b95b5ca
6c226f9
b95b5ca
6c226f9
 
 
 
 
 
 
 
7097513
 
609dcbe
7097513
6c226f9
 
 
b95b5ca
6c226f9
b95b5ca
 
6c226f9
 
 
 
 
33ac9e4
6c226f9
 
 
33ac9e4
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
import torch
import gradio as gr
import yt_dlp
from transformers import pipeline
import os 

MODEL_NAME = "openai/whisper-large-v2"
BATCH_SIZE = 8
FILE_LIMIT_MB = 1000
YT_ATTEMPT_LIMIT = 3

device = 0 if torch.cuda.is_available() else "cpu"

pipe = pipeline(
    task="automatic-speech-recognition",
    model=MODEL_NAME,
    chunk_length_s=30,
    device=device,
)


all_special_ids = pipe.tokenizer.all_special_ids
transcribe_token_id = all_special_ids[-5]
translate_token_id = all_special_ids[-6]


def transcribe(microphone, file_upload, task):
    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):
        raise gr.Error("You have to either use the microphone or upload an audio file")

    file_size_mb = os.stat(file_upload).st_size / (1024 * 1024)
    if file_size_mb > FILE_LIMIT_MB:
        raise gr.Error(
                f"File size exceeds file size limit. Got file of size {file_size_mb:.2f}MB for a limit of {FILE_LIMIT_MB}MB."
        )

    file = microphone if microphone is not None else file_upload

    pipe.model.config.forced_decoder_ids = [[2, transcribe_token_id if task=="transcribe" else translate_token_id]]

    text = pipe(file, batch_size=BATCH_SIZE)["text"]

    return warn_output + text


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 yt_transcribe(yt_url, task, max_filesize=FILE_LIMIT_MB):
    with yt_dlp.YoutubeDL({'format': 'bestaudio/best'}) as ydl:
        try:
            info_dict = ydl.extract_info(yt_url, download=True)
            a = ydl.prepare_filename(info_dict)
        except Exception as e:
            raise gr.Error(f"Error downloading YouTube video: {str(e)}")
    html_embed_str = _return_yt_html_embed(yt_url)
    if os.stat(a).st_size / (1024 * 1024) > max_filesize:
        raise gr.Error(f"Maximum YouTube file size is {max_filesize}MB, got {os.stat(a).st_size / (1024 * 1024):.2f}MB.")
    pipe.model.config.forced_decoder_ids = [[2, transcribe_token_id if task=="transcribe" else translate_token_id]]
    text = pipe(a, batch_size=BATCH_SIZE)["text"]
    os.remove(a)
    return html_embed_str, text


demo = gr.Blocks()

mf_transcribe = gr.Interface(
    fn=transcribe,
    inputs=[
        gr.inputs.Audio(source="microphone", type="filepath", optional=True),
        gr.inputs.Audio(source="upload", type="filepath", optional=True),
        gr.inputs.Radio(["transcribe", "translate"], label="Task", default="transcribe"),
    ],
    outputs="text",
    layout="horizontal",
    theme="huggingface",
    title="Whisper Large V2: Transcribe Audio",
    description=(
        "Transcribe long-form microphone or audio inputs with the click of a button! Demo uses the"
        f" checkpoint [{MODEL_NAME}](https://huggingface.co/{MODEL_NAME}) and 🤗 Transformers to transcribe audio files"
        " of arbitrary length."
    ),
    allow_flagging="never",
)

yt_transcribe = gr.Interface(
    fn=yt_transcribe,
    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")
    ],
    outputs=["html", "text"],
    layout="horizontal",
    theme="huggingface",
    title="Whisper Large V2: Transcribe YouTube",
    description=(
        "Transcribe long-form YouTube videos with the click of a button! Demo uses the checkpoint"
        f" [{MODEL_NAME}](https://huggingface.co/{MODEL_NAME}) and 🤗 Transformers to transcribe video files of"
        " arbitrary length."
    ),
    allow_flagging="never",
)


with demo:
    gr.TabbedInterface([mf_transcribe, yt_transcribe], ["Transcribe Audio", "Transcribe YouTube"])

demo.launch(enable_queue=True)