Spaces:
Running
Running
File size: 3,473 Bytes
6ff26b5 ff5c88b 6ff26b5 | 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 | """Nawah-ASR-50M demo.
The browser does all audio handling: decodeAudioData accepts wav/mp3/m4a/webm, resamples to
16 kHz mono and posts raw float32 samples. So this server needs no ffmpeg, no soundfile, and no
format guessing -- it only ever sees the array the model wants.
The audio token count is per-clip, not fixed: `<audio>` is repeated exactly
`_get_feat_extract_output_lengths(...)` times, which is 25 per second. A 3-second clip costs 75
tokens rather than the 375 a fixed 30-second window would spend on silence.
"""
import os, time
import numpy as np
import torch
from fastapi import FastAPI
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from transformers import (AutoTokenizer, Qwen2AudioForConditionalGeneration,
WhisperFeatureExtractor)
REPO = os.environ.get("ASR_MODEL", "oddadmix/Nawah-ASR-50M-v2")
TOKEN = os.environ.get("MODEL_HF_TOKEN") or os.environ.get("HF_TOKEN")
SR, MAX_SEC = 16000, 30
print(f"[*] loading {REPO}", flush=True)
tok = AutoTokenizer.from_pretrained(REPO, token=TOKEN)
fe = WhisperFeatureExtractor.from_pretrained(REPO, token=TOKEN)
model = Qwen2AudioForConditionalGeneration.from_pretrained(REPO, token=TOKEN,
dtype=torch.float32).eval()
torch.set_num_threads(int(os.environ.get("OMP_NUM_THREADS", 4)))
CV = tok.convert_tokens_to_ids
AUDIO, A_START, A_END = CV("<audio>"), CV("<|audio_start|>"), CV("<|audio_end|>")
PARAMS = sum(p.numel() for p in model.parameters())
print(f"[+] {PARAMS/1e6:.1f}M params ready", flush=True)
app = FastAPI()
class Req(BaseModel):
samples: list[float] = []
max_new: int = 96
@torch.no_grad()
def run(wav: np.ndarray, max_new: int):
wav = wav[: SR * MAX_SEC]
feats = fe([wav], sampling_rate=SR, return_attention_mask=True, return_tensors="pt")
_, n = model.model.audio_tower._get_feat_extract_output_lengths(feats.attention_mask.sum(-1))
n = int(n[0])
ids = torch.tensor([[tok.bos_token_id, A_START] + [AUDIO] * n + [A_END]])
t0 = time.perf_counter()
out = model.generate(input_ids=ids, attention_mask=torch.ones_like(ids),
input_features=feats.input_features,
feature_attention_mask=feats.attention_mask,
max_new_tokens=max_new, do_sample=False, num_beams=1,
pad_token_id=tok.pad_token_id, eos_token_id=tok.eos_token_id)
ms = (time.perf_counter() - t0) * 1000
text = tok.decode(out[0, ids.shape[1]:], skip_special_tokens=True).strip()
dur = len(wav) / SR
return {"text": text, "audio_tokens": n, "duration": round(dur, 2),
"ms": round(ms), "rtf": round(dur * 1000 / max(ms, 1), 1),
"new_tokens": int(out.shape[1] - ids.shape[1])}
@app.get("/")
def index():
return FileResponse("static/index.html")
@app.get("/api/ready")
def ready():
return {"ready": True, "model": REPO, "params": PARAMS, "max_sec": MAX_SEC}
@app.post("/api/asr")
def asr(req: Req):
if not req.samples:
return JSONResponse({"error": "no audio"}, status_code=400)
wav = np.asarray(req.samples, dtype=np.float32)
if wav.size < SR // 4:
return JSONResponse({"error": "clip too short"}, status_code=400)
return JSONResponse(run(wav, max(8, min(req.max_new, 200))))
app.mount("/static", StaticFiles(directory="static"), name="static")
|