A2R-30B-A3B speaker-attributed reasoning demo

.gitattributes

@@ -35,3 +35,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 assets/ami_meeting_2speakers.wav filter=lfs diff=lfs merge=lfs -text
 assets/ami_meeting_4speakers.wav filter=lfs diff=lfs merge=lfs -text
+examples/ami_four_speakers.wav filter=lfs diff=lfs merge=lfs -text
+examples/ami_three_speakers.wav filter=lfs diff=lfs merge=lfs -text
+examples/ami_two_speakers.wav filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,64 +1,35 @@
 ---
-title: A2R Speaker-Attributed Reasoning
-emoji: 🔊
+title: A2R-30B-A3B Speaker-Attributed Reasoning
+emoji: 🗣️
 colorFrom: purple
 colorTo: green
 sdk: gradio
 sdk_version: 6.26.0
 app_file: app.py
-short_description: Ask WHO is speaking in multi-party audio with A2R
 python_version: "3.12"
+short_description: Ask who said what in a multi-speaker recording
 startup_duration_timeout: 1h
-models:
-  - PleasedPenguin/A2R-30B-A3B
-tags:
-  - audio
-  - speaker-attribution
-  - speech-reasoning
 ---
 
-# A2R-30B-A3B · Speaker-Attributed Reasoning
+# A2R-30B-A3B — who said what?
 
