File size: 3,584 Bytes
18fb4f3
 
 
 
 
 
 
 
 
 
 
 
a72cb3f
18fb4f3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a72cb3f
18fb4f3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a72cb3f
 
18fb4f3
 
 
 
 
 
 
 
 
 
a72cb3f
 
18fb4f3
 
 
 
 
 
 
 
 
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
import os
from math import floor
from typing import Optional

import spaces
import torch
import gradio as gr
from transformers import pipeline
from transformers.pipelines.audio_utils import ffmpeg_read


# configuration
MODEL_NAME = "kotoba-tech/kotoba-whisper-bilingual-v1.0"
BATCH_SIZE = 16
CHUNK_LENGTH_S = 15
# device setting
if torch.cuda.is_available():
    torch_dtype = torch.bfloat16
    device = "cuda"
    model_kwargs = {'attn_implementation': 'sdpa'}
else:
    torch_dtype = torch.float32
    device = "cpu"
    model_kwargs = {}

# define the pipeline
pipe = pipeline(
    model=MODEL_NAME,
    chunk_length_s=CHUNK_LENGTH_S,
    batch_size=BATCH_SIZE,
    torch_dtype=torch_dtype,
    device=device,
    model_kwargs=model_kwargs,
    trust_remote_code=True
)


def format_time(start: Optional[float], end: Optional[float]):

    def _format_time(seconds: Optional[float]):
        if seconds is None:
            return "complete    "
        minutes = floor(seconds / 60)
        hours = floor(seconds / 3600)
        seconds = seconds - hours * 3600 - minutes * 60
        m_seconds = floor(round(seconds - floor(seconds), 3) * 10 ** 3)
        seconds = floor(seconds)
        return f'{hours:02}:{minutes:02}:{seconds:02}.{m_seconds:03}'

    return f"[{_format_time(start)}-> {_format_time(end)}]:"


@spaces.GPU
def get_prediction(inputs, task: str, language: Optional[str]):
    generate_kwargs = {"task": task, "language": language}
    prediction = pipe(inputs, return_timestamps=True, generate_kwargs=generate_kwargs)
    text = "".join([c['text'] for c in prediction['chunks']])
    text_timestamped = "\n".join([
        f"{format_time(*c['timestamp'])} {c['text']}" for c in prediction['chunks']
    ])
    return text, text_timestamped


def transcribe(inputs: str, task: str, language: str):
    if inputs is None:
        raise gr.Error("No audio file submitted! Please upload or record an audio file before submitting your request.")
    with open(inputs, "rb") as f:
        inputs = f.read()
    inputs = ffmpeg_read(inputs, pipe.feature_extractor.sampling_rate)
    inputs = {"array": inputs, "sampling_rate": pipe.feature_extractor.sampling_rate}
    return get_prediction(inputs, task, language)


demo = gr.Blocks()
description = (f"Transcribe long-form microphone or audio inputs with the click of a button! Demo uses Kotoba-Whisper "
               f"checkpoint [{MODEL_NAME}](https://huggingface.co/{MODEL_NAME}) and 🤗 Transformers to transcribe audio"
               f" files of arbitrary length.")
title = f"Transcribe Audio with {os.path.basename(MODEL_NAME)}"
mf_transcribe = gr.Interface(
    fn=transcribe,
    inputs=[
        gr.Audio(sources="microphone", type="filepath"),
        gr.Radio(["transcribe", "translate"], label="Task", value="transcribe"),
        gr.Radio(["ja", "en"], label="Output Language", value="ja")
    ],
    outputs=["text", "text"],
    title=title,
    description=description,
    allow_flagging="never",
)
file_transcribe = gr.Interface(
    fn=transcribe,
    inputs=[
        gr.Audio(sources="upload", type="filepath", label="Audio file"),
        gr.Radio(["transcribe", "translate"], label="Task", value="transcribe"),
        gr.Radio(["ja", "en"], label="Output Language", value="ja")
    ],
    outputs=["text", "text"],
    title=title,
    description=description,
    allow_flagging="never",
)
with demo:
    gr.TabbedInterface([mf_transcribe, file_transcribe], ["Microphone", "Audio file"])
demo.queue(api_open=False, default_concurrency_limit=40).launch(show_api=False, show_error=True)