cohere-asr-ja-diffusion

Experimental single step block diffusion model for speculative decoding similar to Orthrus/DFlash. Like Orthrus it conditions on the AR KV cache but like DFlash only the embedding layer is shared. Training and inference code is included based on pure pytorch/flex attention with torch compile.

Notes

• Diffusion for fun, not very practical (increase batching instead)
• torch.compile involves more warmup time
• Pad inputs to static shapes for compile efficiency
• diffusion=False ~10% faster than nano if encoder isn't compiled
• Benchmark sequences short, average of jsut/reazon is below block size

from huggingface_hub import snapshot_download
model_dir = 'diff'
snapshot_download('efwkjn/cohere-asr-ja-diffusion', cache_dir=model_dir, local_dir=model_dir)

import importlib
import subprocess

import numpy as np
import torch
from tokenizers import Tokenizer

cohere_asr = importlib.import_module(model_dir + '.cohere_asr')
file = 'audio.wav'
device = 'cuda'

cmd = [
    'ffmpeg',
    '-threads', '1',
    '-nostdin',
    '-i', file,
    '-f', 's16le',
    '-c:a', 'pcm_s16le',
    '-ar', '16000',
    '-ac', '1',
    '-'
]
with subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL) as p:
    audio = np.frombuffer(p.stdout.read(), dtype=np.int16).astype(np.float32) / 0x8000

fe = cohere_asr.CohereAsrFeatureExtractor(model_dir)
model = cohere_asr.CohereAsr.from_pretrained(model_dir, device)
tokenizer: Tokenizer = Tokenizer.from_file(model_dir + '/tokenizer.json')
prompt = tokenizer.encode(
    '<|startofcontext|><|startoftranscript|><|emo:undefined|>'
    '<|ja|><|ja|><|pnc|><|noitn|><|notimestamp|><|nodiarize|>',
    add_special_tokens=False
).ids

features, lengths = fe([audio[:480000]])
input_ids = torch.tensor(prompt, device=device)[None, :]
features = features.to(device, torch.bfloat16)
lengths = lengths.to(device)
ids = model.generate(input_ids, features, lengths, compile=False, diffusion=True)
print(tokenizer.decode_batch(ids.tolist())[0])