File size: 4,700 Bytes
bc8cab0
79c3205
bc8cab0
 
 
 
 
 
79c3205
 
 
 
bc8cab0
 
 
a7cf274
 
 
bc8cab0
 
 
 
 
 
07da7ec
bc8cab0
 
 
 
 
 
 
 
 
4503426
bc8cab0
 
 
 
 
 
 
4503426
bc8cab0
 
 
 
 
 
 
 
07da7ec
 
bc8cab0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3e0e1ab
 
 
bc8cab0
4503426
3e0e1ab
bc8cab0
 
 
a7cf274
bc8cab0
 
 
 
 
 
3e0e1ab
 
4503426
 
 
3e0e1ab
 
bc8cab0
 
 
 
a7cf274
bc8cab0
 
 
 
 
 
3fb4d4b
 
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
import logging
import warnings

import gradio as gr
import pytube as pt
import torch
from huggingface_hub import model_info
from transformers import pipeline
from transformers.utils.logging import disable_progress_bar

warnings.filterwarnings("ignore")
disable_progress_bar()

DEFAULT_MODEL_NAME = "bhuang/whisper-medium-cv11-french-case-punctuation"
MODEL_NAMES = [
    "openai/whisper-small",
    "openai/whisper-medium",
    "openai/whisper-large-v2",
    "bhuang/whisper-small-cv11-french",
    "bhuang/whisper-small-cv11-french-case-punctuation",
    "bhuang/whisper-medium-cv11-french",
    "bhuang/whisper-medium-cv11-french-case-punctuation",
]
CHUNK_LENGTH_S = 30
MAX_NEW_TOKENS = 225

logging.basicConfig(
    format="%(asctime)s [%(levelname)s] [%(name)s] %(message)s",
    datefmt="%Y-%m-%dT%H:%M:%SZ",
)
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

device = 0 if torch.cuda.is_available() else "cpu"
logger.info(f"Model will be loaded on device {device}")

cached_models = {}

def maybe_load_cached_pipeline(model_name):
    pipe = cached_models.get(model_name)
    if pipe is None:
        # load pipeline
        # todo: set decoding option for pipeline
        pipe = pipeline(
            task="automatic-speech-recognition",
            model=model_name,
            chunk_length_s=CHUNK_LENGTH_S,
            device=device,
        )
        # set forced_decoder_ids
        pipe.model.config.forced_decoder_ids = pipe.tokenizer.get_decoder_prompt_ids(language="fr", task="transcribe")
        # limit genneration max length
        pipe.model.config.max_length = MAX_NEW_TOKENS + 1

        logger.info(f"`{model_name}` pipeline has been initialized")

        cached_models[model_name] = pipe
    return pipe


def transcribe(microphone, file_upload, model_name):
    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

    pipe = maybe_load_cached_pipeline(model_name)
    text = pipe(file)["text"]

    logger.info(f"Transcription: {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, model_name):
    yt = pt.YouTube(yt_url)
    html_embed_str = _return_yt_html_embed(yt_url)
    stream = yt.streams.filter(only_audio=True)[0]
    stream.download(filename="audio.mp3")

    pipe = maybe_load_cached_pipeline(model_name)
    text = pipe("audio.mp3")["text"]

    logger.info(f"Transcription: {text}")

    return html_embed_str, text


# load default model
maybe_load_cached_pipeline(DEFAULT_MODEL_NAME)

demo = gr.Blocks()

mf_transcribe = gr.Interface(
    fn=transcribe,
    inputs=[
        gr.inputs.Audio(source="microphone", type="filepath", optional=True, label="Record"),
        gr.inputs.Audio(source="upload", type="filepath", optional=True, label="Upload File"),
        gr.inputs.Dropdown(choices=MODEL_NAMES, default=DEFAULT_MODEL_NAME, label="Whisper Model"),
    ],
    # outputs="text",
    outputs=gr.outputs.Textbox(label="Transcription"),
    layout="horizontal",
    theme="huggingface",
    title="Whisper Demo: Transcribe Audio",
    description="Transcribe long-form microphone or audio inputs with the click of a button!",
    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.Dropdown(choices=MODEL_NAMES, default=DEFAULT_MODEL_NAME, label="Whisper Model"),
    ],
    # outputs=["html", "text"],
    outputs=[
        gr.outputs.HTML(label="YouTube Page"),
        gr.outputs.Textbox(label="Transcription"),
    ],
    layout="horizontal",
    theme="huggingface",
    title="Whisper Demo: Transcribe YouTube",
    description="Transcribe long-form YouTube videos with the click of a button!",
    allow_flagging="never",
)

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

# demo.launch(server_name="0.0.0.0", debug=True, share=True)
demo.launch(enable_queue=True)