Spaces:
Running on Zero
Running on Zero
| """AudioJev on ZeroGPU: audio in, a calibrated probability per option out. | |
| AudioJev (shlv/AudioJev) is a Qwen2.5-Omni-3B fine-tune that maps audio plus a | |
| multiple-choice question directly to an answer distribution in ONE forward pass — | |
| no text is generated. The last position's logits are read over the candidate | |
| label tokens and softmaxed, and confidence is reported as distribution | |
| concentration (1 - normalized entropy). | |
| The inference path ports the official reference implementation | |
| (github.com/SihanLv/AudioJev-Inference) 1:1: soundfile read -> mono -> 16 kHz | |
| resample_poly -> one extra STFT window of padding -> chat template -> processor | |
| call with per-call audio_kwargs -> feature trim to actual frames -> lm_head | |
| pre-hook that keeps only the final position -> label logits -> softmax. | |
| The only deviations are the ZeroGPU-required ones: the model is placed with | |
| `.to("cuda")` at module scope (no indexed device_map), and `import spaces` | |
| precedes every torch import. | |
| """ | |
| import math | |
| import re | |
| import time | |
| import spaces # ZeroGPU: must precede torch | |
| import torch | |
| from huggingface_hub import snapshot_download | |
| import numpy as np | |
| import soundfile as sf | |
| from scipy.signal import resample_poly | |
| import gradio as gr | |
| from transformers import AutoConfig, Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor | |
| REPO = "shlv/AudioJev" | |
| DEVICE = "cuda" | |
| MAX_PROMPT_TOKENS = 4096 | |
| MAX_AUDIO_BYTES = 20 * 1024 * 1024 | |
| LABELS = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ" | |
| SYSTEM_PROMPT = ( | |
| "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, " | |
| "capable of perceiving auditory and visual inputs, as well as generating text and speech." | |
| ) | |
| # Options whose wording depends on list order ("none of the above", "option C", ...) | |
| # must not be alphabetized — the reference keeps input order for those. | |
| _RELATIVE = re.compile( | |
| r"\b(?:all|none|both|neither|any)\s+of\s+the\s+(?:above|below)\b|" | |
| r"\b(?:option|choice)\s+[A-Z0-9]\b|\b(?:former|latter)\b", re.I | |
| ) | |
| MODES = ["choice", "noul", "score"] | |
| def _prepare_options(mode, instructions, options_text, score_levels_text): | |
| """Port of the reference `_prepare`: validate, then sort by UTF-8 bytes. | |
| The reference sorts choice/noul option texts alphabetically (UTF-8 byte | |
| order) when all texts are distinct and none is order-relative, because the | |
| model was trained on that canonical ordering. Score levels keep list order. | |
| """ | |
| instructions = (instructions or "").strip() | |
| if not instructions: | |
| raise gr.Error("Give a question or instruction for the decision.") | |
| if mode == "choice": | |
| options = [(p.split("=", 1)[0].strip() or p.split("=", 1)[0].strip(), t.strip()) | |
| for p, t in | |
| (("", line) if "=" not in line else (line.split("=", 1)[0], line.split("=", 1)[1]) | |
| for line in options_text.splitlines())] | |
| options = [(k, t.strip()) for k, t in options if t] | |
| if len(options) < 2: | |
| raise gr.Error("Choice mode needs at least 2 options, one per line " | |
| "(optionally 'key=description').") | |
| if len(options) > 36: | |
| raise gr.Error("Choice mode supports at most 36 options.") | |
| keys = [] | |
| for i, (k, _) in enumerate(options): | |
| keys.append(k if k else LABELS[i]) | |
| if len(set(keys)) != len(keys): | |
| raise gr.Error("Duplicate option keys — leave keys empty or make them unique.") | |
| texts = [t for _, t in options] | |
| if len(set(texts)) == len(texts) and not any(_RELATIVE.search(t) for t in texts): | |
| order = sorted(range(len(options)), key=lambda i: texts[i].encode("utf-8")) | |
| options = [options[i] for i in order] | |
| return "choice", instructions, options | |
| if mode == "noul": | |
| options = [("false", "No; the statement is false"), ("true", "Yes; the statement is true")] | |
| texts = [t for _, t in options] | |
| order = sorted(range(len(options)), key=lambda i: texts[i].encode("utf-8")) | |
| options = [options[i] for i in order] | |
| return "noul", instructions, options | |
| levels = [l.strip() for l in score_levels_text.splitlines() if l.strip()] | |
| if not 2 <= len(levels) <= 10: | |
| raise gr.Error("Score mode needs 2-10 ordered levels, one per line (worst to best).") | |
| options = [(str(i), l) for i, l in enumerate(levels)] | |
| return "score", instructions, options | |
| def _render_prompt(instructions, options): | |
| """Verbatim port of the reference `render_prompt`.""" | |
| lines = ["You are a structured decision module.", | |
| "Use the audio and the option descriptions as evidence.", | |
| "Output exactly one listed answer label token and nothing else.", | |
| "Do not output an explanation, punctuation, another word, or a label not listed below.", | |
| "There is no implicit other/none option; choose among the listed options.", | |
| instructions, "", "Options:"] | |
| lines.extend(f"{LABELS[i]}. {text}" for i, (_, text) in enumerate(options)) | |
| lines.extend(["", f"Reply with exactly one option label ({', '.join(LABELS[:len(options)])}).", "Answer:"]) | |
| return "\n".join(lines) | |
| def _concentration(probabilities): | |
| """1 - normalized entropy: distribution concentration, NOT correctness probability.""" | |
| entropy = -sum(p * math.log(p) for p in probabilities if p > 0) | |
| return max(0.0, min(1.0, 1.0 - entropy / math.log(len(probabilities)))) | |
| def _waveform(audio): | |
| """Reference `_waveform`: mono-mix, resample to 16 kHz, validate.""" | |
| if audio is None: | |
| raise gr.Error("Upload or record an audio clip first.") | |
| try: | |
| samples, rate = sf.read(audio, dtype="float32", always_2d=True) | |
| except Exception as exc: | |
| raise gr.Error("That file is not a decodable audio file.") from exc | |
| samples = samples.mean(axis=1) | |
| if not isinstance(rate, int) or rate <= 0 or not np.isfinite(samples).all(): | |
| raise gr.Error("Invalid audio: check the sample rate and content.") | |
| if rate != 16000: | |
| divisor = math.gcd(rate, 16000) | |
| samples = resample_poly(samples, 16000 // divisor, rate // divisor) | |
| samples = np.ascontiguousarray(samples, dtype=np.float32) | |
| if len(samples) < 400: | |
| raise gr.Error("Audio is too short — give it at least 25 ms.") | |
| return samples | |
| def _padded_length(samples, extractor): | |
| """Reference `_padded_length`: one STFT window past the audio, capped at 300 s.""" | |
| hop = extractor.hop_length | |
| return min(-(-(samples + extractor.n_fft) // hop) * hop, extractor.n_samples) | |
| # ---- load once, at startup, onto (emulated) cuda ---- | |
| t0 = time.time() | |
| local = snapshot_download(REPO, allow_patterns=[ | |
| "config.json", "generation_config.json", "model*.safetensors", | |
| "model.safetensors.index.json", "tokenizer*.json", "tokenizer.model", | |
| "vocab.json", "merges.txt", "added_tokens.json", "special_tokens_map.json", | |
| "*preprocessor_config.json", "processor_config.json", "chat_template*", "spk_dict.pt", | |
| ]) | |
| PROCESSOR = Qwen2_5OmniProcessor.from_pretrained(local, use_fast=False) | |
| _token_ids = [PROCESSOR.tokenizer.encode(label, add_special_tokens=False) for label in LABELS] | |
| if any(len(ids) != 1 for ids in _token_ids) or len({ids[0] for ids in _token_ids}) != 36: | |
| raise RuntimeError("model tokenizer does not have 36 distinct single-token labels") | |
| LABEL_IDS = [ids[0] for ids in _token_ids] | |
| config = AutoConfig.from_pretrained(local) | |
| config.enable_audio_output = False # thinker only: talker/token2wav never get built | |
| MODEL = Qwen2_5OmniForConditionalGeneration.from_pretrained( | |
| local, config=config, use_safetensors=True, dtype=torch.bfloat16, | |
| attn_implementation="sdpa").eval().to(DEVICE) | |
| THINKER = MODEL.thinker | |
| THINKER.requires_grad_(False) | |
| print(f"audiojev loaded in {time.time() - t0:.0f}s", flush=True) | |
| def _duration(mode, question, options_text, score_levels_text, audio): | |
| """ZeroGPU budget sized to measured cost: ~5.8 s worst cold call, | |
| <2.5 s warm; 10 s leaves margin for longer clips without waste.""" | |
| return 10 | |
| def decide(mode, question, options_text, score_levels_text, audio): | |
| """Run one AudioJev decision over the audio clip. | |
| Args: | |
| mode: choice (pick one of N), noul (yes/no statement), or score (ordered levels). | |
| question: the question or instruction the model decides on. | |
| options_text: choice options, one per line, optionally 'key=description'. | |
| score_levels_text: score levels, one per line, worst to best. | |
| audio: uploaded/recorded audio clip. | |
| Returns: | |
| Probabilities per option (gr.Label) and a summary line (Markdown). | |
| """ | |
| import os | |
| if audio and os.path.exists(audio) and os.path.getsize(audio) > MAX_AUDIO_BYTES: | |
| raise gr.Error("Audio must be under 20 MB.") | |
| kind, instructions, options = _prepare_options(mode, question, options_text, score_levels_text) | |
| waveform = _waveform(audio) | |
| padded = _padded_length(len(waveform), PROCESSOR.feature_extractor) | |
| prompt = _render_prompt(instructions, options) | |
| conversation = [ | |
| {"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]}, | |
| {"role": "user", "content": [{"type": "audio", "audio": "audio.wav"}, | |
| {"type": "text", "text": prompt}]}, | |
| ] | |
| text = PROCESSOR.apply_chat_template(conversation, add_generation_prompt=True, tokenize=False) | |
| # With audio_kwargs, top-level audio options are ignored and the processor | |
| # empties the dict — build it per call, exactly like the reference backend. | |
| batch = PROCESSOR(text=[text], audio=[waveform], return_tensors="pt", padding=True, | |
| use_audio_in_video=False, | |
| audio_kwargs={"sampling_rate": 16000, "max_length": padded}) | |
| if batch["input_ids"].shape[1] > MAX_PROMPT_TOKENS: | |
| raise gr.Error("Prompt is too long; shorten the question or options.") | |
| frames = int(batch["feature_attention_mask"].sum(-1).max()) | |
| batch["input_features"] = batch["input_features"][:, :, :frames] | |
| batch["feature_attention_mask"] = batch["feature_attention_mask"][:, :frames] | |
| batch = {name: (value.to(DEVICE) if torch.is_tensor(value) else value) | |
| for name, value in batch.items()} | |
| # The tested research readout projects only the final decision position; | |
| # this avoids materializing the full sequence-vocabulary logits matrix. | |
| hook = THINKER.lm_head.register_forward_pre_hook( | |
| lambda _module, args: (args[0][:, -1:, :],)) | |
| try: | |
| logits = THINKER(**batch, use_cache=False, return_dict=True).logits[0, -1].float() | |
| finally: | |
| hook.remove() | |
| selected = logits[LABEL_IDS[:len(options)]] | |
| values = torch.softmax(selected, dim=-1) | |
| if not bool(torch.isfinite(values).all()): | |
| raise gr.Error("Model produced nonfinite probabilities — try again.") | |
| probs = values.cpu().tolist() | |
| values_map = {key: float(v) for (key, _), v in zip(options, probs)} | |
| if kind == "noul": | |
| labels = {f"Yes": values_map["true"], "No": values_map["false"]} | |
| top = max(labels, key=labels.get) | |
| return labels, f"**{top}** · p = {labels[top]*100:.1f}%" | |
| confidence = _concentration(probs) | |
| if kind == "choice": | |
| winner = max(range(len(options)), key=lambda i: probs[i]) | |
| label_probs = {options[i][1]: probs[i] for i in range(len(options))} | |
| return (label_probs, | |
| f"**{options[winner][1]}** · p = {probs[winner]*100:.1f}% · confidence {confidence*100:.0f}%") | |
| # score | |
| score = sum(i * p for i, p in enumerate(probs)) | |
| legend = {f"{i} ({text})": probs[i] for i, (_, text) in enumerate(options)} | |
| return (legend, | |
| f"**score {score:.2f} / {len(options)-1}** · confidence {confidence*100:.0f}%") | |
| EXAMPLE_QUERIES = [ | |
| ("choice", "What is the dominant sound event in the audio?", "examples/carhorn_16k.wav", | |
| "A dog barking\nA person speaking\nA vehicle horn honking\nRain falling"), | |
| ("choice", "Which sound is present in the audio?", "examples/rain_16k.wav", | |
| "Rain falling\nWind blowing\nA person speaking\nTraffic noise"), | |
| ("choice", "What animal makes this sound?", "examples/bark_16k.wav", | |
| "A dog barking\nA cat meowing\nA bird singing\nA cow mooing"), | |
| ] | |
| CSS = """ | |
| #col-container { max-width: 1100px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="AudioJev") as demo: | |
| gr.Markdown( | |
| "# AudioJev — direct audio decisions with order-calibrated probabilities\n" | |
| "Upload or record a clip, ask a question, and list the answer options. " | |
| "AudioJev returns a probability for **every** option in a single forward pass — " | |
| "it never generates text, it only scores the listed candidates." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| audio_in = gr.Audio(label="Audio clip", sources=["upload", "microphone"], type="filepath") | |
| mode = gr.Radio(MODES, value="choice", label="Decision type", | |
| info="choice: pick one of N options · noul: is the statement true? · score: rate on ordered levels") | |
| question = gr.Textbox(label="Question / instruction", lines=2, | |
| value="What is the dominant sound event in the audio?") | |
| options_text = gr.Textbox(label="Options, one per line (choice)", | |
| lines=4, value="A dog barking\nA person speaking\nA vehicle horn honking\nRain falling") | |
| score_levels_text = gr.Textbox(label="Score levels, worst to best (score mode)", | |
| lines=4, value="poor\nfair\ngood\nexcellent", visible=False) | |
| go = gr.Button("Decide", variant="primary") | |
| with gr.Column(): | |
| verdict = gr.Markdown() | |
| probs = gr.Label(label="Option probabilities", num_top_classes=10) | |
| def _toggle_mode(m): | |
| return (gr.update(visible=(m == "choice")), gr.update(visible=(m == "score"))) | |
| mode.change(_toggle_mode, [mode], [options_text, score_levels_text]) | |
| go.click(decide, [mode, question, options_text, score_levels_text, audio_in], [probs, verdict]) | |
| gr.Examples( | |
| examples=[[m, q, o, "poor\nfair\ngood\nexcellent", f] for (m, q, f, o) in EXAMPLE_QUERIES], | |
| inputs=[mode, question, options_text, score_levels_text, audio_in], | |
| outputs=[probs, verdict], | |
| fn=decide, cache_examples=True, cache_mode="lazy", | |
| ) | |
| gr.Markdown( | |
| "### Notes\n" | |
| "* **choice** scores 2–36 options; the model was trained on alphabetically " | |
| "sorted option lists, so distinct non-relative options are re-sorted the same way " | |
| "before scoring (use 'key=description' to control the displayed key).\n" | |
| "* **noul** answers a true/false statement about the audio.\n" | |
| "* **score** spreads probability over ordered levels and reports their expectation.\n" | |
| "* Confidence is distribution concentration (1 − normalized entropy), not the " | |
| "probability of being correct.\n" | |
| "* Audio is capped at 20 MB and 300 s; best under 30 s. Weights: " | |
| "[shlv/AudioJev](https://huggingface.co/shlv/AudioJev) (Qwen Research License). " | |
| "Example clips are CC0 / public-domain audio from Wikimedia Commons." | |
| ) | |
| demo.queue().launch(mcp_server=True) |