w2v-bert2-czech / app.py
mikr's picture
fix
7df6e8c
raw
history blame
1.93 kB
import gradio as gr
import soundfile as sf
import torch
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, pipeline
MODEL_NAME = "mikr/w2v-bert-2.0-czech-colab-cv16"
lang = "cs"
device = 0 if torch.cuda.is_available() else "cpu"
model = Wav2Vec2ForCTC.from_pretrained(MODEL_NAME).to(device)
processor = Wav2Vec2Processor.from_pretrained(MODEL_NAME)
pipe = pipeline(
model=MODEL_NAME,
)
def transcribe(file_upload):
warn_output = ""
if (file_upload is None):
return "ERROR: You have to either use the microphone or upload an audio file"
file = file_upload
text = pipe(file)["text"]
return warn_output + text
def readwav(a_f):
wav, sr = sf.read(a_f, dtype=np.float32)
if len(wav.shape) == 2:
wav = wav.mean(1)
if sr != 16000:
wlen = int(wav.shape[0] / sr * 16000)
wav = signal.resample(wav, wlen)
return wav
def transcribe2(file_upload):
wav = readwav(file_upload)
with torch.inference_mode():
input_values = processor(wav, sampling_rate=16000).input_values[0]
input_values = torch.tensor(input_values, device=device).unsqueeze(0)
logits = model(input_values).logits
pred_ids = torch.argmax(logits, dim=-1)
xcp = processor.batch_decode(pred_ids)
return xcp[0]
iface = gr.Interface(
fn=transcribe2,
inputs=[
gr.File(type="binary", label="Upload Audio File"), # Audio file upload
],
outputs="text",
theme="huggingface",
title="Wav2Vec2-Bert demo - transcribe Czech Audio",
description=(
"Transcribe long-form microphone or audio inputs with the click of a button! Demo uses the fine-tuned"
f" checkpoint [{MODEL_NAME}](https://huggingface.co/{MODEL_NAME}) from Whisper Fine Tuning Sprint Event 2022 "
"and 🤗 Transformers to transcribe audio files of arbitrary length."
),
allow_flagging="never",
)
iface.launch()