|
|
from datetime import datetime |
|
|
import gradio as gr |
|
|
import torch |
|
|
import torchaudio |
|
|
from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq |
|
|
import spaces |
|
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
model_name = "ibm-granite/granite-speech-3.3-8b" |
|
|
|
|
|
processor = AutoProcessor.from_pretrained(model_name) |
|
|
tokenizer = processor.tokenizer |
|
|
model = AutoModelForSpeechSeq2Seq.from_pretrained( |
|
|
model_name, device_map=device, torch_dtype=torch.bfloat16 |
|
|
) |
|
|
|
|
|
|
|
|
def _load_audio_mono_16k(file_path: str) -> torch.Tensor: |
|
|
wav, sr = torchaudio.load(file_path, normalize=True) |
|
|
if wav.shape[0] > 1: |
|
|
wav = torch.mean(wav, dim=0, keepdim=True) |
|
|
if sr != 16000: |
|
|
wav = torchaudio.functional.resample(wav, sr, 16000) |
|
|
return wav |
|
|
|
|
|
@spaces.GPU |
|
|
def process_audio(audio_path: str, instruction: str) -> str: |
|
|
if not audio_path: |
|
|
return "Please upload an audio file." |
|
|
|
|
|
wav = _load_audio_mono_16k(audio_path) |
|
|
|
|
|
date_string = datetime.now().strftime("%B %d, %Y") |
|
|
|
|
|
system_prompt = ( |
|
|
"Knowledge Cutoff Date: April 2024.\n" |
|
|
f"Today's Date: {date_string}.\n" |
|
|
"You are Granite, developed by IBM. You are a helpful AI assistant" |
|
|
) |
|
|
user_prompt = f"<|audio|>{instruction.strip()}" |
|
|
chat = [ |
|
|
{"role": "system", "content": system_prompt}, |
|
|
{"role": "user", "content": user_prompt}, |
|
|
] |
|
|
prompt = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True) |
|
|
|
|
|
model_inputs = processor(prompt, wav, device=device, return_tensors="pt").to(device) |
|
|
outputs = model.generate( |
|
|
**model_inputs, |
|
|
max_new_tokens=4096, |
|
|
do_sample=False, |
|
|
num_beams=1, |
|
|
) |
|
|
|
|
|
num_input_tokens = model_inputs["input_ids"].shape[-1] |
|
|
new_tokens = torch.unsqueeze(outputs[0, num_input_tokens:], dim=0) |
|
|
text = tokenizer.batch_decode(new_tokens, add_special_tokens=False, skip_special_tokens=True)[0] |
|
|
return text |
|
|
|
|
|
|
|
|
with gr.Blocks(title="Granite Speech Demo") as demo: |
|
|
gr.Markdown("# Granite Speech-to-Text Demo") |
|
|
gr.Markdown("Upload audio and transcribe with IBM Granite.") |
|
|
|
|
|
with gr.Row(): |
|
|
with gr.Column(): |
|
|
audio_input = gr.Audio(type="filepath", label="Upload Audio") |
|
|
instruction = gr.Textbox( |
|
|
label="Instruction", |
|
|
value="can you transcribe the speech into a written format?", |
|
|
) |
|
|
submit_btn = gr.Button("Transcribe", variant="primary") |
|
|
with gr.Column(): |
|
|
output_text = gr.Textbox(label="Output", lines=12) |
|
|
|
|
|
submit_btn.click(process_audio, [audio_input, instruction], output_text) |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
demo.queue().launch(share=False, ssr_mode=False) |
|
|
|
|
|
|
|
|
|