-Gradio ZeroGPU demo for
-[**PleasedPenguin/A2R-30B-A3B**](https://huggingface.co/PleasedPenguin/A2R-30B-A3B)
-— the merged 30B (3B active, MoE) speech-language model from
-**HEAR Who Said What: Unlocking Speaker-Attributed Reasoning via Counterfactual
-Voice Grounding** (EMNLP 2026).
+Demo of [`PleasedPenguin/A2R-30B-A3B`](https://huggingface.co/PleasedPenguin/A2R-30B-A3B),
+the model from **"HEAR Who Said What: Unlocking Speaker-Attributed Reasoning via
+Counterfactual Voice Grounding"** ([paper](https://huggingface.co/papers/2608.29120),
+[code](https://github.com/dwsmart32/HEAR),
+[project page](https://attributetoreason.github.io/AttributeToReason/)).
 
-A2R is `Qwen/Qwen3-Omni-30B-A3B-Instruct` fine-tuned with GRPO on CASH-60K
-counterfactual voice-grounding queries, so it answers questions about **who** is
-speaking in multi-party audio — speaker counting, turn order, voice identity,
-overlap — rather than only what is said. It reasons before answering and closes
-with a structured block:
+A2R is `Qwen/Qwen3-Omni-30B-A3B-Instruct` post-trained with GRPO on counterfactual
+audio containing speaker-level hard negatives, so it grounds answers in *vocal
+identity* rather than in the transcript.
 
-```
-<reasoning> … </reasoning><answer>{"Answer": "B"}</answer>
-```
+Prompt formatting, system prompt and sampling defaults (`temperature=0.6`,
+`top_p=0.95`) follow the authors' evaluation harness.
 
-## How it works
+## Example audio attribution
 
-- Upload or record a **multi-speaker clip**, ask a question about the speakers
-  (optionally with answer options).
-- The demo feeds the audio as a single 16 kHz waveform followed by the
-  question, matching the prompt format of the paper's
-  [evaluation harness](https://github.com/dwsmart32/HEAR) (Qwen default system
-  prompt, temperature 0.6, top-p 0.95).
-- Both the parsed answer and the reasoning trace are shown.
-
-The thinker runs in bf16 on a ZeroGPU `xlarge` slice (the model is
-text-output only — `enable_audio_output: false` — so the speech talker is not
-loaded).
-
-## Example clips
-
-The bundled examples are individual-headset-microphone segments from the
-[AMI Meeting Corpus](https://groups.inf.ed.ac.uk/ami/corpus/), © AMI Project,
-licensed under **CC BY 4.0**. AMI is one of the source corpora of the HEAR
-benchmark, so these are faithful example inputs. The gated HEAR synthetic
-audio itself is non-redistributable and is not used here.
-
-Links:
-- Paper: https://huggingface.co/papers/2608.29120
-- Model: https://huggingface.co/PleasedPenguin/A2R-30B-A3B
-- Code: https://github.com/dwsmart32/HEAR
-- Benchmark: https://huggingface.co/datasets/PleasedPenguin/HEAR
+The bundled example clips are excerpts of the
+[AMI Meeting Corpus](https://huggingface.co/datasets/diarizers-community/ami),
+licensed **CC BY 4.0** — one of the source corpora behind the HEAR benchmark.
+The HEAR benchmark audio itself is not redistributable and is therefore not
+bundled here.

app.py

@@ -1,365 +1,283 @@
-"""A2R-30B-A3B — speaker-attributed reasoning over multi-party audio (ZeroGPU).
-
-A2R (HEAR Who Said What, EMNLP 2026) is Qwen3-Omni-30B-A3B-Instruct trained with
-GRPO on CASH-60K to answer questions about WHO is speaking in multi-party audio,
-not only what is said. This demo loads the merged checkpoint's thinker
-(the model is text-output only: `enable_audio_output: false`) in bf16 on a
-ZeroGPU `xlarge` slice and answers speaker-attribution questions about an
-uploaded clip.
-
-The prompt follows the paper's evaluation harness (dwsmart32/HEAR): the audio is
-presented first as a single waveform, then the question + options, with Qwen's
-default system prompt, temperature 0.6 / top_p 0.95. A2R closes with
-`<reasoning>…</reasoning><answer>{"Answer": "X"}</answer>`; we surface both the
-reasoning trace and the parsed answer.
-"""
-
 import os
 
 os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
-os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
-
-import json
-import re
-import time
-from pathlib import Path
-
-import spaces  # must precede torch / transformers
-import gradio as gr
-import numpy as np
-import torch
-import soundfile as sf
-from transformers import AutoProcessor, Qwen3OmniMoeThinkerForConditionalGeneration
-
-MODEL_ID = "PleasedPenguin/A2R-30B-A3B"
-PAPER_URL = "https://huggingface.co/papers/2608.29120"
-MODEL_URL = f"https://huggingface.co/{MODEL_ID}"
-CODE_URL = "https://github.com/dwsmart32/HEAR"
-
-# Qwen's official default system prompt — what the paper's harness applies to a
-# non-baseline Qwen3-Omni checkpoint.
-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."
+
+import spaces  # noqa: E402  (must precede torch / CUDA-touching imports)
+
+import json  # noqa: E402
+import re  # noqa: E402
+import shutil  # noqa: E402
+import threading  # noqa: E402
+import time  # noqa: E402
+
+import gradio as gr  # noqa: E402
+import librosa  # noqa: E402
+import torch  # noqa: E402
+from transformers import (  # noqa: E402
+    AutoProcessor,
+    Qwen3OmniMoeForConditionalGeneration,
+    TextIteratorStreamer,
 )
 
-APP_DIR = Path(__file__).resolve().parent
-ASSET_DIR = APP_DIR / "assets"
-MAX_AUDIO_SECONDS = 120
-SAMPLE_RATE = 16000
-
-# --------------------------------------------------------------------------- #
-# Module-scope model load (bf16 thinker only, ~63 GB — needs a 96GB xlarge).
-# ZeroGPU packs the weights to disk at startup and streams them into VRAM on
-# the first @spaces.GPU call.
-# --------------------------------------------------------------------------- #
-print(f"Loading processor for {MODEL_ID} ...", flush=True)
-PROCESSOR = AutoProcessor.from_pretrained(MODEL_ID)
-
-print("Loading A2R thinker (30B-A3B MoE, bf16) ...", flush=True)
-MODEL = Qwen3OmniMoeThinkerForConditionalGeneration.from_pretrained(
-    MODEL_ID,
+# --------------------------------------------------------------------------------------
+# Model
+# --------------------------------------------------------------------------------------
+HUB_ID = "PleasedPenguin/A2R-30B-A3B"
+# The 70 GB checkpoint is attached to this Space as a read-only volume so it never has to
+# be copied onto the (small) ephemeral Space disk.
+MOUNT = "/models/a2r"
+MODEL_PATH = MOUNT if os.path.isfile(os.path.join(MOUNT, "config.json")) else HUB_ID
+
+_total, _used, _free = shutil.disk_usage("/")
+print(f"[boot] disk: total={_total / 2**30:.1f}GB used={_used / 2**30:.1f}GB free={_free / 2**30:.1f}GB")
+print(f"[boot] loading {MODEL_PATH} ...")
+
+_t0 = time.perf_counter()
+processor = AutoProcessor.from_pretrained(MODEL_PATH)
+model = Qwen3OmniMoeForConditionalGeneration.from_pretrained(
+    MODEL_PATH,
     dtype=torch.bfloat16,
-    low_cpu_mem_usage=True,
     attn_implementation="sdpa",
-).to("cuda")
-MODEL.eval()
-print("A2R thinker loaded and moved to cuda.", flush=True)
-
-
-# --------------------------------------------------------------------------- #
-# Audio helpers
-# --------------------------------------------------------------------------- #
-def load_audio(audio_path: str, target_sr: int = SAMPLE_RATE):
-    """Load any user-supplied audio file to mono float32 at 16 kHz."""
-    audio, sr = sf.read(audio_path, dtype="float32", always_2d=False)
-    audio = np.asarray(audio)
-    if audio.ndim == 2:  # stereo -> mono
-        if audio.shape[1] <= audio.shape[0]:
-            audio = audio.mean(axis=1)
-        else:
-            audio = audio.mean(axis=0)
-    if sr != target_sr:
-        import librosa
-
-        audio = librosa.resample(audio, orig_sr=sr, target_sr=target_sr)
-    return audio.astype(np.float32), target_sr
-
-
-def _seconds(audio: np.ndarray) -> float:
-    return len(audio) / SAMPLE_RATE
-
-
-# --------------------------------------------------------------------------- #
-# Prompt construction (mirrors the HEAR harness prompt builder)
-# --------------------------------------------------------------------------- #
-def build_instruction(question: str, options: str) -> str:
-    """Assemble the question text the way the paper's harness does.
-
-    `options` is a free-form textarea; if the user typed option lines we append
-    them as "(A) ... (B) ..." is left verbatim, otherwise the question is asked
-    open-ended. The harness always asks for a JSON answer.
-    """
-    opts = (options or "").strip()
-    body = question.strip()
-    if opts:
-        body = f"{body}\n\nOptions:\n{opts}"
-    return f"{body}\n\nRespond in JSON format, e.g. {{\"Answer\": \"A\"}}"
+).eval()
+model = model.to("cuda")
+print(f"[boot] model loaded in {time.perf_counter() - _t0:.1f}s (talker: {model.has_talker})")
 
+_total, _used, _free = shutil.disk_usage("/")
+print(f"[boot] disk after load: used={_used / 2**30:.1f}GB free={_free / 2**30:.1f}GB")
 
-def build_conversation(audio: np.ndarray, instruction: str):
-    return [
-        {"role": "system", "content": SYSTEM_PROMPT},
-        {
-            "role": "user",
-            "content": [
-                {"type": "audio", "audio": audio},
-                {"type": "text", "text": instruction},
-            ],
-        },
-    ]
+# Qwen's official Omni system prompt. The paper's evaluation harness
+# (github.com/dwsmart32/HEAR) applies exactly this to A2R.
+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."
+)
 
+MAX_AUDIO_SECONDS = 120
+TARGET_SR = 16000
 
-ANSWER_RE = re.compile(r"<answer>\s*(.*?)\s*</answer>", re.DOTALL | re.IGNORECASE)
-REASONING_RE = re.compile(
-    r"<reasoning>\s*(.*?)\s*</reasoning>", re.DOTALL | re.IGNORECASE
-)
+ANSWER_RE = re.compile(r"<answer>(.*?)</answer>", re.S)
+REASONING_RE = re.compile(r"<reasoning>(.*?)(?:</reasoning>|$)", re.S)
 
 
-def parse_output(text: str):
-    """Split the model's `<reasoning>`/`<answer>` blocks, per the model card.
+def _split_output(text: str):
+    """Split A2R's `<reasoning>...</reasoning><answer>{...}</answer>` output."""
+    reasoning_match = REASONING_RE.search(text)
+    reasoning = reasoning_match.group(1).strip() if reasoning_match else ""
 
-    Mirrors the HEAR harness's extraction: prefer the JSON `{"Answer": "X"}`
-    inside the last `<answer>` block; fall back to the raw block text, then to
-    a standalone-letter scan of the whole response.
-    """
-    reasoning_matches = REASONING_RE.findall(text)
-    answer_matches = ANSWER_RE.findall(text)
-    reasoning = (reasoning_matches[-1] if reasoning_matches else "").strip()
-    answer_raw = (answer_matches[-1] if answer_matches else "").strip()
-    answer_letter = ""
-    if answer_raw:
+    answer_blocks = ANSWER_RE.findall(text)
+    answer = answer_blocks[-1].strip() if answer_blocks else ""
+
+    if not reasoning and not answer:
+        # Model answered without the structured wrapper (common for open-ended questions).
+        return "", text.strip()
+
+    if answer:
         try:
-            obj = json.loads(answer_raw)
-            answer_letter = str(obj.get("Answer", "")).strip()
-        except json.JSONDecodeError:
-            answer_letter = answer_raw.strip()
-    if not answer_letter:
-        m = re.search(r'"Answer"\s*:\s*"?([^",\}]+)"?', text)
-        if m:
-            answer_letter = m.group(1).strip()
-    if not answer_letter:
-        # MCQ fallback: a parenthesized or bolded option letter outside tags.
-        m = re.search(r"[（(]([A-E])[)）]|\*\*([A-E])\*\*", text)
-        if m:
-            answer_letter = m.group(1) or m.group(2)
-    return reasoning, answer_letter, answer_raw
-
-
-# --------------------------------------------------------------------------- #
-# Inference
-# --------------------------------------------------------------------------- #
-def _estimate_duration(
-    audio_path: str | None,
-    question: str,
-    options: str = "",
-    max_new_tokens: int = 1024,
-    temperature: float = 0.6,
-    top_p: float = 0.95,
-    *_args,
-    **_kwargs,
-) -> int:
-    """ZeroGPU duration: weight streaming on cold start + A3B generation time."""
-    return min(420, 150 + int(max_new_tokens) // 6)
+            parsed = json.loads(answer)
+            if isinstance(parsed, dict) and "Answer" in parsed:
+                answer = f"**{parsed['Answer']}**"
+        except Exception:
+            pass
+    return reasoning, answer
+
+
+def _estimate_duration(audio, question, max_new_tokens=768, *args, **kwargs):
+    try:
+        n = int(max_new_tokens)
+    except Exception:
+        n = 768
+    return int(min(360, 60 + n * 0.3))
 
 
 @spaces.GPU(duration=_estimate_duration, size="xlarge")
-def answer_question(
-    audio_path: str,
+def analyze(
+    audio: str,
     question: str,
-    options: str = "",
-    max_new_tokens: int = 1024,
+    max_new_tokens: int = 768,
     temperature: float = 0.6,
     top_p: float = 0.95,
-    progress=gr.Progress(),
 ):
     """Answer a speaker-attribution question about a multi-speaker audio clip.
 
     Args:
-        audio_path: path to the audio clip (multi-party speech works best).
-        question: the question about who speaks / speaker identity.
-        options: optional answer options, one per line, e.g. "(A) Alice".
+        audio: path to an audio file containing one or more speakers.
+        question: the question to ask about who is speaking, optionally with
+            multiple-choice options.
         max_new_tokens: maximum number of tokens A2R may generate.
-        temperature: sampling temperature (the paper uses 0.6).
-        top_p: nucleus sampling threshold (the paper uses 0.95).
+        temperature: sampling temperature.
+        top_p: nucleus sampling probability mass.
 
     Returns:
-        (answer, reasoning, raw_output)
+        The model's answer and its speaker-attribution reasoning trace.
     """
-    t0 = time.time()
-    if not audio_path:
-        raise gr.Error("Please upload or record an audio clip first.")
-    if not (question or "").strip():
-        raise gr.Error("Please enter a question about the speakers.")
-
-    audio, sr = load_audio(audio_path)
-    dur = _seconds(audio)
-    if dur < 0.5:
-        raise gr.Error("The audio clip is too short to reason about speakers.")
-    if dur > MAX_AUDIO_SECONDS:
-        audio = audio[: MAX_AUDIO_SECONDS * SAMPLE_RATE]
-        dur = MAX_AUDIO_SECONDS
-
-    instruction = build_instruction(question, options)
-    conversation = build_conversation(audio, instruction)
-    text = PROCESSOR.apply_chat_template(
-        conversation, add_generation_prompt=True, tokenize=False
+    if not audio:
+        raise gr.Error("Please provide an audio clip.")
+    if not question or not question.strip():
+        raise gr.Error("Please ask a question about the audio.")
+
+    waveform, _ = librosa.load(audio, sr=TARGET_SR, mono=True)
+    if waveform.shape[0] > MAX_AUDIO_SECONDS * TARGET_SR:
+        waveform = waveform[: MAX_AUDIO_SECONDS * TARGET_SR]
+        gr.Info(f"Audio truncated to the first {MAX_AUDIO_SECONDS} seconds.")
+
+    messages = [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {
+            "role": "user",
+            "content": [
+                {"type": "audio", "audio": audio},
+                {"type": "text", "text": question.strip()},
+            ],
+        },
+    ]
+    text = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+    inputs = processor(
+        text=[text],
+        audio=[waveform],
+        sampling_rate=TARGET_SR,
+        return_tensors="pt",
+    ).to("cuda")
+
+    streamer = TextIteratorStreamer(
+        processor.tokenizer, skip_prompt=True, skip_special_tokens=True
     )
-    inputs = PROCESSOR(
-        text=text, audio=[audio], return_tensors="pt", padding=True
+    kwargs = dict(
+        **inputs,
+        streamer=streamer,
+        max_new_tokens=int(max_new_tokens),
+        do_sample=True,
+        temperature=float(temperature),
+        top_p=float(top_p),
     )
-    inputs = {k: (v.to("cuda") if hasattr(v, "to") else v) for k, v in inputs.items()}
-
-    with torch.inference_mode():
-        out_ids = MODEL.generate(
-            **inputs,
-            max_new_tokens=int(max_new_tokens),
-            do_sample=float(temperature) > 0,
-            temperature=float(temperature) if float(temperature) > 0 else None,
-            top_p=float(top_p),
-        )
-    gen_ids = out_ids[:, inputs["input_ids"].shape[1]:]
-    response = PROCESSOR.batch_decode(
-        gen_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
-    )[0].strip()
 
-    reasoning, answer_letter, answer_raw = parse_output(response)
-    elapsed = time.time() - t0
-    stats = f"⏱ {elapsed:.1f}s · {dur:.0f}s audio · {len(gen_ids[0])} tokens"
-    display_answer = answer_letter if answer_letter else response
-    return display_answer, reasoning, response, stats
+    thread = threading.Thread(target=model.thinker.generate, kwargs=kwargs, daemon=True)
+    t0 = time.perf_counter()
+    thread.start()
+
+    acc = ""
+    n_chunks = 0
+    yield gr.update(value="_Listening…_"), gr.update(value="", visible=True)
+    for chunk in streamer:
+        acc += chunk
+        n_chunks += 1
+        if n_chunks % 4 == 0:
+            reasoning, answer = _split_output(acc)
+            yield (
+                gr.update(value=answer or "_Reasoning…_"),
+                gr.update(value=reasoning or acc),
+            )
+    thread.join()
+
+    elapsed = time.perf_counter() - t0
+    reasoning, answer = _split_output(acc)
+    print(f"[infer] {n_chunks} chunks in {elapsed:.1f}s")
+    yield (
+        gr.update(value=answer or "_(no answer parsed)_"),
+        gr.update(value=reasoning or acc),
+    )
 
 
-# --------------------------------------------------------------------------- #
+# --------------------------------------------------------------------------------------
 # UI
-# --------------------------------------------------------------------------- #
+# --------------------------------------------------------------------------------------
+MCQ_TAIL = '\n\nRespond in JSON format, e.g. {"Answer": "A"}'
+
 EXAMPLES = [
     [
-        str(ASSET_DIR / "ami_meeting_4speakers.wav"),
-        "Four people are talking. Who speaks first, and who speaks immediately after them?",
-        "",
+        "examples/ami_two_speakers.wav",
+        "Listen to the main audio and answer the following multiple-choice question."
+        "\n\nHow many distinct speakers are there in the audio?"
+        "\n\nOptions:\n(A) 1\n(B) 2\n(C) 3\n(D) 4" + MCQ_TAIL,
     ],
     [
-        str(ASSET_DIR / "ami_meeting_4speakers.wav"),
-        "How many distinct speakers can you hear in this recording?",
-        "",
+        "examples/ami_three_speakers.wav",
+        "Walk through this recording turn by turn. How many different people speak, and "
+        "in what order do they take turns? Base your answer on the voices you hear, not "
+        "on what is being said.",
     ],
     [
-        str(ASSET_DIR / "ami_meeting_2speakers.wav"),
-        "Two speakers alternate in this clip. Who spoke last — the same person who spoke first, or the other one?",
-        "(A) The same person who spoke first\n(B) The other speaker",
+        "examples/ami_four_speakers.wav",
+        "Listen to the main audio and answer the following multiple-choice question."
+        "\n\nDoes the person who speaks first in the recording speak again later in the "
+        "recording?"
+        "\n\nOptions:\n(A) Yes, the first speaker's voice returns later\n"
+        "(B) No, the first speaker never speaks again" + MCQ_TAIL,
     ],
 ]
 
 CSS = """
-#hero {text-align: center; margin-bottom: 6px;}
-#hero h1 {font-size: 2rem; margin-bottom: 4px;}
-.dark .gradio-container {color: var(--body-text-color);}
+#col-container { max-width: 1100px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
 """
 
-with gr.Blocks(title="A2R Speaker-Attributed Reasoning") as demo:
-    with gr.Column():
-        gr.HTML(
-            """
-            <div id="hero">
-              <h1>🔊 A2R-30B-A3B · Speaker-Attributed Reasoning</h1>
-              <p style="font-size:1.05rem; opacity:.85;">
-                <i>HEAR Who Said What: Unlocking Speaker-Attributed Reasoning via
-                Counterfactual Voice Grounding</i> — ask <b>who</b> is speaking
-                in multi-party audio, not only what is said.
-              </p>
-            </div>
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="A2R — speaker-attributed reasoning") as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
             """
-        )
+# A2R-30B-A3B — who said what?
 
-    with gr.Row(equal_height=False):
-        with gr.Column(scale=5):
-            audio_in = gr.Audio(
-                label="Multi-speaker audio clip",
-                sources=["upload", "microphone"],
-                type="filepath",
-                format="wav",
-            )
-            question_in = gr.Textbox(
-                label="Question about the speakers",
-                placeholder="e.g. Who speaks second? / Who spoke the longest? / How many speakers do you hear?",
-                lines=2,
-                value="Who speaks first, and who speaks immediately after them?",
-            )
-            options_in = gr.Textbox(
-                label="Answer options (optional, one per line)",
-                placeholder="(A) The first speaker\n(B) The second speaker\n(C) Someone else",
-                lines=3,
-            )
-            run_btn = gr.Button("Reason about the speakers", variant="primary")
+Ask questions about **who** is speaking in a multi-party recording, not only about what is
+said. [A2R](https://huggingface.co/PleasedPenguin/A2R-30B-A3B) is `Qwen3-Omni-30B-A3B-Instruct`
+trained with GRPO on counterfactual audio with speaker-level hard negatives, so it grounds its
+answer in vocal cues instead of the transcript.
 
-        with gr.Column(scale=4):
-            answer_out = gr.Textbox(label="Answer", buttons=["copy"])
-            reasoning_out = gr.Markdown(label="Reasoning trace")
-            stats_out = gr.Markdown()
-            with gr.Accordion("Raw model output", open=False):
-                raw_box = gr.Textbox(label="Raw output", lines=6, buttons=["copy"])
+[Paper](https://huggingface.co/papers/2608.29120) ·
+[Project page](https://attributetoreason.github.io/AttributeToReason/) ·
+[Code](https://github.com/dwsmart32/HEAR) ·
+[Model](https://huggingface.co/PleasedPenguin/A2R-30B-A3B)
+"""
+        )
 
-    with gr.Accordion("Advanced settings", open=False):
         with gr.Row():
-            max_new_tokens_in = gr.Slider(
-                128, 2048, value=1024, step=128,
-                label="Max new tokens",
-                info="A2R reasons before answering; 1024 covers most questions.",
-            )
-            temperature_in = gr.Slider(
-                0.0, 1.5, value=0.6, step=0.05,
-                label="Temperature",
-                info="The paper's evaluation uses 0.6.",
+            with gr.Column(scale=1):
+                audio_in = gr.Audio(
+                    label="Multi-speaker audio", type="filepath", sources=["upload", "microphone"]
+                )
+                question_in = gr.Textbox(
+                    label="Question",
+                    lines=7,
+                    placeholder=(
+                        "e.g. How many different people speak in this recording, and in "
+                        "what order do they take turns?"
+                    ),
+                )
+                run = gr.Button("Analyse", variant="primary")
+            with gr.Column(scale=1):
+                gr.Markdown("### Answer")
+                answer_out = gr.Markdown(value="", min_height=90)
+                reasoning_out = gr.Textbox(
+                    label="Reasoning trace", lines=16, show_copy_button=True
+                )
+
+        with gr.Accordion("Advanced settings", open=False):
+            max_new_tokens = gr.Slider(
+                128, 2048, value=768, step=64, label="Max new tokens"
             )
-            top_p_in = gr.Slider(
-                0.1, 1.0, value=0.95, step=0.05, label="Top-p"
-            )
-
-    gr.Examples(
-        examples=EXAMPLES,
-        inputs=[audio_in, question_in, options_in],
-        fn=answer_question,
-        outputs=[answer_out, reasoning_out, raw_box, stats_out],
-        cache_examples=True,
-        cache_mode="lazy",
-        label="Try an example (real meeting audio from the AMI corpus, CC BY 4.0)",
-    )
+            temperature = gr.Slider(0.0, 1.5, value=0.6, step=0.05, label="Temperature")
+            top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.01, label="Top-p")
+
+        gr.Examples(
+            examples=EXAMPLES,
+            inputs=[audio_in, question_in],
+            outputs=[answer_out, reasoning_out],
+            fn=analyze,
+            cache_examples=True,
+            cache_mode="lazy",
+            label="Examples (AMI Meeting Corpus, CC BY 4.0)",
+        )
 
-    gr.Markdown(
-        f"""
-        **About.** A2R is `Qwen/Qwen3-Omni-30B-A3B-Instruct` fine-tuned with GRPO on
-        CASH-60K counterfactual voice-grounding queries, so it answers *who* said
-        what rather than just transcribing. It emits a reasoning trace followed by
-        `<answer>{{"Answer": "…"}}</answer>` — both are shown above.
-        Example clips are individual-headset-mic segments from the
-        [AMI Meeting Corpus](https://groups.inf.ed.ac.uk/ami/corpus/) (CC BY 4.0),
-        the same source corpus family the HEAR benchmark is built on.
-
-        📄 [Paper](https://huggingface.co/papers/2608.29120) ·
-        🤗 [Model](https://huggingface.co/PleasedPenguin/A2R-30B-A3B) ·
-        💻 [Code](https://github.com/dwsmart32/HEAR) ·
-        📊 [HEAR benchmark](https://huggingface.co/datasets/PleasedPenguin/HEAR)
-        """
-    )
+        gr.Markdown(
+            "Example clips are excerpts of the "
+            "[AMI Meeting Corpus](https://huggingface.co/datasets/diarizers-community/ami) "
+            "(CC BY 4.0), one of the source corpora behind the HEAR benchmark. The HEAR "
+            "benchmark audio itself is not redistributable and is therefore not bundled here."
+        )
 
-    run_btn.click(
-        fn=answer_question,
-        inputs=[audio_in, question_in, options_in,
-                max_new_tokens_in, temperature_in, top_p_in],
-        outputs=[answer_out, reasoning_out, raw_box, stats_out],
+    run.click(
+        analyze,
+        inputs=[audio_in, question_in, max_new_tokens, temperature, top_p],
+        outputs=[answer_out, reasoning_out],
+        api_name="analyze",
     )
 
-if __name__ == "__main__":
-    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)
+demo.queue(max_size=12).launch(mcp_server=True)

examples/ami_four_speakers.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f347b457d95437b44ff7a62a05666e076a4e36e36326e487a7acac8f843d6f30
+size 984364

examples/ami_three_speakers.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:573d95505309e40c339080c78a3b0b02c44cdc1213e0cb1eaf20451caf295ffd
+size 1093482

examples/ami_two_speakers.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d42c49e5f34393905bd3acc3e4870d66b3b2c777df950a0f1989fd253934dd6f
+size 1090282

requirements.txt

@@ -1,8 +1,6 @@
 transformers==5.16.1
-accelerate>=1.8.0
-safetensors>=0.4.5
-soundfile>=0.13.0
-librosa>=0.11.0
-numpy>=1.26
-hf-transfer>=0.1.4
-torchvision
+accelerate
+torchvision
+librosa
+soundfile
+numpy