baari 0.1.0__py3-none-any.whl

baari/__init__.py

@@ -0,0 +1,5 @@
+"""End-of-turn detection for Urdu voice agents."""
+
+from baari.detector import TurnDetector, model_path
+
+__all__ = ["TurnDetector", "model_path"]

baari/cli.py

@@ -0,0 +1,242 @@
+import json
+import random
+import zlib
+from collections import defaultdict
+from pathlib import Path
+from typing import Annotated
+
+try:
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+    import typer
+
+    from baari import clips, indicvoices, labels
+    from baari.settings import Settings
+except ImportError as e:
+    raise SystemExit(
+        f'The baari command needs the pipeline extras: pip install "baari[pipeline]" ({e})'
+    ) from e
+
+app = typer.Typer(no_args_is_help=True)
+
+
+def load_index(s: Settings, split: str) -> list[dict]:
+    rows = pq.read_table(s.index_dir / f"{split}.parquet").to_pylist()
+    if split == "valid":
+        return rows
+    return indicvoices.without_test(rows, pq.read_table(s.index_dir / "valid.parquet").to_pylist())
+
+
+@app.command()
+def scan(split: str = "valid") -> None:
+    """Index the conversation segments of one IndicVoices Urdu split (text only, no audio)."""
+    s = Settings()
+    rows = [
+        r
+        for shard in indicvoices.shards(split, s.hf_token)
+        for r in indicvoices.index_shard(shard, s.hf_token)
+    ]
+    s.index_dir.mkdir(parents=True, exist_ok=True)
+    pq.write_table(pa.Table.from_pylist(rows), s.index_dir / f"{split}.parquet")
+    hours = sum(r["duration"] for r in rows) / 3600
+    typer.echo(
+        f"{split}: {len(rows)} segments, {len({r['session'] for r in rows})} sessions, "
+        f"{len({r['speaker_id'] for r in rows})} speakers, {hours:.1f} h"
+    )
+
+
+@app.command()
+def batches(split: str = "valid", max_segments: int = 350) -> None:
+    """Write label batches for the LLM agents."""
+    s = Settings()
+    paths = labels.write_batches(load_index(s, split), s.labels_dir / split / "batches", max_segments)
+    typer.echo(f"wrote {len(paths)} batches to {paths[0].parent}")
+
+
+def split_of(split: str, speaker_id: str, agree: bool) -> str:
+    """valid becomes the test set; any other split is cut by speaker into train and a 5% dev set."""
+    if split == "valid":
+        return "test" if agree else "test_ambiguous"
+    return "dev" if zlib.crc32(speaker_id.encode()) % 20 == 0 else "train"
+
+
+@app.command()
+def build(split: str = "valid", labelers: str = "a,b", seed: int = 0, delete_shards: bool = False) -> None:
+    """Cut 8 s clips for every labelled segment, one shard at a time, and write them to data/clips/<split>.parquet.
+
+    For valid, agreeing labels go to 'test' and the rest to 'test_ambiguous'. Labeler order matters: the
+    first labeler's label is the clip's label and sets the context gaps.
+    """
+    s = Settings()
+    index = load_index(s, split)
+    names = labelers.split(",")
+    results = {n: labels.load_results(s.labels_dir / split / "results" / n, index) for n in names}
+    sessions = labels.sessions(index)
+    rows_of, need, in_shard = defaultdict(set), defaultdict(set), defaultdict(set)
+    for r in index:
+        rows_of[r["shard"]].add(r["row"])
+        need[r["session"]].add(r["shard"])
+        in_shard[r["shard"]].add(r["session"])
+    rng, audio, counts, writer = random.Random(seed), {}, defaultdict(int), None
+    path = s.clips_dir / f"{split}.parquet"
+    path.parent.mkdir(parents=True, exist_ok=True)
+    for n, shard in enumerate(sorted(rows_of), 1):
+        got = indicvoices.shard_audio(shard, rows_of[shard], s.hf_token, s.hf_cache, delete=delete_shards)
+        audio |= {(shard, i): x for i, x in got.items()}
+        out = []
+        # A few sessions span two shards; cut them once all their audio is in.
+        for sid in [
+            sid for sid in sessions if sid in in_shard[shard]
+        ]:  # index order keeps the gaps reproducible
+            need[sid].discard(shard)
+            if need[sid]:
+                continue
+            chunks = sessions[sid]
+            wavs = {
+                c["chunk"]: audio.pop((c["shard"], c["row"]))
+                for c in chunks
+                if (c["shard"], c["row"]) in audio
+            }
+            for c, clip in zip(chunks, clips.build(chunks, wavs, results[names[0]], rng), strict=True):
+                if clip is None:
+                    continue
+                votes = {n: results[n][(sid, c["chunk"])] for n in names}
+                row = {
+                    "audio": clip,
+                    "endpoint_bool": votes[names[0]].complete,
+                    "split": split_of(split, c["speaker_id"], len({v.complete for v in votes.values()}) == 1),
+                    "session": sid,
+                    "chunk": c["chunk"],
+                    "speaker_id": c["speaker_id"],
+                    "spoken_text": c["text"],
+                    **{f"label_{n}": v.complete for n, v in votes.items()},
+                    **{f"confidence_{n}": v.confidence for n, v in votes.items()},
+                }
+                counts[row["split"], row["endpoint_bool"]] += 1
+                out.append(row)
+        if out:
+            table = clips.table(out)
+            writer = writer or pq.ParquetWriter(path, table.schema)
+            writer.write_table(table, row_group_size=500)
+        typer.echo(f"[{n}/{len(rows_of)}] {shard}: {len(out)} clips", err=True)
+    if writer:
+        writer.close()
+    typer.echo(
+        f"{sum(counts.values())} clips: "
+        + json.dumps({f"{k[0]}/{'finished' if k[1] else 'unfinished'}": v for k, v in sorted(counts.items())})
+    )
+
+
+@app.command()
+def spotcheck(
+    split: str = "valid",
+    per_class: int = 25,
+    unfinished_only: bool = False,
+    skip: Path | None = None,
+    name: str = "spotcheck",
+    seed: int = 0,
+) -> None:
+    """Write a local listening page of clean test clips: a random sample of each label, or every
+    unfinished clip. Clips already answered in the `skip` file are left out."""
+    from importlib.resources import files
+
+    s = Settings()
+    done = set(json.loads(skip.read_text())) if skip else set()
+    rows = [
+        r
+        for r in pq.read_table(s.clips_dir / f"{split}.parquet").to_pylist()
+        if r["split"] == "test" and f"{r['session']}|{r['chunk']}" not in done
+    ]
+    rng = random.Random(seed)
+    if unfinished_only:
+        picked = [r for r in rows if not r["endpoint_bool"]]
+    else:
+        picked = [
+            r
+            for label in (False, True)
+            for r in rng.sample([r for r in rows if r["endpoint_bool"] is label], per_class)
+        ]
+    rng.shuffle(picked)
+    out = s.data_dir / name
+    out.mkdir(parents=True, exist_ok=True)
+    clips_meta = []
+    for r in picked:
+        wav = f"{r['session']}_{r['chunk']}.wav"
+        (out / wav).write_bytes(r["audio"]["bytes"])
+        clips_meta.append(
+            {
+                "id": f"{r['session']}|{r['chunk']}",
+                "file": wav,
+                "finished": r["endpoint_bool"],
+                "text": r["spoken_text"],
+            }
+        )
+    page = (
+        files("baari")
+        .joinpath("spotcheck.html")
+        .read_text()
+        .replace("__CLIPS__", json.dumps(clips_meta, ensure_ascii=False))
+        .replace("__NAME__", name)
+    )
+    (out / "index.html").write_text(page)
+    typer.echo(f"open {out / 'index.html'}")
+
+
+@app.command()
+def evaluate(split: str = "valid", subset: str = "test") -> None:
+    """Score the existing turn detectors on the clips."""
+    from baari import evaluate as ev
+
+    s = Settings()
+    report = ev.run(s.clips_dir / f"{split}.parquet", subset, s.hf_cache)
+    s.reports_dir.mkdir(parents=True, exist_ok=True)
+    (s.reports_dir / f"{split}-{subset}.md").write_text(report)
+    typer.echo(report)
+
+
+@app.command()
+def train(
+    clips: Path | None = None,
+    test: Path | None = None,
+    extra: Annotated[
+        list[Path] | None, typer.Option(help="More clip files; their 'train' rows are mixed in.")
+    ] = None,
+    init: Annotated[str, typer.Option(help="fp32 ONNX to start from: a path or repo:file")] = (
+        "adimyth/indic-smart-turn:indic-smart-turn-tiny-fp32.onnx"
+    ),
+    out: Path = Path("data/models/run"),
+    epochs: float = 3,
+    lr: float = 2e-5,
+    batch: int = 64,
+    workers: int = 4,
+    max_steps: int = 0,
+    seed: int = 0,
+) -> None:
+    """Fine-tune on Urdu clips (picked on the dev split), export fp32 and int8 ONNX, score on the test clips."""
+    from huggingface_hub import hf_hub_download
+
+    from baari import train as tr
+
+    s = Settings()
+    start = (
+        Path(init)
+        if Path(init).exists()
+        else Path(hf_hub_download(*init.split(":", 1), cache_dir=s.hf_cache))
+    )
+    log = tr.run(
+        clips or s.clips_dir / "train.parquet",
+        test or s.clips_dir / "valid.parquet",
+        start,
+        out,
+        extra or [],
+        epochs=epochs,
+        lr=lr,
+        batch=batch,
+        seed=seed,
+        workers=workers,
+        cache=s.hf_cache,
+        max_steps=max_steps,
+    )
+    typer.echo(
+        json.dumps({k: log[k] for k in ("sizes", "start", "onnx_fp32_test", "onnx_int8_test")}, indent=1)
+    )

