|
|
import os |
|
|
from math import floor |
|
|
from typing import Optional |
|
|
|
|
|
import spaces |
|
|
import torch |
|
|
import gradio as gr |
|
|
import numpy as np |
|
|
from transformers import pipeline |
|
|
from transformers.pipelines.audio_utils import ffmpeg_read |
|
|
|
|
|
|
|
|
model_name = "kotoba-tech/kotoba-whisper-v2.2" |
|
|
example_file = "sample_diarization_japanese.mp3" |
|
|
|
|
|
|
|
|
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 = {} |
|
|
|
|
|
|
|
|
pipe = pipeline( |
|
|
model=model_name, |
|
|
chunk_length_s=15, |
|
|
batch_size=16, |
|
|
torch_dtype=torch_dtype, |
|
|
device=device, |
|
|
model_kwargs=model_kwargs, |
|
|
trust_remote_code=True |
|
|
) |
|
|
sampling_rate = pipe.feature_extractor.sampling_rate |
|
|
|
|
|
|
|
|
def format_time(start: Optional[float], end: Optional[float]): |
|
|
|
|
|
def _format_time(seconds: Optional[float]): |
|
|
if seconds is None: |
|
|
return "[no timestamp available]" |
|
|
minutes = floor(seconds / 60) |
|
|
hours = floor(seconds / 3600) |
|
|
seconds = seconds - hours * 3600 - minutes * 60 |
|
|
m_seconds = floor(round(seconds - floor(seconds), 1) * 10) |
|
|
seconds = floor(seconds) |
|
|
return f'{minutes:02}:{seconds:02}.{m_seconds:01}' |
|
|
|
|
|
return f"[{_format_time(start)} -> {_format_time(end)}]:" |
|
|
|
|
|
|
|
|
@spaces.GPU |
|
|
def get_prediction(inputs): |
|
|
return pipe(inputs, generate_kwargs={"language": "ja", "task": "transcribe"}) |
|
|
|
|
|
|
|
|
def transcribe(inputs: 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, sampling_rate) |
|
|
array_pad = np.zeros(int(pipe.feature_extractor.sampling_rate * 0.5)) |
|
|
inputs = np.concatenate([array_pad, inputs, array_pad]) |
|
|
prediction = get_prediction({"array": inputs, "sampling_rate": sampling_rate}) |
|
|
output = "" |
|
|
for n, s in enumerate(prediction["speakers"]): |
|
|
text_timestamped = "\n".join([f"- **{format_time(*c['timestamp'])}** {c['text']}" for c in prediction[f"chunks/{s}"]]) |
|
|
output += f'### Speaker {n+1} \n{text_timestamped}\n' |
|
|
return output |
|
|
|
|
|
|
|
|
description = (f"Transcribe and diarize long-form microphone or audio inputs with the click of a button! Demo uses " |
|
|
f"Kotoba-Whisper [{model_name}](https://huggingface.co/{model_name}).") |
|
|
title = f"Audio Transcription and Diarization with {os.path.basename(model_name)}" |
|
|
shared_config = {"fn": transcribe, "title": title, "description": description, "allow_flagging": "never", "examples": [example_file]} |
|
|
o_upload = gr.Markdown() |
|
|
o_mic = gr.Markdown() |
|
|
i_upload = gr.Interface( |
|
|
inputs=[gr.Audio(sources="upload", type="filepath", label="Audio file")], outputs=gr.Markdown(), **shared_config |
|
|
) |
|
|
i_mic = gr.Interface( |
|
|
inputs=[gr.Audio(sources="microphone", type="filepath", label="Microphone input")], outputs=gr.Markdown(), **shared_config |
|
|
) |
|
|
with gr.Blocks() as demo: |
|
|
gr.TabbedInterface([i_upload, i_mic], ["Audio file", "Microphone"]) |
|
|
demo.queue(api_open=False, default_concurrency_limit=40).launch(show_api=False, show_error=True) |
|
|
|