baari/clips.py

@@ -0,0 +1,91 @@
+"""Turn labelled segments into Smart Turn style clips: the last 8 s of 16 kHz audio.
+
+Adapted from indic_turn/build.py in Indic Smart Turn (https://github.com/adimyth/indic-smart-turn).
+Copyright (c) 2026, adimyth. BSD-2-Clause; see LICENSE.indic-smart-turn.
+
+A segment's clip is the segment plus the preceding consecutive segments of the same speaker,
+joined with a short gap after an unfinished segment and a long one after a finished segment
+(the other person would have spoken in between), with trailing silence trimmed to 0.2 s.
+"""
+
+import io
+import random
+
+import numpy as np
+import pyarrow as pa
+import soundfile as sf
+
+from baari.indicvoices import SR
+from baari.labels import Key, SegmentLabel
+
+WINDOW = 8 * SR
+TRAIL = int(0.2 * SR)
+
+
+def frame_energy(x: np.ndarray, win: int = 400, hop: int = 160) -> np.ndarray:
+    n = max(0, (len(x) - win) // hop + 1)
+    if n == 0:
+        return np.zeros(0)
+    idx = np.arange(win)[None, :] + hop * np.arange(n)[:, None]
+    return np.sqrt((x[idx] ** 2).mean(axis=1) + 1e-12)
+
+
+def trim_trailing_silence(
+    x: np.ndarray, rel: float = 0.03, floor: float = 1e-3, hop: int = 160
+) -> np.ndarray:
+    e = frame_energy(x, hop=hop)
+    if len(e) == 0:
+        return x
+    voiced = np.flatnonzero(e > max(floor, rel * e.max()))
+    if len(voiced) == 0:
+        return x
+    return x[: min(len(x), (voiced[-1] + 1) * hop + 400)]
+
+
+def fit_window(x: np.ndarray) -> np.ndarray:
+    x = np.concatenate([trim_trailing_silence(x), np.zeros(TRAIL, np.float32)])
+    return x[-WINDOW:]
+
+
+def build(
+    chunks: list[dict], audio: dict[int, np.ndarray], labels: dict[Key, SegmentLabel], rng: random.Random
+) -> list[np.ndarray | None]:
+    """One clip per chunk of one session, in order; None where the chunk's audio is missing."""
+    clips = []
+    for i, c in enumerate(chunks):
+        if c["chunk"] not in audio:
+            clips.append(None)
+            continue
+        parts, total, j = [audio[c["chunk"]]], len(audio[c["chunk"]]), i - 1
+        while (
+            total < WINDOW
+            and j >= 0
+            and chunks[j]["chunk"] == chunks[j + 1]["chunk"] - 1
+            and chunks[j]["chunk"] in audio
+        ):
+            prev = chunks[j]
+            finished = labels[(prev["session"], prev["chunk"])].complete
+            gap = np.zeros(
+                int((rng.uniform(1.5, 4.0) if finished else rng.uniform(0.15, 0.6)) * SR), np.float32
+            )
+            parts = [audio[prev["chunk"]], gap, *parts]
+            total += len(gap) + len(audio[prev["chunk"]])
+            j -= 1
+        clips.append(fit_window(np.concatenate(parts)))
+    return clips
+
+
+def wav_bytes(x: np.ndarray) -> bytes:
+    buf = io.BytesIO()
+    sf.write(buf, x, SR, format="WAV", subtype="PCM_16")
+    return buf.getvalue()
+
+
+def table(rows: list[dict]) -> pa.Table:
+    """rows carry 'audio' as float32 arrays; stored as 16-bit WAV bytes, as Hugging Face Audio expects."""
+    return pa.Table.from_pylist(
+        [
+            {**r, "audio": {"bytes": wav_bytes(r["audio"]), "path": f"{r['session']}_{r['chunk']}.wav"}}
+            for r in rows
+        ]
+    )

baari/detector.py

@@ -0,0 +1,62 @@
+"""Run baari on audio. Needs only numpy and onnxruntime."""
+
+import os
+import urllib.request
+from pathlib import Path
+
+import numpy as np
+import onnxruntime as ort
+
+from baari.features import compute_whisper_log_mel_features
+
+SAMPLE_RATE = 16000
+MODEL_URL = "https://huggingface.co/samadpls/baari/resolve/main/baari-int8.onnx"
+
+
+def model_path() -> str:
+    """Path to baari's int8 ONNX file, downloaded once into ~/.cache/baari (or $BAARI_CACHE)."""
+    path = Path(os.environ.get("BAARI_CACHE", Path.home() / ".cache" / "baari")) / "baari-int8.onnx"
+    if not path.exists():
+        path.parent.mkdir(parents=True, exist_ok=True)
+        part = path.with_suffix(".part")
+        try:
+            urllib.request.urlretrieve(MODEL_URL, part)
+        except OSError as e:
+            raise RuntimeError(
+                f"Could not download the baari model from {MODEL_URL} ({e}). "
+                "Pass TurnDetector(path=...) to use a local Smart Turn style ONNX file."
+            ) from e
+        part.replace(path)
+    return str(path)
+
+
+class TurnDetector:
+    """Decides, at a pause, whether an Urdu speaker has finished their turn.
+
+    threads: CPU threads for one call. 1 suits a server handling many calls at once; 4 makes a single
+    call about three times faster.
+    """
+
+    def __init__(self, path: str | None = None, threads: int = 1):
+        options = ort.SessionOptions()
+        options.intra_op_num_threads = threads
+        options.inter_op_num_threads = 1
+        self.session = ort.InferenceSession(
+            path or model_path(), sess_options=options, providers=["CPUExecutionProvider"]
+        )
+
+    def probability(self, audio: np.ndarray) -> float:
+        """Probability that the speaker has finished. `audio` is 16 kHz mono, float or int16, ending at
+        the pause. Only the last 8 seconds are used."""
+        x = np.asarray(audio)
+        if x.ndim != 1:
+            raise ValueError(f"expected mono audio as a 1-D array, got shape {x.shape}")
+        x = x / 32768.0 if x.dtype == np.int16 else x
+        n = 8 * SAMPLE_RATE
+        x = x[-n:] if len(x) > n else np.pad(x, (n - len(x), 0))
+        features = compute_whisper_log_mel_features(x.astype(np.float32))[None]
+        return float(np.asarray(self.session.run(None, {"input_features": features})[0]).ravel()[0])
+
+    def is_finished(self, audio: np.ndarray, threshold: float = 0.5) -> bool:
+        """True when the speaker has finished and the agent should reply."""
+        return self.probability(audio) > threshold

baari/evaluate.py

@@ -0,0 +1,152 @@
+"""Score existing end-of-turn models on the Urdu clips.
+
+Smart Turn style models take an 80 x 800 log-mel of the last 8 s and return P(finished).
+PuristanLabs1/urdu-turn-v2 is a PyTorch Whisper-tiny encoder with attention pooling, trained on
+3 s clips; its documented package is not on PyPI, so it is rebuilt here from its model.py.
+"""
+
+import io
+from pathlib import Path
+
+import numpy as np
+import onnxruntime as ort
+import pyarrow.parquet as pq
+import soundfile as sf
+from huggingface_hub import hf_hub_download
+from sklearn.metrics import roc_auc_score
+from transformers import WhisperFeatureExtractor
+
+from baari.features import compute_whisper_log_mel_features
+from baari.indicvoices import SR
+
+ONNX_MODELS = {
+    "Smart Turn v3.2 (CPU, int8)": ("pipecat-ai/smart-turn-v3", "smart-turn-v3.2-cpu.onnx"),
+    "Smart Turn v3.2 (GPU, fp32)": ("pipecat-ai/smart-turn-v3", "smart-turn-v3.2-gpu.onnx"),
+    "Indic Smart Turn tiny (int8)": ("adimyth/indic-smart-turn", "indic-smart-turn-tiny-int8.onnx"),
+    "Indic Smart Turn tiny (fp32)": ("adimyth/indic-smart-turn", "indic-smart-turn-tiny-fp32.onnx"),
+    "Indic Smart Turn base (fp32)": ("adimyth/indic-smart-turn", "indic-smart-turn-base-fp32.onnx"),
+}
+PURISTAN = "PuristanLabs1/urdu-turn-v2 (Urdu)"
+
+
+def last(x: np.ndarray, seconds: int) -> np.ndarray:
+    n = seconds * SR
+    return x[-n:] if len(x) > n else np.pad(x, (n - len(x), 0))
+
+
+class OnnxModel:
+    """Runs a Smart Turn style ONNX file the way Pipecat does: its numpy log-mel features, not
+    transformers'. int8 models can change their answer on about 1 clip in 100 between the two, which
+    differ by about 1e-5, so scores must use the features people will run."""
+
+    def __init__(self, path: str):
+        self.sess = ort.InferenceSession(path, providers=["CPUExecutionProvider"])
+        self.input = self.sess.get_inputs()[0].name
+
+    def __call__(self, wavs: list[np.ndarray]) -> np.ndarray:
+        feats = np.stack([compute_whisper_log_mel_features(last(w, 8).astype(np.float32)) for w in wavs])
+        p = np.asarray(self.sess.run(None, {self.input: feats})[0], dtype=np.float64).reshape(len(wavs), -1)[
+            :, -1
+        ]
+        return p if 0 <= p.min() and p.max() <= 1 else 1 / (1 + np.exp(-p))
+
+
+class PuristanModel:
+    def __init__(self, cache: Path):
+        import torch
+        from torch import nn
+        from transformers import WhisperConfig, WhisperModel
+
+        class WhisperTurnDetector(
+            nn.Module
+        ):  # same layout as the repo's model.py, so the state dict loads strictly
+            def __init__(self):
+                super().__init__()
+                self.whisper = WhisperModel(
+                    WhisperConfig.from_pretrained("openai/whisper-tiny", cache_dir=cache)
+                )
+                h = self.whisper.config.d_model
+                self.attn = nn.Sequential(nn.Linear(h, 128), nn.Tanh(), nn.Linear(128, 1))
+                self.silence_embedding = nn.Parameter(torch.zeros(h))
+                self.classifier = nn.Sequential(
+                    nn.Linear(h, 128), nn.ReLU(), nn.Dropout(0.1), nn.Linear(128, 2)
+                )
+
+            def forward(self, feats, mask):
+                hidden = self.whisper.encoder(feats).last_hidden_state
+                weights = self.attn(hidden).squeeze(-1).masked_fill(mask[:, ::2] == 0, -1e9)
+                pooled = torch.bmm(torch.softmax(weights, dim=1).unsqueeze(1), hidden).squeeze(1)
+                return self.classifier(pooled)
+
+        self.torch = torch
+        self.model = WhisperTurnDetector().eval()
+        weights = hf_hub_download("PuristanLabs1/urdu-turn-v2", "whisper_turn_model.pth", cache_dir=cache)
+        self.model.load_state_dict(torch.load(weights, map_location="cpu", weights_only=True))
+        self.fe = WhisperFeatureExtractor.from_pretrained("openai/whisper-tiny", cache_dir=cache)
+
+    def __call__(self, wavs: list[np.ndarray], seconds: int = 3) -> np.ndarray:
+        f = self.fe(
+            [w[-seconds * SR :] for w in wavs],
+            sampling_rate=SR,
+            return_tensors="pt",
+            return_attention_mask=True,
+        )
+        with self.torch.no_grad():
+            logits = self.model(f.input_features, f.attention_mask)
+        return self.torch.softmax(logits, dim=1)[:, 1].numpy().astype(np.float64)
+
+
+def load_clips(path: Path, subset: str) -> tuple[list[np.ndarray], np.ndarray]:
+    t = pq.read_table(path).to_pydict()
+    keep = [i for i, s in enumerate(t["split"]) if subset == "all" or s == subset]
+    wavs = [sf.read(io.BytesIO(t["audio"][i]["bytes"]), dtype="float32")[0] for i in keep]
+    return wavs, np.array([int(t["endpoint_bool"][i]) for i in keep])
+
+
+def score(model, wavs: list[np.ndarray], batch: int = 32) -> np.ndarray:
+    return np.concatenate([model(wavs[i : i + batch]) for i in range(0, len(wavs), batch)])
+
+
+def balanced(pred: np.ndarray, y: np.ndarray) -> float:
+    return ((pred[y == 0] == 0).mean() + (pred[y == 1] == 1).mean()) / 2
+
+
+def metrics(p: np.ndarray, y: np.ndarray, resamples: int = 1000, seed: int = 0) -> dict:
+    """Balanced accuracy is the headline: most clips are finished, so plain accuracy rewards always saying 'finished'."""
+    pred = (p > 0.5).astype(int)
+    rng = np.random.default_rng(seed)
+    boot = [balanced(pred[i], y[i]) for i in (rng.integers(0, len(y), len(y)) for _ in range(resamples))]
+    return {
+        "balanced": balanced(pred, y),
+        "ci": np.quantile(boot, [0.025, 0.975]),
+        "acc": (pred == y).mean(),
+        "waits": (pred[y == 0] == 0).mean(),
+        "replies": (pred[y == 1] == 1).mean(),
+        "auc": roc_auc_score(y, p) if p.std() > 0 else 0.5,
+    }
+
+
+def run(clips: Path, subset: str, cache: Path) -> str:
+    wavs, y = load_clips(clips, subset)
+    models = {
+        name: OnnxModel(hf_hub_download(repo, f, cache_dir=cache)) for name, (repo, f) in ONNX_MODELS.items()
+    }
+    models[PURISTAN] = PuristanModel(cache)
+    lines = [
+        f"# Existing models on Urdu clips ({clips.stem}, {subset})",
+        "",
+        f"{len(y)} clips: {int(y.sum())} finished, {int(len(y) - y.sum())} unfinished. Threshold 0.5.",
+        "Balanced accuracy is the mean of the two middle columns.",
+        "",
+        "| Model | Balanced accuracy (95% CI) | Waits when unfinished | Replies when finished | Plain accuracy | AUC |",
+        "|---|---|---|---|---|---|",
+    ]
+    scores = {"Always says finished (reference)": np.ones(len(y))}
+    scores |= {name: score(model, wavs) for name, model in models.items()}
+    for name, p in scores.items():
+        m = metrics(p, y)
+        lines.append(
+            f"| {name} | {100 * m['balanced']:.1f}% ({100 * m['ci'][0]:.1f}–{100 * m['ci'][1]:.1f}) | "
+            f"{100 * m['waits']:.1f}% | {100 * m['replies']:.1f}% | {100 * m['acc']:.1f}% | {m['auc']:.3f} |"
+        )
+    return "\n".join(lines) + "\n"