From 4d52a0c162ce11f902cbdd90ee5c5b5779674e95 Mon Sep 17 00:00:00 2001 From: Ankush Malaker <43288948+AnkushMalaker@users.noreply.github.com> Date: Fri, 17 Jul 2026 14:10:18 +0000 Subject: [PATCH 01/16] feat(backend): add audit, memory, and speaker workflows --- backends/advanced/Dockerfile | 26 + backends/advanced/chronicle-data.sh | 15 + backends/advanced/docker-compose.yml | 7 + backends/advanced/init.py | 74 +- backends/advanced/pyproject.toml | 23 + .../scripts/audit_annotation_dataset.py | 630 +++++++++ .../scripts/backfill_cluster_embeddings.py | 41 + .../scripts/backfill_empty_segments.py | 90 ++ .../scripts/collapse_device_registry.js | 86 ++ .../scripts/evaluate_memory_executor.py | 186 +++ .../scripts/purge_orphaned_deferred_jobs.py | 46 + .../advanced/scripts/regen_titles_batch.py | 123 ++ .../scripts/reprocess_conversation.py | 96 ++ .../advanced/scripts/scope_blank_segments.py | 68 + .../advanced/scripts/settle_stuck_status.py | 92 ++ .../controllers/audio_controller.py | 20 +- .../controllers/data_audit_controller.py | 173 ++- .../controllers/drift_controller.py | 23 +- .../guided_annotation_controller.py | 352 +++++ .../guided_enrollment_controller.py | 1028 +++++++++++++++ .../controllers/queue_controller.py | 25 +- .../src/advanced_omi_backend/heartbeat.py | 78 +- .../src/advanced_omi_backend/llm_client.py | 54 +- .../advanced_omi_backend/model_registry.py | 11 +- .../models/conversation.py | 18 + .../models/memory_audit.py | 5 +- .../advanced_omi_backend/prompt_defaults.py | 15 + .../routers/modules/annotation_routes.py | 51 +- .../routers/modules/audio_routes.py | 49 +- .../routers/modules/conversation_routes.py | 24 + .../routers/modules/data_audit_routes.py | 272 +++- .../routers/modules/health_routes.py | 36 +- .../routers/modules/queue_routes.py | 1 + .../services/audio_stream/session_store.py | 53 +- .../services/data_archive.py | 944 +++++++++++++ .../services/memory/agent/__init__.py | 3 + .../services/memory/agent/codex_agent.py | 448 +++++++ .../services/memory/agent/memory_agent.py | 28 +- .../services/memory/agent/vault_tools.py | 101 +- .../services/memory/audit.py | 3 + .../services/memory/base.py | 7 + .../services/memory/config.py | 6 + .../services/memory/conversation_note.py | 197 +++ .../services/memory/providers/chronicle.py | 263 +++- .../services/memory/rebuild.py | 399 ++++++ .../services/memory/vault_lock.py | 33 + .../services/observability/health_poller.py | 65 +- .../services/transcription/__init__.py | 107 +- .../transcription/streaming_consumer.py | 18 +- .../speaker_recognition_client.py | 293 ++++- .../utils/annotation_import.py | 312 +++++ .../utils/silence_condense.py | 216 +++ .../workers/conversation_jobs.py | 32 + .../workers/data_audit_jobs.py | 1 + .../workers/drift_jobs.py | 41 + .../workers/memory_jobs.py | 130 +- .../workers/orchestrator/health_monitor.py | 64 +- .../workers/rq_worker_entry.py | 8 +- .../workers/speaker_benchmark_jobs.py | 392 ++++++ .../workers/speaker_discovery_jobs.py | 243 ++++ .../workers/speaker_mining_jobs.py | 201 +++ .../workers/transcription_jobs.py | 106 +- .../advanced/src/scripts/chronicle_data.py | 308 +++++ .../scripts/graph-validation/sample_data.py | 190 +++ .../scripts/graph-validation/setup_schema.py | 139 ++ .../graph-validation/test_conversation_doc.py | 236 ++++ .../graph-validation/test_entity_graph.py | 161 +++ .../graph-validation/test_fulltext_search.py | 215 +++ .../graph-validation/test_hybrid_search.py | 199 +++ .../graph-validation/test_vector_search.py | 174 +++ .../tests/scripts/graph-validation/utils.py | 292 +++++ .../advanced/tests/test_annotation_export.py | 78 ++ .../advanced/tests/test_annotation_import.py | 107 ++ .../test_audio_persistence_singleflight.py | 76 ++ .../advanced/tests/test_client_lifecycle.py | 56 + .../advanced/tests/test_codex_executor.py | 174 +++ .../tests/test_conversation_audio_salvage.py | 27 + backends/advanced/tests/test_data_archive.py | 389 ++++++ .../tests/test_data_audit_dataset_filter.py | 59 + .../tests/test_leading_silence_trim.py | 53 + .../tests/test_leading_silence_trim_db.py | 129 ++ .../tests/test_memory_agent_completion.py | 257 ++++ .../advanced/tests/test_memory_exclusion.py | 32 + .../advanced/tests/test_memory_rebuild.py | 345 +++++ .../tests/test_memory_transcript_input.py | 52 + .../tests/test_model_registry_reasoning.py | 26 + .../tests/test_openai_compat_routes.py | 36 + backends/advanced/tests/test_session_store.py | 479 +++++++ .../tests/test_smallest_diarization_config.py | 26 + .../advanced/tests/test_speaker_benchmark.py | 30 + .../advanced/tests/test_speaker_discovery.py | 61 + .../test_speaker_identification_policy.py | 31 + .../tests/test_streaming_clock_offset.py | 180 +++ .../tests/test_streaming_provider_health.py | 37 + .../advanced/tests/test_transcript_slicing.py | 112 ++ .../tests/test_vault_rename_person.py | 115 ++ .../advanced/tests/test_vault_scaffold.py | 97 ++ .../tests/test_worker_fleet_health.py | 175 +++ backends/advanced/webui/src/App.tsx | 8 + .../webui/src/components/MemoryAuditCard.tsx | 19 +- .../src/components/SpeakerInlineInput.tsx | 31 +- .../src/components/SpeakerNameDropdown.tsx | 45 +- .../src/components/audio/PlayheadWaveform.tsx | 6 + .../src/components/audio/WaveformDisplay.tsx | 79 +- .../components/audio/WaveformRegionEditor.tsx | 23 +- .../components/dataAudit/AuditFilterBar.tsx | 7 +- .../src/components/dataAudit/DriftPanel.tsx | 76 +- .../components/dataAudit/GuidedEnrollment.tsx | 1163 +++++++++++++++++ .../src/components/dataAudit/filters.tsx | 37 + .../transcript/InsertSegmentForm.tsx | 3 + .../transcript/TranscriptEditor.tsx | 492 ++++++- .../webui/src/pages/ConversationDetail.tsx | 39 +- .../advanced/webui/src/pages/DataAudit.tsx | 48 +- .../webui/src/pages/SpeakerEnrollment.tsx | 5 + backends/advanced/webui/src/pages/Upload.tsx | 230 +++- backends/advanced/webui/src/services/api.ts | 316 ++++- .../advanced/webui/src/utils/speakerLabels.ts | 10 + backends/advanced/worker_healthcheck.py | 11 +- 118 files changed, 16054 insertions(+), 299 deletions(-) create mode 100755 backends/advanced/chronicle-data.sh create mode 100644 backends/advanced/scripts/audit_annotation_dataset.py create mode 100644 backends/advanced/scripts/backfill_cluster_embeddings.py create mode 100644 backends/advanced/scripts/backfill_empty_segments.py create mode 100644 backends/advanced/scripts/collapse_device_registry.js create mode 100644 backends/advanced/scripts/evaluate_memory_executor.py create mode 100644 backends/advanced/scripts/purge_orphaned_deferred_jobs.py create mode 100644 backends/advanced/scripts/regen_titles_batch.py create mode 100644 backends/advanced/scripts/reprocess_conversation.py create mode 100644 backends/advanced/scripts/scope_blank_segments.py create mode 100644 backends/advanced/scripts/settle_stuck_status.py create mode 100644 backends/advanced/src/advanced_omi_backend/controllers/guided_annotation_controller.py create mode 100644 backends/advanced/src/advanced_omi_backend/controllers/guided_enrollment_controller.py create mode 100644 backends/advanced/src/advanced_omi_backend/services/data_archive.py create mode 100644 backends/advanced/src/advanced_omi_backend/services/memory/agent/codex_agent.py create mode 100644 backends/advanced/src/advanced_omi_backend/services/memory/conversation_note.py create mode 100644 backends/advanced/src/advanced_omi_backend/services/memory/rebuild.py create mode 100644 backends/advanced/src/advanced_omi_backend/utils/annotation_import.py create mode 100644 backends/advanced/src/advanced_omi_backend/utils/silence_condense.py create mode 100644 backends/advanced/src/advanced_omi_backend/workers/drift_jobs.py create mode 100644 backends/advanced/src/advanced_omi_backend/workers/speaker_benchmark_jobs.py create mode 100644 backends/advanced/src/advanced_omi_backend/workers/speaker_discovery_jobs.py create mode 100644 backends/advanced/src/advanced_omi_backend/workers/speaker_mining_jobs.py create mode 100644 backends/advanced/src/scripts/chronicle_data.py create mode 100644 backends/advanced/tests/scripts/graph-validation/sample_data.py create mode 100644 backends/advanced/tests/scripts/graph-validation/setup_schema.py create mode 100644 backends/advanced/tests/scripts/graph-validation/test_conversation_doc.py create mode 100644 backends/advanced/tests/scripts/graph-validation/test_entity_graph.py create mode 100644 backends/advanced/tests/scripts/graph-validation/test_fulltext_search.py create mode 100644 backends/advanced/tests/scripts/graph-validation/test_hybrid_search.py create mode 100644 backends/advanced/tests/scripts/graph-validation/test_vector_search.py create mode 100644 backends/advanced/tests/scripts/graph-validation/utils.py create mode 100644 backends/advanced/tests/test_annotation_export.py create mode 100644 backends/advanced/tests/test_annotation_import.py create mode 100644 backends/advanced/tests/test_audio_persistence_singleflight.py create mode 100644 backends/advanced/tests/test_client_lifecycle.py create mode 100644 backends/advanced/tests/test_codex_executor.py create mode 100644 backends/advanced/tests/test_conversation_audio_salvage.py create mode 100644 backends/advanced/tests/test_data_archive.py create mode 100644 backends/advanced/tests/test_data_audit_dataset_filter.py create mode 100644 backends/advanced/tests/test_leading_silence_trim.py create mode 100644 backends/advanced/tests/test_leading_silence_trim_db.py create mode 100644 backends/advanced/tests/test_memory_agent_completion.py create mode 100644 backends/advanced/tests/test_memory_exclusion.py create mode 100644 backends/advanced/tests/test_memory_rebuild.py create mode 100644 backends/advanced/tests/test_memory_transcript_input.py create mode 100644 backends/advanced/tests/test_model_registry_reasoning.py create mode 100644 backends/advanced/tests/test_openai_compat_routes.py create mode 100644 backends/advanced/tests/test_session_store.py create mode 100644 backends/advanced/tests/test_smallest_diarization_config.py create mode 100644 backends/advanced/tests/test_speaker_benchmark.py create mode 100644 backends/advanced/tests/test_speaker_discovery.py create mode 100644 backends/advanced/tests/test_speaker_identification_policy.py create mode 100644 backends/advanced/tests/test_streaming_clock_offset.py create mode 100644 backends/advanced/tests/test_streaming_provider_health.py create mode 100644 backends/advanced/tests/test_transcript_slicing.py create mode 100644 backends/advanced/tests/test_vault_rename_person.py create mode 100644 backends/advanced/tests/test_vault_scaffold.py create mode 100644 backends/advanced/tests/test_worker_fleet_health.py create mode 100644 backends/advanced/webui/src/components/dataAudit/GuidedEnrollment.tsx create mode 100644 backends/advanced/webui/src/pages/SpeakerEnrollment.tsx create mode 100644 backends/advanced/webui/src/utils/speakerLabels.ts diff --git a/backends/advanced/Dockerfile b/backends/advanced/Dockerfile index 2a867bec..7d376f98 100644 --- a/backends/advanced/Dockerfile +++ b/backends/advanced/Dockerfile @@ -41,11 +41,33 @@ RUN git clone ${NOTESMD_REPO} /src \ && CGO_ENABLED=0 go build -trimpath -mod=vendor -o /out/notesmd-cli . +# ============================================ +# Codex CLI fetcher - optional memory-agent executor +# ============================================ +# Static musl binary used when config.yml sets memory.agent_executor: codex (vault +# recording on a ChatGPT subscription). Auth is NOT baked in — it comes from the +# CODEX_HOME volume mount (see docker-compose.yml / the setup wizard). +FROM debian:bookworm-slim AS codex-fetcher +RUN apt-get update && apt-get install -y --no-install-recommends curl ca-certificates \ + && rm -rf /var/lib/apt/lists/* +ARG CODEX_VERSION=0.144.4 +ARG TARGETARCH +RUN case "${TARGETARCH:-amd64}" in \ + arm64) triple=aarch64-unknown-linux-musl ;; \ + *) triple=x86_64-unknown-linux-musl ;; \ + esac \ + && curl -fsSL "https://github.com/openai/codex/releases/download/rust-v${CODEX_VERSION}/codex-${triple}.tar.gz" \ + | tar -xz -C /tmp \ + && install -m 0755 "/tmp/codex-${triple}" /usr/local/bin/codex + + # ============================================ # Production stage # ============================================ FROM python:3.12-slim-bookworm AS prod +COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/ + RUN apt-get update && apt-get install -y --no-install-recommends \ libsndfile1 \ curl \ @@ -56,6 +78,8 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ # notesmd-cli on PATH so the memory agent auto-detects it (shutil.which). COPY --from=notesmd-builder /out/notesmd-cli /usr/local/bin/notesmd-cli +# Codex CLI on PATH for the optional codex memory-agent executor. +COPY --from=codex-fetcher /usr/local/bin/codex /usr/local/bin/codex WORKDIR /app @@ -104,6 +128,8 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ # notesmd-cli on PATH so the memory agent auto-detects it (shutil.which). COPY --from=notesmd-builder /out/notesmd-cli /usr/local/bin/notesmd-cli +# Codex CLI on PATH for the optional codex memory-agent executor. +COPY --from=codex-fetcher /usr/local/bin/codex /usr/local/bin/codex WORKDIR /app diff --git a/backends/advanced/chronicle-data.sh b/backends/advanced/chronicle-data.sh new file mode 100755 index 00000000..fae06bd3 --- /dev/null +++ b/backends/advanced/chronicle-data.sh @@ -0,0 +1,15 @@ +#!/usr/bin/env bash +set -euo pipefail + +cd "$(dirname "$0")" + +if [[ -n "${COMPOSE_CMD:-}" ]]; then + read -r -a compose_command <<<"${COMPOSE_CMD}" +elif [[ "${CONTAINER_ENGINE:-docker}" == "podman" ]]; then + compose_command=(podman-compose) +else + compose_command=(docker compose) +fi + +exec "${compose_command[@]}" run --rm --no-deps chronicle-backend \ + uv run --offline --no-sync python3 src/scripts/chronicle_data.py "$@" diff --git a/backends/advanced/docker-compose.yml b/backends/advanced/docker-compose.yml index 18d3386a..f7749b47 100644 --- a/backends/advanced/docker-compose.yml +++ b/backends/advanced/docker-compose.yml @@ -23,7 +23,11 @@ services: - ../../plugins:/app/plugins # External plugins directory - ../../discovery.py:/app/discovery.py:ro # Service discovery module - /var/run/tailscale/tailscaled.sock:/var/run/tailscale/tailscaled.sock:ro # Tailscale socket for minidisc + # Codex CLI auth for the optional codex memory-agent executor. The wizard points + # CODEX_HOME_DIR at the host's ~/.codex; rw because codex rotates its tokens. + - ${CODEX_HOME_DIR:-./data/codex-home}:/codex-home environment: + - CODEX_HOME=/codex-home - DEEPGRAM_API_KEY=${DEEPGRAM_API_KEY} - PARAKEET_ASR_URL=${PARAKEET_ASR_URL} - OLLAMA_BASE_URL=${OLLAMA_BASE_URL} @@ -94,7 +98,10 @@ services: - ../../plugins:/app/plugins # External plugins directory - ../../discovery.py:/app/discovery.py:ro # Service discovery module - /var/run/tailscale/tailscaled.sock:/var/run/tailscale/tailscaled.sock:ro # Tailscale socket for minidisc + # Codex CLI auth for the optional codex memory-agent executor (memory jobs run here). + - ${CODEX_HOME_DIR:-./data/codex-home}:/codex-home environment: + - CODEX_HOME=/codex-home - DEEPGRAM_API_KEY=${DEEPGRAM_API_KEY} - PARAKEET_ASR_URL=${PARAKEET_ASR_URL} - OPENAI_API_KEY=${OPENAI_API_KEY} diff --git a/backends/advanced/init.py b/backends/advanced/init.py index e7b9fc02..054c77ef 100644 --- a/backends/advanced/init.py +++ b/backends/advanced/init.py @@ -1240,12 +1240,84 @@ def setup_memory(self): Chronicle's agentic Markdown vault is currently the only memory provider, so there is no provider choice to make — we just ensure config.yml/.env - record it. + record it, then choose how the memory agent executes. """ self.config_manager.update_memory_config({"provider": "chronicle"}) self.console.print( "[green][SUCCESS][/green] Memory: Chronicle agentic vault (config.yml + .env)" ) + self.setup_memory_executor() + + def setup_memory_executor(self): + """Choose how the memory agent runs: the built-in LLM tool loop (metered + API calls) or the OpenAI Codex CLI on a ChatGPT subscription.""" + self.print_section("Memory agent executor") + self.console.print( + "[blue][INFO][/blue] The memory agent records each conversation into " + "your vault. It can run through the configured LLM (per-call API usage) " + "or through the OpenAI Codex CLI, which bills against a ChatGPT " + "subscription instead of API keys." + ) + self.console.print() + + existing = ( + self.config_manager.get_memory_config().get("agent_executor") or "direct" + ) + choices = { + "1": "Direct LLM tool loop (uses the configured LLM's API)", + "2": "Codex CLI (uses your ChatGPT subscription)", + } + choice = self.prompt_choice( + "How should the memory agent run?", + choices, + "2" if str(existing).lower() == "codex" else "1", + ) + + if choice != "2": + self.config_manager.update_memory_config({"agent_executor": "direct"}) + self.console.print( + "[green][SUCCESS][/green] Memory agent: direct LLM loop " + "(memory.agent_executor: direct)" + ) + return + + codex_home = Path( + os.environ.get("CODEX_HOME") or (Path.home() / ".codex") + ).expanduser() + auth_file = codex_home / "auth.json" + if not shutil.which("codex"): + self.console.print( + "[yellow][WARNING][/yellow] No `codex` CLI found on this host. The " + "containers ship their own binary, but you still need subscription " + "auth: install Codex and run `codex login` (ChatGPT sign-in), then " + "re-run init." + ) + if auth_file.is_file(): + self.console.print( + f"[green]✅[/green] Found Codex subscription auth at {auth_file}" + ) + else: + self.console.print( + f"[yellow][WARNING][/yellow] No Codex auth at {auth_file} — until " + "`codex login` has been run on this host, the backend automatically " + "falls back to the direct LLM loop." + ) + # The compose files mount CODEX_HOME_DIR at /codex-home inside the backend + # and workers containers (CODEX_HOME env) — read-write, because codex + # rotates the refresh tokens in auth.json. + self.config["CODEX_HOME_DIR"] = str(codex_home) + # danger-full-access: codex's own sandbox (bubblewrap) cannot run nested + # inside the rootless-podman containers; the container is the boundary. + self.config_manager.update_memory_config( + { + "agent_executor": "codex", + "codex": {"sandbox_mode": "danger-full-access"}, + } + ) + self.console.print( + "[green][SUCCESS][/green] Memory agent: Codex CLI " + f"(memory.agent_executor: codex; {codex_home} mounted into the containers)" + ) def setup_optional_services(self): """Configure optional services""" diff --git a/backends/advanced/pyproject.toml b/backends/advanced/pyproject.toml index 31116040..91b6fcac 100644 --- a/backends/advanced/pyproject.toml +++ b/backends/advanced/pyproject.toml @@ -72,6 +72,7 @@ robotframework = "^6.1.1" [tool.pytest.ini_options] minversion = "8.0" testpaths = ["tests"] +norecursedirs = ["tests/scripts"] python_files = ["test_*.py", "*_test.py"] python_classes = ["Test*", "*Tests"] python_functions = ["test_*"] @@ -106,6 +107,28 @@ filterwarnings = [ "ignore::PendingDeprecationWarning", ] +[tool.coverage.run] +branch = true +relative_files = true +source = ["advanced_omi_backend"] + +[tool.coverage.paths] +source = [ + "src/advanced_omi_backend", + "*/site-packages/advanced_omi_backend", +] + +[tool.coverage.report] +precision = 1 +show_missing = true +skip_empty = true + +[tool.coverage.xml] +output = "coverage-reports/coverage.xml" + +[tool.coverage.html] +directory = "coverage-reports/html" + [dependency-groups] dev = [ "black>=25.1.0", diff --git a/backends/advanced/scripts/audit_annotation_dataset.py b/backends/advanced/scripts/audit_annotation_dataset.py new file mode 100644 index 00000000..a9f2b3e5 --- /dev/null +++ b/backends/advanced/scripts/audit_annotation_dataset.py @@ -0,0 +1,630 @@ +#!/usr/bin/env python3 +"""AI-assisted audio review of a Chronicle annotation dataset (read-only). + +Listens to every clip in an exported annotation dataset with an audio-capable +model (Gemini) and flags where the machine transcript disagrees with the audio. +The point is to *review a dataset once before sending it to human annotators* — so +the reviewer sees the likely-wrong spots instead of re-transcribing from scratch. + +This tool is deliberately standalone: it reads the export ZIP + writes a sidecar +report. It makes NO database, transcript, or memory changes. Staging the findings +as editor suggestions is a separate, later step (see ``--stage-suggestions`` in the +follow-up design — intentionally not implemented here until output quality is +reviewed). + +Why Gemini and not GPT: Gemini can directly analyse audio (transcription, +timing, diarization); OpenAI's text models can only judge transcript coherence, +they cannot listen. See https://ai.google.dev/gemini-api/docs/audio + +Usage +----- +Run on 3 clips first, inspect ``ai_audit.jsonl``, then run the whole set: + + uv run --with google-genai python scripts/audit_annotation_dataset.py \ + --dataset annotation_20260628_180119_adaf --limit 3 + + uv run --with google-genai python scripts/audit_annotation_dataset.py \ + --dataset annotation_20260628_180119_adaf + +The run is resumable: clips already present in the output file are skipped, so a +crash or Ctrl-C mid-run costs nothing. Pass ``--overwrite`` to start fresh. + +Requires ``GOOGLE_API_KEY`` (read from the environment, else from +``extras/ml-experiments/.env`` or ``backends/advanced/.env``). +""" + +from __future__ import annotations + +import argparse +import io +import json +import os +import sys +import tempfile +import time +import zipfile +from pathlib import Path +from typing import Any, Optional + +# --------------------------------------------------------------------------- +# Paths / config +# --------------------------------------------------------------------------- + +# scripts/ -> backends/advanced -> repo root +BACKEND_DIR = Path(__file__).resolve().parent.parent +REPO_ROOT = BACKEND_DIR.parent.parent +DEFAULT_EXPORTS_DIR = BACKEND_DIR / "data" / "exports" + +DEFAULT_MODEL = "gemini-3.1-pro-preview" +OUTPUT_NAME = "ai_audit.jsonl" + +# Headroom so a dense clip's thinking + JSON output can't hit the cap and truncate +# the reply into invalid JSON (findings-only output is otherwise small). +MAX_OUTPUT_TOKENS = 65536 +UPLOAD_PROCESSING_TIMEOUT_S = 180 +REQUEST_TIMEOUT_MS = 300_000 +# Retry a clip whose reply doesn't parse — the pro model is occasionally +# nondeterministic about emitting strictly valid JSON. +MAX_ATTEMPTS = 3 + + +ISSUE_TYPES = [ + "mistranscription", # wrong word(s) — the transcript says X, the audio says Y + "missing_speech", # audible speech absent from the transcript + "hallucinated_text", # transcript text with no corresponding audio + "wrong_speaker", # speech attributed to the wrong speaker + "punctuation", # punctuation/casing only; wording is correct + "script_mix", # Hindi/English written in the wrong script (romanized vs Devanagari) + "other", +] + +AUDIT_PROMPT = """\ +You are auditing an automatic speech-recognition (ASR) transcript against its \ +source audio. This transcript will be used to fine-tune an STT model, so \ +verbatim accuracy matters more than readability. + +You are given: +1. The audio clip. +2. The current machine transcript, split into numbered segments. + +Listen to the ENTIRE clip and compare it to each segment. Report every place the \ +transcript disagrees with what is actually said. Focus on content errors that \ +would teach the model wrong things: wrong words, missed speech, invented text, \ +and code-switch script errors. Ignore trivial stylistic choices. + +Rules for suggested_text: +* Give the corrected VERBATIM text for that segment (what is actually spoken). +* Keep Hindi in Devanagari and English in Roman/Latin script; do not transliterate. +* Preserve filler words, false starts, repetitions exactly as spoken. +* If a segment is already correct, do NOT include it in findings. + +Current transcript segments: +{segments_block} + +Respond with a single JSON object of exactly this shape (no markdown, no prose \ +outside the JSON): +{{ + "transcript_quality": one of ["good", "minor_issues", "major_issues", "unusable"], + "overall_notes": "one or two sentences on the clip's overall transcript quality", + "findings": [ + {{ + "segment_index": , + "original_text": "", + "suggested_text": "", + "issue_type": one of {issue_types}, + "confidence": , + "explanation": "" + }} + ] +}} +If the whole transcript is faithful, return "findings": []. +""" + + +def _load_api_key() -> str: + key = os.environ.get("GOOGLE_API_KEY") + if key: + return key + # Fall back to the two .env files that may hold it. + try: + from dotenv import dotenv_values + except ImportError: + dotenv_values = None + if dotenv_values is not None: + for env_path in ( + REPO_ROOT / "extras" / "ml-experiments" / ".env", + BACKEND_DIR / ".env", + ): + if env_path.exists(): + val = dotenv_values(env_path).get("GOOGLE_API_KEY") + if val: + return val + print( + "ERROR: GOOGLE_API_KEY not found in environment or " + "extras/ml-experiments/.env / backends/advanced/.env", + file=sys.stderr, + ) + sys.exit(2) + + +def _resolve_dataset_zip(args: argparse.Namespace) -> Path: + if args.dataset_path: + p = Path(args.dataset_path) + if p.is_dir(): + p = p / "dataset.zip" + if not p.exists(): + print(f"ERROR: dataset ZIP not found: {p}", file=sys.stderr) + sys.exit(2) + return p + exports_dir = Path(args.exports_dir) if args.exports_dir else DEFAULT_EXPORTS_DIR + zip_path = exports_dir / args.dataset / "dataset.zip" + if not zip_path.exists(): + print( + f"ERROR: dataset ZIP not found: {zip_path}\n" + f" (looked under exports dir {exports_dir})", + file=sys.stderr, + ) + sys.exit(2) + return zip_path + + +def _read_manifest(zf: zipfile.ZipFile) -> list[dict[str, Any]]: + records = [] + for line in zf.read("manifest.jsonl").decode("utf-8").splitlines(): + if line.strip(): + records.append(json.loads(line)) + return records + + +def _segments_block(record: dict[str, Any]) -> str: + """Numbered segment listing for the prompt. Falls back to the whole text.""" + segments = record.get("segments") or [] + if not segments: + return f"[0] (Unknown Speaker) {record.get('text', '').strip()}" + lines = [] + for i, seg in enumerate(segments): + speaker = seg.get("identified_as") or seg.get("speaker") or "Unknown Speaker" + start = seg.get("start", 0.0) + end = seg.get("end", 0.0) + text = (seg.get("text") or "").strip() + lines.append(f"[{i}] ({speaker} | {start:.1f}-{end:.1f}s) {text}") + return "\n".join(lines) + + +def _parse_json_response(raw: str) -> dict[str, Any]: + """Parse a model JSON reply, tolerating ```json fences and stray prose.""" + text = raw.strip() + if text.startswith("```"): + text = text.split("```", 2)[1] + if text.lstrip().startswith("json"): + text = text.lstrip()[4:] + text = text.strip() + # Grab the outermost {...} if the model wrapped it in prose. + if not text.startswith("{"): + start = text.find("{") + end = text.rfind("}") + if start != -1 and end != -1: + text = text[start : end + 1] + return json.loads(text) + + +def _normalize_findings( + parsed: dict[str, Any], record: dict[str, Any] +) -> dict[str, Any]: + segments = record.get("segments") or [] + findings = [] + for raw in parsed.get("findings") or []: + if not isinstance(raw, dict): + continue + try: + seg_index = int(raw.get("segment_index", -1)) + except (TypeError, ValueError): + seg_index = -1 + issue = str(raw.get("issue_type", "other")) + if issue not in ISSUE_TYPES: + issue = "other" + try: + confidence = float(raw.get("confidence", 0.0)) + except (TypeError, ValueError): + confidence = 0.0 + confidence = max(0.0, min(1.0, confidence)) + original = raw.get("original_text") + if original is None and 0 <= seg_index < len(segments): + original = segments[seg_index].get("text", "") + findings.append( + { + "segment_index": seg_index, + "original_text": original or "", + "suggested_text": raw.get("suggested_text") or "", + "issue_type": issue, + "confidence": round(confidence, 3), + "explanation": (raw.get("explanation") or "").strip(), + } + ) + quality = str(parsed.get("transcript_quality", "unknown")) + if quality not in {"good", "minor_issues", "major_issues", "unusable"}: + quality = "unknown" + return { + "transcript_quality": quality, + "overall_notes": (parsed.get("overall_notes") or "").strip(), + "findings": findings, + } + + +def _audit_clip( + client, + types, + *, + model: str, + record: dict[str, Any], + audio_bytes: bytes, + error_dir: Optional[Path] = None, +) -> dict[str, Any]: + """Upload one clip + transcript to Gemini and return the normalized audit.""" + prompt = AUDIT_PROMPT.format( + segments_block=_segments_block(record), + issue_types=json.dumps(ISSUE_TYPES), + ) + + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: + tmp.write(audio_bytes) + tmp_path = tmp.name + + uploaded = None + try: + uploaded = client.files.upload(file=tmp_path) + deadline = time.time() + UPLOAD_PROCESSING_TIMEOUT_S + while uploaded.state == "PROCESSING": + if time.time() > deadline: + raise TimeoutError("Gemini file stuck in PROCESSING") + time.sleep(2) + uploaded = client.files.get(name=uploaded.name) + if uploaded.state == "FAILED": + raise RuntimeError("Gemini rejected the audio upload") + + last_error: Exception | None = None + last_raw = "" + for _attempt in range(MAX_ATTEMPTS): + resp = client.models.generate_content( + model=model, + contents=[uploaded, prompt], + config=types.GenerateContentConfig( + max_output_tokens=MAX_OUTPUT_TOKENS, + response_mime_type="application/json", + ), + ) + last_raw = resp.text or "" + finish = str(resp.candidates[0].finish_reason) if resp.candidates else "?" + try: + parsed = _parse_json_response(last_raw) + except json.JSONDecodeError as exc: + last_error = RuntimeError( + f"unparseable JSON (finish_reason={finish}): {exc}" + ) + continue + audit = _normalize_findings(parsed, record) + usage = resp.usage_metadata + audit["tokens"] = { + "input": getattr(usage, "prompt_token_count", 0) or 0, + "output": getattr(usage, "candidates_token_count", 0) or 0, + "thinking": getattr(usage, "thoughts_token_count", 0) or 0, + } + return audit + + # All attempts failed to parse — persist the last raw reply for inspection + # rather than throwing it away, then surface the error. + if error_dir is not None: + error_dir.mkdir(parents=True, exist_ok=True) + (error_dir / f"{record['clip_id']}.raw.txt").write_text(last_raw) + raise last_error or RuntimeError("audit failed") + finally: + if uploaded is not None: + try: + client.files.delete(name=uploaded.name) + except Exception: + pass + try: + os.unlink(tmp_path) + except OSError: + pass + + +async def _stage_suggestions(args: argparse.Namespace) -> int: + """Read ai_audit.jsonl and stage its findings as PENDING model suggestions. + + Findings are matched to the *imported* conversations (the memory-excluded + copies created by Upload → Annotation workspace) via ``external_source_id`` + computed by the very same parser the importer used, so segment indices and + ids line up exactly. Nothing is applied to a transcript — each finding becomes + a suggestion the editor renders for accept/edit/reject. + """ + from beanie import init_beanie + + from advanced_omi_backend.database import db + from advanced_omi_backend.models.annotation import ( + Annotation, + AnnotationSource, + AnnotationStatus, + AnnotationType, + ) + from advanced_omi_backend.models.conversation import Conversation + from advanced_omi_backend.models.user import User + from advanced_omi_backend.utils.annotation_import import parse_annotation_dataset + + zip_path = _resolve_dataset_zip(args) + out_path = Path(args.out) if args.out else zip_path.parent / OUTPUT_NAME + if not out_path.exists(): + print( + f"ERROR: no audit report at {out_path}. Run the audit (without " + f"--stage-suggestions) first.", + file=sys.stderr, + ) + return 2 + + audits: dict[str, dict[str, Any]] = {} + for line in out_path.read_text().splitlines(): + if line.strip(): + row = json.loads(line) + audits[row["clip_id"]] = row + + # Same parser the import endpoint used → identical dataset_id + clip_ids. + dataset = parse_annotation_dataset(zip_path.read_bytes()) + + await init_beanie(database=db, document_models=[User, Conversation, Annotation]) + + staged = 0 + skipped_existing = 0 + skipped_conf = 0 + skipped_bounds = 0 + missing_conv = 0 + touched: set[str] = set() + + for clip in dataset.clips: + row = audits.get(clip.clip_id) + if not row: + continue + external_source_id = f"{dataset.dataset_id}:{clip.clip_id}" + conv = await Conversation.find_one( + Conversation.external_source_type == "annotation_dataset", + Conversation.external_source_id == external_source_id, + ) + if not conv: + missing_conv += 1 + continue + + active = conv.active_transcript + segments = active.segments if active and active.segments else [] + seg_count = len(segments) + + if args.replace: + await Annotation.find( + Annotation.conversation_id == conv.conversation_id, + Annotation.annotation_type == AnnotationType.TRANSCRIPT, + Annotation.source == AnnotationSource.MODEL_SUGGESTION, + Annotation.status == AnnotationStatus.PENDING, + ).delete() + + for finding in row.get("findings", []): + if float(finding.get("confidence", 0.0)) < args.min_confidence: + skipped_conf += 1 + continue + seg_idx = finding.get("segment_index") + if seg_idx is None or seg_idx < 0 or seg_idx >= seg_count: + skipped_bounds += 1 + continue + + if not args.replace: + existing = await Annotation.find_one( + Annotation.conversation_id == conv.conversation_id, + Annotation.segment_index == seg_idx, + Annotation.source == AnnotationSource.MODEL_SUGGESTION, + Annotation.status == AnnotationStatus.PENDING, + ) + if existing: + skipped_existing += 1 + continue + + annotation = Annotation( + annotation_type=AnnotationType.TRANSCRIPT, + user_id=conv.user_id, + conversation_id=conv.conversation_id, + segment_index=seg_idx, + # Truthful "before": the segment's current text, matching the + # convention used when the editor creates transcript annotations. + original_text=segments[seg_idx].text, + corrected_text=finding.get("suggested_text", ""), + source=AnnotationSource.MODEL_SUGGESTION, + status=AnnotationStatus.PENDING, + ) + await annotation.save() + staged += 1 + touched.add(conv.conversation_id) + + print( + f"Dataset: {dataset.dataset_id}\n" + f"Staged {staged} suggestion(s) across {len(touched)} conversation(s).\n" + f" skipped: {skipped_existing} already-pending, " + f"{skipped_conf} below --min-confidence={args.min_confidence}, " + f"{skipped_bounds} out-of-bounds segment_index\n" + f" {missing_conv} clip(s) had no imported conversation " + f"(import the dataset via Upload → Annotation workspace first)\n" + f"\nReview them at: https://localhost/data-audit?dataset={dataset.dataset_id}" + ) + return 0 + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + src = ap.add_mutually_exclusive_group(required=True) + src.add_argument( + "--dataset", + help="Export id under data/exports/, e.g. annotation_20260628_180119_adaf", + ) + src.add_argument( + "--dataset-path", + help="Explicit path to a dataset.zip (or a directory containing one)", + ) + ap.add_argument("--model", default=DEFAULT_MODEL, help=f"default: {DEFAULT_MODEL}") + ap.add_argument( + "--limit", + type=int, + default=None, + help="Only audit the first N not-yet-audited clips (use 3 for a first look)", + ) + ap.add_argument( + "--out", + default=None, + help=f"Output path (default: {OUTPUT_NAME} beside the dataset)", + ) + ap.add_argument( + "--overwrite", + action="store_true", + help="Ignore existing output and re-audit every clip", + ) + ap.add_argument( + "--exports-dir", default=None, help="Override data/exports location" + ) + ap.add_argument( + "--stage-suggestions", + action="store_true", + help=( + "Instead of auditing, read an existing ai_audit.jsonl and stage its " + "findings as PENDING model suggestions on the imported conversations " + "so they appear in the transcript editor. Requires the dataset to have " + "been imported (Upload → Annotation workspace) first." + ), + ) + ap.add_argument( + "--min-confidence", + type=float, + default=0.0, + help="When staging, only stage findings at or above this confidence (0.0-1.0)", + ) + ap.add_argument( + "--replace", + action="store_true", + help="When staging, clear existing pending model suggestions on each " + "conversation first (else duplicate segments are skipped)", + ) + args = ap.parse_args() + + if args.stage_suggestions: + import asyncio + + return asyncio.run(_stage_suggestions(args)) + + api_key = _load_api_key() + zip_path = _resolve_dataset_zip(args) + out_path = Path(args.out) if args.out else zip_path.parent / OUTPUT_NAME + + # Resume: collect clip_ids already in the output file. + done: set[str] = set() + if out_path.exists() and not args.overwrite: + for line in out_path.read_text().splitlines(): + if line.strip(): + try: + done.add(json.loads(line)["clip_id"]) + except (json.JSONDecodeError, KeyError): + pass + elif args.overwrite and out_path.exists(): + out_path.unlink() + + archive_bytes = zip_path.read_bytes() + with zipfile.ZipFile(io.BytesIO(archive_bytes)) as zf: + records = _read_manifest(zf) + audio_cache = { + r["clip_id"]: zf.read(r["audio_path"]) + for r in records + if r.get("audio_path") in zf.namelist() + } + + pending = [r for r in records if r["clip_id"] not in done] + if args.limit is not None: + pending = pending[: args.limit] + + print( + f"Dataset: {zip_path}\n" + f"Model: {args.model}\n" + f"Clips: {len(records)} total, {len(done)} already audited, " + f"{len(pending)} to do this run\n" + f"Output: {out_path}\n", + flush=True, + ) + if not pending: + print("Nothing to do. (Use --overwrite to re-audit.)") + return 0 + + from google import genai + from google.genai import types + + client = genai.Client( + api_key=api_key, + http_options=types.HttpOptions(timeout=REQUEST_TIMEOUT_MS), + ) + + ok = 0 + failed = 0 + quality_counts: dict[str, int] = {} + total_findings = 0 + t_start = time.time() + + with out_path.open("a", encoding="utf-8") as out_f: + for i, record in enumerate(pending, start=1): + clip_id = record["clip_id"] + audio_bytes = audio_cache.get(clip_id) + label = f"[{i}/{len(pending)}] {clip_id}" + if not audio_bytes: + print(f"{label} — SKIP (audio missing in ZIP)", flush=True) + failed += 1 + continue + t0 = time.time() + try: + audit = _audit_clip( + client, + types, + model=args.model, + record=record, + audio_bytes=audio_bytes, + error_dir=out_path.parent / "ai_audit_errors", + ) + except Exception as exc: # noqa: BLE001 - report + continue + print(f"{label} — ERROR: {exc}", flush=True) + failed += 1 + continue + + row = { + "clip_id": clip_id, + "conversation_id": record.get("conversation_id"), + "conversation_title": record.get("conversation_title"), + "audio_path": record.get("audio_path"), + "duration_seconds": record.get("duration_seconds"), + "model": args.model, + **audit, + } + out_f.write(json.dumps(row, ensure_ascii=False) + "\n") + out_f.flush() + + ok += 1 + n_find = len(audit["findings"]) + total_findings += n_find + quality_counts[audit["transcript_quality"]] = ( + quality_counts.get(audit["transcript_quality"], 0) + 1 + ) + print( + f"{label} — {audit['transcript_quality']}, " + f"{n_find} finding(s), {time.time() - t0:.1f}s", + flush=True, + ) + + print( + f"\nDone in {time.time() - t_start:.0f}s. " + f"audited={ok} failed={failed} total_findings={total_findings}\n" + f"quality: {quality_counts}\n" + f"Report: {out_path}", + flush=True, + ) + return 0 if failed == 0 else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/backends/advanced/scripts/backfill_cluster_embeddings.py b/backends/advanced/scripts/backfill_cluster_embeddings.py new file mode 100644 index 00000000..77dac1e1 --- /dev/null +++ b/backends/advanced/scripts/backfill_cluster_embeddings.py @@ -0,0 +1,41 @@ +"""One-time backfill of per-cluster speaker centroids for existing conversations. + +Embeds one centroid per diarized speaker from each conversation's already-stored segment +boundaries (no re-diarization) and saves it to +``TranscriptVersion.metadata["cluster_centroids"]``. This lets the "reprocess impact" +finder re-identify past conversations against the current gallery with pure vector math. + +GPU-bound (runs the wespeaker embedder via the speaker service). Idempotent: skips +conversations that already have stored centroids unless you flip only_missing. + +Run inside the backend container: + python3 /app/scripts/backfill_cluster_embeddings.py # all missing + python3 /app/scripts/backfill_cluster_embeddings.py 5 # first 5 (smoke test) +""" + +import asyncio +import sys + +from beanie import init_beanie + +from advanced_omi_backend.controllers.drift_controller import ( + backfill_cluster_embeddings, +) +from advanced_omi_backend.database import db +from advanced_omi_backend.models.audio_chunk import AudioChunkDocument +from advanced_omi_backend.models.conversation import Conversation +from advanced_omi_backend.models.user import User + + +async def main() -> None: + limit = int(sys.argv[1]) if len(sys.argv) > 1 else None + # AudioChunkDocument is needed because backfill reconstructs conversation audio. + await init_beanie( + database=db, document_models=[User, Conversation, AudioChunkDocument] + ) + result = await backfill_cluster_embeddings(limit=limit, only_missing=True) + print(f"Backfill result: {result}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/backends/advanced/scripts/backfill_empty_segments.py b/backends/advanced/scripts/backfill_empty_segments.py new file mode 100644 index 00000000..3c64cdd3 --- /dev/null +++ b/backends/advanced/scripts/backfill_empty_segments.py @@ -0,0 +1,90 @@ +"""Backfill renderable segments for transcript versions that have text but no segments. + +The streaming transcription path used to store a transcript version with text + words +but an EMPTY segments list when the provider didn't diarize (and speaker recognition was +off/unavailable). The web UI only renders the segment list, so those conversations showed +"No transcript segments available" despite having a transcript. + +This script finds every transcript version where `transcript` has text but `segments` is +empty, and builds a single fallback "Speaker 0" segment from the version's words (or the +full text). It is idempotent — versions that already have segments are left untouched. + +Run inside the backend/worker container: + uv run python3 scripts/backfill_empty_segments.py # dry run (default) + uv run python3 scripts/backfill_empty_segments.py --apply # write changes +""" + +import argparse +import asyncio +import sys + +from beanie import init_beanie + +from advanced_omi_backend.database import db +from advanced_omi_backend.models.conversation import Conversation +from advanced_omi_backend.models.user import User + + +def _build_fallback_segment(version) -> "Conversation.SpeakerSegment": + """Build a single full-span fallback segment from a version's words/text.""" + words = version.words or [] + start = words[0].start if words else 0.0 + end = words[-1].end if words else 0.0 + return Conversation.SpeakerSegment( + speaker="Speaker 0", + start=start, + end=end, + text=version.transcript or "", + words=list(words), + ) + + +async def main(apply: bool) -> None: + await init_beanie(database=db, document_models=[User, Conversation]) + + scanned = 0 + fixed_versions = 0 + fixed_convs = 0 + + async for conv in Conversation.find_all(): + scanned += 1 + changed = False + for version in conv.transcript_versions: + has_text = bool((version.transcript or "").strip()) + if has_text and not version.segments: + seg = _build_fallback_segment(version) + version.segments = [seg] + # Mark provenance so it's distinguishable from real diarization + version.diarization_source = version.diarization_source or None + if isinstance(version.metadata, dict): + version.metadata.setdefault( + "segments_created_by", "backfill_fallback" + ) + fixed_versions += 1 + changed = True + print( + f" conv={conv.conversation_id[:12]} version={version.version_id} " + f"-> 1 fallback segment ({len(seg.text)} chars, {len(seg.words)} words)" + ) + if changed: + fixed_convs += 1 + if apply: + await conv.save() + + print( + f"\nScanned {scanned} conversations. " + f"{'Fixed' if apply else 'Would fix'} {fixed_versions} versions " + f"across {fixed_convs} conversations." + ) + if not apply and fixed_versions: + print("Dry run — re-run with --apply to write changes.") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--apply", action="store_true", help="Write changes (default: dry run)" + ) + args = parser.parse_args() + asyncio.run(main(args.apply)) + sys.exit(0) diff --git a/backends/advanced/scripts/collapse_device_registry.js b/backends/advanced/scripts/collapse_device_registry.js new file mode 100644 index 00000000..418b48b7 --- /dev/null +++ b/backends/advanced/scripts/collapse_device_registry.js @@ -0,0 +1,86 @@ +// One-time migration: collapse the bloated `registered_clients` map. +// +// Before stable client_ids existed, every reconnect of one device minted a new +// counter-suffixed client_id (havpe, havpe-2, … havpe-286), so the registry +// accumulated hundreds of rows for a handful of physical devices. This script +// rewrites each user's `registered_clients` keyed by the STABLE client_id +// (user_suffix + sanitized device_name) — one row per device — merging the +// counter duplicates and preserving any user-set friendly name + the earliest +// first_seen / latest last_seen. +// +// IDEMPOTENT: re-running yields the same result. Run it AFTER the stable-id +// backend code is deployed (otherwise the old counter immediately re-pollutes). +// +// docker exec mongosh chronicle --file /path/collapse_device_registry.js +// DRY_RUN: set env DRYRUN=1 (default) to preview; DRYRUN=0 to apply. + +// mongosh exposes process.env +const APPLY = (typeof process !== "undefined" && process.env && process.env.DRYRUN === "0"); + +// Mirror backend generate_client_id sanitization: +// lowercase, keep [a-z0-9-], first 10 chars. +function sanitizeDevice(d) { + return String(d || "") + .toLowerCase() + .split("") + .filter((c) => /[a-z0-9-]/.test(c)) + .join("") + .slice(0, 10); +} + +function stableClientId(userIdHex, deviceName) { + return userIdHex.slice(-6) + "-" + sanitizeDevice(deviceName); +} + +let usersTouched = 0; +let totalBefore = 0; +let totalAfter = 0; + +db.users.find({ "registered_clients": { $exists: true, $ne: {} } }).forEach((u) => { + const idHex = u._id.toString(); + const rc = u.registered_clients || {}; + const keys = Object.keys(rc); + if (keys.length === 0) return; + + const collapsed = {}; + keys.forEach((k) => { + const e = rc[k] || {}; + // Without a device_name we can't derive a stable id — keep the row as-is. + const stableId = e.device_name ? stableClientId(idHex, e.device_name) : k; + + const cur = collapsed[stableId]; + if (!cur) { + collapsed[stableId] = { + client_id: stableId, + device_name: e.device_name || null, + name: e.name || e.device_name || stableId, + first_seen: e.first_seen || e.last_seen || new Date(), + last_seen: e.last_seen || e.first_seen || new Date(), + }; + } else { + // Merge duplicates: keep a real user-set name, widen the time range. + if ((!cur.name || cur.name === cur.device_name) && e.name && e.name !== e.device_name) { + cur.name = e.name; + } + if (e.first_seen && e.first_seen < cur.first_seen) cur.first_seen = e.first_seen; + if (e.last_seen && e.last_seen > cur.last_seen) cur.last_seen = e.last_seen; + } + }); + + const before = keys.length; + const after = Object.keys(collapsed).length; + totalBefore += before; + totalAfter += after; + if (after === before) return; // nothing to collapse for this user + + usersTouched += 1; + print(`${u.email}: ${before} -> ${after} [${Object.keys(collapsed).join(", ")}]`); + + if (APPLY) { + db.users.updateOne({ _id: u._id }, { $set: { registered_clients: collapsed } }); + } +}); + +print(""); +print(`${APPLY ? "APPLIED" : "DRY RUN (set DRYRUN=0 to apply)"}: ` + + `users touched=${usersTouched}, entries ${totalBefore} -> ${totalAfter}`); diff --git a/backends/advanced/scripts/evaluate_memory_executor.py b/backends/advanced/scripts/evaluate_memory_executor.py new file mode 100644 index 00000000..e32912a1 --- /dev/null +++ b/backends/advanced/scripts/evaluate_memory_executor.py @@ -0,0 +1,186 @@ +#!/usr/bin/env python3 +"""Replay a transcript canary set into a fresh vault with one memory executor. + +This is intentionally separate from the rebuild path: it never touches Mongo, queues, +the live vault, or the audit ledger. It is a small reproducible quality experiment for +comparing the direct tool-calling agent with the Codex CLI agent. +""" + +from __future__ import annotations + +import argparse +import asyncio +import contextlib +import json +import sys +import time +from pathlib import Path + +DEFAULT_PREFIXES = ( + "5c0b2333", + "99e2a020", + "7620c7b4", + "ea282e40", + "a4ed37ac", + "19f5a281", + "d0de4521", + "fd3f7f7d", +) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--executor", choices=("codex", "direct"), required=True) + parser.add_argument("--dataset", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--model", help="Codex model override recorded in the manifest") + parser.add_argument( + "--reasoning-effort", + choices=("none", "minimal", "low", "medium", "high", "xhigh", "max", "ultra"), + help="Codex reasoning-effort override recorded in the manifest", + ) + parser.add_argument( + "--conversation", + action="append", + dest="prefixes", + help="conversation id or unique prefix; repeatable (defaults to audited 8)", + ) + parser.add_argument( + "--keep-output", + action="store_true", + help="reuse the output vault; default refuses a non-empty destination", + ) + return parser.parse_args() + + +def _load_rows(dataset: Path, prefixes: tuple[str, ...]) -> list[dict]: + rows = [ + json.loads(line) for line in dataset.read_text().splitlines() if line.strip() + ] + selected = [] + for prefix in prefixes: + matches = [row for row in rows if row["conversation_id"].startswith(prefix)] + if len(matches) != 1: + raise SystemExit( + f"{prefix!r} matched {len(matches)} conversations (expected 1)" + ) + selected.append(matches[0]) + return selected + + +def _prepare_output(output: Path, keep: bool) -> None: + if output.exists() and any(output.iterdir()) and not keep: + raise SystemExit( + f"refusing non-empty output {output}; remove it or pass --keep-output" + ) + output.mkdir(parents=True, exist_ok=True) + + +@contextlib.contextmanager +def _isolated_vault_lock(*_args, **_kwargs): + """No Redis dependency: this process is the sole writer of a throwaway vault.""" + yield + + +async def _run(args: argparse.Namespace) -> int: + # The script lives at backends/advanced/scripts; make src imports work when invoked + # directly without requiring an editable install. + backend = Path(__file__).resolve().parents[1] + sys.path.insert(0, str(backend / "src")) + + from advanced_omi_backend.services.memory import vault_lock + from advanced_omi_backend.services.memory.agent import codex_agent + from advanced_omi_backend.services.memory.agent.codex_agent import CodexMemoryAgent + from advanced_omi_backend.services.memory.agent.memory_agent import MemoryAgent + from advanced_omi_backend.services.memory.vault_scaffold import seed_vault_scaffold + + # Production uses Redis to serialize writers. This evaluator owns a unique fresh + # directory and runs sequentially, so requiring the application stack would add no + # safety and would make the experiment needlessly hard to reproduce. + vault_lock.vault_run_lock = _isolated_vault_lock + vault_lock.vault_note_lock = _isolated_vault_lock + configured = codex_agent._codex_settings() if args.executor == "codex" else {} + if args.executor == "codex" and (args.model or args.reasoning_effort): + configured.update( + { + key: value + for key, value in { + "model": args.model, + "reasoning_effort": args.reasoning_effort, + }.items() + if value + } + ) + codex_agent._codex_settings = lambda: configured + + # Codex receives the vault as both subprocess cwd and --cd. Resolve once so a + # caller's relative path is not interpreted relative to itself by the CLI. + args.output = args.output.resolve() + args.dataset = args.dataset.resolve() + prefixes = tuple(args.prefixes or DEFAULT_PREFIXES) + rows = _load_rows(args.dataset, prefixes) + _prepare_output(args.output, args.keep_output) + seed_vault_scaffold(args.output) + agent_type = CodexMemoryAgent if args.executor == "codex" else MemoryAgent + + manifest = { + "executor": args.executor, + "model": args.model + or (configured.get("model") if args.executor == "codex" else None), + "reasoning_effort": args.reasoning_effort + or (configured.get("reasoning_effort") if args.executor == "codex" else None), + "dataset": str(args.dataset.resolve()), + "output": str(args.output.resolve()), + "started_at_epoch": time.time(), + "runs": [], + } + manifest_path = args.output / "evaluation-manifest.json" + failures = 0 + for index, row in enumerate(rows, 1): + conversation_id = row["conversation_id"] + print( + f"[{index}/{len(rows)}] {args.executor}: {conversation_id} ({row['n_chars']} chars)", + flush=True, + ) + started = time.perf_counter() + result = await agent_type(args.output).run( + row["transcript"], + conversation_id, + date=row.get("created_at"), + duration_minutes=(row.get("duration_s") or 0) / 60, + title=row.get("title"), + ) + note = args.output / "Conversations" / f"{conversation_id}.md" + ok = note.is_file() and not result.truncated and not result.stalled + failures += int(not ok) + manifest["runs"].append( + { + "conversation_id": conversation_id, + "ok": ok, + "elapsed_seconds": round(time.perf_counter() - started, 3), + "rounds": result.rounds, + "tool_calls": result.tool_calls, + "touched": result.touched, + "removed": result.removed, + "errors": result.errors, + "truncated": result.truncated, + "stalled": result.stalled, + "summary": result.summary, + } + ) + manifest_path.write_text(json.dumps(manifest, indent=2) + "\n") + + manifest["finished_at_epoch"] = time.time() + manifest["failures"] = failures + manifest_path.write_text(json.dumps(manifest, indent=2) + "\n") + print(f"vault: {args.output.resolve()}\nmanifest: {manifest_path.resolve()}") + return int(bool(failures)) + + +def main() -> int: + args = _parse_args() + return asyncio.run(_run(args)) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/backends/advanced/scripts/purge_orphaned_deferred_jobs.py b/backends/advanced/scripts/purge_orphaned_deferred_jobs.py new file mode 100644 index 00000000..c52e2f3e --- /dev/null +++ b/backends/advanced/scripts/purge_orphaned_deferred_jobs.py @@ -0,0 +1,46 @@ +"""Manual cleanup: delete orphaned deferred RQ jobs. + +A deferred job is "orphaned" when nothing will ever promote it: none of its +dependencies is still pending (every dependency is either missing from Redis — +evicted/deleted — or already terminal). The detection + deletion logic lives in +:mod:`advanced_omi_backend.services.job_reaper` and is shared with the periodic +backstop reaper (services/reaper.py); this is just a CLI wrapper for an on-demand +sweep. + +Usage (run where the backend package is importable): + python scripts/purge_orphaned_deferred_jobs.py # dry run + python scripts/purge_orphaned_deferred_jobs.py --delete # actually delete +""" + +import sys + +from advanced_omi_backend.services.job_reaper import ( + find_orphaned_deferred_jobs, + reap_orphaned_deferred_jobs, +) + + +def main() -> int: + if "--delete" not in sys.argv: + orphans = find_orphaned_deferred_jobs() + print(f"Found {len(orphans)} orphaned deferred job(s):") + for queue_name, job_id, conv, reason in orphans: + print(f" [{queue_name}] {job_id} conv={conv} :: {reason}") + print( + "\nDRY RUN — re-run with --delete to remove them (deletion cascades " + "through dependent chains)." + ) + return 0 + + result = reap_orphaned_deferred_jobs() + for d in result["details"]: + print(f" deleted [{d['queue']}] {d['job_id']} conv={d['conversation_id']}") + print(f"\nDeleted {result['deleted']} orphaned deferred job(s) total.") + remaining = find_orphaned_deferred_jobs() + if remaining: + print(f"WARNING: {len(remaining)} orphan(s) still remain.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/backends/advanced/scripts/regen_titles_batch.py b/backends/advanced/scripts/regen_titles_batch.py new file mode 100644 index 00000000..6e8d597d --- /dev/null +++ b/backends/advanced/scripts/regen_titles_batch.py @@ -0,0 +1,123 @@ +"""Batch-regenerate titles/summaries for completed conversations stuck with a +placeholder title (e.g. "Recording...", "Audio Recording (Transcription Failed)"). + +These conversations have a real transcript but never got an LLM title (completed via +a path that skipped generate_title_summary, or the title job failed). This enqueues +generate_title_summary_job (reads the existing transcript — NO re-transcription) in +batches of BATCH_SIZE, polling until each title changes off the placeholder. + +Run inside the backend/worker container: + python3 /app/regen_titles_batch.py # dry run (lists targets) + python3 /app/regen_titles_batch.py --apply # enqueue + poll +""" + +import argparse +import asyncio +import sys +import time + +from beanie import init_beanie + +from advanced_omi_backend.database import db +from advanced_omi_backend.models.conversation import Conversation +from advanced_omi_backend.models.user import User + +BATCH_SIZE = 5 +POLL_SECS = 20 +MAX_WAIT_SECS = 8 * 60 + +PLACEHOLDER_TITLES = ("Reprocessing...", "Recording...", "Transcribing...") + + +def _is_placeholder_title(title: str) -> bool: + return "Audio Recording (" in title or title in PLACEHOLDER_TITLES + + +def _is_placeholder_conv(c: "Conversation") -> bool: + t = c.title or "" + return bool(t) and _is_placeholder_title(t) + + +async def _targets() -> list: + out = [] + async for c in Conversation.find_all(): + if ( + not c.deleted + and c.processing_status == Conversation.ConversationStatus.COMPLETED.value + and _is_placeholder_conv(c) + ): + out.append(c.conversation_id) + return out + + +async def _title_of(cid: str): + c = await Conversation.find_one(Conversation.conversation_id == cid) + return (c.title or "") if c else "" + + +async def main(apply: bool) -> None: + await init_beanie(database=db, document_models=[User, Conversation]) + ids = await _targets() + print(f"=== {len(ids)} completed convs with placeholder titles ===", flush=True) + for cid in ids: + print(f" {cid[:12]} title={await _title_of(cid)!r}", flush=True) + if not apply: + print("Dry run — re-run with --apply to enqueue.", flush=True) + return + + # Import the queue lazily (module-level side effects: redis/queue setup). + from advanced_omi_backend.controllers.queue_controller import default_queue + from advanced_omi_backend.workers.conversation_jobs import ( + generate_title_summary_job, + ) + + done = {} + nbatches = (len(ids) + BATCH_SIZE - 1) // BATCH_SIZE + for b in range(nbatches): + batch = ids[b * BATCH_SIZE : (b + 1) * BATCH_SIZE] + print(f"--- BATCH {b+1}/{nbatches}: {[c[:8] for c in batch]} ---", flush=True) + for cid in batch: + default_queue.enqueue( + generate_title_summary_job, + cid, + job_id=f"regen_title_{cid[:12]}", + job_timeout=300, + description=f"Regenerate title/summary for {cid[:8]}", + ) + print(f" enqueued {cid[:8]}", flush=True) + + deadline = time.time() + MAX_WAIT_SECS + while True: + await asyncio.sleep(POLL_SECS) + pending = [] + for cid in batch: + t = await _title_of(cid) + if (not t) or _is_placeholder_title(t): + pending.append(cid) + else: + done[cid] = t + print( + f" poll: {len(batch)-len(pending)}/{len(batch)} retitled; " + f"pending={[c[:8] for c in pending]}", + flush=True, + ) + if not pending or time.time() > deadline: + if pending: + print( + f" batch {b+1}: TIMEOUT, pending={[c[:8] for c in pending]}", + flush=True, + ) + break + + print("=== FINAL ===", flush=True) + for cid in ids: + print(f" {cid[:12]} -> {done.get(cid, '(unchanged)')!r}", flush=True) + print(f"Retitled {len(done)}/{len(ids)}.", flush=True) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--apply", action="store_true") + args = parser.parse_args() + asyncio.run(main(args.apply)) + sys.exit(0) diff --git a/backends/advanced/scripts/reprocess_conversation.py b/backends/advanced/scripts/reprocess_conversation.py new file mode 100644 index 00000000..8a39185d --- /dev/null +++ b/backends/advanced/scripts/reprocess_conversation.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 +"""Reprocess a conversation's transcript via the backend API. + +This triggers POST /api/conversations/{id}/reprocess-transcript, which re-runs +batch transcription using the *currently configured* default batch STT provider +(config/config.yml -> defaults.stt) and then the full post-conversation chain +(speaker recognition -> memory -> title/summary -> dispatch complete). + +There is no per-reprocess provider override in the backend, so the "config" is +whatever defaults.stt points at. To reprocess with smallest.ai Pulse (hi), set + defaults.stt: stt-smallest +in config/config.yml and restart backend + workers, then run this script. + +Usage: + uv run python3 scripts/reprocess_conversation.py [ ...] + +Env (optional, defaults shown): + BACKEND_URL=http://localhost:8000 + ADMIN_EMAIL / ADMIN_PASSWORD (else read from backends/advanced/.env) +""" + +import json +import os +import sys +import urllib.error +import urllib.parse +import urllib.request +from pathlib import Path + +BACKEND_URL = os.environ.get("BACKEND_URL", "http://localhost:8000").rstrip("/") +ENV_FILE = Path(__file__).resolve().parent.parent / ".env" + + +def _read_env_file(path: Path) -> dict: + values = {} + if not path.exists(): + return values + for line in path.read_text().splitlines(): + line = line.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, _, val = line.partition("=") + values[key.strip()] = val.strip().strip("'").strip('"') + return values + + +def _creds() -> tuple[str, str]: + env = _read_env_file(ENV_FILE) + email = os.environ.get("ADMIN_EMAIL") or env.get("ADMIN_EMAIL") + password = os.environ.get("ADMIN_PASSWORD") or env.get("ADMIN_PASSWORD") + if not email or not password: + sys.exit("ERROR: ADMIN_EMAIL / ADMIN_PASSWORD not found (env or .env).") + return email, password + + +def login(email: str, password: str) -> str: + data = urllib.parse.urlencode({"username": email, "password": password}).encode() + req = urllib.request.Request( + f"{BACKEND_URL}/auth/jwt/login", + data=data, + headers={"Content-Type": "application/x-www-form-urlencoded"}, + method="POST", + ) + with urllib.request.urlopen(req, timeout=30) as resp: + token = json.loads(resp.read())["access_token"] + print(f"[auth] logged in as {email}") + return token + + +def reprocess(conversation_id: str, token: str) -> None: + req = urllib.request.Request( + f"{BACKEND_URL}/api/conversations/{conversation_id}/reprocess-transcript", + data=b"", + headers={"Authorization": f"Bearer {token}"}, + method="POST", + ) + try: + with urllib.request.urlopen(req, timeout=60) as resp: + body = json.loads(resp.read()) + print(f"[ok] {conversation_id}: {json.dumps(body, indent=2)}") + except urllib.error.HTTPError as e: + detail = e.read().decode(errors="replace") + print(f"[FAIL] {conversation_id}: HTTP {e.code} -> {detail}") + + +def main() -> None: + ids = sys.argv[1:] + if not ids: + sys.exit(__doc__) + token = login(*_creds()) + for conversation_id in ids: + reprocess(conversation_id, token) + + +if __name__ == "__main__": + main() diff --git a/backends/advanced/scripts/scope_blank_segments.py b/backends/advanced/scripts/scope_blank_segments.py new file mode 100644 index 00000000..0bbd63bf --- /dev/null +++ b/backends/advanced/scripts/scope_blank_segments.py @@ -0,0 +1,68 @@ +"""Read-only scope: find conversations whose ACTIVE transcript version has words/text +but an empty segments list (so the WebUI renders blank). Breaks them down by +diarization_source to distinguish pyannote-wiped from provider/never-diarized. + +Run inside the backend/worker container (read-only, no writes): + python3 /app/scope_blank_segments.py +""" + +import asyncio +import sys +from collections import Counter + +from beanie import init_beanie + +from advanced_omi_backend.database import db +from advanced_omi_backend.models.conversation import Conversation +from advanced_omi_backend.models.user import User + + +async def main() -> None: + await init_beanie(database=db, document_models=[User, Conversation]) + + by_diar = Counter() + rows = [] + async for c in Conversation.find_all(): + if c.deleted: + continue + av = c.active_transcript + if not av: + continue + has_text = bool((av.transcript or "").strip()) + words = av.words or [] + segs = av.segments or [] + if (has_text or words) and not segs: + diar = av.diarization_source or "(none)" + by_diar[diar] += 1 + word_spk = sorted({w.speaker for w in words if w.speaker is not None}) + rows.append( + ( + c.conversation_id, + c.processing_status, + c.client_id, + round(c.audio_total_duration or 0, 1), + len(words), + word_spk, + diar, + av.provider, + (c.title or "")[:40], + ) + ) + + print( + f"=== {len(rows)} non-deleted convs: active version has words/text but 0 segments ===" + ) + print(f"by diarization_source: {dict(by_diar)}\n") + rows.sort(key=lambda r: -r[3]) + print( + f"{'conv':10} {'status':10} {'client':20} {'dur':>7} {'wrds':>4} {'wordspk':12} {'diar':10} {'provider':10} title" + ) + for cid, st, client, dur, nw, wspk, diar, prov, title in rows: + print( + f"{cid[:8]:10} {str(st):10} {str(client):20} {dur:>7} {nw:>4} {str(wspk):12} {str(diar):10} {str(prov):10} {title!r}" + ) + + +if __name__ == "__main__": + asyncio.run(main()) + sys.exit(0) diff --git a/backends/advanced/scripts/settle_stuck_status.py b/backends/advanced/scripts/settle_stuck_status.py new file mode 100644 index 00000000..d0870e07 --- /dev/null +++ b/backends/advanced/scripts/settle_stuck_status.py @@ -0,0 +1,92 @@ +"""Settle conversations stuck in a non-terminal processing_status. + +The status reconciler (services/status_reconciler.py) only scans non-deleted +conversations, and the no-speech/dead-end paths used to leave processing_status at +"active" (with a stale "Reprocessing..."/"Audio Recording (...)" title) — see the +fixes in mark_conversation_deleted and reprocess_speakers. That left a backlog of +conversations (mostly soft-deleted) stuck "active"/None/legacy "transcription_failed". + +This one-off applies the SAME fact-derived logic (Conversation.apply_status, the single +owner of the field) to EVERY conversation — including deleted ones — and clears stale +placeholder titles. Idempotent: terminal, correctly-titled conversations are untouched. + +Run inside the backend/worker container: + uv run python3 scripts/settle_stuck_status.py # dry run (default) + uv run python3 scripts/settle_stuck_status.py --apply # write changes +""" + +import argparse +import asyncio +import sys + +from beanie import init_beanie + +from advanced_omi_backend.database import db +from advanced_omi_backend.models.conversation import Conversation +from advanced_omi_backend.models.user import User + +TERMINAL = { + Conversation.ConversationStatus.COMPLETED.value, + Conversation.ConversationStatus.FAILED.value, +} +PLACEHOLDER_TITLES = ("Reprocessing...", "Recording...", "Transcribing...") + + +def _is_placeholder_title(title: str) -> bool: + # "Audio Recording (Processing...)" / "(Batch Transcription...)" / "(Transcription + # Failed)" — require the parenthetical so real LLM titles like "Audio Recording + # Issues" are NOT matched. + return "Audio Recording (" in title or title in PLACEHOLDER_TITLES + + +async def main(apply: bool) -> None: + await init_beanie(database=db, document_models=[User, Conversation]) + + scanned = settled = titled = 0 + async for conv in Conversation.find_all(): + scanned += 1 + changed = False + + # Only touch non-terminal conversations (active / None / legacy strings) — + # leave already-settled completed/failed ones (and their titles) alone. + if conv.processing_status not in TERMINAL: + # Settle status from facts. apply_status: has transcript -> completed; + # settled & none -> failed; else active. + before = conv.processing_status + if conv.apply_status(settled=True): + settled += 1 + changed = True + print( + f" conv={conv.conversation_id[:12]} deleted={conv.deleted} " + f"status {before!r} -> {conv.processing_status!r}" + + (f" stage={conv.failure_stage}" if conv.failure_stage else "") + ) + + # Clear stale in-flight placeholder titles on these dead-end convs. + title = conv.title or "" + if title and _is_placeholder_title(title): + conv.title = None + titled += 1 + changed = True + print(f" conv={conv.conversation_id[:12]} cleared title {title!r}") + + if changed and apply: + await conv.save() + + verb = "Settled" if apply else "Would settle" + print( + f"\nScanned {scanned} conversations. {verb} {settled} statuses, " + f"cleared {titled} placeholder titles." + ) + if not apply and (settled or titled): + print("Dry run — re-run with --apply to write changes.") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--apply", action="store_true", help="Write changes (default: dry run)" + ) + args = parser.parse_args() + asyncio.run(main(args.apply)) + sys.exit(0) diff --git a/backends/advanced/src/advanced_omi_backend/controllers/audio_controller.py b/backends/advanced/src/advanced_omi_backend/controllers/audio_controller.py index 4e10408e..1dcfaec9 100644 --- a/backends/advanced/src/advanced_omi_backend/controllers/audio_controller.py +++ b/backends/advanced/src/advanced_omi_backend/controllers/audio_controller.py @@ -14,6 +14,7 @@ from fastapi import UploadFile from fastapi.responses import JSONResponse +from rq import Retry from advanced_omi_backend.controllers.queue_controller import ( JOB_RESULT_TTL, @@ -49,6 +50,7 @@ async def upload_and_process_audio_files( files: list[UploadFile], device_name: str = "upload", source: str = "upload", + annotation_only: bool = False, ) -> dict: """ Upload audio files and process them directly. @@ -63,6 +65,7 @@ async def upload_and_process_audio_files( files: List of uploaded audio files device_name: Device identifier source: Source of the upload (e.g., 'upload', 'gdrive') + annotation_only: Create editable transcription records without memory extraction """ try: if not files: @@ -159,9 +162,18 @@ async def upload_and_process_audio_files( user_id=user.user_id, client_id=client_id, title=title, - summary="Processing uploaded audio file...", + summary=( + "Processing annotation-only audio file..." + if annotation_only + else "Processing uploaded audio file..." + ), external_source_id=external_source_id, external_source_type=external_source_type, + data_purpose="annotation" if annotation_only else None, + memory_excluded=annotation_only, + memory_exclusion_reason=( + "annotation_only_upload" if annotation_only else None + ), ) await conversation.insert() conversation_id = ( @@ -229,6 +241,9 @@ async def upload_and_process_audio_files( job_timeout=-1, result_ttl=JOB_RESULT_TTL, job_id=transcribe_job_id, + # Bulk uploads can trip provider rate limits (HTTP 429); + # spread retries out so the batch drains instead of failing. + retry=Retry(max=4, interval=[60, 300, 900, 1800]), description=f"Transcribe uploaded file {conversation_id[:8]}", meta={ "conversation_id": conversation_id, @@ -251,12 +266,14 @@ async def upload_and_process_audio_files( transcript_version_id=version_id, # Pass the version_id from transcription job depends_on_job=transcription_job, # Wait for transcription to complete (or None) client_id=client_id, # Pass client_id for UI tracking + skip_memory_extraction=annotation_only, ) file_result = { "filename": filename, "status": "started", # RQ standard: job has been enqueued "conversation_id": conversation_id, + "annotation_only": annotation_only, "transcript_job_id": ( transcription_job.id if transcription_job else None ), @@ -309,6 +326,7 @@ async def upload_and_process_audio_files( response_body = { "message": f"Uploaded and processing {len(successful_files)} file(s)", "client_id": client_id, + "annotation_only": annotation_only, "files": processed_files, "summary": { "total": len(files), diff --git a/backends/advanced/src/advanced_omi_backend/controllers/data_audit_controller.py b/backends/advanced/src/advanced_omi_backend/controllers/data_audit_controller.py index e4c66d90..96aeb7c8 100644 --- a/backends/advanced/src/advanced_omi_backend/controllers/data_audit_controller.py +++ b/backends/advanced/src/advanced_omi_backend/controllers/data_audit_controller.py @@ -9,6 +9,7 @@ import json import logging +import re import shutil import statistics import uuid @@ -40,10 +41,19 @@ new_export_id, validate_export_id, ) +from advanced_omi_backend.utils.annotation_import import ( + AnnotationDatasetError, + parse_annotation_dataset, +) from advanced_omi_backend.utils.audio_chunk_utils import ( audio_cache_duration_matches, + convert_audio_to_chunks, reconstruct_audio_segment, ) +from advanced_omi_backend.utils.audio_utils import ( + AudioValidationError, + validate_and_prepare_audio, +) from advanced_omi_backend.utils.transcript_slicing import ( build_transcript_text, shift_segments, @@ -88,6 +98,8 @@ "archive_reason": 1, "processing_status": 1, "failure_stage": 1, + "external_source_id": 1, + "external_source_type": 1, "vad_analysis": 1, "derived_from": 1, "active_transcript_version": 1, @@ -196,6 +208,7 @@ async def list_for_audit( created_before: Optional[datetime] = None, include_speakers: Optional[List[str]] = None, exclude_speakers: Optional[List[str]] = None, + dataset_id: Optional[str] = None, archived_only: bool = False, hide_failed: bool = False, hide_reviewed: bool = False, @@ -234,6 +247,10 @@ async def list_for_audit( "$ne": Conversation.ConversationStatus.FAILED.value } + if dataset_id: + base["external_source_type"] = "annotation_dataset" + base["external_source_id"] = {"$regex": f"^{re.escape(dataset_id)}:"} + # Date range goes into the Mongo query (not the Python predicate) so # it narrows the MAX_SCAN working set instead of competing with it. if created_after or created_before: @@ -253,6 +270,28 @@ async def list_for_audit( raw_docs = await cursor.to_list(length=MAX_SCAN) scan_capped = len(raw_docs) >= MAX_SCAN + dataset_base: dict = {} if user.is_superuser else {"user_id": str(user.user_id)} + dataset_base.update( + { + "external_source_type": "annotation_dataset", + "audio_archived": {"$ne": True}, + "deleted": {"$ne": True}, + "audio_chunks_count": {"$gt": 0}, + } + ) + dataset_docs = await ( + collection.find(dataset_base, {"external_source_id": 1}) + .sort("created_at", -1) + .limit(MAX_SCAN) + ).to_list(length=MAX_SCAN) + available_datasets = list( + dict.fromkeys( + source_id.rsplit(":", 1)[0] + for doc in dataset_docs + if (source_id := doc.get("external_source_id")) and ":" in source_id + ) + ) + # Match threshold the pipeline used + a small comfort margin: an # identification within this band of the cutoff is a weak/suspect match # (the "low-confidence" review signal), computed from stored confidence. @@ -389,6 +428,7 @@ async def list_for_audit( "marginal_margin": marginal_margin, "unanalyzed_count": unanalyzed_count, "speakers": sorted(available_speakers), + "datasets": available_datasets, } except Exception as e: @@ -1458,10 +1498,141 @@ async def merge_conversations(user: User, conversation_ids: List[str]): # --------------------------------------------------------------------------- -# Annotation dataset export +# Annotation dataset import / export # --------------------------------------------------------------------------- +async def import_annotation_dataset(user: User, archive_bytes: bytes): + """Import an export-compatible ZIP as isolated, editor-ready conversations.""" + try: + dataset = parse_annotation_dataset(archive_bytes) + except AnnotationDatasetError as exc: + return JSONResponse(status_code=422, content={"error": str(exc)}) + + client_id = f"{str(user.id)[-6:]}-annotation-import" + results = [] + for clip in dataset.clips: + external_source_id = f"{dataset.dataset_id}:{clip.clip_id}" + existing = await Conversation.find_one( + Conversation.user_id == user.user_id, + Conversation.external_source_type == "annotation_dataset", + Conversation.external_source_id == external_source_id, + ) + if existing: + results.append( + { + "clip_id": clip.clip_id, + "status": "skipped", + "reason": "already_imported", + "conversation_id": existing.conversation_id, + } + ) + continue + + conversation = None + try: + audio_data, sample_rate, sample_width, channels, duration = ( + await validate_and_prepare_audio( + audio_data=clip.audio_bytes, + expected_sample_rate=16000, + convert_to_mono=True, + auto_resample=True, + ) + ) + segments = [ + Conversation.SpeakerSegment(**segment) for segment in clip.segments + ] + conversation = create_conversation( + user_id=user.user_id, + client_id=client_id, + title=clip.conversation_title, + summary="Imported annotation dataset; excluded from user memory.", + external_source_id=external_source_id, + external_source_type="annotation_dataset", + data_purpose="annotation", + memory_excluded=True, + memory_exclusion_reason="annotation_dataset_import", + ) + version_id = str(uuid.uuid4()) + conversation.add_transcript_version( + version_id=version_id, + transcript=clip.transcript, + segments=segments, + provider="annotation-import", + model=f"chronicle-dataset-v{dataset.schema_version}", + metadata={ + "dataset_id": dataset.dataset_id, + "clip_id": clip.clip_id, + "source_conversation_id": clip.source_conversation_id, + "source_client_id": clip.source_client_id, + "source_audio_path": clip.audio_path, + "annotation_notes": clip.notes, + }, + set_as_active=True, + ) + conversation.apply_status(settled=bool(clip.transcript.strip())) + await conversation.insert() + + chunk_count = await convert_audio_to_chunks( + conversation_id=conversation.conversation_id, + audio_data=audio_data, + sample_rate=sample_rate, + channels=channels, + sample_width=sample_width, + ) + results.append( + { + "clip_id": clip.clip_id, + "status": "imported", + "conversation_id": conversation.conversation_id, + "duration_seconds": round(duration, 2), + "chunk_count": chunk_count, + "transcript_source": clip.transcript_source, + } + ) + except (AudioValidationError, ValueError) as exc: + logger.warning(f"Could not import annotation clip {clip.clip_id}: {exc}") + if conversation and conversation.id: + await AudioChunkDocument.find( + AudioChunkDocument.conversation_id == conversation.conversation_id + ).delete() + await conversation.delete() + results.append( + {"clip_id": clip.clip_id, "status": "error", "error": str(exc)} + ) + except Exception as exc: + logger.exception(f"Could not import annotation clip {clip.clip_id}") + if conversation and conversation.id: + await AudioChunkDocument.find( + AudioChunkDocument.conversation_id == conversation.conversation_id + ).delete() + await conversation.delete() + results.append( + {"clip_id": clip.clip_id, "status": "error", "error": str(exc)} + ) + + imported = sum(result["status"] == "imported" for result in results) + skipped = sum(result["status"] == "skipped" for result in results) + failed = sum(result["status"] == "error" for result in results) + response = { + "dataset_id": dataset.dataset_id, + "schema_version": dataset.schema_version, + "message": f"Imported {imported} annotation clip(s)", + "results": results, + "summary": { + "total": len(results), + "imported": imported, + "skipped": skipped, + "failed": failed, + }, + } + if failed == len(results): + return JSONResponse(status_code=400, content=response) + if failed: + return JSONResponse(status_code=207, content=response) + return response + + async def start_screening( user: User, conversation_ids: List[str], diff --git a/backends/advanced/src/advanced_omi_backend/controllers/drift_controller.py b/backends/advanced/src/advanced_omi_backend/controllers/drift_controller.py index d803c1d5..5a27261a 100644 --- a/backends/advanced/src/advanced_omi_backend/controllers/drift_controller.py +++ b/backends/advanced/src/advanced_omi_backend/controllers/drift_controller.py @@ -13,7 +13,7 @@ import logging from collections import Counter -from typing import Optional +from typing import Callable, Optional from advanced_omi_backend.config import get_diarization_settings from advanced_omi_backend.models.conversation import Conversation @@ -116,26 +116,33 @@ async def find_drift_conversations() -> dict: async def backfill_cluster_embeddings( - limit: Optional[int] = None, only_missing: bool = True + limit: Optional[int] = None, + only_missing: bool = True, + progress_callback: Optional[Callable[[int, int, int, int, int], None]] = None, ) -> dict: """One-time: embed per-cluster centroids for conversations that lack them. Reconstructs each conversation's audio and pools one centroid per existing diarized speaker (no re-diarization) via ``/v1/embed-clusters``, storing the result keyed by - the segments' display labels. GPU-bound (runs the embedder); intended to be invoked - from a script inside the backend container. + the segments' display labels. GPU-bound (runs the embedder); intended to run as a + background job or from the maintenance script inside the backend container. """ client = SpeakerRecognitionClient() convs = await Conversation.find({"deleted": {"$ne": True}}).to_list() done = skipped = failed = 0 + total = len(convs) for conv in convs: version = conv.active_transcript if not version or not version.segments: skipped += 1 + if progress_callback: + progress_callback(done + skipped + failed, total, done, skipped, failed) continue if only_missing and (version.metadata or {}).get("cluster_centroids"): skipped += 1 + if progress_callback: + progress_callback(done + skipped + failed, total, done, skipped, failed) continue speech = _speech_segments(version) @@ -145,6 +152,8 @@ async def backfill_cluster_embeddings( ] if not diar: skipped += 1 + if progress_callback: + progress_callback(done + skipped + failed, total, done, skipped, failed) continue max_end = max(d["end"] for d in diar) @@ -157,6 +166,8 @@ async def backfill_cluster_embeddings( "audio reconstruct failed for %s: %s", conv.conversation_id[:8], e ) failed += 1 + if progress_callback: + progress_callback(done + skipped + failed, total, done, skipped, failed) continue resp = await client.embed_clusters(audio, diar) @@ -167,6 +178,8 @@ async def backfill_cluster_embeddings( resp.get("error") or resp.get("message"), ) failed += 1 + if progress_callback: + progress_callback(done + skipped + failed, total, done, skipped, failed) continue if not version.metadata: @@ -179,6 +192,8 @@ async def backfill_cluster_embeddings( conv.conversation_id[:8], len(resp["clusters"]), ) + if progress_callback: + progress_callback(done + skipped + failed, total, done, skipped, failed) if limit and done >= limit: break diff --git a/backends/advanced/src/advanced_omi_backend/controllers/guided_annotation_controller.py b/backends/advanced/src/advanced_omi_backend/controllers/guided_annotation_controller.py new file mode 100644 index 00000000..aaf92330 --- /dev/null +++ b/backends/advanced/src/advanced_omi_backend/controllers/guided_annotation_controller.py @@ -0,0 +1,352 @@ +"""Guided annotation — active-learning selection of conversation windows to +ground-truth for speaker identity and boundaries. + +The guided-enrollment sibling: enrollment picks *clips for one speaker's +gallery*; this picks *conversation windows for human annotation* (speaker +boundaries + identities in the transcript editor). The goal is the most +model-information per minute of annotation effort, so windows are ranked by a +per-segment informativeness sum normalized by window duration. + +Per-segment informativeness (stored transcript data only — no audio pass): + * label uncertainty — unlabeled segments carry maximal label entropy; + attributed segments score by proximity of the stored cosine confidence to + the operating threshold (the decision boundary, where a human label + resolves the most uncertainty), + * shortness — sub-4s segments are the pipeline's dominant failure regime + (short-duration verification error grows ~2.4x at 2s), so confirming them + is worth more than confirming long clear turns, + * overlap — segments overlapping a different speaker mark boundary regions + where diarization ground truth is scarcest. + +Windows already covered by human annotations are down-weighted, decided +windows are never re-suggested (``annotation_review_targets``), and batches +round-robin across speaker signatures so one frequent pair cannot monopolize +the queue (session/pair diversity beats more-of-the-same, as in enrollment). +""" + +import logging +from datetime import datetime, timezone +from typing import List, Optional + +from advanced_omi_backend.config import get_diarization_settings +from advanced_omi_backend.models.conversation import Conversation +from advanced_omi_backend.users import User + +logger = logging.getLogger(__name__) + +MIN_WINDOW_SECONDS = 30.0 +MAX_WINDOW_SECONDS = 300.0 +WINDOW_STEP_SECONDS = 30.0 +MAX_WINDOWS_PER_CONVERSATION = 2 +UNCERTAINTY_BAND = 0.25 +SHORT_SEGMENT_CAP = 4.0 +MIN_WINDOW_SEGMENTS = 3 + +W_LABEL, W_SHORT, W_OVERLAP = 0.55, 0.25, 0.20 + +HUMAN_ANNOTATION_TYPES = ["diarization", "timing", "insert", "deletion"] + + +def _targets_collection(): + return Conversation.get_pymongo_collection().database["annotation_review_targets"] + + +def _annotations_collection(): + return Conversation.get_pymongo_collection().database["annotations"] + + +def _target_key(conversation_id: str, window_start: float) -> str: + return f"{conversation_id}:{round(window_start, 1)}" + + +def _active_segments(doc: dict) -> list: + versions = doc.get("transcript_versions") or [] + active_id = doc.get("active_transcript_version") + active = next((v for v in versions if v.get("version_id") == active_id), None) + if active is None and versions: + active = versions[-1] + return (active or {}).get("segments") or [] + + +def _segment_info(seg: dict, segments: list, threshold: float) -> Optional[dict]: + """Informativeness of confirming one segment's speaker, or None for + segments a human label teaches us nothing new about.""" + if seg.get("segment_type") not in (None, "speech"): + return None + start = float(seg.get("start") or 0.0) + end = float(seg.get("end") or 0.0) + duration = end - start + if duration <= 0: + return None + + identified = seg.get("identified_as") + confidence = seg.get("confidence") + if identified is None or confidence is None: + label_info = 1.0 # unlabeled: maximal label entropy + else: + label_info = 1.0 - min(1.0, abs(confidence - threshold) / UNCERTAINTY_BAND) + + shortness = 1.0 - min(duration, SHORT_SEGMENT_CAP) / SHORT_SEGMENT_CAP + + overlap = 0.0 + for other in segments: + if other is seg or other.get("segment_type") not in (None, "speech"): + continue + if other.get("identified_as") == identified and identified is not None: + continue + if other.get("start", 0) < end and start < other.get("end", 0): + overlap = 1.0 + break + + score = W_LABEL * label_info + W_SHORT * shortness + W_OVERLAP * overlap + return { + "start": start, + "end": end, + "duration": duration, + "identified_as": identified, + "score": score, + "unlabeled": identified is None, + "uncertain": identified is not None and label_info >= 0.6, + "short": duration < 2.0, + "overlap": overlap > 0, + } + + +def _conversation_windows(doc: dict, threshold: float, window_seconds: float) -> list: + """Best non-overlapping annotation windows of one conversation, scored by + informativeness per minute of annotation effort.""" + segments = _active_segments(doc) + infos = [ + info + for seg in segments + if (info := _segment_info(seg, segments, threshold)) is not None + ] + if len(infos) < MIN_WINDOW_SEGMENTS: + return [] + audio_end = max( + float(doc.get("audio_total_duration") or 0.0), + max(i["end"] for i in infos), + ) + + candidates = [] + start = 0.0 + while start < audio_end: + end = min(start + window_seconds, audio_end) + inside = [i for i in infos if i["start"] < end and i["end"] > start] + if len(inside) >= MIN_WINDOW_SEGMENTS: + minutes = max((end - start) / 60.0, 0.25) + speakers = sorted( + {i["identified_as"] for i in inside if i["identified_as"]} + ) + candidates.append( + { + "window_start": round(start, 1), + "window_end": round(end, 1), + "score": sum(i["score"] for i in inside) / minutes, + "n_segments": len(inside), + "n_unlabeled": sum(i["unlabeled"] for i in inside), + "n_uncertain": sum(i["uncertain"] for i in inside), + "n_short": sum(i["short"] for i in inside), + "n_overlap": sum(i["overlap"] for i in inside), + "speakers": speakers, + } + ) + if end >= audio_end: + break + start += WINDOW_STEP_SECONDS + + candidates.sort(key=lambda w: w["score"], reverse=True) + picked: list = [] + for window in candidates: + if len(picked) >= MAX_WINDOWS_PER_CONVERSATION: + break + if any( + window["window_start"] < p["window_end"] + and p["window_start"] < window["window_end"] + for p in picked + ): + continue + picked.append(window) + return picked + + +def _window_reasons(window: dict) -> List[str]: + reasons = [] + if window["n_unlabeled"]: + reasons.append(f"{window['n_unlabeled']} unlabeled segments") + if window["n_uncertain"]: + reasons.append(f"{window['n_uncertain']} segments near the decision threshold") + if window["n_short"]: + reasons.append(f"{window['n_short']} short (<2s) segments — the failure regime") + if window["n_overlap"]: + reasons.append(f"{window['n_overlap']} cross-speaker overlaps") + if window.get("prior_annotations"): + reasons.append( + f"{window['prior_annotations']} human annotations already in this window" + ) + return reasons + + +async def _prior_annotation_points(conversation_ids: List[str]) -> dict: + """Per conversation, the time points a human has already annotated.""" + points: dict = {cid: [] for cid in conversation_ids} + async for row in _annotations_collection().find( + { + "conversation_id": {"$in": conversation_ids}, + "annotation_type": {"$in": HUMAN_ANNOTATION_TYPES}, + "source": "user", + }, + { + "conversation_id": 1, + "segment_start_time": 1, + "new_start": 1, + "insert_start": 1, + }, + ): + for field in ("segment_start_time", "new_start", "insert_start"): + value = row.get(field) + if value is not None: + points[row["conversation_id"]].append(float(value)) + break + return points + + +def _diverse_batch(windows: list, batch_size: int) -> list: + """Round-robin across speaker signatures so one pair can't fill the batch.""" + by_signature: dict = {} + for window in windows: + signature = "+".join(window["speakers"]) or "(unknown only)" + by_signature.setdefault(signature, []).append(window) + for group in by_signature.values(): + group.sort(key=lambda w: w["score"], reverse=True) + signatures = sorted( + by_signature, key=lambda s: by_signature[s][0]["score"], reverse=True + ) + + batch: list = [] + while len(batch) < batch_size and any(by_signature.values()): + for signature in signatures: + group = by_signature[signature] + if group: + batch.append(group.pop(0)) + if len(batch) >= batch_size: + break + return batch + + +async def suggest_annotation_targets( + user: User, + batch_size: int = 6, + window_seconds: float = 120.0, +): + """Ranked conversation windows whose ground-truth annotation is expected + to teach the speaker pipeline the most per minute of effort.""" + batch_size = max(1, min(batch_size, 20)) + window_seconds = max(MIN_WINDOW_SECONDS, min(window_seconds, MAX_WINDOW_SECONDS)) + threshold = get_diarization_settings().get("similarity_threshold", 0.5) + + decided = { + _target_key(r["conversation_id"], r["window_start"]) + async for r in _targets_collection().find( + {}, {"conversation_id": 1, "window_start": 1} + ) + } + + query = { + "deleted": {"$ne": True}, + "audio_archived": {"$ne": True}, + "audio_chunks_count": {"$gt": 0}, + } + if not user.is_superuser: + query["user_id"] = str(user.user_id) + + collection = Conversation.get_pymongo_collection() + windows: list = [] + scanned = 0 + async for doc in collection.find( + query, + { + "conversation_id": 1, + "title": 1, + "created_at": 1, + "client_id": 1, + "audio_total_duration": 1, + "active_transcript_version": 1, + "transcript_versions.version_id": 1, + "transcript_versions.segments.start": 1, + "transcript_versions.segments.end": 1, + "transcript_versions.segments.text": 1, + "transcript_versions.segments.identified_as": 1, + "transcript_versions.segments.confidence": 1, + "transcript_versions.segments.segment_type": 1, + }, + ): + scanned += 1 + for window in _conversation_windows(doc, threshold, window_seconds): + if _target_key(doc["conversation_id"], window["window_start"]) in decided: + continue + windows.append( + { + **window, + "conversation_id": doc["conversation_id"], + "conversation_title": doc.get("title"), + "conversation_date": str(doc.get("created_at") or ""), + "client_id": doc.get("client_id"), + } + ) + + prior = await _prior_annotation_points( + list({w["conversation_id"] for w in windows}) + ) + for window in windows: + n_prior = sum( + 1 + for t in prior.get(window["conversation_id"], []) + if window["window_start"] <= t <= window["window_end"] + ) + window["prior_annotations"] = n_prior + # A partially annotated window still needs finishing but resolves less + # new uncertainty per minute than an untouched one. + window["score"] = round(window["score"] / (1.0 + n_prior), 4) + + windows.sort(key=lambda w: w["score"], reverse=True) + batch = _diverse_batch(windows, batch_size) + for window in batch: + window["reasons"] = _window_reasons(window) + + return { + "threshold": threshold, + "window_seconds": window_seconds, + "batch": batch, + "conversations_scanned": scanned, + "pool_windows": len(windows), + "decided_total": len(decided), + } + + +async def decide_annotation_targets(user: User, decisions: List[dict]): + """Record window outcomes so finished/skipped windows leave the queue.""" + targets = _targets_collection() + recorded = 0 + for decision in decisions: + conversation_id = decision.get("conversation_id") + window_start = decision.get("window_start") + if not conversation_id or window_start is None: + continue + await targets.update_one( + { + "conversation_id": conversation_id, + "window_start": round(float(window_start), 1), + }, + { + "$set": { + "window_end": decision.get("window_end"), + "decision": decision.get("decision"), + "decided_by": str(user.user_id), + "decided_at": datetime.now(timezone.utc), + } + }, + upsert=True, + ) + recorded += 1 + return {"recorded": recorded, "status": "ok"} diff --git a/backends/advanced/src/advanced_omi_backend/controllers/guided_enrollment_controller.py b/backends/advanced/src/advanced_omi_backend/controllers/guided_enrollment_controller.py new file mode 100644 index 00000000..b3b1aa65 --- /dev/null +++ b/backends/advanced/src/advanced_omi_backend/controllers/guided_enrollment_controller.py @@ -0,0 +1,1028 @@ +"""Guided speaker enrollment — active-learning clip selection for one speaker. + +The user picks an enrolled speaker; we scan the corpus for candidate segments, +score a shortlist against the speaker's gallery on the speaker service, and +serve small batches (3-5) of the most *informative* clips for a human yes/no. +Accepted clips are appended to the speaker's voiceprint; every decision is +recorded in the ``enrollment_reviews`` collection so a clip is never re-shown. + +Selection follows the enrollment literature: total net speech helps up to +~30-60s then flattens; clips from *different* sessions/acoustic conditions beat +more-of-the-same; and human confirmation adds the most on clips the system is +uncertain about. Ranking therefore combines + * novelty — 1 - max cosine to the speaker's existing per-clip gallery, + * uncertainty — proximity of the centroid cosine to the operating threshold + (confirmed "hard positives" expand coverage the most), + * duration — capped so one long clip can't dominate, +with a plausibility gate (too-low cosine, or a different speaker scoring +clearly higher, wastes the user's time) and ≤2 clips per conversation for +session diversity. +""" + +import asyncio +import logging +from datetime import datetime, timezone +from typing import List, Optional + +from fastapi.responses import JSONResponse +from rq.exceptions import NoSuchJobError +from rq.job import Job + +from advanced_omi_backend.config import get_diarization_settings +from advanced_omi_backend.controllers.queue_controller import ( + JOB_RESULT_TTL, + default_queue, +) +from advanced_omi_backend.models.conversation import Conversation +from advanced_omi_backend.speaker_recognition_client import SpeakerRecognitionClient +from advanced_omi_backend.users import User +from advanced_omi_backend.utils.audio_chunk_utils import reconstruct_audio_segment +from advanced_omi_backend.workers.speaker_benchmark_jobs import ( + run_speaker_benchmark_job, +) +from advanced_omi_backend.workers.speaker_discovery_jobs import ( + discover_speaker_candidates_job, +) + +logger = logging.getLogger(__name__) + +MIN_CLIP_SECONDS = 3.0 +MAX_CLIP_SECONDS = 30.0 +MIN_PLAUSIBLE_SIM = 0.30 +OTHER_SPEAKER_MARGIN = 0.07 +UNCERTAINTY_BAND = 0.25 +MAX_PER_CONVERSATION = 2 +SCORE_CONCURRENCY = 4 + +W_NOVELTY, W_UNCERTAINTY, W_DURATION = 0.40, 0.35, 0.25 + + +def _reviews_collection(): + return Conversation.get_pymongo_collection().database["enrollment_reviews"] + + +def _batches_collection(): + return Conversation.get_pymongo_collection().database["enrollment_batches"] + + +def _discovery_collection(): + return Conversation.get_pymongo_collection().database["speaker_corpus_matches"] + + +def _discovery_runs_collection(): + return Conversation.get_pymongo_collection().database["speaker_discovery_runs"] + + +def _job_status(job_id: Optional[str]) -> Optional[str]: + if not job_id: + return None + try: + status = Job.fetch(job_id, connection=default_queue.connection).get_status( + refresh=True + ) + return status.value if hasattr(status, "value") else str(status) + except NoSuchJobError: + return "expired" + + +def _clip_key(conversation_id: str, start: float) -> str: + return f"{conversation_id}:{round(start, 2)}" + + +def _active_segments(doc: dict) -> list: + versions = doc.get("transcript_versions") or [] + active_id = doc.get("active_transcript_version") + active = next((v for v in versions if v.get("version_id") == active_id), None) + if active is None and versions: + active = versions[-1] + return (active or {}).get("segments") or [] + + +def _overlaps_other_speaker(seg: dict, segments: list) -> bool: + for other in segments: + if other is seg or other.get("segment_type") not in (None, "speech"): + continue + if other.get("speaker") == seg.get("speaker"): + continue + if other.get("start", 0) < seg.get("end", 0) and seg.get( + "start", 0 + ) < other.get("end", 0): + return True + return False + + +async def _gallery_stats( + speaker_client: SpeakerRecognitionClient, speaker_name: str +) -> Optional[dict]: + speaker = await speaker_client.get_speaker_by_name(speaker_name) + if not speaker: + return None + return { + "speaker_id": speaker["id"], + "speaker_name": speaker["name"], + "n_clips": speaker.get("audio_sample_count"), + "total_duration_s": speaker.get("total_audio_duration"), + } + + +async def _gallery_health( + speaker_client: SpeakerRecognitionClient, speaker_id: str +) -> Optional[dict]: + report = await speaker_client.get_enrollment_health(user_id=1) + if report.get("error"): + logger.warning("Guided enrollment health audit failed: %s", report) + return None + speaker = next( + ( + item + for item in report.get("speakers", []) + if item["speaker_id"] == speaker_id + ), + None, + ) + if not speaker: + return None + n_clips = speaker["n_clips"] + return { + "n_clips": n_clips, + "median_self": speaker.get("median_self"), + "n_flagged": speaker["n_flagged"], + "flagged_rate": round(speaker["n_flagged"] / n_clips, 4) if n_clips else 0.0, + "verdict": speaker["verdict"], + } + + +async def _candidate_pool(user: User, speaker_name: str, reviewed: set) -> list: + """All unreviewed candidate clips for a speaker, with cheap priors only.""" + query = { + "deleted": {"$ne": True}, + "audio_archived": {"$ne": True}, + "audio_chunks_count": {"$gt": 0}, + } + if not user.is_superuser: + query["user_id"] = str(user.user_id) + + collection = Conversation.get_pymongo_collection() + pool = [] + async for doc in collection.find( + query, + { + "conversation_id": 1, + "title": 1, + "created_at": 1, + "audio_total_duration": 1, + "active_transcript_version": 1, + "transcript_versions.version_id": 1, + "transcript_versions.segments.start": 1, + "transcript_versions.segments.end": 1, + "transcript_versions.segments.text": 1, + "transcript_versions.segments.speaker": 1, + "transcript_versions.segments.identified_as": 1, + "transcript_versions.segments.confidence": 1, + "transcript_versions.segments.segment_type": 1, + }, + ): + segments = _active_segments(doc) + speaker_present = any(s.get("identified_as") == speaker_name for s in segments) + if not speaker_present: + continue + audio_duration = doc.get("audio_total_duration") or 0.0 + for index, seg in enumerate(segments): + if seg.get("segment_type") not in (None, "speech"): + continue + identified = seg.get("identified_as") + # Attributed segments are candidates at any confidence; unknown + # segments only in conversations where the speaker appears (they + # are the likeliest missed hard positives). + if identified is not None and identified != speaker_name: + continue + start = float(seg.get("start") or 0.0) + end = float(seg.get("end") or 0.0) + if audio_duration: + end = min(end, audio_duration) + duration = end - start + if duration < MIN_CLIP_SECONDS: + continue + end = min(end, start + MAX_CLIP_SECONDS) + if _clip_key(doc["conversation_id"], start) in reviewed: + continue + if _overlaps_other_speaker(seg, segments): + continue + pool.append( + { + "conversation_id": doc["conversation_id"], + "conversation_title": doc.get("title"), + "conversation_date": str(doc.get("created_at") or ""), + "conversation_duration": round(float(audio_duration), 3), + "segment_index": index, + "start": round(start, 3), + "end": round(end, 3), + "duration": round(end - start, 3), + "text": (seg.get("text") or "")[:300], + "current_label": identified, + "stored_confidence": seg.get("confidence"), + } + ) + return pool + + +def _prior(clip: dict, threshold: float) -> float: + """Cheap pre-score ordering before the expensive embedding pass.""" + dur = min(clip["duration"], 10.0) / 10.0 + conf = clip["stored_confidence"] + if conf is None or clip["current_label"] is None: + band = 0.6 # unknown segments in the speaker's conversations: promising + else: + band = 1.0 - min(1.0, abs(conf - threshold) / UNCERTAINTY_BAND) + return 0.5 * dur + 0.5 * band + + +def _shortlist(pool: list, threshold: float, max_scan: int) -> list: + ranked = sorted(pool, key=lambda c: _prior(c, threshold), reverse=True) + picked: list = [] + per_conv: dict = {} + for clip in ranked: + if len(picked) >= max_scan: + break + cid = clip["conversation_id"] + if per_conv.get(cid, 0) >= 3: + continue + per_conv[cid] = per_conv.get(cid, 0) + 1 + picked.append(clip) + return picked + + +async def _score_clip( + speaker_client: SpeakerRecognitionClient, + sem: asyncio.Semaphore, + clip: dict, + speaker_id: str, +) -> Optional[dict]: + async with sem: + try: + wav = await reconstruct_audio_segment( + clip["conversation_id"], clip["start"], clip["end"] + ) + except Exception as e: + logger.warning( + "Guided enrollment: reconstruction failed for %s [%s-%s]: %s", + clip["conversation_id"], + clip["start"], + clip["end"], + e, + ) + return None + scores = await speaker_client.score_enrollment_candidate(wav, speaker_id) + if scores.get("error") or scores.get("sim_centroid") is None: + return None + return {**clip, "scores": scores} + + +def _information_score(clip: dict, threshold: float) -> Optional[dict]: + """Gate on plausibility, then rank by marginal information.""" + s = clip["scores"] + sim = s["sim_centroid"] + best_other = s.get("best_other") or {} + if sim < MIN_PLAUSIBLE_SIM: + return None + if best_other.get("score", 0.0) >= sim + OTHER_SPEAKER_MARGIN: + return None # probably the other speaker — not worth the user's time + + novelty = 1.0 - ( + s.get("max_clip_sim") if s.get("max_clip_sim") is not None else 0.0 + ) + uncertainty = 1.0 - min(1.0, abs(sim - threshold) / UNCERTAINTY_BAND) + dur = min(clip["duration"], 10.0) / 10.0 + score = W_NOVELTY * novelty + W_UNCERTAINTY * uncertainty + W_DURATION * dur + + reasons = [] + if novelty >= 0.5: + reasons.append("new acoustic condition for the gallery") + if abs(sim - threshold) <= 0.1: + reasons.append("near the decision boundary — confirmation helps most") + if sim >= threshold + 0.1: + reasons.append("confident match") + if clip["current_label"] is None: + reasons.append("currently unlabeled in the transcript") + elif clip.get("speaker_name") and clip["current_label"] != clip["speaker_name"]: + reasons.append(f'currently labeled {clip["current_label"]} — possible mismatch') + if clip["duration"] >= 8: + reasons.append("long clip") + + return { + **clip, + "info_score": round(score, 4), + "novelty": round(novelty, 3), + "uncertainty": round(uncertainty, 3), + "reasons": reasons, + } + + +async def suggest_clips( + user: User, + speaker_name: str, + batch_size: int = 4, + max_scan: int = 24, + order: str = "informative", +): + """Next batch of candidate clips for one speaker. + + ``order="informative"`` ranks by marginal information (novelty + boundary + uncertainty + duration) — best for teaching the model. ``order="confidence"`` + ranks by raw similarity to the gallery — best for finding the speaker fast + when the gallery is small and most candidates are low-similarity noise. + """ + batch_size = max(1, min(batch_size, 8)) + max_scan = max(batch_size, min(max_scan, 48)) + + speaker_client = SpeakerRecognitionClient() + if not speaker_client.enabled: + return JSONResponse( + status_code=503, content={"error": "Speaker recognition is not enabled"} + ) + gallery = await _gallery_stats(speaker_client, speaker_name) + if not gallery: + return JSONResponse( + status_code=404, + content={"error": f"No enrolled speaker named '{speaker_name}'"}, + ) + reviewed = { + _clip_key(r["conversation_id"], r["segment_start"]) + async for r in _reviews_collection().find( + {"speaker_name": speaker_name}, + {"conversation_id": 1, "segment_start": 1}, + ) + } + + threshold = get_diarization_settings().get("similarity_threshold", 0.5) + pool = await _candidate_pool(user, speaker_name, reviewed) + shortlist = _shortlist(pool, threshold, max_scan) + + sem = asyncio.Semaphore(SCORE_CONCURRENCY) + scored = await asyncio.gather( + *( + _score_clip(speaker_client, sem, clip, gallery["speaker_id"]) + for clip in shortlist + ) + ) + ranked = sorted( + filter( + None, + (_information_score(c, threshold) for c in scored if c is not None), + ), + key=lambda c: c["info_score"], + reverse=True, + ) + + discovery_query = { + "requested_by": str(user.user_id), + "speaker_id": gallery["speaker_id"], + "review_key": {"$nin": list(reviewed)}, + "$or": [ + {"human_label": None}, + {"human_label": gallery["speaker_name"]}, + ], + } + discovery_rows = ( + await _discovery_collection().find(discovery_query, {"_id": 0}).to_list() + ) + discovery_count = len(discovery_rows) + combined = { + _clip_key(candidate["conversation_id"], candidate["start"]): candidate + for candidate in ranked + } + for candidate in filter( + None, (_information_score(row, threshold) for row in discovery_rows) + ): + key = _clip_key(candidate["conversation_id"], candidate["start"]) + if key not in combined or candidate["info_score"] > combined[key]["info_score"]: + combined[key] = candidate + rank_key = ( + (lambda candidate: candidate["scores"]["sim_centroid"]) + if order == "confidence" + else (lambda candidate: candidate["info_score"]) + ) + ranked = sorted(combined.values(), key=rank_key, reverse=True) + + batch: list = [] + per_conv: dict = {} + for clip in ranked: + if len(batch) >= batch_size: + break + cid = clip["conversation_id"] + if per_conv.get(cid, 0) >= MAX_PER_CONVERSATION: + continue + per_conv[cid] = per_conv.get(cid, 0) + 1 + batch.append(clip) + + return { + "speaker": gallery, + "threshold": threshold, + "batch": batch, + "scanned": len(shortlist), + "gated_out": sum(1 for c in scored if c is not None) - len(ranked), + "pool_remaining": max(0, len(pool) - len(shortlist)), + "reviewed_total": len(reviewed), + "discovery_indexed": discovery_count > 0, + "discovery_candidates": discovery_count, + } + + +async def enqueue_corpus_discovery( + user: User, speaker_name: str, include_deleted: bool = False +): + """Index all corpus speech once, then score it against one live gallery.""" + client = SpeakerRecognitionClient() + gallery = await _gallery_stats(client, speaker_name) + if not gallery: + return JSONResponse( + status_code=404, + content={"error": f"No enrolled speaker named '{speaker_name}'"}, + ) + run_key = { + "requested_by": str(user.user_id), + "speaker_id": gallery["speaker_id"], + } + existing = await _discovery_runs_collection().find_one(run_key) + existing_status = _job_status(existing.get("job_id") if existing else None) + if existing_status in {"queued", "started", "deferred", "scheduled"}: + return { + "job_id": existing["job_id"], + "status": existing_status, + "reused": True, + } + job = default_queue.enqueue( + discover_speaker_candidates_job, + requested_by=str(user.user_id), + speaker_id=gallery["speaker_id"], + speaker_name=gallery["speaker_name"], + include_all_users=bool(user.is_superuser), + include_deleted=include_deleted, + job_timeout=14400, + result_ttl=JOB_RESULT_TTL, + description=f"Speaker corpus discovery: {gallery['speaker_name']}", + ) + await _discovery_runs_collection().update_one( + run_key, + { + "$set": { + "speaker_name": gallery["speaker_name"], + "job_id": job.id, + "queued_at": datetime.now(timezone.utc), + } + }, + upsert=True, + ) + return {"job_id": job.id, "status": "queued", "reused": False} + + +async def mine_uploaded_files(user: User, speaker_name: str, files: list): + """Ingest uploaded audio files as a mining corpus for one speaker. + + Files become annotation-only conversations (audio chunks + batch + transcription + speaker identification, no memory extraction); a + corpus-discovery job is chained behind the transcription jobs so mined + speech is scored against the speaker's gallery as soon as it has segments. + """ + # Lazy import: audio_controller pulls in the transcription stack. + from advanced_omi_backend.controllers.audio_controller import ( + upload_and_process_audio_files, + ) + from advanced_omi_backend.workers.speaker_mining_jobs import ( + MINING_DEVICE_NAME, + _parse_upload_response, + enqueue_discovery_after, + ) + + client = SpeakerRecognitionClient() + if not client.enabled: + return JSONResponse( + status_code=503, content={"error": "Speaker recognition is not enabled"} + ) + gallery = await _gallery_stats(client, speaker_name) + if not gallery: + return JSONResponse( + status_code=404, + content={"error": f"No enrolled speaker named '{speaker_name}'"}, + ) + + body = _parse_upload_response( + await upload_and_process_audio_files( + user, files, device_name=MINING_DEVICE_NAME, annotation_only=True + ) + ) + started = [f for f in body.get("files", []) if f.get("status") == "started"] + transcript_job_ids = [ + f["transcript_job_id"] for f in started if f.get("transcript_job_id") + ] + + discovery_job_id = None + if started: + discovery_job_id = await enqueue_discovery_after( + str(user.user_id), + gallery["speaker_id"], + gallery["speaker_name"], + transcript_job_ids, + bool(user.is_superuser), + ) + + return { + "speaker_name": gallery["speaker_name"], + "ingested": len(started), + "failed": [ + {"filename": f.get("filename"), "error": f.get("error")} + for f in body.get("files", []) + if f.get("status") == "error" + ], + "transcription_jobs": len(transcript_job_ids), + "transcription_available": bool(transcript_job_ids) or not started, + "discovery_job_id": discovery_job_id, + } + + +async def enqueue_local_mining(user: User, speaker_name: str, paths: List[str]): + """Queue server-side corpus mining (admin): ingest files already on the + backend's data volume — e.g. backup WAVs of purged conversations — and + chain discovery for the speaker.""" + from advanced_omi_backend.workers.speaker_mining_jobs import mine_local_corpus_job + + client = SpeakerRecognitionClient() + gallery = await _gallery_stats(client, speaker_name) + if not gallery: + return JSONResponse( + status_code=404, + content={"error": f"No enrolled speaker named '{speaker_name}'"}, + ) + if not paths: + return JSONResponse(status_code=400, content={"error": "No paths provided"}) + if len(paths) > 1000: + return JSONResponse( + status_code=400, content={"error": "Too many paths (max 1000)"} + ) + job = default_queue.enqueue( + mine_local_corpus_job, + requested_by=str(user.user_id), + speaker_id=gallery["speaker_id"], + speaker_name=gallery["speaker_name"], + paths=paths, + include_all_users=bool(user.is_superuser), + job_timeout=14400, + result_ttl=JOB_RESULT_TTL, + description=f"Speaker mining ingest: {len(paths)} files for {gallery['speaker_name']}", + ) + return {"job_id": job.id, "status": "queued", "files": len(paths)} + + +async def corpus_discovery_state(user: User, speaker_name: str): + client = SpeakerRecognitionClient() + gallery = await _gallery_stats(client, speaker_name) + if not gallery: + return JSONResponse( + status_code=404, + content={"error": f"No enrolled speaker named '{speaker_name}'"}, + ) + key = { + "requested_by": str(user.user_id), + "speaker_id": gallery["speaker_id"], + } + run = await _discovery_runs_collection().find_one(key, {"_id": 0}) + job_id = run.get("job_id") if run else None + return { + "speaker_name": gallery["speaker_name"], + "job_id": job_id, + "status": _job_status(job_id), + "matched_segments": await _discovery_collection().count_documents(key), + } + + +async def decide_clips(user: User, speaker_name: str, decisions: List[dict]): + """Record review decisions and enroll clips with a confirmed identity.""" + speaker_client = SpeakerRecognitionClient() + if not speaker_client.enabled: + return JSONResponse( + status_code=503, content={"error": "Speaker recognition is not enabled"} + ) + gallery = await _gallery_stats(speaker_client, speaker_name) + if not gallery: + return JSONResponse( + status_code=404, + content={"error": f"No enrolled speaker named '{speaker_name}'"}, + ) + health_before = await _gallery_health(speaker_client, gallery["speaker_id"]) + + reviews = _reviews_collection() + enrolled, reassigned, rejected, skipped, bad_clips, multiple_speakers, errors = ( + 0, + 0, + 0, + 0, + 0, + 0, + [], + ) + for decision in decisions: + conversation_id = decision.get("conversation_id") + start = decision.get("start") + end = decision.get("end") + original_start = decision.get("original_start") + original_end = decision.get("original_end") + review_decision = decision.get("decision") + actual_speaker = decision.get("actual_speaker") + if ( + not conversation_id + or start is None + or end is None + or end <= start + or original_start is None + or original_end is None + or original_end <= original_start + ): + errors.append({"clip": decision, "error": "invalid clip bounds"}) + continue + + if review_decision == "another_speaker" and not actual_speaker: + errors.append({"clip": decision, "error": "actual_speaker is required"}) + continue + + enroll_error = None + enrollment_target = ( + speaker_name if review_decision == "accept" else actual_speaker + ) + if enrollment_target: + try: + target_gallery = await _gallery_stats(speaker_client, enrollment_target) + if not target_gallery: + raise ValueError(f"No enrolled speaker named '{enrollment_target}'") + wav = await reconstruct_audio_segment(conversation_id, start, end) + result = await speaker_client.append_to_speaker( + target_gallery["speaker_id"], wav + ) + if result.get("error"): + enroll_error = result["error"] + else: + enrolled += 1 + if enrollment_target != speaker_name: + reassigned += 1 + except Exception as e: + enroll_error = str(e) + if enroll_error: + errors.append({"clip": decision, "error": enroll_error}) + elif review_decision == "reject": + rejected += 1 + elif review_decision == "skip": + skipped += 1 + elif review_decision == "multiple_speakers": + multiple_speakers += 1 + elif review_decision == "bad_clip": + bad_clips += 1 + + await reviews.update_one( + { + "speaker_name": speaker_name, + "conversation_id": conversation_id, + "segment_start": round(float(original_start), 3), + }, + { + "$set": { + "speaker_id": gallery["speaker_id"], + "segment_end": round(float(original_end), 3), + "selected_start": round(float(start), 3), + "selected_end": round(float(end), 3), + "decision": review_decision, + "actual_speaker": enrollment_target, + "enrolled": enrollment_target is not None and enroll_error is None, + "enroll_error": enroll_error, + "scores": decision.get("scores"), + "reviewed_by": str(user.user_id), + "reviewed_at": datetime.now(timezone.utc), + } + }, + upsert=True, + ) + + speaker_after = await _gallery_stats(speaker_client, speaker_name) + health_after = await _gallery_health(speaker_client, gallery["speaker_id"]) + accepted_novelties = [ + 1.0 - decision["scores"]["max_clip_sim"] + for decision in decisions + if decision.get("decision") == "accept" + and decision.get("scores", {}).get("max_clip_sim") is not None + ] + coverage = { + "accepted_novelty_mean": ( + round(sum(accepted_novelties) / len(accepted_novelties), 3) + if accepted_novelties + else None + ) + } + snapshot = { + "speaker_id": gallery["speaker_id"], + "speaker_name": speaker_name, + "reviewed_by": str(user.user_id), + "created_at": datetime.now(timezone.utc), + "health_before": health_before, + "health_after": health_after, + "coverage": coverage, + "decisions": { + "enrolled": enrolled, + "reassigned": reassigned, + "rejected": rejected, + "skipped": skipped, + "multiple_speakers": multiple_speakers, + "bad_clips": bad_clips, + }, + } + await _batches_collection().insert_one(snapshot) + benchmark_job_id = None + discovery_job_id = None + if enrolled > 0: + benchmark_job = default_queue.enqueue( + run_speaker_benchmark_job, + user_id=str(user.user_id), + job_timeout=7200, + result_ttl=JOB_RESULT_TTL, + description="Speaker enhancement: post-enrollment cross-validation", + ) + benchmark_job_id = benchmark_job.id + discovery_response = await enqueue_corpus_discovery(user, speaker_name) + if isinstance(discovery_response, dict): + discovery_job_id = discovery_response.get("job_id") + + return { + "speaker": speaker_after, + "health_before": health_before, + "health_after": health_after, + "coverage": coverage, + "benchmark_job_id": benchmark_job_id, + "discovery_job_id": discovery_job_id, + "enrolled": enrolled, + "reassigned": reassigned, + "rejected": rejected, + "skipped": skipped, + "multiple_speakers": multiple_speakers, + "bad_clips": bad_clips, + "errors": errors, + "status": "ok" if not errors else "partial", + } + + +async def gallery_clips(user: User, speaker_name: str): + """List a speaker's enrolled clips with per-clip contamination flags. + + Powers the gallery-management panel: each clip carries the audit's + self-similarity, closest-other-speaker score, and mislabel/junk/weak flags + so the user can spot and remove bad enrollments. + """ + speaker_client = SpeakerRecognitionClient() + if not speaker_client.enabled: + return JSONResponse( + status_code=503, content={"error": "Speaker recognition is not enabled"} + ) + gallery = await _gallery_stats(speaker_client, speaker_name) + if not gallery: + return JSONResponse( + status_code=404, + content={"error": f"No enrolled speaker named '{speaker_name}'"}, + ) + report = await speaker_client.get_enrollment_health(user_id=1) + if report.get("error"): + return JSONResponse( + status_code=503, + content={"error": "Unable to audit the speaker's enrolled clips"}, + ) + speaker = next( + ( + item + for item in report.get("speakers", []) + if item["speaker_id"] == gallery["speaker_id"] + ), + None, + ) + return { + "speaker": gallery, + "verdict": speaker["verdict"] if speaker else None, + "median_self": speaker.get("median_self") if speaker else None, + "clips": speaker["clips"] if speaker else [], + "thresholds": report.get("thresholds"), + } + + +async def delete_gallery_clip( + user: User, speaker_name: str, segment_id: int, hard: bool = False +): + """Remove one enrolled clip from a speaker's voiceprint. + + The clip must belong to the named speaker (guards against stale UI state + deleting another speaker's clip). Quarantined by default so it's + recoverable; the speaker service recomputes the centroid either way. + """ + speaker_client = SpeakerRecognitionClient() + if not speaker_client.enabled: + return JSONResponse( + status_code=503, content={"error": "Speaker recognition is not enabled"} + ) + gallery = await _gallery_stats(speaker_client, speaker_name) + if not gallery: + return JSONResponse( + status_code=404, + content={"error": f"No enrolled speaker named '{speaker_name}'"}, + ) + report = await speaker_client.get_enrollment_health(user_id=1) + speaker = next( + ( + item + for item in report.get("speakers", []) + if item["speaker_id"] == gallery["speaker_id"] + ), + None, + ) + if not speaker or not any(c["segment_id"] == segment_id for c in speaker["clips"]): + return JSONResponse( + status_code=404, + content={ + "error": f"Clip {segment_id} is not enrolled for '{speaker_name}'" + }, + ) + result = await speaker_client.delete_enrollment_segment(segment_id, hard=hard) + if result.get("error"): + return JSONResponse(status_code=502, content=result) + logger.info( + "Guided enrollment: removed clip %s from %s (%s)", + segment_id, + speaker_name, + "hard" if hard else "quarantined", + ) + return { + **result, + "speaker": await _gallery_stats(speaker_client, speaker_name), + "health": await _gallery_health(speaker_client, gallery["speaker_id"]), + } + + +async def reset_speaker_state( + user: User, speaker_name: str, purge_gallery: bool = False +): + """Forget all guided-enrollment state for a speaker name. + + Deleting a speaker on the speaker service does not touch the backend's + review ledger, so a re-enrolled speaker with the same name inherits stale + 'already reviewed' exclusions and old session history. This clears the + review decisions, session snapshots, and corpus-discovery matches recorded + under the name (across old speaker ids), so every clip becomes suggestible + again. With ``purge_gallery`` the speaker's voiceprint and enrollment audio + are also deleted from the speaker service, leaving a truly blank slate. + """ + scope = {"speaker_name": speaker_name} + if not user.is_superuser: + reviews_scope = {**scope, "reviewed_by": str(user.user_id)} + discovery_scope = {**scope, "requested_by": str(user.user_id)} + else: + reviews_scope = scope + discovery_scope = scope + + deleted = { + "reviews": ( + await _reviews_collection().delete_many(reviews_scope) + ).deleted_count, + "sessions": ( + await _batches_collection().delete_many(reviews_scope) + ).deleted_count, + "discovery_matches": ( + await _discovery_collection().delete_many(discovery_scope) + ).deleted_count, + "discovery_runs": ( + await _discovery_runs_collection().delete_many(discovery_scope) + ).deleted_count, + } + + gallery_deleted = False + if purge_gallery: + speaker_client = SpeakerRecognitionClient() + if not speaker_client.enabled: + return JSONResponse( + status_code=503, + content={ + "error": "Speaker recognition is not enabled", + "deleted": deleted, + }, + ) + gallery = await _gallery_stats(speaker_client, speaker_name) + if gallery: + result = await speaker_client.delete_speaker( + gallery["speaker_id"], delete_audio=True + ) + if result.get("error"): + return JSONResponse( + status_code=502, content={**result, "deleted": deleted} + ) + gallery_deleted = True + + logger.info( + "Guided enrollment reset for '%s' by %s: %s%s", + speaker_name, + user.user_id, + deleted, + " + gallery purged" if gallery_deleted else "", + ) + return { + "speaker_name": speaker_name, + "deleted": deleted, + "gallery_deleted": gallery_deleted, + "status": "ok", + } + + +async def enrollment_history(user: User, speaker_name: str, limit: int = 50): + """Return dated gallery-quality snapshots for one speaker, newest first.""" + query = {"speaker_name": speaker_name} + if not user.is_superuser: + query["reviewed_by"] = str(user.user_id) + rows = [] + async for row in ( + _batches_collection() + .find(query, {"_id": 0}) + .sort("created_at", -1) + .limit(max(1, min(limit, 200))) + ): + created_at = row.get("created_at") + if isinstance(created_at, datetime): + row["created_at"] = created_at.isoformat() + rows.append(row) + return {"speaker_name": speaker_name, "sessions": rows} + + +async def enqueue_benchmark(user: User): + job = default_queue.enqueue( + run_speaker_benchmark_job, + user_id=str(user.user_id), + job_timeout=7200, + result_ttl=JOB_RESULT_TTL, + description="Speaker enhancement: grouped cross-validation", + ) + return {"job_id": job.id, "status": "queued"} + + +async def latest_benchmark(user: User): + row = ( + await _batches_collection() + .database["speaker_benchmark_runs"] + .find_one({"user_id": str(user.user_id)}, {"_id": 0}, sort=[("created_at", -1)]) + ) + if row and isinstance(row.get("created_at"), datetime): + row["created_at"] = row["created_at"].isoformat() + return {"report": row} + + +async def reconstructed_baseline(user: User): + first_review = await _reviews_collection().find_one( + {"reviewed_by": str(user.user_id)}, + {"reviewed_at": 1}, + sort=[("reviewed_at", 1)], + ) + cutoff = first_review.get("reviewed_at") if first_review else None + if cutoff is None: + return {"cutoff": None, "speakers": [], "status": "no_guided_reviews"} + + client = SpeakerRecognitionClient() + baseline = await client.get_enrollment_health(user_id=1, before=cutoff) + current = await client.get_enrollment_health(user_id=1) + if baseline.get("error") or current.get("error"): + return JSONResponse( + status_code=503, + content={"error": "Unable to compute speaker gallery baseline"}, + ) + current_by_id = {speaker["speaker_id"]: speaker for speaker in current["speakers"]} + speakers = [] + for before in baseline["speakers"]: + after = current_by_id.get(before["speaker_id"]) + speakers.append( + { + "speaker_id": before["speaker_id"], + "name": before["name"], + "baseline": { + "n_clips": before["n_clips"], + "median_self": before["median_self"], + "n_flagged": before["n_flagged"], + "verdict": before["verdict"], + }, + "current": ( + { + "n_clips": after["n_clips"], + "median_self": after["median_self"], + "n_flagged": after["n_flagged"], + "verdict": after["verdict"], + } + if after + else None + ), + } + ) + return { + "cutoff": cutoff.isoformat(), + "speakers": speakers, + "status": "reconstructed", + "limitations": ( + "Uses surviving clip rows and their current speaker assignment. Clips deleted " + "or relabeled before reconstruction cannot be restored to their prior state." + ), + } diff --git a/backends/advanced/src/advanced_omi_backend/controllers/queue_controller.py b/backends/advanced/src/advanced_omi_backend/controllers/queue_controller.py index 385c7d5e..6507ff6c 100644 --- a/backends/advanced/src/advanced_omi_backend/controllers/queue_controller.py +++ b/backends/advanced/src/advanced_omi_backend/controllers/queue_controller.py @@ -23,6 +23,11 @@ from advanced_omi_backend.config import get_misc_settings from advanced_omi_backend.config_loader import get_service_config +from advanced_omi_backend.heartbeat import ( + FLEET_HEALTH_KEY, + evaluate_fleet_health, + is_rq_worker_fresh, +) from advanced_omi_backend.redis_factory import create_sync_redis from advanced_omi_backend.services.memory.audit import MemoryCause, UpdateStrategy from advanced_omi_backend.services.sse_publisher import publish_sse_event @@ -792,6 +797,7 @@ def start_post_conversation_jobs( client_id: Optional[str] = None, end_reason: str = "file_upload", skip_speaker_recognition: bool = False, + skip_memory_extraction: bool = False, memory_cause: MemoryCause = MemoryCause.AUTO_EXTRACTION, memory_strategy: UpdateStrategy = UpdateStrategy.FULL, ) -> Dict[str, str]: @@ -816,6 +822,8 @@ def start_post_conversation_jobs( end_reason: Reason conversation ended (e.g., 'file_upload', 'websocket_disconnect', 'user_stopped') skip_speaker_recognition: Skip the speaker step even when enabled — used by split/merge, whose transcripts already carry speaker labels + skip_memory_extraction: Skip memory extraction even when globally enabled — + used for annotation/training datasets that should not enter user memory Returns: Dict with job IDs for speaker_recognition, memory, title_summary, event_dispatch @@ -897,6 +905,11 @@ def start_post_conversation_jobs( memory_enabled = memory_config.get( "enabled", True ) # Default to True for backward compatibility + if memory_enabled and skip_memory_extraction: + logger.info( + f"⏭️ Memory extraction skipped by caller for conversation {conversation_id[:8]}" + ) + memory_enabled = False memory_job = None if memory_enabled: @@ -1059,12 +1072,18 @@ def get_queue_health() -> Dict[str, Any]: "total_workers": 0, "active_workers": 0, "idle_workers": 0, + "worker_fleet": { + "healthy": False, + "status": "unknown", + "detail": "Worker fleet health has not been checked", + }, } # Check Redis connection try: redis_conn.ping() health["redis_connection"] = "healthy" + health["worker_fleet"] = evaluate_fleet_health(redis_conn.get(FLEET_HEALTH_KEY)) except Exception as e: health["redis_connection"] = f"unhealthy: {e}" return health @@ -1080,7 +1099,11 @@ def get_queue_health() -> Dict[str, Any]: } # Check workers - workers = Worker.all(connection=redis_conn) + workers = [ + worker + for worker in Worker.all(connection=redis_conn) + if is_rq_worker_fresh(worker) + ] health["total_workers"] = len(workers) for worker in workers: diff --git a/backends/advanced/src/advanced_omi_backend/heartbeat.py b/backends/advanced/src/advanced_omi_backend/heartbeat.py index 6b026ad2..0df49d98 100644 --- a/backends/advanced/src/advanced_omi_backend/heartbeat.py +++ b/backends/advanced/src/advanced_omi_backend/heartbeat.py @@ -1,15 +1,17 @@ -"""Liveness heartbeats for the custom stream-consumer workers. +"""Liveness heartbeats for the worker fleet and custom stream consumers. RQ workers register and heartbeat with RQ itself (observable via ``Worker.all()``), -so a dead/wedged RQ worker drops out of the registration count. The custom -stream-consumer workers (``streaming-stt``, ``windowed-batch``, +but registrations can outlive their containers and must be freshness-checked. The +custom stream-consumer workers (``streaming-stt``, ``windowed-batch``, ``wakeword-dispatch``) are plain processes, where "process is alive" does NOT prove the main loop is still turning. Each beats once per main-loop iteration; the workers-container healthcheck (``worker_healthcheck.py``) flags a stale heartbeat so a wedged-but-alive consumer stops reporting healthy. """ +import json import time +from typing import Any HEARTBEAT_KEY_PREFIX = "worker:heartbeat:" @@ -18,6 +20,76 @@ # wedge (loop stops beating) is detectable long before that. HEARTBEAT_TTL_SECONDS = 3600 +# The orchestrator owns this heartbeat. Unlike RQ registrations, it disappears or +# becomes stale when the entire workers container is absent, which lets the backend +# detect an outage that cannot be observed by the in-container self-healer. +FLEET_HEALTH_KEY = "worker:fleet:health" +FLEET_HEARTBEAT_TTL_SECONDS = 300 +FLEET_HEARTBEAT_MAX_AGE_SECONDS = 30 + + +def evaluate_fleet_health( + raw: str | bytes | None, + *, + now: float | None = None, + max_age_seconds: float = FLEET_HEARTBEAT_MAX_AGE_SECONDS, +) -> dict[str, Any]: + """Decode an orchestrator heartbeat into a stable health result.""" + checked_at = time.time() if now is None else now + if raw is None: + return { + "healthy": False, + "status": "missing", + "detail": "No worker fleet heartbeat has been published", + } + + if isinstance(raw, bytes): + raw = raw.decode("utf-8", errors="replace") + + try: + payload = json.loads(raw) + timestamp = float(payload["timestamp"]) + except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc: + return { + "healthy": False, + "status": "invalid", + "detail": f"Worker fleet heartbeat is invalid: {exc}", + } + + age_seconds = max(0.0, checked_at - timestamp) + result = {**payload, "age_seconds": age_seconds} + if age_seconds > max_age_seconds: + result.update( + healthy=False, + status="stale", + detail=( + f"Worker fleet heartbeat is {age_seconds:.1f}s old " + f"(maximum {max_age_seconds:.0f}s)" + ), + ) + return result + + status = str(payload.get("status") or "invalid") + result["status"] = status + result["healthy"] = status == "healthy" + if not result["healthy"] and not result.get("detail"): + result["detail"] = f"Worker fleet reported status '{status}'" + return result + + +def is_rq_worker_fresh(worker: Any, *, now: float | None = None) -> bool: + """Return whether an RQ registration has heartbeated within its own TTL.""" + last_heartbeat = getattr(worker, "last_heartbeat", None) + if last_heartbeat is None: + return False + try: + heartbeat_at = last_heartbeat.timestamp() + worker_ttl = float(getattr(worker, "worker_ttl", 420)) + except (AttributeError, TypeError, ValueError): + return False + checked_at = time.time() if now is None else now + return checked_at - heartbeat_at <= worker_ttl + async def beat(redis_client, worker_name: str) -> None: """Record that ``worker_name``'s main loop just iterated. Best-effort. diff --git a/backends/advanced/src/advanced_omi_backend/llm_client.py b/backends/advanced/src/advanced_omi_backend/llm_client.py index e13e5377..5ec98239 100644 --- a/backends/advanced/src/advanced_omi_backend/llm_client.py +++ b/backends/advanced/src/advanced_omi_backend/llm_client.py @@ -244,6 +244,36 @@ def reset_llm_client(): ) +def _is_context_length_error(exc: Exception) -> bool: + """Return whether a provider 400 specifically reports context overflow.""" + if not isinstance(exc, openai.BadRequestError): + return False + + values: list[str] = [] + + def collect(value: Any) -> None: + if isinstance(value, dict): + for key, item in value.items(): + if key in {"code", "type", "message"}: + values.append(str(item).lower()) + collect(item) + elif isinstance(value, (list, tuple)): + for item in value: + collect(item) + + collect(getattr(exc, "body", None)) + details = " ".join(values) + return any( + marker in details + for marker in ( + "exceed_context_size_error", + "context_length_exceeded", + "exceeds the available context size", + "maximum context length", + ) + ) + + def _get_fallback_model_def(): """ModelDef named by defaults.fallback_llm, or None when unset/invalid or identical to defaults.llm (retrying the same model is pointless).""" @@ -384,12 +414,14 @@ async def async_chat_with_tools( model: str | None = None, temperature: float | None = None, operation: str | None = None, + force_fallback: bool = False, ): """Async wrapper for chat completion with tool calling. When ``operation`` is provided, parameters are resolved from config. - Unreachable-primary calls retry once against ``defaults.fallback_llm`` - (see :func:`async_generate`). + Unreachable-primary and context-overflow calls retry once against + ``defaults.fallback_llm``. ``force_fallback`` is used for a semantic retry after + a provider returned a syntactically valid but incomplete result. Tracing is handled automatically by the OTEL instrumentor. """ @@ -409,9 +441,25 @@ async def _chat_once(op, model_override): registry = get_models_registry() if registry: op = registry.get_llm_operation(operation) + if force_fallback: + fb_op = registry.get_fallback_llm_operation(operation, primary=op) + if fb_op is None: + raise RuntimeError( + f"No fallback LLM is configured for operation {operation!r}" + ) + logger.warning( + "Using fallback LLM %r for semantic retry of operation %r", + fb_op.model_name, + operation, + ) + return await _chat_once(fb_op, None) try: return await _chat_once(op, model) - except _FALLBACK_EXCEPTIONS as e: + except Exception as e: + if not isinstance( + e, _FALLBACK_EXCEPTIONS + ) and not _is_context_length_error(e): + raise fb_op = registry.get_fallback_llm_operation(operation, primary=op) if fb_op is None: raise diff --git a/backends/advanced/src/advanced_omi_backend/model_registry.py b/backends/advanced/src/advanced_omi_backend/model_registry.py index c5f5af86..dd25fb89 100644 --- a/backends/advanced/src/advanced_omi_backend/model_registry.py +++ b/backends/advanced/src/advanced_omi_backend/model_registry.py @@ -10,6 +10,7 @@ from __future__ import annotations import logging +import re import time from pathlib import Path from typing import Any, Dict, List, Optional @@ -306,11 +307,11 @@ def to_api_params(self) -> Dict[str, Any]: if self.reasoning_effort: if openai_reasoning: effort = self.reasoning_effort.strip().lower() - # "none" is only accepted from gpt-5.1 on; earlier reasoning - # models (gpt-5-nano, o3, …) reject it — "minimal" is their floor. - if effort in ("none", "off", "0") and not model_name.lower().startswith( - "gpt-5.1" - ): + # "none" is accepted by versioned GPT-5.1+ models. Earlier GPT-5 + # variants (gpt-5, gpt-5-mini/nano) require "minimal" instead. + version_match = re.match(r"^gpt-5\.(\d+)", model_name.lower()) + supports_none = bool(version_match and int(version_match.group(1)) >= 1) + if effort in ("none", "off", "0") and not supports_none: effort = "minimal" params["reasoning_effort"] = effort elif self.model_def.thinking: diff --git a/backends/advanced/src/advanced_omi_backend/models/conversation.py b/backends/advanced/src/advanced_omi_backend/models/conversation.py index 9cbc87a1..46807994 100644 --- a/backends/advanced/src/advanced_omi_backend/models/conversation.py +++ b/backends/advanced/src/advanced_omi_backend/models/conversation.py @@ -189,6 +189,18 @@ class DerivedFrom(BaseModel): external_source_type: Optional[str] = Field( None, description="Type of external source (gdrive, dropbox, s3, etc.)" ) + data_purpose: Optional[str] = Field( + None, + description="Operational purpose for this conversation, e.g. annotation or normal_capture", + ) + memory_excluded: bool = Field( + False, + description="When true, this conversation must not create or reprocess memories", + ) + memory_exclusion_reason: Optional[str] = Field( + None, + description="Why memory processing is disabled for this conversation", + ) # MongoDB chunk-based audio storage (new system) audio_chunks_count: Optional[int] = Field( @@ -536,6 +548,9 @@ def create_conversation( segments: Optional[List["Conversation.SpeakerSegment"]] = None, external_source_id: Optional[str] = None, external_source_type: Optional[str] = None, + data_purpose: Optional[str] = None, + memory_excluded: bool = False, + memory_exclusion_reason: Optional[str] = None, ) -> Conversation: """ Factory function to create a new conversation. @@ -565,6 +580,9 @@ def create_conversation( "active_transcript_version": None, "external_source_id": external_source_id, "external_source_type": external_source_type, + "data_purpose": data_purpose, + "memory_excluded": memory_excluded, + "memory_exclusion_reason": memory_exclusion_reason, } # Only set conversation_id if provided, otherwise let the model auto-generate it diff --git a/backends/advanced/src/advanced_omi_backend/models/memory_audit.py b/backends/advanced/src/advanced_omi_backend/models/memory_audit.py index 0f3163a4..a259f0b2 100644 --- a/backends/advanced/src/advanced_omi_backend/models/memory_audit.py +++ b/backends/advanced/src/advanced_omi_backend/models/memory_audit.py @@ -38,8 +38,9 @@ class MemoryAuditEntry(Document): cause: Optional[str] = Field( None, description="Why the memory changed (provenance), one of MemoryCause: " - "auto_extraction, memory_replay, transcript_reprocess, speaker_reprocess, " - "annotation_apply, obsidian_sync, delete_all. See services/memory/audit.py.", + "auto_extraction, memory_replay, memory_rebuild, transcript_reprocess, " + "speaker_reprocess, annotation_apply, obsidian_sync, delete_all. " + "See services/memory/audit.py.", ) strategy: Optional[str] = Field( None, diff --git a/backends/advanced/src/advanced_omi_backend/prompt_defaults.py b/backends/advanced/src/advanced_omi_backend/prompt_defaults.py index 11ae9005..22acb927 100644 --- a/backends/advanced/src/advanced_omi_backend/prompt_defaults.py +++ b/backends/advanced/src/advanced_omi_backend/prompt_defaults.py @@ -702,3 +702,18 @@ def register_all_defaults(registry: PromptRegistry) -> None: category="memory", variables=["vault_summary"], ) + + from advanced_omi_backend.services.memory.agent.codex_agent import ( + DEFAULT_CODEX_AGENT_SYSTEM_PROMPT, + ) + + registry.register_default( + "memory.codex_agent_system", + template=DEFAULT_CODEX_AGENT_SYSTEM_PROMPT, + name="Codex Memory Agent System Prompt", + description="Vault-aware instructions for the Codex CLI memory-agent executor " + "(memory.agent_executor: codex) — same conventions as the direct agent, phrased " + "for direct file editing. Supports a {{vault_summary}} slot.", + category="memory", + variables=["vault_summary"], + ) diff --git a/backends/advanced/src/advanced_omi_backend/routers/modules/annotation_routes.py b/backends/advanced/src/advanced_omi_backend/routers/modules/annotation_routes.py index 4ddba293..9e2c1c19 100644 --- a/backends/advanced/src/advanced_omi_backend/routers/modules/annotation_routes.py +++ b/backends/advanced/src/advanced_omi_backend/routers/modules/annotation_routes.py @@ -60,6 +60,11 @@ def _apply_diarization_label(segment, corrected_speaker: str) -> None: segment.segment_type = Conversation.SegmentType.EVENT.value +def _should_reprocess_memory(conversation: Conversation) -> bool: + """False for annotation/training imports that are explicitly excluded from memory.""" + return not getattr(conversation, "memory_excluded", False) + + @router.get("/suggestions") async def get_pending_suggestions( current_user: User = Depends(current_active_user), @@ -1090,15 +1095,21 @@ async def apply_diarization_annotations( annotation.processed_by = "apply" await annotation.save() - # Chain memory reprocessing. Diarization-only edits change speaker - # attribution, so use the same speaker-diff strategy as a speaker - # reprocess (it falls back to a full re-extraction if no diff applies). - enqueue_memory_processing( - conversation_id=conversation_id, - priority=JobPriority.NORMAL, - cause=MemoryCause.ANNOTATION_APPLY, - strategy=UpdateStrategy.SPEAKER_DIFF, - ) + # Chain memory reprocessing unless this is an annotation/training import. + # Diarization-only edits change speaker attribution, so use the same + # speaker-diff strategy as a speaker reprocess (it falls back to full + # re-extraction if no diff applies). + if _should_reprocess_memory(conversation): + enqueue_memory_processing( + conversation_id=conversation_id, + priority=JobPriority.NORMAL, + cause=MemoryCause.ANNOTATION_APPLY, + strategy=UpdateStrategy.SPEAKER_DIFF, + ) + else: + logger.info( + f"Skipping memory reprocessing for memory-excluded conversation {conversation_id[:8]}" + ) return JSONResponse( content={ @@ -1337,14 +1348,20 @@ async def apply_all_annotations( annotation.status = AnnotationStatus.ACCEPTED await annotation.save() - # Trigger memory reprocessing (once for all changes). Combined apply may - # change transcript text as well as speakers, so re-extract in full. - enqueue_memory_processing( - conversation_id=conversation_id, - priority=JobPriority.NORMAL, - cause=MemoryCause.ANNOTATION_APPLY, - strategy=UpdateStrategy.FULL, - ) + # Trigger memory reprocessing (once for all changes) unless this is an + # annotation/training import. Combined apply may change transcript text as + # well as speakers, so re-extract in full for normal conversations. + if _should_reprocess_memory(conversation): + enqueue_memory_processing( + conversation_id=conversation_id, + priority=JobPriority.NORMAL, + cause=MemoryCause.ANNOTATION_APPLY, + strategy=UpdateStrategy.FULL, + ) + else: + logger.info( + f"Skipping memory reprocessing for memory-excluded conversation {conversation_id[:8]}" + ) return JSONResponse( content={ diff --git a/backends/advanced/src/advanced_omi_backend/routers/modules/audio_routes.py b/backends/advanced/src/advanced_omi_backend/routers/modules/audio_routes.py index 3ce4328b..f5c39da0 100644 --- a/backends/advanced/src/advanced_omi_backend/routers/modules/audio_routes.py +++ b/backends/advanced/src/advanced_omi_backend/routers/modules/audio_routes.py @@ -60,6 +60,10 @@ async def upload_audio_from_drive_folder( ), current_user: User = Depends(current_superuser), device_name: str = Query(default="upload"), + annotation_only: bool = Query( + default=False, + description="Create editable annotation records without memory extraction", + ), ): try: files = await download_audio_files_from_drive(gdrive_folder_id, current_user.id) @@ -67,7 +71,26 @@ async def upload_audio_from_drive_folder( raise HTTPException(status_code=400, detail=str(e)) return await audio_controller.upload_and_process_audio_files( - current_user, files, device_name, source="gdrive" + current_user, + files, + device_name, + source="gdrive", + annotation_only=annotation_only, + ) + + +@router.post("/upload_audio_from_gdrive/annotation") +async def upload_audio_from_drive_folder_for_annotation( + gdrive_folder_id: str = Query(..., description="Google Drive folder ID"), + current_user: User = Depends(current_superuser), + device_name: str = Query(default="annotation-import"), +): + """Import a Drive folder into the annotation workspace without memory writes.""" + return await upload_audio_from_drive_folder( + gdrive_folder_id=gdrive_folder_id, + current_user=current_user, + device_name=device_name, + annotation_only=True, ) @@ -444,6 +467,10 @@ async def upload_audio_files( device_name: str = Query( default="upload", description="Device name for uploaded files" ), + annotation_only: bool = Query( + default=False, + description="Create editable annotation records without memory extraction", + ), ): """ Upload and process audio files. Admin only. @@ -456,5 +483,23 @@ async def upload_audio_files( - Summary of enqueued vs failed uploads """ return await audio_controller.upload_and_process_audio_files( - current_user, files, device_name + current_user, files, device_name, annotation_only=annotation_only + ) + + +@router.post("/upload/annotation") +async def upload_audio_files_for_annotation( + files: list[UploadFile] = File(...), + current_user: User = Depends(current_superuser), + device_name: str = Query( + default="annotation-upload", + description="Device name for annotation workspace uploads", + ), +): + """Transcribe raw audio into the annotation workspace without memory writes.""" + return await audio_controller.upload_and_process_audio_files( + current_user, + files, + device_name, + annotation_only=True, ) diff --git a/backends/advanced/src/advanced_omi_backend/routers/modules/conversation_routes.py b/backends/advanced/src/advanced_omi_backend/routers/modules/conversation_routes.py index 08290bc2..030e0125 100644 --- a/backends/advanced/src/advanced_omi_backend/routers/modules/conversation_routes.py +++ b/backends/advanced/src/advanced_omi_backend/routers/modules/conversation_routes.py @@ -13,6 +13,10 @@ from advanced_omi_backend.auth import current_active_user, current_superuser from advanced_omi_backend.controllers import conversation_controller from advanced_omi_backend.controllers.drift_controller import find_drift_conversations +from advanced_omi_backend.controllers.queue_controller import ( + JOB_RESULT_TTL, + default_queue, +) from advanced_omi_backend.models.conversation import Conversation from advanced_omi_backend.models.waveform import WaveformData from advanced_omi_backend.users import User @@ -21,6 +25,7 @@ get_trimmed_opus_for_time_range, reconstruct_audio_segment, ) +from advanced_omi_backend.workers.drift_jobs import cluster_embedding_backfill_job from advanced_omi_backend.workers.waveform_jobs import generate_waveform_data logger = logging.getLogger(__name__) @@ -96,6 +101,25 @@ async def identify_drift(current_user: User = Depends(current_superuser)): return await find_drift_conversations() +@router.post("/drift/backfill-cluster-embeddings") +async def backfill_drift_cluster_embeddings( + current_user: User = Depends(current_superuser), +): + """Queue the GPU-backed embedding pass needed to analyze older conversations.""" + job = default_queue.enqueue( + cluster_embedding_backfill_job, + job_timeout=7200, + result_ttl=JOB_RESULT_TTL, + description="Backfill speaker cluster embeddings for drift analysis", + ) + logger.info( + "Enqueued cluster-embedding backfill job %s for admin %s", + job.id, + current_user.user_id, + ) + return {"job_id": job.id, "status": "queued"} + + @router.get("/{conversation_id}") async def get_conversation_detail( conversation_id: str, current_user: User = Depends(current_active_user) diff --git a/backends/advanced/src/advanced_omi_backend/routers/modules/data_audit_routes.py b/backends/advanced/src/advanced_omi_backend/routers/modules/data_audit_routes.py index 8e019caa..20a395b2 100644 --- a/backends/advanced/src/advanced_omi_backend/routers/modules/data_audit_routes.py +++ b/backends/advanced/src/advanced_omi_backend/routers/modules/data_audit_routes.py @@ -11,16 +11,20 @@ from datetime import datetime from typing import Dict, List, Optional -from fastapi import APIRouter, Depends, Query -from fastapi.responses import JSONResponse +from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, UploadFile +from fastapi.responses import JSONResponse, Response from pydantic import BaseModel, Field from advanced_omi_backend.auth import ( current_active_user, current_active_user_optional, + current_superuser, get_user_from_token_param, ) -from advanced_omi_backend.controllers import data_audit_controller +from advanced_omi_backend.controllers import ( + data_audit_controller, + guided_enrollment_controller, +) from advanced_omi_backend.users import User logger = logging.getLogger(__name__) @@ -148,6 +152,11 @@ async def list_conversations( exclude_speakers: Optional[str] = Query( None, description="Comma-separated speakers a conversation must contain none of" ), + dataset_id: Optional[str] = Query( + None, + max_length=200, + description="Only conversations imported from this annotation dataset", + ), archived_only: bool = Query( False, description="List archived metadata stubs instead of active conversations", @@ -180,6 +189,7 @@ def _csv(v: Optional[str]) -> Optional[list]: created_before=created_before, include_speakers=_csv(include_speakers), exclude_speakers=_csv(exclude_speakers), + dataset_id=dataset_id, archived_only=archived_only, hide_failed=hide_failed, hide_reviewed=hide_reviewed, @@ -278,6 +288,241 @@ async def speakers_confidence(current_user: User = Depends(current_active_user)) return await data_audit_controller.speaker_confidence_overview(current_user) +class GuidedSuggestRequest(BaseModel): + speaker_name: str = Field(..., description="Enrolled speaker to improve") + batch_size: int = Field(5, ge=3, le=5, description="Maximum clips per review batch") + max_scan: int = Field( + 24, ge=4, le=48, description="Shortlist size scored per request" + ) + include_deleted: bool = Field( + False, + description="Discovery only: also index speech from soft-deleted " + "conversations whose audio chunks are still present", + ) + order: str = Field( + "informative", + pattern="^(informative|confidence)$", + description="Batch ranking: 'informative' (novelty + boundary " + "uncertainty) or 'confidence' (highest similarity first)", + ) + + +class GuidedDecisionClip(BaseModel): + conversation_id: str + start: float = Field(..., ge=0.0) + end: float = Field(..., gt=0.0) + original_start: float = Field(..., ge=0.0) + original_end: float = Field(..., gt=0.0) + decision: str = Field( + ..., pattern="^(accept|reject|skip|bad_clip|multiple_speakers|another_speaker)$" + ) + actual_speaker: Optional[str] = None + scores: Optional[Dict] = Field( + None, description="Scores shown at review time, kept for provenance" + ) + + +class GuidedDecideRequest(BaseModel): + speaker_name: str + decisions: List[GuidedDecisionClip] = Field(..., min_length=1, max_length=16) + + +@router.post("/enrollment/guided/suggest") +async def guided_enrollment_suggest( + body: GuidedSuggestRequest, + current_user: User = Depends(current_active_user), +): + """Next batch of highest-information candidate clips for one enrolled speaker: + corpus scan → embedding scores vs the speaker's gallery → ranked by novelty + + boundary uncertainty + duration, max 2 clips per conversation.""" + return await guided_enrollment_controller.suggest_clips( + current_user, body.speaker_name, body.batch_size, body.max_scan, body.order + ) + + +@router.post("/enrollment/guided/decide") +async def guided_enrollment_decide( + body: GuidedDecideRequest, + current_user: User = Depends(current_active_user), +): + """Record review decisions; confirmed clips are appended to the selected + speaker's voiceprint, and reviewed clips are not re-suggested.""" + return await guided_enrollment_controller.decide_clips( + current_user, + body.speaker_name, + [d.model_dump() for d in body.decisions], + ) + + +@router.post("/enrollment/guided/discover") +async def guided_enrollment_discover( + body: GuidedSuggestRequest, + current_user: User = Depends(current_active_user), +): + """Queue reusable corpus-speech indexing and selected-gallery matching.""" + return await guided_enrollment_controller.enqueue_corpus_discovery( + current_user, body.speaker_name, body.include_deleted + ) + + +@router.post("/enrollment/guided/mine") +async def guided_enrollment_mine( + speaker_name: str = Form(..., description="Enrolled speaker to mine for"), + files: List[UploadFile] = File(..., description="Unlabelled audio corpus"), + current_user: User = Depends(current_superuser), +): + """Upload an unlabelled audio corpus and mine it for one speaker's voice. + + Files are ingested as annotation-only conversations (no memory extraction) + and corpus discovery is chained behind their transcription jobs, so mined + clips surface in guided enrollment automatically.""" + return await guided_enrollment_controller.mine_uploaded_files( + current_user, speaker_name, files + ) + + +class GuidedMineLocalRequest(BaseModel): + speaker_name: str = Field(..., description="Enrolled speaker to mine for") + paths: List[str] = Field( + ..., + min_length=1, + max_length=1000, + description="Absolute paths under the backend data directory (/app/data)", + ) + + +@router.post("/enrollment/guided/mine-local") +async def guided_enrollment_mine_local( + body: GuidedMineLocalRequest, + current_user: User = Depends(current_superuser), +): + """Queue server-side corpus mining over files already on the data volume + (e.g. backup WAVs of purged conversations). Admin only.""" + return await guided_enrollment_controller.enqueue_local_mining( + current_user, body.speaker_name, body.paths + ) + + +@router.get("/enrollment/guided/discover") +async def guided_enrollment_discovery_state( + speaker_name: str, + current_user: User = Depends(current_active_user), +): + """Return the persisted corpus-discovery job so refreshed pages can reattach.""" + return await guided_enrollment_controller.corpus_discovery_state( + current_user, speaker_name + ) + + +@router.get("/enrollment/guided/history") +async def guided_enrollment_history( + speaker_name: str, + limit: int = 50, + current_user: User = Depends(current_active_user), +): + """Dated before/after gallery-health snapshots for enrollment sessions.""" + return await guided_enrollment_controller.enrollment_history( + current_user, speaker_name, limit + ) + + +@router.get("/enrollment/guided/gallery") +async def guided_enrollment_gallery( + speaker_name: str, + current_user: User = Depends(current_active_user), +): + """A speaker's enrolled clips with per-clip contamination flags + (self-similarity, closest other speaker, mislabel/junk/weak).""" + return await guided_enrollment_controller.gallery_clips(current_user, speaker_name) + + +class GalleryClipDeleteRequest(BaseModel): + speaker_name: str = Field(..., description="Speaker the clip must belong to") + hard: bool = Field( + False, description="Permanently delete the audio instead of quarantining" + ) + + +@router.post("/enrollment/guided/gallery/segments/{segment_id}/delete") +async def guided_enrollment_gallery_delete( + segment_id: int, + body: GalleryClipDeleteRequest, + current_user: User = Depends(current_active_user), +): + """Remove one enrolled clip from the speaker's voiceprint; the speaker + service recomputes the centroid. Quarantined (recoverable) by default.""" + return await guided_enrollment_controller.delete_gallery_clip( + current_user, body.speaker_name, segment_id, body.hard + ) + + +@router.get("/enrollment/guided/gallery/segments/{segment_id}/audio") +async def guided_enrollment_gallery_audio( + segment_id: int, + token: Optional[str] = Query( + default=None, description="JWT token for audio element access" + ), + current_user: Optional[User] = Depends(current_active_user_optional), +): + """Stream one enrolled clip for playback in the gallery panel.""" + if not current_user and token: + current_user = await get_user_from_token_param(token) + if not current_user: + raise HTTPException(status_code=401, detail="Authentication required") + from advanced_omi_backend.speaker_recognition_client import SpeakerRecognitionClient + + audio = await SpeakerRecognitionClient().get_enrollment_segment_audio(segment_id) + if audio is None: + raise HTTPException(status_code=404, detail="Clip audio not available") + return Response(content=audio, media_type="audio/wav") + + +class GuidedResetRequest(BaseModel): + speaker_name: str = Field(..., description="Speaker profile to clean") + purge_gallery: bool = Field( + False, + description="Also delete the speaker's voiceprint and enrollment audio " + "from the speaker service", + ) + + +@router.post("/enrollment/guided/reset") +async def guided_enrollment_reset( + body: GuidedResetRequest, + current_user: User = Depends(current_active_user), +): + """Forget all guided-enrollment state recorded under a speaker name + (review decisions, session history, corpus-discovery matches) so clips + become suggestible again after a delete/re-enroll. Optionally also purges + the voiceprint gallery on the speaker service.""" + return await guided_enrollment_controller.reset_speaker_state( + current_user, body.speaker_name, body.purge_gallery + ) + + +@router.post("/enrollment/benchmark") +async def run_enrollment_benchmark( + current_user: User = Depends(current_active_user), +): + """Queue five-fold conversation-grouped evaluation over human-labeled clips.""" + return await guided_enrollment_controller.enqueue_benchmark(current_user) + + +@router.get("/enrollment/benchmark/latest") +async def latest_enrollment_benchmark( + current_user: User = Depends(current_active_user), +): + return await guided_enrollment_controller.latest_benchmark(current_user) + + +@router.get("/enrollment/baseline") +async def enrollment_baseline( + current_user: User = Depends(current_active_user), +): + """Reconstruct all speaker galleries immediately before guided review began.""" + return await guided_enrollment_controller.reconstructed_baseline(current_user) + + @router.get("/triage/pending") async def triage_pending(current_user: User = Depends(current_active_user)): """Count of unapplied speaker-triage decisions and conversations they span.""" @@ -359,6 +604,27 @@ async def start_export( ) +@router.post("/import") +async def import_dataset( + dataset: UploadFile = File(..., description="Chronicle annotation dataset ZIP"), + current_user: User = Depends(current_active_user), +): + """Import WAV clips and existing transcripts from an annotation dataset ZIP. + + Imported clips are immediately available in the transcript editor and are + permanently excluded from memory processing. + """ + filename = dataset.filename or "" + if not filename.lower().endswith(".zip"): + return JSONResponse( + status_code=422, + content={"error": "Annotation dataset must be a .zip file"}, + ) + return await data_audit_controller.import_annotation_dataset( + current_user, await dataset.read() + ) + + @router.get("/exports") async def list_exports(current_user: User = Depends(current_active_user)): """List completed annotation exports with their summaries.""" diff --git a/backends/advanced/src/advanced_omi_backend/routers/modules/health_routes.py b/backends/advanced/src/advanced_omi_backend/routers/modules/health_routes.py index aa7f472f..353f6e80 100644 --- a/backends/advanced/src/advanced_omi_backend/routers/modules/health_routes.py +++ b/backends/advanced/src/advanced_omi_backend/routers/modules/health_routes.py @@ -16,7 +16,11 @@ from advanced_omi_backend import __version__ as package_version from advanced_omi_backend.client_manager import get_client_manager -from advanced_omi_backend.controllers.queue_controller import get_queue_health +from advanced_omi_backend.controllers.queue_controller import ( + get_queue_health, + redis_conn, +) +from advanced_omi_backend.heartbeat import FLEET_HEALTH_KEY, evaluate_fleet_health from advanced_omi_backend.llm_client import ( async_health_check, async_health_check_fallback, @@ -189,6 +193,7 @@ async def health_check(): worker_count = queue_health.get("total_workers", 0) active_workers = queue_health.get("active_workers", 0) idle_workers = queue_health.get("idle_workers", 0) + worker_fleet = queue_health.get("worker_fleet", {}) if redis_healthy: health_status["services"]["redis"] = { @@ -200,6 +205,22 @@ async def health_check(): "idle_workers": idle_workers, "queues": queue_health.get("queues", {}), } + + fleet_healthy = worker_fleet.get("healthy", False) + health_status["services"]["workers"] = { + **worker_fleet, + "status": ( + "✅ Worker fleet healthy" + if fleet_healthy + else f"❌ {worker_fleet.get('detail', 'Worker fleet unavailable')}" + ), + "fleet_status": worker_fleet.get("status", "unknown"), + "healthy": fleet_healthy, + "critical": True, + } + if not fleet_healthy: + overall_healthy = False + critical_services_healthy = False else: health_status["services"]["redis"] = { "status": f"❌ Connection Failed: {queue_health.get('redis_connection')}", @@ -542,8 +563,19 @@ async def readiness_check(): # Only check critical services for readiness try: - # Quick MongoDB ping to ensure we can serve requests + # MongoDB and the worker fleet are both required for accepting audio. A + # backend with no workers can receive recordings but cannot persist or + # transcribe them, so reporting ready would permit silent data loss. await asyncio.wait_for(mongo_client.admin.command("ping"), timeout=2.0) + await asyncio.wait_for(asyncio.to_thread(redis_conn.ping), timeout=2.0) + fleet = await asyncio.wait_for( + asyncio.to_thread( + lambda: evaluate_fleet_health(redis_conn.get(FLEET_HEALTH_KEY)) + ), + timeout=2.0, + ) + if not fleet["healthy"]: + raise RuntimeError(fleet.get("detail") or "Worker fleet unavailable") return JSONResponse( content={"status": "ready", "timestamp": int(time.time())}, status_code=200 ) diff --git a/backends/advanced/src/advanced_omi_backend/routers/modules/queue_routes.py b/backends/advanced/src/advanced_omi_backend/routers/modules/queue_routes.py index ddd81b08..c95582ec 100644 --- a/backends/advanced/src/advanced_omi_backend/routers/modules/queue_routes.py +++ b/backends/advanced/src/advanced_omi_backend/routers/modules/queue_routes.py @@ -486,6 +486,7 @@ async def get_queue_worker_details(current_user: User = Depends(current_active_u }, "queues": queue_health.get("queues", {}), "redis_connection": queue_health.get("redis_connection", "unknown"), + "worker_fleet": queue_health.get("worker_fleet", {}), } return status diff --git a/backends/advanced/src/advanced_omi_backend/services/audio_stream/session_store.py b/backends/advanced/src/advanced_omi_backend/services/audio_stream/session_store.py index becef711..ab752405 100644 --- a/backends/advanced/src/advanced_omi_backend/services/audio_stream/session_store.py +++ b/backends/advanced/src/advanced_omi_backend/services/audio_stream/session_store.py @@ -115,6 +115,10 @@ class SessionView: # audio seconds sent to the streaming transcription provider (session-relative # clock; persisted so word-timestamp offsets survive provider reconnects) transcription_seconds_sent: float = 0.0 + transcription_provider_status: str = "" + transcription_provider_connected_at: float = 0.0 + transcription_last_audio_sent_at: float = 0.0 + transcription_last_message_at: float = 0.0 # job ids speech_detection_job_id: str = "" audio_persistence_job_id: str = "" @@ -202,6 +206,16 @@ def fjson(key: str, default): speech_detected_at=s("speech_detected_at"), chunks_published=fint("chunks_published"), transcription_seconds_sent=ffloat("transcription_seconds_sent") or 0.0, + transcription_provider_status=s("transcription_provider_status"), + transcription_provider_connected_at=( + ffloat("transcription_provider_connected_at") or 0.0 + ), + transcription_last_audio_sent_at=( + ffloat("transcription_last_audio_sent_at") or 0.0 + ), + transcription_last_message_at=( + ffloat("transcription_last_message_at") or 0.0 + ), speech_detection_job_id=s("speech_detection_job_id"), audio_persistence_job_id=s("audio_persistence_job_id"), websocket_connected=s("websocket_connected") == "true", @@ -289,6 +303,10 @@ async def init_session( # Connection-scoped — reset on every (re)connect (see docstring). "transcription_error": "", "transcription_seconds_sent": "0", + "transcription_provider_status": "disconnected", + "transcription_provider_connected_at": "0", + "transcription_last_audio_sent_at": "0", + "transcription_last_message_at": "0", "completion_reason": "", }, ) @@ -383,7 +401,40 @@ async def set_markers(self, session_id: str, markers: list) -> None: await self._redis.hset(self._key(session_id), "markers", json.dumps(markers)) async def set_transcription_error(self, session_id: str, message: str) -> None: - await self._redis.hset(self._key(session_id), "transcription_error", message) + await self._redis.hset( + self._key(session_id), + mapping={ + "transcription_error": message, + "transcription_provider_status": "error", + }, + ) + + async def mark_transcription_provider_connected(self, session_id: str) -> None: + now = str(time.time()) + await self._redis.hset( + self._key(session_id), + mapping={ + "transcription_provider_status": "connected", + "transcription_provider_connected_at": now, + "transcription_error": "", + }, + ) + + async def mark_transcription_audio_sent(self, session_id: str) -> None: + await self._redis.hset( + self._key(session_id), "transcription_last_audio_sent_at", str(time.time()) + ) + + async def mark_transcription_provider_message(self, session_id: str) -> None: + await self._redis.hset( + self._key(session_id), "transcription_last_message_at", str(time.time()) + ) + + async def mark_transcription_provider_disconnected(self, session_id: str) -> None: + key = self._key(session_id) + status = await self._redis.hget(key, "transcription_provider_status") + if _to_str(status) != "error": + await self._redis.hset(key, "transcription_provider_status", "disconnected") async def set_transcription_seconds(self, session_id: str, seconds: float) -> None: """Persist the streaming provider's session-relative audio clock.""" diff --git a/backends/advanced/src/advanced_omi_backend/services/data_archive.py b/backends/advanced/src/advanced_omi_backend/services/data_archive.py new file mode 100644 index 00000000..e3bf1659 --- /dev/null +++ b/backends/advanced/src/advanced_omi_backend/services/data_archive.py @@ -0,0 +1,944 @@ +"""Portable, checksummed archives for Chronicle's durable user data. + +MongoDB collections are stored as concatenated BSON documents so types and +compressed audio bytes round-trip without JSON coercion. Filesystem-backed +vault and legacy audio files are included as ordinary ZIP members. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import json +import logging +import os +import shutil +import zipfile +from collections import defaultdict +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path, PurePosixPath +from typing import Any, BinaryIO, Iterable, Iterator, Optional +from urllib.parse import quote + +from bson import BSON +from pymongo import ReplaceOne + +from advanced_omi_backend.utils.audio_chunk_utils import decode_opus_to_pcm + +ARCHIVE_FORMAT = "chronicle-data-archive" +ARCHIVE_VERSION = 1 +ARCHIVE_SUFFIX = ".chronicle" +DATABASE_PREFIX = "database/" +FILES_PREFIX = "files/" +MANIFEST_PATH = "manifest.json" +FILE_ROOTS = ("conversation_docs", "memory_md", "audio_chunks") +DERIVED_MEMORY_COLLECTIONS = frozenset({"memory_audit"}) +SYNC_MARKERS = frozenset({".stfolder", ".stignore"}) + +logger = logging.getLogger(__name__) + + +class ArchiveError(RuntimeError): + """Raised when an archive is invalid or cannot be restored safely.""" + + +@dataclass(frozen=True) +class ArchiveSummary: + path: Path + collections: int + documents: int + files: int + bytes_written: int + + +@dataclass(frozen=True) +class ImportSummary: + collections: int + documents: int + files: int + skipped_collections: tuple[str, ...] + duplicate_audio_warnings: tuple["DuplicateAudioWarning", ...] + duplicate_chunk_warnings: tuple["DuplicateChunkWarning", ...] = () + + +@dataclass(frozen=True) +class DuplicateAudioWarning: + kept_conversation_id: str + skipped_conversation_id: str + fingerprint: str + kept_source: str + + +@dataclass(frozen=True) +class DuplicateChunkWarning: + conversation_id: str + chunk_index: int + kept_chunk_id: str + skipped_chunk_id: str + kept_source: str + + +class _DigestWriter: + """Track the digest and size of bytes written to another binary stream.""" + + def __init__(self, stream: BinaryIO): + self.stream = stream + self.digest = hashlib.sha256() + self.size = 0 + + def write(self, data: bytes) -> int: + written = self.stream.write(data) + if written != len(data): + raise OSError(f"Short archive write: expected {len(data)}, wrote {written}") + self.digest.update(data) + self.size += written + return written + + @property + def sha256(self) -> str: + return self.digest.hexdigest() + + +def _utc_now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _safe_member_path(member: str) -> PurePosixPath: + path = PurePosixPath(member) + if ( + path.is_absolute() + or not path.parts + or any(part in ("", ".", "..") for part in path.parts) + ): + raise ArchiveError(f"Unsafe archive member path: {member!r}") + return path + + +def _collection_member(collection_name: str) -> str: + return f"{DATABASE_PREFIX}{quote(collection_name, safe='')}.bson" + + +def _iter_regular_files(data_dir: Path) -> Iterator[tuple[Path, str]]: + for root_name in FILE_ROOTS: + root = data_dir / root_name + if not root.is_dir(): + continue + for path in sorted(root.rglob("*")): + if path.is_symlink() or not path.is_file(): + continue + relative = path.relative_to(data_dir).as_posix() + yield path, f"{FILES_PREFIX}{relative}" + + +async def create_data_archive( + database: Any, + output_path: Path, + *, + data_dir: Path, + overwrite: bool = False, +) -> ArchiveSummary: + """Export all Mongo collections and durable filesystem data to one archive.""" + output_path = output_path.expanduser().resolve() + if output_path.suffix != ARCHIVE_SUFFIX: + output_path = output_path.with_name(output_path.name + ARCHIVE_SUFFIX) + if output_path.exists() and not overwrite: + raise ArchiveError(f"Archive already exists: {output_path}") + output_path.parent.mkdir(parents=True, exist_ok=True) + temp_path = output_path.with_name(f".{output_path.name}.partial-{os.getpid()}") + if temp_path.exists(): + temp_path.unlink() + + manifest: dict[str, Any] = { + "format": ARCHIVE_FORMAT, + "schema_version": ARCHIVE_VERSION, + "created_at": _utc_now(), + "database": getattr(database, "name", None), + "file_roots": list(FILE_ROOTS), + "collections": {}, + "files": {}, + } + total_documents = 0 + total_files = 0 + + try: + with zipfile.ZipFile( + temp_path, + mode="w", + compression=zipfile.ZIP_DEFLATED, + compresslevel=6, + allowZip64=True, + ) as archive: + collection_names = sorted( + name + for name in await database.list_collection_names() + if not name.startswith("system.") + ) + for collection_name in collection_names: + member = _collection_member(collection_name) + count = 0 + with archive.open(member, mode="w", force_zip64=True) as stream: + writer = _DigestWriter(stream) + cursor = database[collection_name].find({}) + if collection_name == "audio_chunks": + cursor = cursor.sort( + [ + ("conversation_id", 1), + ("chunk_index", 1), + ("created_at", 1), + ("_id", 1), + ] + ) + async for document in cursor: + writer.write(BSON.encode(document)) + count += 1 + manifest["collections"][collection_name] = { + "member": member, + "documents": count, + } + manifest["files"][member] = { + "sha256": writer.sha256, + "size": writer.size, + "kind": "collection", + } + total_documents += count + + for source_path, member in _iter_regular_files(data_dir): + _safe_member_path(member) + with source_path.open("rb") as source, archive.open( + member, mode="w", force_zip64=True + ) as stream: + writer = _DigestWriter(stream) + for chunk in iter(lambda: source.read(1024 * 1024), b""): + writer.write(chunk) + manifest["files"][member] = { + "sha256": writer.sha256, + "size": writer.size, + "kind": "data_file", + } + total_files += 1 + + archive.writestr( + MANIFEST_PATH, + json.dumps(manifest, indent=2, sort_keys=True).encode("utf-8"), + ) + + temp_path.replace(output_path) + except Exception: + temp_path.unlink(missing_ok=True) + raise + + return ArchiveSummary( + path=output_path, + collections=len(manifest["collections"]), + documents=total_documents, + files=total_files, + bytes_written=output_path.stat().st_size, + ) + + +def _load_manifest(archive: zipfile.ZipFile) -> dict[str, Any]: + names = archive.namelist() + if len(names) != len(set(names)): + raise ArchiveError("Archive contains duplicate member names") + try: + manifest = json.loads(archive.read(MANIFEST_PATH)) + except KeyError as exc: + raise ArchiveError("Archive has no manifest.json") from exc + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise ArchiveError("Archive manifest is not valid UTF-8 JSON") from exc + + if manifest.get("format") != ARCHIVE_FORMAT: + raise ArchiveError(f"Unsupported archive format: {manifest.get('format')!r}") + if manifest.get("schema_version") != ARCHIVE_VERSION: + raise ArchiveError( + f"Unsupported archive schema version: {manifest.get('schema_version')!r}" + ) + if not isinstance(manifest.get("collections"), dict) or not isinstance( + manifest.get("files"), dict + ): + raise ArchiveError("Archive manifest is missing collections or files") + + expected = {MANIFEST_PATH, *manifest["files"].keys()} + actual = set(names) + if actual != expected: + missing = sorted(expected - actual) + extra = sorted(actual - expected) + raise ArchiveError(f"Archive member mismatch; missing={missing}, extra={extra}") + for member in manifest["files"]: + _safe_member_path(member) + return manifest + + +def verify_data_archive(archive_path: Path) -> dict[str, Any]: + """Validate structure, CRCs, sizes, and SHA-256 hashes before import.""" + try: + with zipfile.ZipFile(archive_path, mode="r", allowZip64=True) as archive: + manifest = _load_manifest(archive) + for member, expected in manifest["files"].items(): + digest = hashlib.sha256() + size = 0 + with archive.open(member, mode="r") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + size += len(chunk) + if size != expected.get("size"): + raise ArchiveError( + f"Archive size mismatch for {member}: expected " + f"{expected.get('size')}, got {size}" + ) + if digest.hexdigest() != expected.get("sha256"): + raise ArchiveError(f"Archive checksum mismatch for {member}") + return manifest + except zipfile.BadZipFile as exc: + raise ArchiveError(f"Not a valid Chronicle archive: {archive_path}") from exc + + +def _iter_bson(stream: BinaryIO) -> Iterator[dict[str, Any]]: + while True: + length_bytes = stream.read(4) + if not length_bytes: + return + if len(length_bytes) != 4: + raise ArchiveError("Truncated BSON document length") + length = int.from_bytes(length_bytes, byteorder="little", signed=True) + if length < 5 or length > 256 * 1024 * 1024: + raise ArchiveError(f"Invalid BSON document length: {length}") + remainder = stream.read(length - 4) + if len(remainder) != length - 4: + raise ArchiveError("Truncated BSON document") + yield BSON(length_bytes + remainder).decode() + + +def _iter_unique_audio_chunks( + documents: Iterable[dict[str, Any]], +) -> Iterator[tuple[dict[str, Any], Optional[DuplicateChunkWarning]]]: + """Yield the first chunk for each conversation/index and report later copies.""" + kept: dict[tuple[str, int], str] = {} + for document in documents: + conversation_id = str(document.get("conversation_id", "")) + chunk_index = int(document.get("chunk_index", 0)) + key = (conversation_id, chunk_index) + chunk_id = str(document.get("_id", "")) + if key in kept: + yield document, DuplicateChunkWarning( + conversation_id=conversation_id, + chunk_index=chunk_index, + kept_chunk_id=kept[key], + skipped_chunk_id=chunk_id, + kept_source="archive", + ) + continue + kept[key] = chunk_id + yield document, None + + +def _chunk_digest(document: dict[str, Any]) -> bytes: + audio_data = document.get("audio_data") + if not isinstance(audio_data, (bytes, bytearray)): + raise ArchiveError("Audio chunk has no binary audio_data") + digest = hashlib.sha256() + digest.update(int(document.get("chunk_index", 0)).to_bytes(8, "big", signed=False)) + digest.update(int(document.get("sample_rate", 16000)).to_bytes(4, "big")) + digest.update(int(document.get("channels", 1)).to_bytes(2, "big")) + digest.update(len(audio_data).to_bytes(8, "big")) + digest.update(audio_data) + return digest.digest() + + +def _finalize_audio_fingerprints( + chunks: dict[str, list[tuple[int, bytes]]], +) -> dict[str, str]: + fingerprints: dict[str, str] = {} + for conversation_id, chunk_digests in chunks.items(): + digest = hashlib.sha256() + digest.update(len(chunk_digests).to_bytes(8, "big")) + for chunk_index, chunk_digest in sorted(chunk_digests): + digest.update(chunk_index.to_bytes(8, "big", signed=False)) + digest.update(chunk_digest) + fingerprints[conversation_id] = digest.hexdigest() + return fingerprints + + +def _finalize_audio_structures( + chunks: dict[str, list[tuple[int, int, int, int]]], +) -> dict[str, str]: + structures: dict[str, str] = {} + for conversation_id, chunk_metadata in chunks.items(): + digest = hashlib.sha256() + digest.update(len(chunk_metadata).to_bytes(8, "big")) + for chunk_index, original_size, sample_rate, channels in sorted(chunk_metadata): + digest.update(chunk_index.to_bytes(8, "big", signed=False)) + digest.update(original_size.to_bytes(8, "big", signed=False)) + digest.update(sample_rate.to_bytes(4, "big", signed=False)) + digest.update(channels.to_bytes(2, "big", signed=False)) + structures[conversation_id] = digest.hexdigest() + return structures + + +def _archive_audio_fingerprints( + archive: zipfile.ZipFile, manifest: dict[str, Any] +) -> dict[str, str]: + metadata = manifest["collections"].get("audio_chunks") + if not metadata: + return {} + chunks: dict[str, list[tuple[int, bytes]]] = {} + with archive.open(metadata["member"], mode="r") as stream: + for document, duplicate in _iter_unique_audio_chunks(_iter_bson(stream)): + if duplicate: + continue + conversation_id = document.get("conversation_id") + if not conversation_id: + raise ArchiveError("Audio chunk has no conversation_id") + chunk_index = int(document.get("chunk_index", 0)) + chunks.setdefault(str(conversation_id), []).append( + (chunk_index, _chunk_digest(document)) + ) + return _finalize_audio_fingerprints(chunks) + + +def _archive_audio_structures( + archive: zipfile.ZipFile, manifest: dict[str, Any] +) -> dict[str, str]: + metadata = manifest["collections"].get("audio_chunks") + if not metadata: + return {} + chunks: dict[str, list[tuple[int, int, int, int]]] = {} + with archive.open(metadata["member"], mode="r") as stream: + for document, duplicate in _iter_unique_audio_chunks(_iter_bson(stream)): + if duplicate: + continue + conversation_id = document.get("conversation_id") + if not conversation_id: + raise ArchiveError("Audio chunk has no conversation_id") + chunks.setdefault(str(conversation_id), []).append( + ( + int(document.get("chunk_index", 0)), + int(document.get("original_size", 0)), + int(document.get("sample_rate", 16000)), + int(document.get("channels", 1)), + ) + ) + return _finalize_audio_structures(chunks) + + +async def _database_audio_fingerprints(database: Any) -> dict[str, str]: + chunks: dict[str, list[tuple[int, bytes]]] = {} + cursor = database["audio_chunks"].find( + {}, + projection={ + "conversation_id": 1, + "chunk_index": 1, + "sample_rate": 1, + "channels": 1, + "audio_data": 1, + }, + ) + async for document in cursor: + conversation_id = document.get("conversation_id") + if not conversation_id: + continue + chunk_index = int(document.get("chunk_index", 0)) + chunks.setdefault(str(conversation_id), []).append( + (chunk_index, _chunk_digest(document)) + ) + return _finalize_audio_fingerprints(chunks) + + +async def _database_audio_structures(database: Any) -> dict[str, str]: + chunks: dict[str, list[tuple[int, int, int, int]]] = {} + cursor = database["audio_chunks"].find( + {}, + projection={ + "conversation_id": 1, + "chunk_index": 1, + "original_size": 1, + "sample_rate": 1, + "channels": 1, + }, + ) + async for document in cursor: + conversation_id = document.get("conversation_id") + if not conversation_id: + continue + chunks.setdefault(str(conversation_id), []).append( + ( + int(document.get("chunk_index", 0)), + int(document.get("original_size", 0)), + int(document.get("sample_rate", 16000)), + int(document.get("channels", 1)), + ) + ) + return _finalize_audio_structures(chunks) + + +async def _pcm_fingerprints( + chunks: dict[str, list[dict[str, Any]]], +) -> dict[str, str]: + semaphore = asyncio.Semaphore(4) + + async def fingerprint_conversation( + conversation_id: str, documents: list[dict[str, Any]] + ) -> tuple[str, str]: + async with semaphore: + digest = hashlib.sha256() + for document in sorted( + documents, key=lambda item: item.get("chunk_index", 0) + ): + try: + pcm = await decode_opus_to_pcm( + bytes(document["audio_data"]), + int(document.get("sample_rate", 16000)), + int(document.get("channels", 1)), + ) + except Exception as exc: + raise ArchiveError( + f"Could not decode audio for duplicate check: {conversation_id}" + ) from exc + digest.update(pcm) + return conversation_id, digest.hexdigest() + + results = await asyncio.gather( + *( + fingerprint_conversation(conversation_id, documents) + for conversation_id, documents in chunks.items() + ) + ) + return dict(results) + + +async def _archive_pcm_fingerprints( + archive: zipfile.ZipFile, + manifest: dict[str, Any], + conversation_ids: set[str], +) -> dict[str, str]: + if not conversation_ids: + return {} + metadata = manifest["collections"].get("audio_chunks") + if not metadata: + return {} + chunks: dict[str, list[dict[str, Any]]] = defaultdict(list) + with archive.open(metadata["member"], mode="r") as stream: + for document, duplicate in _iter_unique_audio_chunks(_iter_bson(stream)): + if duplicate: + continue + conversation_id = str(document.get("conversation_id", "")) + if conversation_id in conversation_ids: + chunks[conversation_id].append(document) + return await _pcm_fingerprints(chunks) + + +async def _database_pcm_fingerprints( + database: Any, conversation_ids: set[str] +) -> dict[str, str]: + if not conversation_ids: + return {} + chunks: dict[str, list[dict[str, Any]]] = defaultdict(list) + cursor = database["audio_chunks"].find( + {"conversation_id": {"$in": sorted(conversation_ids)}}, + projection={ + "conversation_id": 1, + "chunk_index": 1, + "sample_rate": 1, + "channels": 1, + "audio_data": 1, + }, + ) + async for document in cursor: + conversation_id = str(document.get("conversation_id", "")) + if conversation_id in conversation_ids: + chunks[conversation_id].append(document) + return await _pcm_fingerprints(chunks) + + +def _conversation_sort_key(document: dict[str, Any]) -> tuple[float, str]: + created_at = document.get("created_at") + if isinstance(created_at, datetime): + if created_at.tzinfo is None: + created_at = created_at.replace(tzinfo=timezone.utc) + timestamp = created_at.timestamp() + else: + timestamp = float("inf") + return timestamp, str(document.get("conversation_id", "")) + + +def _archive_conversation_sort_keys( + archive: zipfile.ZipFile, manifest: dict[str, Any] +) -> dict[str, tuple[float, str]]: + metadata = manifest["collections"].get("conversations") + if not metadata: + return {} + sort_keys: dict[str, tuple[float, str]] = {} + with archive.open(metadata["member"], mode="r") as stream: + for document in _iter_bson(stream): + conversation_id = document.get("conversation_id") + if conversation_id: + sort_keys[str(conversation_id)] = _conversation_sort_key(document) + return sort_keys + + +async def _database_conversation_sort_keys( + database: Any, conversation_ids: set[str] +) -> dict[str, tuple[float, str]]: + if not conversation_ids: + return {} + keys: dict[str, tuple[float, str]] = {} + cursor = database["conversations"].find( + {"conversation_id": {"$in": sorted(conversation_ids)}}, + projection={"conversation_id": 1, "created_at": 1}, + ) + async for document in cursor: + conversation_id = document.get("conversation_id") + if conversation_id: + keys[str(conversation_id)] = _conversation_sort_key(document) + return keys + + +async def _duplicate_audio_plan( + database: Any, + archive: zipfile.ZipFile, + manifest: dict[str, Any], + *, + replace: bool, +) -> tuple[set[str], tuple[DuplicateAudioWarning, ...]]: + archive_fingerprints = _archive_audio_fingerprints(archive, manifest) + archive_structures = _archive_audio_structures(archive, manifest) + existing_fingerprints: dict[str, str] = {} + existing_structures: dict[str, str] = {} + if not replace: + existing_fingerprints = await _database_audio_fingerprints(database) + existing_structures = await _database_audio_structures(database) + + conversations_by_structure: dict[str, list[tuple[str, str]]] = defaultdict(list) + for conversation_id, structure in archive_structures.items(): + conversations_by_structure[structure].append(("archive", conversation_id)) + for conversation_id, structure in existing_structures.items(): + conversations_by_structure[structure].append(("existing", conversation_id)) + candidate_archive_ids: set[str] = set() + candidate_existing_ids: set[str] = set() + for candidates in conversations_by_structure.values(): + if len(candidates) < 2 or not any( + source == "archive" for source, _ in candidates + ): + continue + candidate_archive_ids.update( + conversation_id + for source, conversation_id in candidates + if source == "archive" + ) + candidate_existing_ids.update( + conversation_id + for source, conversation_id in candidates + if source == "existing" + ) + + archive_pcm = await _archive_pcm_fingerprints( + archive, manifest, candidate_archive_ids + ) + existing_pcm = await _database_pcm_fingerprints(database, candidate_existing_ids) + for conversation_id, fingerprint in archive_pcm.items(): + archive_fingerprints[conversation_id] = f"pcm:{fingerprint}" + for conversation_id, fingerprint in existing_pcm.items(): + existing_fingerprints[conversation_id] = f"pcm:{fingerprint}" + for conversation_id in archive_fingerprints.keys() - archive_pcm.keys(): + archive_fingerprints[conversation_id] = ( + f"compressed:{archive_fingerprints[conversation_id]}" + ) + for conversation_id in existing_fingerprints.keys() - existing_pcm.keys(): + existing_fingerprints[conversation_id] = ( + f"compressed:{existing_fingerprints[conversation_id]}" + ) + + archive_sort_keys = _archive_conversation_sort_keys(archive, manifest) + by_fingerprint: dict[str, list[str]] = {} + for conversation_id, fingerprint in archive_fingerprints.items(): + by_fingerprint.setdefault(fingerprint, []).append(conversation_id) + + skipped: set[str] = set() + warnings: list[DuplicateAudioWarning] = [] + archive_winners: dict[str, str] = {} + for fingerprint, conversation_ids in by_fingerprint.items(): + ordered = sorted( + conversation_ids, + key=lambda item: archive_sort_keys.get(item, (float("inf"), item)), + ) + winner = ordered[0] + archive_winners[fingerprint] = winner + for duplicate in ordered[1:]: + skipped.add(duplicate) + warnings.append( + DuplicateAudioWarning( + kept_conversation_id=winner, + skipped_conversation_id=duplicate, + fingerprint=fingerprint, + kept_source="archive", + ) + ) + + if not replace: + existing_by_fingerprint: dict[str, set[str]] = {} + for conversation_id, fingerprint in existing_fingerprints.items(): + existing_by_fingerprint.setdefault(fingerprint, set()).add(conversation_id) + existing_ids = { + conversation_id + for ids in existing_by_fingerprint.values() + for conversation_id in ids + } + existing_sort_keys = await _database_conversation_sort_keys( + database, existing_ids + ) + for fingerprint, archive_winner in archive_winners.items(): + existing_ids_for_audio = existing_by_fingerprint.get(fingerprint, set()) + if not existing_ids_for_audio or archive_winner in existing_ids_for_audio: + continue + existing_winner = min( + existing_ids_for_audio, + key=lambda item: existing_sort_keys.get(item, (float("inf"), item)), + ) + skipped.add(archive_winner) + warnings.append( + DuplicateAudioWarning( + kept_conversation_id=existing_winner, + skipped_conversation_id=archive_winner, + fingerprint=fingerprint, + kept_source="existing_database", + ) + ) + + for warning in warnings: + logger.warning( + "Duplicate audio skipped during import: conversation %s matches %s %s " + "(fingerprint=%s)", + warning.skipped_conversation_id, + warning.kept_source, + warning.kept_conversation_id, + warning.fingerprint, + ) + return skipped, tuple(warnings) + + +async def _restore_collection( + database: Any, + archive: zipfile.ZipFile, + collection_name: str, + member: str, + expected_count: int, + *, + replace: bool, + skipped_conversation_ids: set[str], + batch_size: int = 500, +) -> tuple[int, tuple[DuplicateChunkWarning, ...]]: + collection = database[collection_name] + if replace: + await collection.delete_many({}) + + restored = 0 + scanned = 0 + chunk_warnings: list[DuplicateChunkWarning] = [] + archive_chunk_keys: dict[tuple[str, int], str] = {} + existing_chunk_keys: dict[tuple[str, int], str] = {} + if collection_name == "audio_chunks" and not replace: + cursor = collection.find( + {}, projection={"conversation_id": 1, "chunk_index": 1} + ) + async for existing in cursor: + key = ( + str(existing.get("conversation_id", "")), + int(existing.get("chunk_index", 0)), + ) + existing_chunk_keys.setdefault(key, str(existing.get("_id", ""))) + operations: list[ReplaceOne] = [] + with archive.open(member, mode="r") as stream: + for document in _iter_bson(stream): + scanned += 1 + if "_id" not in document: + raise ArchiveError(f"Document in {collection_name} has no _id") + if str(document.get("conversation_id", "")) in skipped_conversation_ids: + continue + if collection_name == "audio_chunks": + key = ( + str(document.get("conversation_id", "")), + int(document.get("chunk_index", 0)), + ) + chunk_id = str(document["_id"]) + kept_id = existing_chunk_keys.get(key) or archive_chunk_keys.get(key) + if kept_id is not None: + warning = DuplicateChunkWarning( + conversation_id=key[0], + chunk_index=key[1], + kept_chunk_id=kept_id, + skipped_chunk_id=chunk_id, + kept_source=( + "existing_database" + if key in existing_chunk_keys + else "archive" + ), + ) + chunk_warnings.append(warning) + logger.warning( + "Duplicate audio chunk skipped during import: conversation %s " + "chunk %d (%s kept, skipped _id=%s)", + warning.conversation_id, + warning.chunk_index, + warning.kept_source, + warning.skipped_chunk_id, + ) + continue + archive_chunk_keys[key] = chunk_id + operations.append( + ReplaceOne({"_id": document["_id"]}, document, upsert=True) + ) + if len(operations) >= batch_size: + await collection.bulk_write(operations, ordered=False) + restored += len(operations) + operations.clear() + if operations: + await collection.bulk_write(operations, ordered=False) + restored += len(operations) + if scanned != expected_count: + raise ArchiveError( + f"Document count mismatch for {collection_name}: " + f"expected {expected_count}, scanned {scanned}" + ) + return restored, tuple(chunk_warnings) + + +def _destination_for_data_member(data_dir: Path, member: str) -> Path: + member_path = _safe_member_path(member) + prefix = PurePosixPath(FILES_PREFIX.rstrip("/")) + try: + relative = member_path.relative_to(prefix) + except ValueError as exc: + raise ArchiveError(f"Not a filesystem data member: {member}") from exc + if not relative.parts or relative.parts[0] not in FILE_ROOTS: + raise ArchiveError(f"Unsupported filesystem data root: {member}") + destination = data_dir.joinpath(*relative.parts) + if data_dir.resolve() not in destination.resolve().parents: + raise ArchiveError(f"Filesystem member escapes data directory: {member}") + return destination + + +def clear_vault_contents(user_root: Path) -> int: + """Delete derived vault contents while retaining Syncthing pairing markers.""" + if not user_root.is_dir(): + return 0 + deleted = 0 + for entry in user_root.iterdir(): + if entry.name in SYNC_MARKERS: + continue + if entry.is_dir() and not entry.is_symlink(): + deleted += sum(1 for path in entry.rglob("*") if path.is_file()) + shutil.rmtree(entry) + else: + entry.unlink(missing_ok=True) + deleted += 1 + return deleted + + +def _clear_restored_file_roots(data_dir: Path, root_names: Iterable[str]) -> None: + for root_name in root_names: + if root_name not in FILE_ROOTS: + raise ArchiveError(f"Unsupported filesystem data root: {root_name}") + root = data_dir / root_name + if not root.is_dir(): + continue + if root_name in ("conversation_docs", "memory_md"): + for entry in root.iterdir(): + if entry.is_dir() and not entry.is_symlink(): + clear_vault_contents(entry) + elif entry.name not in SYNC_MARKERS: + entry.unlink(missing_ok=True) + continue + for entry in root.iterdir(): + if entry.is_dir() and not entry.is_symlink(): + shutil.rmtree(entry) + else: + entry.unlink(missing_ok=True) + + +def _restore_data_files( + archive: zipfile.ZipFile, + manifest: dict[str, Any], + data_dir: Path, + *, + replace: bool, +) -> int: + members = [ + member + for member, metadata in manifest["files"].items() + if metadata.get("kind") == "data_file" + ] + if replace: + _clear_restored_file_roots( + data_dir, manifest.get("file_roots", list(FILE_ROOTS)) + ) + restored = 0 + for member in members: + destination = _destination_for_data_member(data_dir, member) + destination.parent.mkdir(parents=True, exist_ok=True) + temp_path = destination.with_name(f".{destination.name}.restore-{os.getpid()}") + try: + with archive.open(member, mode="r") as source, temp_path.open( + "wb" + ) as target: + shutil.copyfileobj(source, target, length=1024 * 1024) + temp_path.replace(destination) + except Exception: + temp_path.unlink(missing_ok=True) + raise + restored += 1 + return restored + + +async def import_data_archive( + database: Any, + archive_path: Path, + *, + data_dir: Path, + replace: bool = False, + restore_files: bool = True, + fresh_memory: bool = False, +) -> ImportSummary: + """Verify and import an archive, optionally excluding all derived memory state.""" + manifest = verify_data_archive(archive_path) + if fresh_memory and restore_files: + raise ArchiveError("fresh_memory cannot be combined with restore_files") + + skipped = set(DERIVED_MEMORY_COLLECTIONS if fresh_memory else ()) + restored_collections = 0 + restored_documents = 0 + restored_files = 0 + duplicate_warnings: tuple[DuplicateAudioWarning, ...] = () + duplicate_chunk_warnings: list[DuplicateChunkWarning] = [] + with zipfile.ZipFile(archive_path, mode="r", allowZip64=True) as archive: + skipped_conversation_ids, duplicate_warnings = await _duplicate_audio_plan( + database, archive, manifest, replace=replace + ) + for collection_name, metadata in manifest["collections"].items(): + if collection_name in skipped: + continue + restored_count, collection_chunk_warnings = await _restore_collection( + database, + archive, + collection_name, + metadata["member"], + metadata["documents"], + replace=replace, + skipped_conversation_ids=skipped_conversation_ids, + ) + restored_documents += restored_count + duplicate_chunk_warnings.extend(collection_chunk_warnings) + restored_collections += 1 + if restore_files: + restored_files = _restore_data_files( + archive, manifest, data_dir, replace=replace + ) + + return ImportSummary( + collections=restored_collections, + documents=restored_documents, + files=restored_files, + skipped_collections=tuple(sorted(skipped)), + duplicate_audio_warnings=duplicate_warnings, + duplicate_chunk_warnings=tuple(duplicate_chunk_warnings), + ) diff --git a/backends/advanced/src/advanced_omi_backend/services/memory/agent/__init__.py b/backends/advanced/src/advanced_omi_backend/services/memory/agent/__init__.py index 55e7fb4b..806522ef 100644 --- a/backends/advanced/src/advanced_omi_backend/services/memory/agent/__init__.py +++ b/backends/advanced/src/advanced_omi_backend/services/memory/agent/__init__.py @@ -1,5 +1,6 @@ """Chronicle memory agent: a tool-calling agent that maintains the markdown vault.""" +from .codex_agent import CodexMemoryAgent, codex_executor_available from .memory_agent import MemoryAgent, MemoryAgentResult, search_vault from .vault_tools import ( VAULT_SEARCH_TOOL_SCHEMAS, @@ -9,6 +10,8 @@ ) __all__ = [ + "CodexMemoryAgent", + "codex_executor_available", "MemoryAgent", "MemoryAgentResult", "search_vault", diff --git a/backends/advanced/src/advanced_omi_backend/services/memory/agent/codex_agent.py b/backends/advanced/src/advanced_omi_backend/services/memory/agent/codex_agent.py new file mode 100644 index 00000000..211759f7 --- /dev/null +++ b/backends/advanced/src/advanced_omi_backend/services/memory/agent/codex_agent.py @@ -0,0 +1,448 @@ +"""Codex CLI memory-agent executor. + +Alternative executor for the Chronicle memory agent: instead of the built-in +tool-calling loop (metered per-call API usage via the model registry), it shells out +to the OpenAI Codex CLI (``codex exec``) working directly inside the user's vault +directory — so vault recording runs on a ChatGPT subscription (``~/.codex/auth.json``, +mounted as ``CODEX_HOME`` in containers) instead of API calls. + +Selected via config.yml ``memory.agent_executor: codex``. Satisfies the same contract +as :class:`MemoryAgent` (constructor + ``run() -> MemoryAgentResult``) so the +chronicle provider's note-guarantee retry, audit recording, and job bookkeeping work +unchanged. Differences from the direct loop: + +- ``touched``/``removed`` are computed from a before/after filesystem diff, never + trusted from the CLI's own reporting. A file rename shows up as a removal (with its + pre-removal content preserved for the ledger) plus a creation — the pair is not + re-associated. +- The per-write ``vault_note_lock`` backstops don't apply (Codex edits files itself), + so the whole run holds the run-scale :func:`vault_run_lock` on the same key, and + the hard rules those tools enforced are stated in the prompt instead. +- ``force_fallback=True`` (the note-guarantee recovery attempt) delegates to the + direct :class:`MemoryAgent` — if a Codex run failed to produce the note, retrying + through a different path beats re-running the same CLI. +""" + +import asyncio +import json +import logging +import os +import shutil +import tempfile +from datetime import datetime, timezone +from pathlib import Path +from typing import Dict, List, Optional + +from ..vault_templates import CONVERSATION_TEMPLATE, PERSON_TEMPLATE, TOPIC_TEMPLATE +from .memory_agent import MemoryAgentResult, _for_prompt, _get_prompt + +logger = logging.getLogger("memory_service.agent.codex") + +CODEX_AGENT_SYSTEM_PROMPT_ID = "memory.codex_agent_system" + +# Fallback timeout when config carries none; the run lock TTL is derived from it. +DEFAULT_RUN_TIMEOUT_SECONDS = 900 +_STDERR_TAIL_CHARS = 2000 + +DEFAULT_CODEX_AGENT_SYSTEM_PROMPT = ( + """\ +You are Chronicle's memory agent. You maintain a personal Obsidian-style markdown VAULT — +the current working directory — by reading and editing its files directly. Given one +transcribed conversation, record it and update what the vault knows about the people, +topics, and things involved — making the SMALLEST edits that capture the new information. +Never regenerate a whole note when an edit will do. + +# Vault layout +- Conversations/.md — one per conversation. +- People/.md — one per person (speakers and named people). +- Topics/.md — one per recurring topic. +- /.md — notes for any OTHER recurring kind of thing (Places, Projects, + Books, Companies…). Each category has a hub note .md and a + Templates/ Template.md describing its shape. +- Templates/ holds note templates and Templates/Bases/ the aggregation views — this is + scaffolding; never write captured content there. + +Notes are aggregated by the `categories` property (a wikilink to the category hub, e.g. +`categories: ["[[People]]"]`), NOT by folder — so always set `categories` correctly. + +# Conventions (this vault follows the Kepano / "file over app" style) +- Link profusely: every person, topic, and thing is a [[wikilink]]. An unresolved link + (no note yet) is fine — it is a breadcrumb for later. +- Category names and property names are PLURAL where applicable and REUSED across + categories (org, role, date, location, topics…) so things stay findable. Prefer an + existing category/property over inventing a near-duplicate. +- Use `list` properties (`["[[A]]", "[[B]]"]`) for anything that may hold more than one value. +- Capture what was actually said; quote key facts verbatim; never invent. + +# Note templates — fill these EXACTLY (they are the schema) +Conversation note — `Conversations/.md`: +``` +""" + + _for_prompt(CONVERSATION_TEMPLATE) + + """``` +Person note — `People/.md`: +``` +""" + + _for_prompt(PERSON_TEMPLATE) + + """``` +Topic note — `Topics/.md`: +``` +""" + + _for_prompt(TOPIC_TEMPLATE) + + """``` +In a template: replace `` with the ISO date and `` with the note's title; +fill the blank properties and bullets. Copy the `![[Conversations.base#…]]` embed line +VERBATIM into every new person/topic/category note — it auto-lists that note's +conversations; never edit or remove it. + +# Organic categories +Most conversations only touch People and Topics. But when something is a substantive, +recurring KIND of thing that is not People/Topics/Conversations (a place, project, book, +company…), mint the category ONCE by hand: create `Templates/<Category> Template.md` +(model it on `Templates/Topic Template.md`, with the few short reusable frontmatter keys +its notes need), a hub note `<Category>.md` (model it on an existing hub), and — if +`Templates/Bases/` holds per-category `.base` files — a matching one copied from an +existing category's with the names substituted. Then file notes at `<Category>/<Name>.md` +with `categories: ["[[<Category>]]"]`. Do NOT over-create categories — only when the +thing will plausibly recur and matters. + +# Required outcome +Inspect the vault with the tools and reading strategy you judge appropriate before editing. +Reuse exact existing note names so [[wikilinks]] resolve. Then: +1. Create the conversation note at `Conversations/<conversation_id>.md` from the + Conversation template; put every identified person in `people:` and every theme in + `topics:` as [[wikilinks]]. +2. For each person/topic/thing: if its note exists, READ it and append only the genuinely + new facts — a bullet under `## About` and a dated line under `## Mentions`. NEVER + rewrite, re-order, or wholesale replace an existing note, never paste template + scaffold (`## About`/`## Conversations`/`## Mentions`) into a note that already has + it, and never duplicate a fact — each `## Section` heading must appear exactly once + per note. If the note does not exist, create it from the matching template. +3. If the conversation re-identifies a speaker (e.g. "Speaker 0" is actually Alice), + rename `People/<old>.md` to `People/<new>.md` AND rewrite every `[[old]]` wikilink + across the vault (`notesmd-cli move "People/<old>.md" "People/<new>.md"` does both in + one shot if installed; otherwise grep for the links and edit each file). If the target + note already exists, merge the old note's fact bullets into it instead, delete the old + note, and rewrite the links. +4. HARD RULES: `Unknown Speaker N` is a diarization placeholder, not a person — never + put it in `people:`, create a note for it, or wikilink it. Hermes is Chronicle's + voice assistant, not a human — link it as the topic [[Hermes]]; never create or + update `People/Hermes.md`. Never write captured content into Templates/ (category + minting is the only Templates/ write). Never touch files outside this directory, + never run git commands, never create scratch/plan files, and never delete a note + except when merging a rename. +5. Work until everything is recorded. Your FINAL message must be only a 1-2 sentence + summary of what you changed. + +Be precise and conservative: capture what was actually said, link things, avoid invention. +{{vault_summary}}""" +) + + +def _codex_settings() -> dict: + """The ``memory.codex`` mapping from config.yml (soft dependency — {} if absent).""" + try: + from advanced_omi_backend.model_registry import get_models_registry + + reg = get_models_registry() + mem = (reg.memory if reg else None) or {} + cfg = mem.get("codex") or {} + return dict(cfg) if isinstance(cfg, dict) else {} + except Exception as e: # noqa: BLE001 — registry optional (tests, host scripts) + logger.debug("model registry unavailable for codex settings (%s)", e) + return {} + + +def _codex_home() -> Path: + return Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex")) + + +def codex_executor_available() -> tuple[bool, str]: + """Whether the Codex CLI executor can run: binary on PATH + subscription auth.""" + binary = shutil.which(os.environ.get("CODEX_BINARY", "codex")) + if not binary: + return False, "codex binary not found on PATH" + auth = _codex_home() / "auth.json" + if not auth.is_file(): + return False, f"no Codex auth at {auth} (run `codex login` / mount CODEX_HOME)" + return True, binary + + +class CodexMemoryAgent: + """Runs one ``codex exec`` over the vault to turn a transcript into vault edits.""" + + def __init__( + self, + vault_root: Path, + operation: str = "memory_agent", + *, + force_fallback: bool = False, + ): + # `operation` is accepted for signature-compatibility with MemoryAgent; the + # Codex executor doesn't use the model registry for its own calls. + self.root = Path(vault_root) + self.operation = operation + self.force_fallback = force_fallback + + async def run( + self, + transcript: str, + conversation_id: str, + *, + date: Optional[str] = None, + duration_minutes: Optional[float] = None, + title: Optional[str] = None, + vault_summary: str = "", + guidance: str = "", + ) -> MemoryAgentResult: + if self.force_fallback: + # Note-guarantee recovery: the Codex run already failed to produce a valid + # conversation note — retry through the direct agent on the fallback LLM + # rather than re-running the same CLI. + from .memory_agent import MemoryAgent + + logger.warning( + "codex agent recovery for conv=%s: delegating to the direct memory " + "agent (fallback LLM)", + conversation_id, + ) + return await MemoryAgent(self.root, force_fallback=True).run( + transcript, + conversation_id, + date=date, + duration_minutes=duration_minutes, + title=title, + vault_summary=vault_summary, + guidance=guidance, + ) + + available, detail = codex_executor_available() + if not available: + return MemoryAgentResult( + conversation_id=conversation_id, + rounds=0, + touched=[], + summary="", + errors=[f"codex executor unavailable: {detail}"], + truncated=True, + ) + binary = detail + + date = date or datetime.now(timezone.utc).isoformat() + system_prompt = await _get_prompt( + CODEX_AGENT_SYSTEM_PROMPT_ID, + DEFAULT_CODEX_AGENT_SYSTEM_PROMPT, + vault_summary, + ) + guidance_block = f"\n\n{guidance}" if guidance else "" + prompt = ( + f"{system_prompt}\n\n" + f"New conversation to record.\n" + f"conversation_id: {conversation_id}\n" + f"date: {date}\n" + f"duration_minutes: {duration_minutes if duration_minutes is not None else 'unknown'}\n\n" + f"source_title: {title or 'unknown'}\n\n" + f"Transcript (speaker-labelled):\n{transcript}" + f"{guidance_block}" + ) + + settings = _codex_settings() + timeout = int(settings.get("timeout_seconds") or DEFAULT_RUN_TIMEOUT_SECONDS) + sandbox_mode = str(settings.get("sandbox_mode") or "workspace-write") + model = str(settings.get("model") or "") + reasoning_effort = str(settings.get("reasoning_effort") or "") + + from ..vault_lock import VaultLockTimeout + + try: + # asyncio.to_thread: the Redis run lock is sync; the subprocess itself is + # driven inside the thread too so lock lifetime and process lifetime match. + return await asyncio.to_thread( + self._run_locked, + binary, + prompt, + conversation_id, + timeout, + sandbox_mode, + model, + reasoning_effort, + ) + except VaultLockTimeout as e: + return MemoryAgentResult( + conversation_id=conversation_id, + rounds=0, + touched=[], + summary="", + errors=[str(e)], + truncated=True, + ) + + # ------------------------------------------------------------------ + # subprocess + diff (sync; runs in a worker thread under the run lock) + # ------------------------------------------------------------------ + + def _run_locked( + self, + binary: str, + prompt: str, + conversation_id: str, + timeout: int, + sandbox_mode: str, + model: str, + reasoning_effort: str, + ) -> MemoryAgentResult: + import subprocess + + from ..vault_lock import vault_run_lock + + with vault_run_lock(self.root.name, ttl_seconds=timeout + 60): + before = self._snapshot() + with tempfile.NamedTemporaryFile( + mode="r", suffix=".txt", prefix="codex-last-msg-", delete=False + ) as last_msg_file: + last_msg_path = Path(last_msg_file.name) + cmd = [ + binary, + "exec", + "--json", + "--skip-git-repo-check", # the vault is not a git repository + "--ephemeral", # don't persist session files into CODEX_HOME + "--cd", + str(self.root), + "--sandbox", + sandbox_mode, + "--output-last-message", + str(last_msg_path), + ] + if model: + cmd += ["-m", model] + if reasoning_effort: + cmd += ["-c", f'model_reasoning_effort="{reasoning_effort}"'] + cmd += ["-"] # prompt on stdin (avoids ARG_MAX / quoting) + + env = {**os.environ, "RUST_LOG": os.environ.get("RUST_LOG", "error")} + errors: List[str] = [] + stdout = "" + timed_out = False + logger.info( + "codex agent starting for conv=%s (sandbox=%s model=%s timeout=%ds)", + conversation_id, + sandbox_mode, + model or "default", + timeout, + ) + try: + proc = subprocess.run( + cmd, + input=prompt, + capture_output=True, + text=True, + timeout=timeout, + env=env, + cwd=str(self.root), + ) + stdout = proc.stdout or "" + if proc.returncode != 0: + errors.append( + f"codex exec exited {proc.returncode}: " + f"{(proc.stderr or '')[-_STDERR_TAIL_CHARS:].strip()}" + ) + except subprocess.TimeoutExpired as e: + timed_out = True + stdout = ( + e.stdout.decode() + if isinstance(e.stdout, bytes) + else (e.stdout or "") + ) + errors.append(f"codex exec timed out after {timeout}s") + except OSError as e: + errors.append(f"codex exec failed to start: {e}") + + command_count, turn_count, event_errors = self._parse_events(stdout) + errors.extend(event_errors) + + summary = "" + try: + summary = last_msg_path.read_text(encoding="utf-8").strip() + except OSError: + pass + finally: + last_msg_path.unlink(missing_ok=True) + + after = self._snapshot() + + touched = sorted( + rel for rel, content in after.items() if before.get(rel) != content + ) + removed = [ + {"old_path": rel, "new_path": "", "before": before[rel]} + for rel in sorted(before) + if rel not in after + ] + failed = timed_out or (not summary and bool(errors)) + result = MemoryAgentResult( + conversation_id=conversation_id, + rounds=max(turn_count, 1), + touched=touched, + summary=summary, + tool_calls=command_count, + removed=removed, + errors=errors, + truncated=failed, + ) + logger.info( + "codex agent done: conv=%s turns=%d commands=%d touched=%d removed=%d " + "errors=%d%s — %s", + conversation_id, + result.rounds, + command_count, + len(touched), + len(removed), + len(errors), + " (FAILED)" if failed else "", + summary[:160], + ) + return result + + def _snapshot(self) -> Dict[str, str]: + """Vault-relative ``*.md`` contents (same shape the provider's audit diff uses).""" + snapshot: Dict[str, str] = {} + if not self.root.exists(): + return snapshot + for path in self.root.rglob("*.md"): + try: + snapshot[path.relative_to(self.root).as_posix()] = path.read_text( + encoding="utf-8" + ) + except OSError: + continue + return snapshot + + @staticmethod + def _parse_events(stdout: str) -> tuple[int, int, List[str]]: + """Tolerantly scan the ``--json`` JSONL stream for counts and errors.""" + commands = 0 + turns = 0 + errors: List[str] = [] + for line in stdout.splitlines(): + line = line.strip() + if not line.startswith("{"): + continue + try: + event = json.loads(line) + except json.JSONDecodeError: + continue + etype = str(event.get("type", "")) + item = event.get("item") or {} + item_type = str(item.get("item_type") or item.get("type") or "") + if etype == "item.completed" and "command" in item_type: + commands += 1 + elif etype == "turn.completed": + turns += 1 + elif etype == "turn.failed": + turns += 1 + failure = event.get("error") or {} + errors.append(f"codex turn failed: {failure.get('message', failure)}") + elif etype == "error": + errors.append(f"codex error: {event.get('message', event)}") + return commands, turns, errors diff --git a/backends/advanced/src/advanced_omi_backend/services/memory/agent/memory_agent.py b/backends/advanced/src/advanced_omi_backend/services/memory/agent/memory_agent.py index 37df3004..4365af5d 100644 --- a/backends/advanced/src/advanced_omi_backend/services/memory/agent/memory_agent.py +++ b/backends/advanced/src/advanced_omi_backend/services/memory/agent/memory_agent.py @@ -106,6 +106,10 @@ def _for_prompt(template: str) -> str: existing category/property over inventing a near-duplicate. - Use `list` properties (`["[[A]]", "[[B]]"]`) for anything that may hold more than one value. - Capture what was actually said; quote key facts verbatim; never invent. +- `Unknown Speaker N` is a diarization placeholder, not a person. Never put it in + `people:`, create a `People/Unknown Speaker N.md` note, or wikilink it. +- Hermes is Chronicle's voice assistant/system, not a human. Link it as the recurring + topic `[[Hermes]]`; never create or update `People/Hermes.md`. # Note templates — fill these EXACTLY (they are the schema) Conversation note — `Conversations/<conversation_id>.md`: @@ -176,6 +180,10 @@ class MemoryAgentResult: touched: List[str] summary: str tool_calls: int = 0 + # Notes retired by a rename/merge this run (VaultTools.removed entries): each is + # {"old_path", "new_path", "before"}. Recorded as ``rename`` audit-ledger entries + # so a note disappearing is never invisible in the ledger. + removed: List[dict] = field(default_factory=list) errors: List[str] = field(default_factory=list) truncated: bool = ( False # loop ended on a truncated/empty LLM response, not a deliberate finish @@ -203,7 +211,13 @@ async def _get_prompt(prompt_id: str, default: str, vault_summary: str = "") -> class MemoryAgent: """Runs the tool loop that turns a transcript into vault edits.""" - def __init__(self, vault_root: Path, operation: str = "memory_agent"): + def __init__( + self, + vault_root: Path, + operation: str = "memory_agent", + *, + force_fallback: bool = False, + ): # `operation` selects the model/params from model_registry. A dedicated # "memory_agent" operation is used (not "memory_extraction", which may force # response_format=json and conflict with tool calling): reasoning models spend @@ -211,6 +225,7 @@ def __init__(self, vault_root: Path, operation: str = "memory_agent"): # operation carries a larger max_tokens budget and a low reasoning_effort. self.tools = VaultTools(vault_root) self.operation = operation + self.force_fallback = force_fallback async def run( self, @@ -219,6 +234,7 @@ async def run( *, date: Optional[str] = None, duration_minutes: Optional[float] = None, + title: Optional[str] = None, vault_summary: str = "", guidance: str = "", ) -> MemoryAgentResult: @@ -233,6 +249,7 @@ async def run( f"conversation_id: {conversation_id}\n" f"date: {date}\n" f"duration_minutes: {duration_minutes if duration_minutes is not None else 'unknown'}\n\n" + f"source_title: {title or 'unknown'}\n\n" f"Transcript (speaker-labelled):\n{transcript}" f"{guidance_block}" ) @@ -248,7 +265,10 @@ async def run( for round_idx in range(MAX_TOOL_ROUNDS): response = await async_chat_with_tools( - messages, tools=VAULT_TOOL_SCHEMAS, operation=self.operation + messages, + tools=VAULT_TOOL_SCHEMAS, + operation=self.operation, + force_fallback=self.force_fallback, ) choice = response.choices[0] msg = choice.message @@ -285,6 +305,7 @@ async def run( touched=sorted(self.tools.touched), summary=summary, tool_calls=tool_calls, + removed=list(self.tools.removed), errors=errors, truncated=True, ) @@ -301,6 +322,7 @@ async def run( touched=sorted(self.tools.touched), summary=summary, tool_calls=tool_calls, + removed=list(self.tools.removed), errors=errors, ) @@ -351,6 +373,7 @@ async def run( touched=sorted(self.tools.touched), summary="(stopped: stalled retrying a failing edit)", tool_calls=tool_calls, + removed=list(self.tools.removed), errors=errors, stalled=True, ) @@ -368,6 +391,7 @@ async def run( touched=sorted(self.tools.touched), summary="(stopped at max rounds)", tool_calls=tool_calls, + removed=list(self.tools.removed), errors=errors, ) diff --git a/backends/advanced/src/advanced_omi_backend/services/memory/agent/vault_tools.py b/backends/advanced/src/advanced_omi_backend/services/memory/agent/vault_tools.py index 4860ee88..c54488e0 100644 --- a/backends/advanced/src/advanced_omi_backend/services/memory/agent/vault_tools.py +++ b/backends/advanced/src/advanced_omi_backend/services/memory/agent/vault_tools.py @@ -77,6 +77,33 @@ def _section_counts(text: str) -> Counter: return Counter(m.group(1).lower() for m in _H2_RE.finditer(text)) +# Sections whose bullets carry accumulated facts and must survive a person-note +# merge. ``Conversations`` is intentionally excluded — it holds a dynamic +# ``![[Conversations.base]]`` embed, not facts to migrate. +_MERGE_SECTIONS = ("About", "Mentions") + + +def _extract_section_bullets(content: str, heading: str) -> List[str]: + """Return the non-empty bullet lines under a ``## heading`` (case-insensitive). + + A bare ``-`` placeholder (an empty template bullet) is skipped so migrating an + untouched section contributes nothing. + """ + want = heading.casefold() + out: List[str] = [] + in_section = False + for line in content.splitlines(): + m = _H2_RE.match(line.rstrip()) + if m: + in_section = m.group(1).casefold() == want + continue + if in_section: + stripped = line.strip() + if stripped.startswith("-") and stripped.lstrip("-").strip(): + out.append(line.rstrip()) + return out + + def _assert_no_new_section_dupes(rel: str, before: str, after: str) -> None: """Reject a mutation that *introduces* a duplicated ``## Section`` heading. @@ -112,6 +139,10 @@ def __init__(self, vault_root: Path): self._rg = shutil.which("rg") self._notesmd = os.getenv("NOTESMD_CLI_BIN") or shutil.which("notesmd-cli") self.touched: set = set() # vault-relative paths created/edited this run + # Notes retired by a rename/merge this run. Each entry is + # {"old_path", "new_path", "before"} — the audit step turns these into + # ``rename`` ledger entries so a note vanishing is never invisible. + self.removed: List[dict] = [] @contextlib.contextmanager def _locked(self) -> Iterator[None]: @@ -319,6 +350,16 @@ def write_note(self, path: str, content: str, overwrite: bool = False) -> str: # one wholesale either duplicates the body or drops accumulated facts. # Force those updates through edit_note. top_folder = Path(rel).parts[0] if len(Path(rel).parts) > 1 else "" + note_stem = Path(rel).stem + if top_folder == "People" and ( + re.fullmatch(r"unknown speaker(?:\s+\d+)?", note_stem, re.IGNORECASE) + or note_stem.casefold() == "hermes" + ): + raise VaultToolError( + "Unknown Speaker diarization placeholders and the Hermes assistant are not people; " + "do not create or link a person note for them. Use Topics/Hermes.md " + "for the Hermes assistant." + ) if existed and overwrite and top_folder in ("People", "Topics"): raise VaultToolError( f"Refusing to overwrite existing note '{rel}'. People/Topics notes " @@ -361,27 +402,79 @@ def rename_person(self, old_name: str, new_name: str) -> str: raise VaultToolError( f"Person note 'People/{old_name}.md' does not exist." ) + # Snapshot the retiring note before it moves/unlinks so the audit ledger + # keeps its final content and the merge never loses facts unrecorded. + old_content = old_fp.read_text(encoding="utf-8") if new_fp.exists(): - # Merge case — a plain move would clobber the target. Rewrite backlinks in - # Python, leave the bodies for the agent to consolidate via edit_note. + # Merge case — a plain move would clobber the target. Migrate the old + # note's facts into the target *before* deleting it (non-lossy by + # construction — never rely on a follow-up edit_note that may not come), + # rewrite backlinks, then remove the old note. + migrated = self._migrate_person_facts(old_content, new_fp, old_rel) n = self._rewrite_backlinks_python(old_name, new_name) old_fp.unlink() + self.touched.add(new_rel) + self._record_removal(old_rel, new_rel, old_content) return ( - f"'{new_name}' already existed — merged: rewrote {n} backlink(s) and " + f"'{new_name}' already existed — merged into People/{new_name}.md: " + f"migrated {migrated} fact bullet(s), rewrote {n} backlink(s), and " f"deleted People/{old_name}.md. Review People/{new_name}.md and use " - f"edit_note to consolidate any duplicated facts." + f"edit_note to de-duplicate any overlapping facts." ) self.touched.add(new_rel) if self._notesmd: try: self._move_cli(old_rel, new_rel) + self._record_removal(old_rel, new_rel, old_content) return f"Renamed People/{old_name} -> People/{new_name} (backlinks rewritten)." except Exception as e: # noqa: BLE001 logger.warning("notesmd-cli move failed (%s); using python", e) n = self._rewrite_backlinks_python(old_name, new_name) old_fp.rename(new_fp) + self._record_removal(old_rel, new_rel, old_content) return f"Renamed People/{old_name} -> People/{new_name} ({n} backlink(s) rewritten)." + def _record_removal(self, old_rel: str, new_rel: str, before: str) -> None: + """Queue a rename/merge removal for the audit ledger and clear any prior + ``touched`` entry for the vanished path (it no longer exists to re-read).""" + self.touched.discard(old_rel) + self.removed.append( + {"old_path": old_rel, "new_path": new_rel, "before": before} + ) + + def _migrate_person_facts( + self, old_content: str, new_fp: Path, old_rel: str + ) -> int: + """Append the retiring note's fact bullets into the merge target so a merge + never silently drops accumulated facts. Bullets land under the matching + ``## About`` / ``## Mentions`` heading; if the target lacks a heading, they go + into a ``## Merged from <old>`` section rather than being lost. Duplicates are + acceptable here — the agent de-duplicates afterward. Returns bullets migrated. + """ + target = new_fp.read_text(encoding="utf-8") + migrated = 0 + orphaned: List[str] = [] + for heading in _MERGE_SECTIONS: + bullets = _extract_section_bullets(old_content, heading) + if not bullets: + continue + block = "\n".join(bullets) + try: + target = apply_section_edit(target, heading, block, "append") + except SectionEditError: + orphaned.extend(bullets) + migrated += len(bullets) + if orphaned: + target = ( + target.rstrip("\n") + + f"\n\n## Merged from {old_rel}\n" + + "\n".join(orphaned) + + "\n" + ) + if migrated: + new_fp.write_text(target, encoding="utf-8") + return migrated + def _move_cli(self, old_rel: str, new_rel: str) -> None: assert self._notesmd is not None # callers gate on truthiness subprocess.run( diff --git a/backends/advanced/src/advanced_omi_backend/services/memory/audit.py b/backends/advanced/src/advanced_omi_backend/services/memory/audit.py index c28063ca..9961e9a0 100644 --- a/backends/advanced/src/advanced_omi_backend/services/memory/audit.py +++ b/backends/advanced/src/advanced_omi_backend/services/memory/audit.py @@ -48,6 +48,7 @@ class MemoryCause(str, Enum): AUTO_EXTRACTION = "auto_extraction" # automatic post-conversation pipeline MEMORY_REPLAY = "memory_replay" # manual re-extract, same inputs + MEMORY_REBUILD = "memory_rebuild" # clean vault replay from durable transcripts TRANSCRIPT_REPROCESS = "transcript_reprocess" # re-ran ASR SPEAKER_REPROCESS = "speaker_reprocess" # re-ran diarization ANNOTATION_APPLY = "annotation_apply" # user applied annotation corrections @@ -67,6 +68,7 @@ class UpdateStrategy(str, Enum): _CAUSE_KIND = { MemoryCause.AUTO_EXTRACTION: "extraction", MemoryCause.MEMORY_REPLAY: "reprocess", + MemoryCause.MEMORY_REBUILD: "reprocess", MemoryCause.TRANSCRIPT_REPROCESS: "reprocess", MemoryCause.SPEAKER_REPROCESS: "reprocess", MemoryCause.ANNOTATION_APPLY: "reprocess", @@ -77,6 +79,7 @@ class UpdateStrategy(str, Enum): _CAUSE_LABEL = { MemoryCause.AUTO_EXTRACTION: "AI extraction", MemoryCause.MEMORY_REPLAY: "Memory replay", + MemoryCause.MEMORY_REBUILD: "Memory rebuild", MemoryCause.TRANSCRIPT_REPROCESS: "Transcript reprocess", MemoryCause.SPEAKER_REPROCESS: "Speaker reprocess", MemoryCause.ANNOTATION_APPLY: "Annotation applied", diff --git a/backends/advanced/src/advanced_omi_backend/services/memory/base.py b/backends/advanced/src/advanced_omi_backend/services/memory/base.py index 9b57c747..dffbf2da 100644 --- a/backends/advanced/src/advanced_omi_backend/services/memory/base.py +++ b/backends/advanced/src/advanced_omi_backend/services/memory/base.py @@ -101,6 +101,10 @@ async def add_memory( user_email: str, allow_update: bool = False, db_helper: Any = None, + *, + source_date: Optional[str] = None, + source_duration_minutes: Optional[float] = None, + source_title: Optional[str] = None, ) -> Tuple[bool, List[str]]: """Add memories extracted from a transcript. @@ -112,6 +116,9 @@ async def add_memory( user_email: User email address allow_update: Whether to allow updating existing memories db_helper: Optional database helper for tracking relationships + source_date: Trusted source conversation timestamp + source_duration_minutes: Trusted source audio duration in minutes + source_title: Trusted source conversation title Returns: Tuple of (success: bool, created_memory_ids: List[str]) diff --git a/backends/advanced/src/advanced_omi_backend/services/memory/config.py b/backends/advanced/src/advanced_omi_backend/services/memory/config.py index e2500fca..64171921 100644 --- a/backends/advanced/src/advanced_omi_backend/services/memory/config.py +++ b/backends/advanced/src/advanced_omi_backend/services/memory/config.py @@ -42,6 +42,10 @@ class MemoryConfig: extraction_prompt: str = None extraction_enabled: bool = True timeout_seconds: int = 1200 + # How the memory agent executes: "direct" = built-in tool-calling loop (metered + # API calls via the model registry); "codex" = OpenAI Codex CLI operating on the + # vault directory (ChatGPT subscription). See agent/codex_agent.py. + agent_executor: str = "direct" def load_config_yml() -> Dict[str, Any]: @@ -140,6 +144,7 @@ def build_memory_config_from_env() -> MemoryConfig: # Timeouts/tunables from registry.memory timeout_seconds = int(mem_settings.get("timeout_seconds", 1200)) + agent_executor = str(mem_settings.get("agent_executor") or "direct").lower() memory_logger.info( f"🔧 Memory config: Provider={memory_provider_enum.value}, " @@ -154,6 +159,7 @@ def build_memory_config_from_env() -> MemoryConfig: extraction_prompt=extraction_prompt, extraction_enabled=extraction_enabled, timeout_seconds=timeout_seconds, + agent_executor=agent_executor, ) except ImportError: diff --git a/backends/advanced/src/advanced_omi_backend/services/memory/conversation_note.py b/backends/advanced/src/advanced_omi_backend/services/memory/conversation_note.py new file mode 100644 index 00000000..a9aa58ac --- /dev/null +++ b/backends/advanced/src/advanced_omi_backend/services/memory/conversation_note.py @@ -0,0 +1,197 @@ +"""Validation and deterministic rendering for generated conversation notes.""" + +from __future__ import annotations + +import json +import re +from pathlib import Path +from typing import Any, Iterable + +from ruamel.yaml import YAML + +_YAML = YAML(typ="safe") +_FRONTMATTER_BOUNDARY = re.compile(r"^---\s*$", re.MULTILINE) +_H2 = re.compile(r"^##\s+(.+?)\s*$", re.MULTILINE) +_H3 = re.compile(r"^###\s+(.+?)\s*$", re.MULTILINE) +_PLACEHOLDERS = {"", "-", "none", "n/a", "unknown", "untitled", "[ ]", "- [ ]"} + + +class ConversationNoteError(ValueError): + """The model output cannot be made into a substantive conversation note.""" + + +def _frontmatter_and_body(content: str) -> tuple[dict[str, Any], str]: + boundaries = list(_FRONTMATTER_BOUNDARY.finditer(content)) + if len(boundaries) < 2: + raise ConversationNoteError("missing YAML frontmatter") + first, second = boundaries[:2] + try: + metadata = _YAML.load(content[first.end() : second.start()]) or {} + except Exception as exc: + raise ConversationNoteError(f"invalid YAML frontmatter: {exc}") from exc + if not isinstance(metadata, dict): + raise ConversationNoteError("frontmatter must be a mapping") + body = f"{content[: first.start()]}\n{content[second.end() :]}".strip() + return metadata, body + + +def _sections(body: str) -> dict[str, str]: + matches = list(_H3.finditer(body)) + sections: dict[str, str] = {} + for index, match in enumerate(matches): + end = matches[index + 1].start() if index + 1 < len(matches) else len(body) + sections[match.group(1).strip().casefold()] = body[match.end() : end].strip() + return sections + + +def _substantive_lines(value: str) -> list[str]: + result: list[str] = [] + seen: set[str] = set() + for line in value.splitlines(): + cleaned = line.strip() + payload = re.sub(r"^-\s*(?:\[[ xX]\]\s*)?", "", cleaned).strip() + if payload.casefold() in _PLACEHOLDERS: + continue + key = cleaned.casefold() + if key not in seen: + seen.add(key) + result.append(cleaned) + return result + + +def _list_property(metadata: dict[str, Any], name: str) -> list[str]: + value = metadata.get(name, []) + if not isinstance(value, list): + return [] + items = list( + dict.fromkeys(str(item).strip() for item in value if str(item).strip()) + ) + if name == "people": + items = [ + item + for item in items + if not re.search(r"\bunknown speaker(?:\s+\d+)?\b", item, re.IGNORECASE) + ] + return items + + +def _render_list(name: str, values: Iterable[str]) -> list[str]: + values = list(values) + if not values: + return [f"{name}: []"] + return [f"{name}:", *(f" - {json.dumps(value)}" for value in values)] + + +def canonicalize_conversation_note( + path: Path, + *, + conversation_id: str, + date: str, + duration_minutes: float | None, + title: str | None, +) -> None: + """Validate model-written content and replace its metadata with trusted values. + + The LLM still extracts the semantic content, but it never controls identity, + chronology, or duration. Exact repeated lines are removed while rendering. + """ + content = path.read_text(encoding="utf-8") + if "```" in content or "\\n" in content: + raise ConversationNoteError("note contains a code fence or escaped newlines") + metadata, body = _frontmatter_and_body(content) + sections = _sections(body) + + summary_lines = _substantive_lines(sections.get("summary", "")) + fact_lines = _substantive_lines(sections.get("key facts", "")) + action_lines = _substantive_lines(sections.get("action items", "")) + if len(" ".join(summary_lines)) < 20: + raise ConversationNoteError("summary is empty or placeholder content") + if not fact_lines: + raise ConversationNoteError("key facts are empty or placeholder content") + + heading = _H2.search(body) + generated_title = heading.group(1).strip() if heading else "" + if generated_title.casefold() in _PLACEHOLDERS: + generated_title = (title or "").strip() + if generated_title.casefold() in _PLACEHOLDERS: + raise ConversationNoteError("title is empty or placeholder content") + + duration = "" if duration_minutes is None else f"{float(duration_minutes):g}" + people = _list_property(metadata, "people") + topics = _list_property(metadata, "topics") + hermes_links = [item for item in people if item.casefold() == "[[hermes]]"] + people = [item for item in people if item.casefold() != "[[hermes]]"] + if hermes_links and not any(item.casefold() == "[[hermes]]" for item in topics): + topics.append("[[Hermes]]") + + lines = [ + "---", + "categories:", + ' - "[[Conversations]]"', + f"conversation_id: {json.dumps(str(conversation_id))}", + f"date: {json.dumps(str(date))}", + *_render_list("people", people), + *_render_list("topics", topics), + f"duration_minutes: {duration}", + "---", + f"## {generated_title}", + "", + "### Summary", + *summary_lines, + "", + "### Key Facts", + *fact_lines, + "", + "### Action Items", + *(action_lines or ["- [ ]"]), + "", + ] + path.write_text("\n".join(lines), encoding="utf-8") + + +def write_source_fallback_conversation_note( + path: Path, + *, + transcript: str, + conversation_id: str, + date: str, + duration_minutes: float | None, + title: str | None, +) -> None: + """Write a minimal lossless note after both semantic LLM attempts fail.""" + excerpt = " ".join(transcript.split())[:500].strip() + if not excerpt: + raise ConversationNoteError( + "cannot create source fallback for empty transcript" + ) + safe_excerpt = excerpt.replace('"', "'") + fallback_title = (title or "").strip() or f"Conversation on {date[:10]}" + duration = "" if duration_minutes is None else f"{float(duration_minutes):g}" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + "\n".join( + [ + "---", + "categories:", + ' - "[[Conversations]]"', + f"conversation_id: {json.dumps(str(conversation_id))}", + f"date: {json.dumps(str(date))}", + "people: []", + "topics: []", + f"duration_minutes: {duration}", + "---", + f"## {fallback_title}", + "", + "### Summary", + f'The source transcript contains this short utterance: "{safe_excerpt}"', + "", + "### Key Facts", + f'- Verbatim source excerpt: "{safe_excerpt}"', + "", + "### Action Items", + "- [ ]", + "", + ] + ), + encoding="utf-8", + ) diff --git a/backends/advanced/src/advanced_omi_backend/services/memory/providers/chronicle.py b/backends/advanced/src/advanced_omi_backend/services/memory/providers/chronicle.py index 5d280fee..adfb5d43 100644 --- a/backends/advanced/src/advanced_omi_backend/services/memory/providers/chronicle.py +++ b/backends/advanced/src/advanced_omi_backend/services/memory/providers/chronicle.py @@ -15,12 +15,18 @@ import logging import time +from datetime import datetime, timezone from pathlib import Path from typing import Any, Iterable, List, Optional, Tuple from ..audit import record_vault_change from ..base import MemoryEntry, MemoryServiceBase from ..config import MemoryConfig +from ..conversation_note import ( + ConversationNoteError, + canonicalize_conversation_note, + write_source_fallback_conversation_note, +) from ..vault_manager import ConvDocVaultManager from ..vault_scaffold import is_scaffold_note, seed_vault_scaffold @@ -52,6 +58,29 @@ async def initialize(self) -> None: self._initialized = True memory_logger.info("✅ Chronicle memory service initialized (agentic vault).") + def _agent_class(self): + """The write-agent executor: the built-in tool loop, or the Codex CLI. + + Falls back to the direct agent (with a warning) when Codex is configured but + the CLI/auth is unavailable, so memory jobs keep flowing. + """ + # Lazy import: circular dependency (agent → memory_agent → llm_client → + # services.memory.config → service_factory → this module) + from ..agent import MemoryAgent + + if (getattr(self.config, "agent_executor", "direct") or "direct") == "codex": + from ..agent.codex_agent import CodexMemoryAgent, codex_executor_available + + available, detail = codex_executor_available() + if available: + return CodexMemoryAgent + memory_logger.warning( + "memory.agent_executor is 'codex' but the executor is unavailable " + "(%s) — using the direct LLM agent", + detail, + ) + return MemoryAgent + # ========================================================================= # ADD MEMORY # ========================================================================= @@ -65,12 +94,30 @@ async def add_memory( user_email: str, allow_update: bool = False, db_helper: Any = None, + *, + source_date: Optional[str] = None, + source_duration_minutes: Optional[float] = None, + source_title: Optional[str] = None, ) -> Tuple[bool, List[str]]: await self._ensure_initialized() - return await self._add_memory_agent(transcript, source_id, user_id) + return await self._add_memory_agent( + transcript, + source_id, + user_id, + source_date=source_date, + source_duration_minutes=source_duration_minutes, + source_title=source_title, + ) async def _add_memory_agent( - self, transcript: str, source_id: str, user_id: str + self, + transcript: str, + source_id: str, + user_id: str, + *, + source_date: Optional[str] = None, + source_duration_minutes: Optional[float] = None, + source_title: Optional[str] = None, ) -> Tuple[bool, List[str]]: """Write path via the tool-calling memory agent. @@ -79,10 +126,6 @@ async def _add_memory_agent( vault-relative note paths stand in for the chunk/memory ids the older index path returned, so the existing job bookkeeping (counts, versions) works unchanged. """ - # Lazy import: circular dependency (agent → memory_agent → llm_client → - # services.memory.config → service_factory → this module) - from ..agent import MemoryAgent - if not transcript or len(transcript.strip()) < 10: memory_logger.info(f"Skipping empty transcript for {source_id}") return True, [] @@ -94,8 +137,15 @@ async def _add_memory_agent( # per-user vault_note_lock (lock-write-unlock, never across LLM calls). seed_vault_scaffold(user_root) # idempotent: .base + hub notes existing_before = self._vault_note_set(user_root) - agent = MemoryAgent(user_root) - result = await agent.run(transcript, source_id) + result = await self._run_agent_with_note_guarantee( + self._agent_class(), + user_root, + transcript, + source_id, + source_date=source_date, + source_duration_minutes=source_duration_minutes, + source_title=source_title, + ) if (result.truncated or result.stalled) and not result.touched: reason = ( "truncated LLM response" if result.truncated else "stalled retry loop" @@ -110,6 +160,21 @@ async def _add_memory_agent( time.perf_counter() - t0, ) return False, [] + await self._record_agent_touches( + user_id, + source_id, + user_root, + result.touched, + existing_before, + removed=result.removed, + ) + expected_note = user_root / "Conversations" / f"{Path(source_id).name}.md" + if not expected_note.is_file(): + memory_logger.error( + "❌ add_memory(agent) %s: required conversation note was not created", + source_id, + ) + return False, result.touched memory_logger.info( "✅ add_memory(agent) %s: touched=%d rounds=%d tools=%d errors=%d (%.2fs) — %s", source_id, @@ -120,11 +185,117 @@ async def _add_memory_agent( time.perf_counter() - t0, result.summary[:160], ) - await self._record_agent_touches( - user_id, source_id, user_root, result.touched, existing_before - ) return True, result.touched + async def _run_agent_with_note_guarantee( + self, + agent_class, + user_root: Path, + transcript: str, + source_id: str, + *, + guidance: str = "", + source_date: Optional[str] = None, + source_duration_minutes: Optional[float] = None, + source_title: Optional[str] = None, + ): + """Retry when the exact conversation note is absent or fails validation.""" + trusted_date = source_date or datetime.now(timezone.utc).isoformat() + result = await agent_class(user_root).run( + transcript, + source_id, + date=trusted_date, + duration_minutes=source_duration_minutes, + title=source_title, + guidance=guidance, + ) + note_name = Path(source_id).name + expected_note = user_root / "Conversations" / f"{note_name}.md" + if self._canonicalize_conversation_note( + expected_note, + source_id, + trusted_date, + source_duration_minutes, + source_title, + ): + return result + + memory_logger.warning( + "Memory agent did not create Conversations/%s.md; retrying with the " + "configured fallback LLM", + note_name, + ) + recovery_guidance = (f"{guidance}\n\n" if guidance else "") + ( + "RECOVERY REQUIREMENT: the previous attempt did not create the required " + f"conversation note. You MUST write it at exactly Conversations/{note_name}.md " + f"using conversation_id {source_id}. Do not alter or abbreviate the ID. " + "The Summary and Key Facts sections MUST contain substantive text. For a " + "short or low-information transcript, summarize the exact utterance rather " + "than leaving either section blank." + ) + recovery = await agent_class(user_root, force_fallback=True).run( + transcript, + source_id, + date=trusted_date, + duration_minutes=source_duration_minutes, + title=source_title, + guidance=recovery_guidance, + ) + recovery.rounds += result.rounds + recovery.tool_calls += result.tool_calls + recovery.touched = list(dict.fromkeys((*result.touched, *recovery.touched))) + recovery.removed = [*result.removed, *recovery.removed] + recovery.errors = [*result.errors, *recovery.errors] + recovery_valid = self._canonicalize_conversation_note( + expected_note, + source_id, + trusted_date, + source_duration_minutes, + source_title, + ) + if not recovery_valid: + memory_logger.warning( + "Both memory-agent attempts produced an invalid note for %s; " + "writing a source-preserving fallback", + source_id, + ) + write_source_fallback_conversation_note( + expected_note, + transcript=transcript, + conversation_id=source_id, + date=trusted_date, + duration_minutes=source_duration_minutes, + title=source_title, + ) + recovery.touched = list( + dict.fromkeys((*recovery.touched, f"Conversations/{note_name}.md")) + ) + return recovery + + @staticmethod + def _canonicalize_conversation_note( + path: Path, + source_id: str, + source_date: str, + source_duration_minutes: Optional[float], + source_title: Optional[str], + ) -> bool: + if not path.is_file(): + return False + try: + canonicalize_conversation_note( + path, + conversation_id=source_id, + date=source_date, + duration_minutes=source_duration_minutes, + title=source_title, + ) + return True + except ConversationNoteError as exc: + memory_logger.warning("Invalid conversation note %s: %s", path, exc) + path.unlink(missing_ok=True) + return False + async def _reprocess_memory_agent( self, transcript: str, @@ -140,10 +311,6 @@ async def _reprocess_memory_agent( backlinks) instead of leaving orphaned ``[[Speaker 0]]`` notes. Person/topic notes are kept and surgically updated — only the conversation note is regenerated. """ - # Lazy import: circular dependency (agent → memory_agent → llm_client → - # services.memory.config → service_factory → this module) - from ..agent import MemoryAgent - if not transcript or len(transcript.strip()) < 10: memory_logger.info(f"Skipping empty transcript for {source_id}") return True, [] @@ -163,8 +330,13 @@ async def _reprocess_memory_agent( if conv_note.exists(): conv_note.unlink() - agent = MemoryAgent(user_root) - result = await agent.run(transcript, source_id, guidance=guidance) + result = await self._run_agent_with_note_guarantee( + self._agent_class(), + user_root, + transcript, + source_id, + guidance=guidance, + ) if result.truncated and not result.touched: memory_logger.error( "❌ reprocess_memory(agent) %s: aborted on truncated LLM response after " @@ -175,6 +347,20 @@ async def _reprocess_memory_agent( time.perf_counter() - t0, ) return False, [] + await self._record_agent_touches( + user_id, + source_id, + user_root, + result.touched, + existing_before, + removed=result.removed, + ) + if not conv_note.is_file(): + memory_logger.error( + "❌ reprocess_memory(agent) %s: required conversation note was not created", + source_id, + ) + return False, result.touched memory_logger.info( "✅ reprocess_memory(agent) %s: touched=%d rounds=%d tools=%d errors=%d (%.2fs) — %s", source_id, @@ -185,16 +371,21 @@ async def _reprocess_memory_agent( time.perf_counter() - t0, result.summary[:160], ) - await self._record_agent_touches( - user_id, source_id, user_root, result.touched, existing_before - ) return True, result.touched - def _vault_note_set(self, user_root: Path) -> set: - """Vault-relative paths of all notes currently on disk (for create/update audit).""" + def _vault_note_set(self, user_root: Path) -> dict[str, str]: + """Snapshot vault-relative note contents for create/update audit records.""" if not user_root.exists(): - return set() - return {p.relative_to(user_root).as_posix() for p in user_root.rglob("*.md")} + return {} + snapshot: dict[str, str] = {} + for path in user_root.rglob("*.md"): + try: + snapshot[path.relative_to(user_root).as_posix()] = path.read_text( + encoding="utf-8" + ) + except OSError: + continue + return snapshot async def _record_agent_touches( self, @@ -202,9 +393,28 @@ async def _record_agent_touches( source_id: str, user_root: Path, touched: Iterable[str], - existing_before: set, + existing_before: dict[str, str], + removed: Optional[Iterable[dict]] = None, ) -> None: - """Record one audit-ledger entry per note the memory agent changed.""" + """Record one audit-ledger entry per note the memory agent changed. + + ``removed`` are notes retired by a rename/merge (``VaultTools.removed``): each + is logged as a ``rename`` entry carrying the pre-removal content, so a note + disappearing from the vault is never invisible in the ledger (the gap that made + a rename look like an unexplained clobber followed later by a fresh ``create``). + """ + for entry in removed or (): + await record_vault_change( + user_id=user_id, + conversation_id=source_id, + operation="rename", + note_path=entry.get("old_path"), + before=entry.get("before"), + after=None, + agent_mode=True, + summary=f"renamed/merged into {entry.get('new_path')}", + new_path=entry.get("new_path"), + ) for rel in sorted(touched): try: after: Optional[str] = (user_root / rel).read_text(encoding="utf-8") @@ -216,6 +426,7 @@ async def _record_agent_touches( conversation_id=source_id, operation="create" if is_new else "update", note_path=rel, + before=existing_before.get(rel), after=after, agent_mode=True, summary=( diff --git a/backends/advanced/src/advanced_omi_backend/services/memory/rebuild.py b/backends/advanced/src/advanced_omi_backend/services/memory/rebuild.py new file mode 100644 index 00000000..d755a086 --- /dev/null +++ b/backends/advanced/src/advanced_omi_backend/services/memory/rebuild.py @@ -0,0 +1,399 @@ +"""Bulk reconstruction of derived Markdown memories from durable transcripts.""" + +from __future__ import annotations + +import os +import tarfile +import uuid +from collections import defaultdict +from dataclasses import dataclass, replace +from datetime import datetime, timezone +from enum import Enum +from pathlib import Path +from typing import Any, Iterable, Optional + +from rq.exceptions import NoSuchJobError +from rq.job import Dependency, Job + +from advanced_omi_backend.controllers.queue_controller import ( + JOB_RESULT_TTL, + default_queue, + memory_queue, + post_conv_enqueue_kwargs, + transcription_queue, +) +from advanced_omi_backend.services.data_archive import clear_vault_contents +from advanced_omi_backend.services.memory.audit import MemoryCause, UpdateStrategy +from advanced_omi_backend.workers.memory_jobs import enqueue_memory_processing +from advanced_omi_backend.workers.speaker_jobs import recognise_speakers_job + +VAULT_ROOTS = ("conversation_docs", "memory_md") +MEMORY_REBUILD_JOB_TIMEOUT = 7200 + +import logging + +logger = logging.getLogger(__name__) + + +class MemoryRebuildError(RuntimeError): + """Raised when a clean memory rebuild cannot start safely.""" + + +class RebuildStage(str, Enum): + """Earliest derived-data stage to rerun during reconstruction.""" + + MEMORY = "memory" + SPEAKERS = "speakers" + + +@dataclass(frozen=True) +class RebuildConversation: + conversation_id: str + user_id: str + created_at: Any + transcript_version_id: str + memory_excluded: bool = False + has_audio: bool = True + active_transcript_version_id: Optional[str] = None + + +@dataclass(frozen=True) +class RebuildPlan: + conversations: tuple[RebuildConversation, ...] + user_ids: tuple[str, ...] + + @property + def count(self) -> int: + return len(self.conversations) + + @property + def memory_count(self) -> int: + return sum(not item.memory_excluded for item in self.conversations) + + @property + def speaker_count(self) -> int: + return sum(item.has_audio for item in self.conversations) + + +@dataclass(frozen=True) +class RebuildResult: + run_id: str + jobs: tuple[str, ...] + speaker_jobs: tuple[str, ...] + skipped_speaker_conversations: tuple[str, ...] + memory_jobs: tuple[str, ...] + from_stage: RebuildStage + user_ids: tuple[str, ...] + deleted_vault_files: int + deleted_audit_entries: int + vault_backup: Optional[Path] + + +def _normalise_user_ids(user_ids: Optional[Iterable[str]]) -> tuple[str, ...]: + if not user_ids: + return () + normalised = tuple(sorted({str(user_id) for user_id in user_ids})) + if any(not user_id or Path(user_id).name != user_id for user_id in normalised): + raise MemoryRebuildError("Invalid user ID") + return normalised + + +def _speaker_source_version(document: dict[str, Any]) -> str: + """Walk back prior speaker-only versions to the underlying ASR transcript.""" + active_id = document["active_transcript_version"] + versions = { + version.get("version_id"): version + for version in document.get("transcript_versions", []) + if isinstance(version, dict) and version.get("version_id") + } + current_id = active_id + visited: set[str] = set() + while current_id not in visited: + visited.add(current_id) + version = versions.get(current_id) or {} + metadata = version.get("metadata") or {} + if metadata.get("reprocessing_type") != "speaker_diarization": + break + source_id = metadata.get("source_version_id") + if not source_id or source_id not in versions: + break + current_id = source_id + return current_id + + +async def build_rebuild_plan( + database: Any, + user_ids: Optional[Iterable[str]] = None, + *, + from_stage: RebuildStage = RebuildStage.MEMORY, +) -> RebuildPlan: + """Select replayable conversations in stable chronological order.""" + requested_users = _normalise_user_ids(user_ids) + query: dict[str, Any] = { + "active_transcript_version": {"$ne": None}, + "deleted": {"$ne": True}, + } + if from_stage is RebuildStage.MEMORY: + query["memory_excluded"] = {"$ne": True} + if requested_users: + query["user_id"] = {"$in": list(requested_users)} + + cursor = ( + database["conversations"] + .find( + query, + projection={ + "conversation_id": 1, + "user_id": 1, + "created_at": 1, + "active_transcript_version": 1, + "transcript_versions.version_id": 1, + "transcript_versions.metadata": 1, + "memory_excluded": 1, + }, + ) + .sort([("user_id", 1), ("created_at", 1), ("conversation_id", 1)]) + ) + conversations_list: list[RebuildConversation] = [] + async for document in cursor: + conversations_list.append( + RebuildConversation( + conversation_id=document["conversation_id"], + user_id=str(document["user_id"]), + created_at=document.get("created_at"), + transcript_version_id=( + _speaker_source_version(document) + if from_stage is RebuildStage.SPEAKERS + else document["active_transcript_version"] + ), + memory_excluded=document.get("memory_excluded", False) is True, + active_transcript_version_id=document["active_transcript_version"], + ) + ) + if from_stage is RebuildStage.SPEAKERS and conversations_list: + conversation_ids = [item.conversation_id for item in conversations_list] + audio_conversation_ids = set( + await database["audio_chunks"].distinct( + "conversation_id", + { + "conversation_id": {"$in": conversation_ids}, + "deleted": {"$ne": True}, + }, + ) + ) + conversations_list = [ + replace(item, has_audio=item.conversation_id in audio_conversation_ids) + for item in conversations_list + ] + conversations = tuple(conversations_list) + selected_users = tuple(sorted({item.user_id for item in conversations})) + if requested_users: + selected_users = requested_users + return RebuildPlan(conversations=conversations, user_ids=selected_users) + + +def _active_rebuild_jobs(conversation_ids: set[str]) -> list[str]: + job_ids: set[str] = set() + for queue in (transcription_queue, memory_queue, default_queue): + job_ids.update(queue.get_job_ids()) + for registry in ( + queue.started_job_registry, + queue.deferred_job_registry, + queue.scheduled_job_registry, + ): + job_ids.update(registry.get_job_ids()) + + active: list[str] = [] + for job_id in job_ids: + try: + job = Job.fetch(job_id, connection=memory_queue.connection) + except NoSuchJobError: + continue + if (job.meta or {}).get("conversation_id") in conversation_ids: + active.append(job_id) + return sorted(active) + + +def _enqueue_speaker_rebuild( + item: RebuildConversation, + *, + run_id: str, + sequence: int, + depends_on: Optional[str], +) -> Job: + """Create a speaker version from the imported active transcript.""" + dependency = Dependency(jobs=depends_on, allow_failure=True) if depends_on else None + target_version_id = str(uuid.uuid4()) + return transcription_queue.enqueue( + recognise_speakers_job, + item.conversation_id, + target_version_id, + "", + None, + item.transcript_version_id, + job_timeout=1200, + result_ttl=JOB_RESULT_TTL, + job_id=(f"speaker_rebuild_{run_id}_{sequence}_" f"{item.conversation_id[:12]}"), + description=f"Rebuild speakers for {item.conversation_id[:8]}", + **post_conv_enqueue_kwargs( + "speaker", + { + "conversation_id": item.conversation_id, + "source_version_id": item.transcript_version_id, + "version_id": target_version_id, + "trigger": "archive_rebuild", + }, + depends_on=dependency, + ), + ) + + +def _backup_user_vaults( + data_dir: Path, user_ids: tuple[str, ...], backup_dir: Path +) -> Optional[Path]: + vaults: list[Path] = [] + for root_name in VAULT_ROOTS: + for user_id in user_ids: + user_root = data_dir / root_name / user_id + if user_root.is_dir(): + vaults.append(user_root) + if not vaults: + return None + + backup_dir.mkdir(parents=True, exist_ok=True) + timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S") + final_path = backup_dir / f"memory_vault_{timestamp}.tar.gz" + temp_path = final_path.with_name(f".{final_path.name}.partial-{os.getpid()}") + try: + with tarfile.open(temp_path, mode="w:gz") as archive: + for vault in vaults: + archive.add(vault, arcname=vault.relative_to(data_dir).as_posix()) + temp_path.replace(final_path) + except Exception: + temp_path.unlink(missing_ok=True) + raise + return final_path + + +async def execute_memory_rebuild( + database: Any, + plan: RebuildPlan, + *, + data_dir: Path, + backup_dir: Optional[Path] = None, + from_stage: RebuildStage = RebuildStage.MEMORY, +) -> RebuildResult: + """Clear derived memory state and enqueue ordered replay chains per user.""" + if not plan.user_ids: + raise MemoryRebuildError("No users matched the rebuild request") + if not plan.conversations: + raise MemoryRebuildError("No conversations with active transcripts matched") + + conversation_ids = {item.conversation_id for item in plan.conversations} + active_jobs = _active_rebuild_jobs(conversation_ids) + if active_jobs: + preview = ", ".join(active_jobs[:5]) + raise MemoryRebuildError( + "Existing speaker or memory jobs target this rebuild set. Wait for them " + "to finish or " + f"stop/cancel them first: {preview}" + ) + + vault_backup = None + if backup_dir is not None: + vault_backup = _backup_user_vaults(data_dir, plan.user_ids, backup_dir) + + deleted_files = 0 + for root_name in VAULT_ROOTS: + for user_id in plan.user_ids: + deleted_files += clear_vault_contents(data_dir / root_name / user_id) + + audit_result = await database["memory_audit"].delete_many( + {"user_id": {"$in": list(plan.user_ids)}} + ) + deleted_audit_entries = int(audit_result.deleted_count) + + if from_stage is RebuildStage.SPEAKERS: + # Memory is allowed to continue after a failed speaker job. Resetting the + # pointer first guarantees that continuation reads the clean ASR source, + # never the previously generated (and potentially polluted) speaker version. + for item in plan.conversations: + if ( + item.active_transcript_version_id + and item.active_transcript_version_id != item.transcript_version_id + ): + await database["conversations"].update_one( + {"conversation_id": item.conversation_id}, + {"$set": {"active_transcript_version": item.transcript_version_id}}, + ) + + run_id = uuid.uuid4().hex[:12] + by_user: dict[str, list[RebuildConversation]] = defaultdict(list) + for item in plan.conversations: + by_user[item.user_id].append(item) + + speaker_jobs: list[str] = [] + skipped_speaker_conversations: list[str] = [] + memory_jobs: list[str] = [] + for user_id in plan.user_ids: + user_conversations = by_user[user_id] + memory_dependency = None + if from_stage is RebuildStage.SPEAKERS: + speaker_dependency = None + speaker_conversations = [ + item for item in user_conversations if item.has_audio + ] + skipped = [item for item in user_conversations if not item.has_audio] + for item in skipped: + logger.warning( + "Speaker rebuild skipped conversation %s: no stored audio chunks", + item.conversation_id, + ) + skipped_speaker_conversations.append(item.conversation_id) + for sequence, item in enumerate(speaker_conversations, start=1): + speaker_job = _enqueue_speaker_rebuild( + item, + run_id=run_id, + sequence=sequence, + depends_on=speaker_dependency, + ) + speaker_jobs.append(speaker_job.id) + speaker_dependency = speaker_job.id + if speaker_dependency: + memory_dependency = Dependency( + jobs=speaker_dependency, + allow_failure=True, + ) + + memory_conversations = [ + item for item in user_conversations if not item.memory_excluded + ] + for sequence, item in enumerate(memory_conversations, start=1): + job = enqueue_memory_processing( + item.conversation_id, + cause=MemoryCause.MEMORY_REBUILD, + strategy=UpdateStrategy.FULL, + depends_on=memory_dependency, + job_timeout=MEMORY_REBUILD_JOB_TIMEOUT, + job_id=( + f"memory_rebuild_{run_id}_{sequence}_" + f"{item.conversation_id[:12]}" + ), + ) + memory_jobs.append(job.id) + memory_dependency = job + + jobs = tuple((*speaker_jobs, *memory_jobs)) + + return RebuildResult( + run_id=run_id, + jobs=jobs, + speaker_jobs=tuple(speaker_jobs), + skipped_speaker_conversations=tuple(skipped_speaker_conversations), + memory_jobs=tuple(memory_jobs), + from_stage=from_stage, + user_ids=plan.user_ids, + deleted_vault_files=deleted_files, + deleted_audit_entries=deleted_audit_entries, + vault_backup=vault_backup, + ) diff --git a/backends/advanced/src/advanced_omi_backend/services/memory/vault_lock.py b/backends/advanced/src/advanced_omi_backend/services/memory/vault_lock.py index 5d70e26d..1b4a67c7 100644 --- a/backends/advanced/src/advanced_omi_backend/services/memory/vault_lock.py +++ b/backends/advanced/src/advanced_omi_backend/services/memory/vault_lock.py @@ -43,6 +43,39 @@ class VaultLockTimeout(Exception): """The per-user vault lock could not be acquired within the wait window.""" +@contextlib.contextmanager +def vault_run_lock(user_id: str, ttl_seconds: int = 1200) -> Iterator[None]: + """Hold the per-user vault lock for a whole external-executor run. + + The Codex CLI executor edits vault files directly for the duration of one agent + run (minutes, not milliseconds), so it takes the SAME ``vault:write:{user_id}`` + key with a run-scale TTL instead of the per-operation one. While it is held, + per-operation writers (``vault_note_lock``) block for their 30s window and then + fail closed with a retryable error — acceptable because memory jobs are already + serialised on one worker, so contention is limited to rare chat-driven writes. + """ + client = create_sync_redis(decode_responses=True) + lock = client.lock( + f"vault:write:{user_id}", + timeout=ttl_seconds, + blocking_timeout=_LOCK_WAIT_SECONDS, + ) + try: + if not lock.acquire(): + raise VaultLockTimeout( + f"vault run lock for user {user_id} not acquired within " + f"{_LOCK_WAIT_SECONDS}s" + ) + try: + yield + finally: + with contextlib.suppress(Exception): + lock.release() + finally: + with contextlib.suppress(Exception): + client.close() + + @contextlib.contextmanager def vault_note_lock(user_id: str) -> Iterator[None]: """Hold the per-user lock around ONE vault-mutating filesystem operation. diff --git a/backends/advanced/src/advanced_omi_backend/services/observability/health_poller.py b/backends/advanced/src/advanced_omi_backend/services/observability/health_poller.py index 5d6f7855..f8578e0f 100644 --- a/backends/advanced/src/advanced_omi_backend/services/observability/health_poller.py +++ b/backends/advanced/src/advanced_omi_backend/services/observability/health_poller.py @@ -11,7 +11,8 @@ ``error`` (``critical`` for a crash loop); returning to healthy → ``info``. 2. **Failed RQ jobs** — newly-failed jobs (hard crashes / timeouts) that the per-site soft-failure taps can't see, deduped by job id. -3. **Config diagnostics** — new configuration *issues* (errors), forgotten when +3. **Worker fleet heartbeat** — missing, stale, or unhealthy supervisor state. +4. **Config diagnostics** — new configuration *issues* (errors), forgotten when resolved so they re-alarm if they recur. Last-known state lives in Redis so it survives a backend restart without @@ -32,21 +33,78 @@ get_config_diagnostics, get_external_services, ) +from advanced_omi_backend.heartbeat import FLEET_HEALTH_KEY, evaluate_fleet_health from advanced_omi_backend.redis_factory import create_async_redis from advanced_omi_backend.services.observability.system_events import record_event logger = logging.getLogger("observability.health_poller") -POLL_INTERVAL_SECS = 30 +POLL_INTERVAL_SECS = 15 # Let services finish booting before the first sample (avoid "starting" noise). -INITIAL_DELAY_SECS = 45 +INITIAL_DELAY_SECS = 20 # Redis keys for last-known state. _HEALTH_KEY = "system:health:last" # hash: "{node}/{service}" -> health _SEEN_FAILED_KEY = "system:health:seen_failed_jobs" # set of job ids _CONFIG_SEEN_KEY = "system:health:config_issues" # set of issue keys +_WORKER_HEALTH_FIELD = "internal/workers-fleet" _SEEN_FAILED_TTL = 7 * 24 * 3600 +_BAD_WORKER_STATES = {"missing", "stale", "invalid", "unhealthy"} + + +async def _poll_worker_fleet(redis, *, now: float | None = None) -> None: + """Record worker fleet outage and recovery transitions.""" + result = evaluate_fleet_health(await redis.get(FLEET_HEALTH_KEY), now=now) + health = result["status"] + previous = await redis.hget(_HEALTH_KEY, _WORKER_HEALTH_FIELD) + + # A fresh orchestrator can legitimately be in its startup grace period. Keep a + # prior outage active until it reports healthy, but do not alarm on startup. + if health == "starting": + if previous is None: + await redis.hset(_HEALTH_KEY, _WORKER_HEALTH_FIELD, health) + return + + if health == previous: + return + + previous_bad = previous in _BAD_WORKER_STATES + current_bad = health in _BAD_WORKER_STATES + await redis.hset(_HEALTH_KEY, _WORKER_HEALTH_FIELD, health) + + metadata = { + "health": health, + "previous": previous, + "heartbeat_age_seconds": round(result.get("age_seconds", 0), 1), + "workers_total": result.get("workers_total"), + "workers_alive": result.get("workers_alive"), + } + if current_bad and not previous_bad: + await record_event( + severity="critical", + category="service", + source="workers", + title="Worker fleet unavailable", + detail=( + f"{result.get('detail')}. Audio persistence, speech detection, " + "transcription, memory, and background jobs may not run." + ), + metadata=metadata, + ) + elif previous_bad and health == "healthy": + await record_event( + severity="info", + category="service", + source="workers", + title="Worker fleet recovered", + detail=( + f"The worker supervisor reports {result.get('workers_alive', 0)}/" + f"{result.get('workers_total', 0)} child processes alive." + ), + metadata=metadata, + ) + def _bad_severity(health: str | None, detail: str) -> str | None: """Return the event severity for a bad health state, or None if it's fine.""" @@ -177,6 +235,7 @@ async def run_health_poller(app=None) -> None: while True: try: await _poll_external_services(redis) + await _poll_worker_fleet(redis) await _poll_failed_jobs(redis) await _poll_config_diagnostics(redis) except asyncio.CancelledError: diff --git a/backends/advanced/src/advanced_omi_backend/services/transcription/__init__.py b/backends/advanced/src/advanced_omi_backend/services/transcription/__init__.py index 9812d477..ae6916a8 100644 --- a/backends/advanced/src/advanced_omi_backend/services/transcription/__init__.py +++ b/backends/advanced/src/advanced_omi_backend/services/transcription/__init__.py @@ -7,9 +7,11 @@ """ import asyncio +import hashlib import json import logging import re +from datetime import datetime, timezone from typing import Optional from urllib.parse import urlencode @@ -263,6 +265,98 @@ async def transcribe( progress_callback=None, priority: bool = False, **kwargs, + ) -> dict: + """Transcribe with a persistent response cache. + + Batch providers are typically paid, per-minute APIs; re-transcribing + identical audio (reprocessing, retries after downstream failures, bulk + speaker mining over backup files) must not bill twice. The normalized + result is stored in Mongo keyed by the audio content hash plus the + provider configuration (minus the API key, so key rotation keeps the + cache). Hot-word/context hints are deliberately NOT part of the key — + a hint tweak isn't worth re-billing the whole corpus. Cache failures + never block transcription. + """ + cache = None + cache_key = None + if self.model.model_provider != "mock": + try: + config = ( + self.model.model_dump() + if hasattr(self.model, "model_dump") + else dict(vars(self.model)) + ) + config.pop("api_key", None) + fingerprint = json.dumps( + { + "provider": self._name, + "config": config, + "diarize": diarize, + "sample_rate": sample_rate, + }, + sort_keys=True, + default=str, + ) + cache_key = { + "audio_sha256": hashlib.sha256(audio_data).hexdigest(), + "request_sha256": hashlib.sha256(fingerprint.encode()).hexdigest(), + } + from advanced_omi_backend.models.conversation import Conversation + + cache = Conversation.get_pymongo_collection().database[ + "transcription_response_cache" + ] + row = await cache.find_one(cache_key, {"result": 1}) + if row and row.get("result") is not None: + logger.info( + f"♻️ Transcription cache hit for '{self._name}' " + f"({len(audio_data)} bytes) — reusing stored response, " + "no provider call" + ) + return row["result"] + except Exception as e: + logger.debug(f"Transcription cache lookup skipped: {e}") + cache = None + + result = await self._transcribe_uncached( + audio_data, + sample_rate, + diarize=diarize, + context_info=context_info, + progress_callback=progress_callback, + priority=priority, + **kwargs, + ) + + if cache is not None and cache_key is not None and result: + try: + # Mongo documents cap at 16MB; skip pathological payloads. + if len(json.dumps(result, default=str)) < 12_000_000: + await cache.update_one( + cache_key, + { + "$set": { + "provider": self._name, + "audio_bytes": len(audio_data), + "result": result, + "created_at": datetime.now(timezone.utc), + } + }, + upsert=True, + ) + except Exception as e: + logger.debug(f"Transcription cache write skipped: {e}") + return result + + async def _transcribe_uncached( + self, + audio_data: bytes, + sample_rate: int, + diarize: bool = False, + context_info: Optional[str] = None, + progress_callback=None, + priority: bool = False, + **kwargs, ) -> dict: # Special handling for mock provider (no HTTP server needed) if self.model.model_provider == "mock": @@ -469,15 +563,6 @@ async def transcribe( msg += f"Service error: {detail}" raise RuntimeError(msg) from e - # DEBUG: Log Deepgram response structure - if "results" in data and "channels" in data.get("results", {}): - channels = data["results"]["channels"] - if channels and "alternatives" in channels[0]: - alt = channels[0]["alternatives"][0] - logger.debug( - f"DEBUG Registry: Deepgram alternative keys: {list(alt.keys())}" - ) - # Extract normalized shape text, words, segments = "", [], [] extract = (op.get("response", {}) or {}).get("extract") or {} @@ -740,6 +825,10 @@ async def process_audio_chunk( except asyncio.TimeoutError: # No message available yet return None + except (ConnectionError, OSError, websockets.exceptions.ConnectionClosed): + # Let the consumer reconnect; a dead transport is different from a + # healthy socket that simply has no transcript available yet. + raise except Exception as e: logger.error(f"Error processing audio chunk result for {client_id}: {e}") return None diff --git a/backends/advanced/src/advanced_omi_backend/services/transcription/streaming_consumer.py b/backends/advanced/src/advanced_omi_backend/services/transcription/streaming_consumer.py index 08f68ff8..a1d4323d 100644 --- a/backends/advanced/src/advanced_omi_backend/services/transcription/streaming_consumer.py +++ b/backends/advanced/src/advanced_omi_backend/services/transcription/streaming_consumer.py @@ -356,7 +356,9 @@ async def start_session_stream( "last_activity": time.time(), "sample_rate": sample_rate, "time_offset": time_offset, + "last_health_persisted_at": 0.0, } + await self.store.mark_transcription_provider_connected(session_id) # Only buffer audio for speaker identification when provider lacks diarization if not self._provider_has_diarization: @@ -506,6 +508,12 @@ async def end_session_stream(self, session_id: str): completion_status = "error" finally: + try: + await self.store.mark_transcription_provider_disconnected(session_id) + except Exception: + logger.warning( + f"Failed to mark transcription provider disconnected for {session_id}" + ) # Cleanup must run on both paths (previously the error path leaked # active_sessions / audio buffer entries). self.active_sessions.pop(session_id, None) @@ -555,14 +563,22 @@ async def process_audio_chunk( # send — a chunk that dies mid-send is re-sent after reconnect. session = self.active_sessions.get(session_id) if session is not None: - session["last_activity"] = time.time() + now = time.time() + session["last_activity"] = now self._session_audio_seconds[session_id] = ( self._session_audio_seconds.get(session_id, 0.0) + len(audio_chunk) / (session.get("sample_rate", 16000) * 2) ) + # Audio chunks can arrive several times per second. Persisting + # this health timestamp at most every five seconds keeps the + # cross-worker signal useful without amplifying Redis traffic. + if now - session.get("last_health_persisted_at", 0.0) >= 5.0: + await self.store.mark_transcription_audio_sent(session_id) + session["last_health_persisted_at"] = now # Provider returns None if no response yet, or a dict with results if result: + await self.store.mark_transcription_provider_message(session_id) is_final = result.get("is_final", False) text = result.get("text", "") words = result.get("words") or [] diff --git a/backends/advanced/src/advanced_omi_backend/speaker_recognition_client.py b/backends/advanced/src/advanced_omi_backend/speaker_recognition_client.py index 7bfc7436..ee03b9d2 100644 --- a/backends/advanced/src/advanced_omi_backend/speaker_recognition_client.py +++ b/backends/advanced/src/advanced_omi_backend/speaker_recognition_client.py @@ -18,6 +18,7 @@ import traceback import uuid import wave +from datetime import datetime from pathlib import Path from typing import Dict, List, Optional @@ -35,6 +36,61 @@ logger = logging.getLogger(__name__) +def _select_label_mappings( + label_votes: Dict[str, List[tuple[str, float]]], + *, + similarity_threshold: float, +) -> Dict[str, tuple[str, float]]: + """Choose conservative, one-to-one names for diarized speaker labels. + + Two agreeing samples are enough. A lone sample must clear a deliberately + stricter threshold, and an enrolled identity may name only one diarized label. + """ + candidates: list[tuple[str, str, int, float]] = [] + single_sample_threshold = max(0.65, similarity_threshold + 0.15) + for label, votes in label_votes.items(): + by_name: Dict[str, List[float]] = {} + for name, confidence in votes: + by_name.setdefault(name, []).append(float(confidence)) + if not by_name: + continue + best_name = max( + by_name, + key=lambda name: ( + len(by_name[name]), + sum(by_name[name]) / len(by_name[name]), + ), + ) + scores = by_name[best_name] + average = sum(scores) / len(scores) + if len(scores) < 2 and average < single_sample_threshold: + logger.info( + "Speaker label %r left unknown: one sample confidence %.3f < %.3f", + label, + average, + single_sample_threshold, + ) + continue + candidates.append((label, best_name, len(scores), average)) + + selected: Dict[str, tuple[str, float]] = {} + used_names: set[str] = set() + for label, name, vote_count, average in sorted( + candidates, key=lambda item: (item[2], item[3]), reverse=True + ): + identity = name.casefold() + if identity in used_names: + logger.info( + "Speaker label %r left unknown: identity %r already assigned", + label, + name, + ) + continue + used_names.add(identity) + selected[label] = (name, average) + return selected + + class SpeakerRecognitionClient: """Client for communicating with the speaker recognition service.""" @@ -525,9 +581,9 @@ async def _identify_one(seg: Dict) -> Optional[Dict]: await asyncio.gather(*all_tasks, return_exceptions=True) # Majority-vote per label - label_mapping: Dict[str, tuple] = {} # label -> (identified_name, confidence) + label_votes: Dict[str, List[tuple[str, float]]] = {} for label, tasks in label_tasks.items(): - name_votes: Dict[str, List[float]] = {} + votes: List[tuple[str, float]] = [] for task in tasks: try: result = task.result() @@ -536,27 +592,17 @@ async def _identify_one(seg: Dict) -> Optional[Dict]: if result and result.get("found"): name = result.get("speaker_name", "Unknown") confidence = result.get("confidence", 0.0) - name_votes.setdefault(name, []).append(confidence) - - if name_votes: - # Pick name with most votes, break ties by average confidence - best_name = max( - name_votes.keys(), - key=lambda n: ( - len(name_votes[n]), - sum(name_votes[n]) / len(name_votes[n]), - ), - ) - avg_confidence = sum(name_votes[best_name]) / len(name_votes[best_name]) - label_mapping[label] = (best_name, avg_confidence) - logger.info( - f"🎤 Label '{label}' -> '{best_name}' " - f"({len(name_votes[best_name])}/{len(tasks)} votes, conf={avg_confidence:.3f})" - ) - else: + votes.append((name, confidence)) + label_votes[label] = votes + if not votes: logger.info( f"🎤 Label '{label}' -> no identification (keeping original)" ) + label_mapping = _select_label_mappings( + label_votes, similarity_threshold=similarity_threshold + ) + for label, (name, confidence) in label_mapping.items(): + logger.info("🎤 Label %r -> %r (conf=%.3f)", label, name, confidence) # Build result segments in same format as diarize_identify_match() # Non-speech segments are kept but not speaker-identified @@ -1312,6 +1358,213 @@ async def append_to_speaker(self, speaker_id: str, audio_data: bytes) -> Dict: logger.error(f"🎤 ❌ Error appending to speaker: {e}") return {"error": "unknown_error", "message": str(e)} + async def score_enrollment_candidate( + self, audio_wav_bytes: bytes, speaker_id: str + ) -> Dict: + """Score a candidate clip's enrollment value for one target speaker. + + POST /enrollment/candidates/score — returns sim_centroid (cosine to the + target's centroid), max_clip_sim (redundancy vs the target's per-clip + gallery), n_gallery_clips, best_other ({speaker_id, name, score} of the + closest other enrolled speaker), and duration. + """ + if not self.enabled: + return {"error": "speaker_recognition_disabled"} + + try: + async with aiohttp.ClientSession() as session: + form_data = aiohttp.FormData() + form_data.add_field( + "file", + audio_wav_bytes, + filename="candidate.wav", + content_type="audio/wav", + ) + form_data.add_field("speaker_id", speaker_id) + + async with session.post( + f"{self.service_url}/enrollment/candidates/score", + data=form_data, + timeout=aiohttp.ClientTimeout(total=30), + ) as response: + if response.status != 200: + response_text = await response.text() + logger.warning( + f"🎤 /enrollment/candidates/score returned {response.status}: {response_text}" + ) + return {"error": "score_failed", "status": response.status} + return await response.json() + + except aiohttp.ClientError as e: + logger.error(f"🎤 Failed to score enrollment candidate: {e}") + return {"error": "connection_failed", "message": str(e)} + + async def get_enrollment_health( + self, user_id: int = 1, before: Optional[datetime] = None + ) -> Dict: + """Return per-clip gallery cohesion and contamination metrics.""" + if not self.enabled: + return {"error": "speaker_recognition_disabled"} + try: + async with aiohttp.ClientSession() as session: + params = {"user_id": user_id} + if before is not None: + params["before"] = before.isoformat() + async with session.get( + f"{self.service_url}/enrollment/health", + params=params, + timeout=aiohttp.ClientTimeout(total=30), + ) as response: + if response.status != 200: + return { + "error": "health_audit_failed", + "status": response.status, + } + return await response.json() + except aiohttp.ClientError as e: + logger.error(f"🎤 Failed to audit enrollment health: {e}") + return {"error": "connection_failed", "message": str(e)} + + async def get_enrollment_segment_audio(self, segment_id: int) -> Optional[bytes]: + """Fetch one enrolled clip's audio for playback (None if unavailable).""" + if not self.enabled: + return None + try: + async with aiohttp.ClientSession() as session: + async with session.get( + f"{self.service_url}/enrollment/segments/{segment_id}/audio", + timeout=aiohttp.ClientTimeout(total=30), + ) as response: + if response.status != 200: + return None + return await response.read() + except aiohttp.ClientError as e: + logger.error(f"🎤 Failed to fetch enrollment segment audio: {e}") + return None + + async def delete_enrollment_segment( + self, segment_id: int, hard: bool = False + ) -> Dict: + """Remove one clip from a speaker's voiceprint (quarantined by default); + the service recomputes the speaker's centroid.""" + if not self.enabled: + return {"error": "speaker_recognition_disabled"} + try: + async with aiohttp.ClientSession() as session: + form_data = aiohttp.FormData() + form_data.add_field("hard", "true" if hard else "false") + async with session.post( + f"{self.service_url}/enrollment/segments/{segment_id}/delete", + data=form_data, + timeout=aiohttp.ClientTimeout(total=30), + ) as response: + if response.status != 200: + response_text = await response.text() + logger.warning( + f"🎤 Segment delete returned {response.status}: {response_text}" + ) + return { + "error": "segment_delete_failed", + "status": response.status, + } + return await response.json() + except aiohttp.ClientError as e: + logger.error(f"🎤 Failed to delete enrollment segment: {e}") + return {"error": "connection_failed", "message": str(e)} + + async def delete_speaker(self, speaker_id: str, delete_audio: bool = True) -> Dict: + """Delete an enrolled speaker (and, by default, their enrollment audio).""" + if not self.enabled: + return {"error": "speaker_recognition_disabled"} + try: + async with aiohttp.ClientSession() as session: + async with session.delete( + f"{self.service_url}/speakers/{speaker_id}", + params={"delete_audio": "true" if delete_audio else "false"}, + timeout=aiohttp.ClientTimeout(total=30), + ) as response: + if response.status != 200: + response_text = await response.text() + logger.warning( + f"🎤 Speaker delete returned {response.status}: {response_text}" + ) + return { + "error": "speaker_delete_failed", + "status": response.status, + } + return await response.json() + except aiohttp.ClientError as e: + logger.error(f"🎤 Failed to delete speaker: {e}") + return {"error": "connection_failed", "message": str(e)} + + async def extract_speaker_embedding(self, audio_wav_bytes: bytes) -> Dict: + """Extract an evaluation embedding without mutating speaker enrollment.""" + if not self.enabled: + return {"error": "speaker_recognition_disabled"} + try: + async with aiohttp.ClientSession() as session: + form_data = aiohttp.FormData() + form_data.add_field( + "file", + audio_wav_bytes, + filename="evaluation.wav", + content_type="audio/wav", + ) + async with session.post( + f"{self.service_url}/enrollment/candidates/embed", + data=form_data, + timeout=aiohttp.ClientTimeout(total=30), + ) as response: + if response.status != 200: + return {"error": "embedding_failed", "status": response.status} + return await response.json() + except aiohttp.ClientError as e: + logger.error(f"🎤 Failed to extract evaluation embedding: {e}") + return {"error": "connection_failed", "message": str(e)} + + async def score_cached_embeddings( + self, speaker_id: str, embeddings: list[list[float]] + ) -> Dict: + """Score cached corpus embeddings against the current live gallery.""" + if not self.enabled: + return {"error": "speaker_recognition_disabled"} + try: + async with aiohttp.ClientSession() as session: + async with session.post( + f"{self.service_url}/enrollment/candidates/score-embeddings", + json={"speaker_id": speaker_id, "embeddings": embeddings}, + timeout=aiohttp.ClientTimeout(total=120), + ) as response: + if response.status != 200: + return { + "error": "embedding_score_failed", + "status": response.status, + "message": await response.text(), + } + return await response.json() + except aiohttp.ClientError as e: + logger.error(f"🎤 Failed to score cached embeddings: {e}") + return {"error": "connection_failed", "message": str(e)} + + async def get_embedding_info(self) -> Dict: + """Return the active speaker embedding model fingerprint.""" + if not self.enabled: + return {"error": "speaker_recognition_disabled"} + try: + async with aiohttp.ClientSession() as session: + async with session.get( + f"{self.service_url}/enrollment/candidates/embedding-info", + timeout=aiohttp.ClientTimeout(total=10), + ) as response: + if response.status != 200: + return { + "error": "embedding_info_failed", + "status": response.status, + } + return await response.json() + except aiohttp.ClientError as e: + return {"error": "connection_failed", "message": str(e)} + async def reidentify_clusters( self, clusters: Dict[str, list], diff --git a/backends/advanced/src/advanced_omi_backend/utils/annotation_import.py b/backends/advanced/src/advanced_omi_backend/utils/annotation_import.py new file mode 100644 index 00000000..341a9c89 --- /dev/null +++ b/backends/advanced/src/advanced_omi_backend/utils/annotation_import.py @@ -0,0 +1,312 @@ +"""Read Chronicle annotation dataset ZIPs without extracting them to disk.""" + +import hashlib +import io +import json +import zipfile +from dataclasses import dataclass +from pathlib import PurePosixPath +from typing import Any + +from advanced_omi_backend.utils.annotation_export import MANIFEST_NAME, META_NAME + +SCHEMA_VERSION = 1 +MAX_ARCHIVE_BYTES = 512 * 1024 * 1024 +MAX_UNCOMPRESSED_BYTES = 2 * 1024 * 1024 * 1024 +MAX_CLIPS = 500 + + +class AnnotationDatasetError(ValueError): + """The uploaded ZIP does not satisfy Chronicle's annotation dataset contract.""" + + +@dataclass(frozen=True) +class AnnotationClip: + clip_id: str + audio_path: str + source_conversation_id: str + conversation_title: str + source_client_id: str + transcript: str + transcript_source: str + segments: list[dict[str, Any]] + audio_bytes: bytes + sample_rate: int + duration_seconds: float + notes: str | None + + +@dataclass(frozen=True) +class AnnotationDataset: + dataset_id: str + schema_version: int + clips: list[AnnotationClip] + + +def _safe_archive_path(value: Any, field: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise AnnotationDatasetError(f"{field} must be a non-empty path") + path = PurePosixPath(value) + if path.is_absolute() or ".." in path.parts or "\\" in value: + raise AnnotationDatasetError(f"Unsafe {field}: {value!r}") + return str(path) + + +def _normalise_segments( + value: Any, *, clip_id: str, duration_seconds: float +) -> list[dict[str, Any]]: + if not isinstance(value, list): + raise AnnotationDatasetError(f"Clip {clip_id!r} segments must be a list") + + segments = [] + for index, raw in enumerate(value): + if not isinstance(raw, dict): + raise AnnotationDatasetError( + f"Clip {clip_id!r} segment {index} must be an object" + ) + try: + start = float(raw["start"]) + end = float(raw["end"]) + except (KeyError, TypeError, ValueError) as exc: + raise AnnotationDatasetError( + f"Clip {clip_id!r} segment {index} has invalid timing" + ) from exc + if start < 0 or end <= start or end > duration_seconds + 0.1: + raise AnnotationDatasetError( + f"Clip {clip_id!r} segment {index} falls outside the audio duration" + ) + text = raw.get("text") + if not isinstance(text, str): + raise AnnotationDatasetError( + f"Clip {clip_id!r} segment {index} text must be a string" + ) + speaker = raw.get("speaker") or "Unknown Speaker" + if not isinstance(speaker, str): + raise AnnotationDatasetError( + f"Clip {clip_id!r} segment {index} speaker must be a string" + ) + identified_as = raw.get("identified_as") + if identified_as is not None and not isinstance(identified_as, str): + raise AnnotationDatasetError( + f"Clip {clip_id!r} segment {index} identified_as must be a string or null" + ) + segments.append( + { + "start": start, + "end": end, + "speaker": speaker, + "identified_as": identified_as, + "text": text, + } + ) + return segments + + +def _active_transcript( + record: dict[str, Any], duration_seconds: float +) -> tuple[str, list, str]: + clip_id = record["clip_id"] + annotation = record.get("annotation") or {} + if not isinstance(annotation, dict): + raise AnnotationDatasetError(f"Clip {clip_id!r} annotation must be an object") + + annotation_text = annotation.get("text") + annotation_segments = annotation.get("segments") + uses_annotation = annotation_text is not None or annotation_segments is not None + + raw_segments = ( + annotation_segments + if annotation_segments is not None + else record.get("segments", []) + ) + segments = _normalise_segments( + raw_segments, clip_id=clip_id, duration_seconds=duration_seconds + ) + + if annotation_text is not None: + if not isinstance(annotation_text, str): + raise AnnotationDatasetError( + f"Clip {clip_id!r} annotation.text must be a string or null" + ) + transcript = annotation_text + if annotation_segments is None: + segments = ( + [ + { + "start": 0.0, + "end": duration_seconds, + "speaker": "Unknown Speaker", + "identified_as": None, + "text": transcript, + } + ] + if transcript.strip() + else [] + ) + elif uses_annotation: + transcript = " ".join(segment["text"].strip() for segment in segments).strip() + else: + transcript = record.get("text", "") + if not isinstance(transcript, str): + raise AnnotationDatasetError(f"Clip {clip_id!r} text must be a string") + + return transcript, segments, "human_annotation" if uses_annotation else "manifest" + + +def parse_annotation_dataset(archive_bytes: bytes) -> AnnotationDataset: + """Validate and read an export-compatible annotation dataset ZIP.""" + if not archive_bytes: + raise AnnotationDatasetError("Annotation dataset is empty") + if len(archive_bytes) > MAX_ARCHIVE_BYTES: + raise AnnotationDatasetError( + "Annotation dataset exceeds the 512 MB upload limit" + ) + + try: + archive = zipfile.ZipFile(io.BytesIO(archive_bytes)) + except zipfile.BadZipFile as exc: + raise AnnotationDatasetError( + "Annotation dataset must be a valid ZIP file" + ) from exc + + with archive: + names = set() + total_size = 0 + for info in archive.infolist(): + safe_name = _safe_archive_path(info.filename, "ZIP entry") + if safe_name in names: + raise AnnotationDatasetError(f"Duplicate ZIP entry: {safe_name}") + names.add(safe_name) + total_size += info.file_size + if total_size > MAX_UNCOMPRESSED_BYTES: + raise AnnotationDatasetError( + "Annotation dataset expands beyond the 2 GB limit" + ) + if MANIFEST_NAME not in names: + raise AnnotationDatasetError( + f"Annotation dataset is missing {MANIFEST_NAME}" + ) + + manifest_bytes = archive.read(MANIFEST_NAME) + metadata = {} + if META_NAME in names: + try: + metadata = json.loads(archive.read(META_NAME)) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise AnnotationDatasetError(f"Invalid {META_NAME}") from exc + if not isinstance(metadata, dict): + raise AnnotationDatasetError(f"{META_NAME} must contain an object") + + schema_version = metadata.get("schema_version", SCHEMA_VERSION) + if schema_version != SCHEMA_VERSION: + raise AnnotationDatasetError( + f"Unsupported annotation dataset schema version: {schema_version}" + ) + dataset_id = metadata.get("export_id") + if not isinstance(dataset_id, str) or not dataset_id.strip(): + dataset_id = ( + f"annotation_import_{hashlib.sha256(manifest_bytes).hexdigest()[:16]}" + ) + + try: + manifest_text = manifest_bytes.decode("utf-8") + except UnicodeDecodeError as exc: + raise AnnotationDatasetError(f"{MANIFEST_NAME} must be UTF-8") from exc + + clips = [] + seen_clip_ids = set() + seen_audio_paths = set() + for line_number, line in enumerate(manifest_text.splitlines(), start=1): + if not line.strip(): + continue + if len(clips) >= MAX_CLIPS: + raise AnnotationDatasetError( + f"Annotation dataset exceeds {MAX_CLIPS} clips" + ) + try: + record = json.loads(line) + except json.JSONDecodeError as exc: + raise AnnotationDatasetError( + f"Invalid JSON on {MANIFEST_NAME} line {line_number}" + ) from exc + if not isinstance(record, dict): + raise AnnotationDatasetError( + f"{MANIFEST_NAME} line {line_number} must contain an object" + ) + + clip_id = record.get("clip_id") + if not isinstance(clip_id, str) or not clip_id.strip(): + raise AnnotationDatasetError( + f"{MANIFEST_NAME} line {line_number} has no clip_id" + ) + if clip_id in seen_clip_ids: + raise AnnotationDatasetError(f"Duplicate clip_id: {clip_id}") + seen_clip_ids.add(clip_id) + + audio_path = _safe_archive_path(record.get("audio_path"), "audio_path") + if audio_path in seen_audio_paths: + raise AnnotationDatasetError(f"Duplicate audio_path: {audio_path}") + seen_audio_paths.add(audio_path) + if audio_path not in names: + raise AnnotationDatasetError( + f"Clip {clip_id!r} is missing audio file {audio_path!r}" + ) + if PurePosixPath(audio_path).suffix.lower() != ".wav": + raise AnnotationDatasetError( + f"Clip {clip_id!r} audio_path must reference a WAV file" + ) + + try: + duration_seconds = float(record["duration_seconds"]) + sample_rate = int(record["sample_rate"]) + except (KeyError, TypeError, ValueError) as exc: + raise AnnotationDatasetError( + f"Clip {clip_id!r} has invalid duration or sample rate" + ) from exc + if duration_seconds <= 0 or sample_rate <= 0: + raise AnnotationDatasetError( + f"Clip {clip_id!r} has invalid duration or sample rate" + ) + + transcript, segments, transcript_source = _active_transcript( + record, duration_seconds + ) + annotation = record.get("annotation") or {} + notes = annotation.get("notes") + if notes is not None and not isinstance(notes, str): + raise AnnotationDatasetError( + f"Clip {clip_id!r} annotation.notes must be a string or null" + ) + source_conversation_id = record.get("conversation_id") or clip_id + if not isinstance(source_conversation_id, str): + raise AnnotationDatasetError( + f"Clip {clip_id!r} conversation_id must be a string" + ) + + clips.append( + AnnotationClip( + clip_id=clip_id, + audio_path=audio_path, + source_conversation_id=source_conversation_id, + conversation_title=str(record.get("conversation_title") or clip_id), + source_client_id=str( + record.get("client_id") or "annotation-import" + ), + transcript=transcript, + transcript_source=transcript_source, + segments=segments, + audio_bytes=archive.read(audio_path), + sample_rate=sample_rate, + duration_seconds=duration_seconds, + notes=notes, + ) + ) + + if not clips: + raise AnnotationDatasetError(f"{MANIFEST_NAME} contains no clips") + + return AnnotationDataset( + dataset_id=dataset_id, + schema_version=schema_version, + clips=clips, + ) diff --git a/backends/advanced/src/advanced_omi_backend/utils/silence_condense.py b/backends/advanced/src/advanced_omi_backend/utils/silence_condense.py new file mode 100644 index 00000000..f914d88b --- /dev/null +++ b/backends/advanced/src/advanced_omi_backend/utils/silence_condense.py @@ -0,0 +1,216 @@ +"""Silence-aware condensing for paid batch transcription. + +Batch STT providers bill by audio duration — silence included. Long-form +recordings (wearables, meeting rooms) are often mostly silence, so before +sending audio to the provider we run the local TEN VAD over the PCM, cut out +silences longer than ``CUT_GAP_SECONDS`` (keeping ``PAD_SECONDS`` of context +around speech), and transcribe the condensed audio instead. Word and segment +timestamps are then mapped back onto the original timeline, splitting any +segment that spans a cut, so everything downstream (speaker identification, +enrollment clips, playback) still refers to real conversation time. + +Condensing is skipped when it wouldn't save at least ``MIN_SAVINGS_FRACTION`` +of the audio, and produces an empty result without any provider call when the +audio contains no speech at all. +""" + +import logging +from typing import List, Optional, Tuple + +import numpy as np + +from advanced_omi_backend.services.vad import get_vad_provider +from advanced_omi_backend.utils.vad_analysis import ( + _pcm_to_mono_int16, + frame_speech_intervals, + merge_speech_regions, +) + +logger = logging.getLogger(__name__) + +MIN_AUDIO_SECONDS = 10.0 # below this, condensing overhead isn't worth it +MIN_SAVINGS_FRACTION = 0.15 # send original unless we cut at least this much +CUT_GAP_SECONDS = 3.0 # only cut silences longer than this +PAD_SECONDS = 0.5 # context kept around each speech region +VAD_SAMPLE_RATE = 16000 + +# (condensed_start_seconds, original_start_seconds, length_seconds) per region +CondenseMap = List[Tuple[float, float, float]] + + +def condense_silence( + pcm_data: bytes, sample_rate: int, channels: int, sample_width: int +) -> Tuple[bytes, Optional[CondenseMap], Optional[float]]: + """Cut long silences out of PCM audio using the local VAD. + + Returns ``(pcm, mapping, speech_seconds)``: + - ``mapping is None`` — audio unchanged (too short, wrong format, VAD + unavailable, or not enough silence to be worth cutting). + - ``mapping == []`` — no speech at all; ``pcm`` is empty and the caller + should skip transcription entirely. + - otherwise — ``pcm`` is the condensed audio and ``mapping`` translates + condensed time back to original time (see :func:`remap_condensed_result`). + """ + if sample_width != 2: + return pcm_data, None, None + frame_bytes = sample_width * channels + total_samples = len(pcm_data) // frame_bytes + if not sample_rate or total_samples <= 0: + return pcm_data, None, None + duration = total_samples / sample_rate + if duration < MIN_AUDIO_SECONDS: + return pcm_data, None, None + + try: + mono = _pcm_to_mono_int16(pcm_data, channels) + if sample_rate != VAD_SAMPLE_RATE: + # Resample for VAD scoring only — the audio we send is untouched. + n16 = int(mono.size * VAD_SAMPLE_RATE / sample_rate) + positions = np.linspace(0, mono.size - 1, n16) + mono = np.interp( + positions, np.arange(mono.size), mono.astype(np.float32) + ).astype(np.int16) + provider = get_vad_provider() + scores = provider.score(mono, VAD_SAMPLE_RATE) + hop_seconds = provider.frame_hop_ms / 1000.0 + except Exception as e: + logger.warning("Silence condensing skipped (VAD failed): %s", e) + return pcm_data, None, None + + raw_intervals = frame_speech_intervals(scores, hop_seconds, 0.0) + regions = merge_speech_regions( + raw_intervals, + duration, + pad_seconds=PAD_SECONDS, + merge_gap_seconds=CUT_GAP_SECONDS, + max_count=100_000, # never coarsen: every region maps to billed audio + ) + if not regions: + return b"", [], 0.0 + + speech_seconds = sum(end - start for start, end in regions) + if speech_seconds >= duration * (1.0 - MIN_SAVINGS_FRACTION): + return pcm_data, None, speech_seconds + + pieces: List[bytes] = [] + mapping: CondenseMap = [] + cursor = 0.0 + for start, end in regions: + byte_start = int(start * sample_rate) * frame_bytes + byte_end = min(int(end * sample_rate) * frame_bytes, len(pcm_data)) + piece = pcm_data[byte_start:byte_end] + if not piece: + continue + length = len(piece) / (sample_rate * frame_bytes) + pieces.append(piece) + mapping.append((cursor, start, length)) + cursor += length + + if not mapping: + return b"", [], 0.0 + logger.info( + "🔇 Condensed %.0fs audio to %.0fs of speech (%d regions, %.0f%% saved)", + duration, + cursor, + len(mapping), + (1.0 - cursor / duration) * 100, + ) + return b"".join(pieces), mapping, speech_seconds + + +def _region_index(t: float, mapping: CondenseMap) -> int: + """Region containing condensed time ``t``; exact boundaries belong to the + NEXT region (a time at a cut is the start of the following speech).""" + for index, (cond_start, _orig, length) in enumerate(mapping): + if t < cond_start + length - 1e-6: + return index + return len(mapping) - 1 + + +def _to_original_in(t: float, region: Tuple[float, float, float]) -> float: + cond_start, orig_start, length = region + return orig_start + max(0.0, min(t - cond_start, length)) + + +def _word_text(word: dict) -> str: + return str( + word.get("punctuated_word") or word.get("word") or word.get("text") or "" + ).strip() + + +def remap_condensed_result(result: dict, mapping: CondenseMap) -> dict: + """Rewrite a transcription result from condensed time to original time. + + Words are pointlike and shift directly. A segment that spans a cut is + split into one sub-segment per speech region — text is redistributed by + word timestamps when the provider returned words, otherwise the full text + stays on the piece with the largest share of the segment. + """ + if not mapping: + return result + + words = result.get("words") or [] + # Group words by region while still in condensed time. + word_regions = [_region_index(w.get("start", 0.0), mapping) for w in words] + + segments = result.get("segments") or [] + new_segments = [] + for segment in segments: + seg_start = float(segment.get("start", 0.0)) + seg_end = float(segment.get("end", seg_start)) + first = _region_index(seg_start, mapping) + last = _region_index(max(seg_end - 1e-3, seg_start), mapping) + if first == last: + new_segments.append( + { + **segment, + "start": round(_to_original_in(seg_start, mapping[first]), 3), + "end": round(_to_original_in(seg_end, mapping[first]), 3), + } + ) + continue + + # Segment spans a cut: one sub-segment per region it touches. + seg_words = [ + (w, r) + for w, r in zip(words, word_regions) + if seg_start - 1e-3 <= w.get("start", 0.0) < seg_end + 1e-3 + ] + pieces = [] + for region_index in range(first, last + 1): + region = mapping[region_index] + cond_start, _orig, length = region + piece_start = max(seg_start, cond_start) + piece_end = min(seg_end, cond_start + length) + if piece_end <= piece_start: + continue + in_piece = [w for w, r in seg_words if r == region_index] + text = " ".join(filter(None, (_word_text(w) for w in in_piece))) + pieces.append( + { + **segment, + "start": round(_to_original_in(piece_start, region), 3), + "end": round(_to_original_in(piece_end, region), 3), + "text": text, + "_overlap": piece_end - piece_start, + } + ) + if pieces and not any(p["text"] for p in pieces): + # No word timestamps: keep the full text on the largest piece. + largest = max(pieces, key=lambda p: p["_overlap"]) + largest["text"] = segment.get("text", "") + for piece in pieces: + piece.pop("_overlap", None) + new_segments.extend(pieces) + + # Words are pointlike: map both ends through the region their start is in, + # so a word never straddles a cut. + for word, region_index in zip(words, word_regions): + region = mapping[region_index] + start = word.get("start", 0.0) + end = word.get("end", start) + word["start"] = round(_to_original_in(start, region), 3) + word["end"] = round(_to_original_in(end, region), 3) + + result["segments"] = new_segments + return result diff --git a/backends/advanced/src/advanced_omi_backend/workers/conversation_jobs.py b/backends/advanced/src/advanced_omi_backend/workers/conversation_jobs.py index 1130d40b..ccee55ca 100644 --- a/backends/advanced/src/advanced_omi_backend/workers/conversation_jobs.py +++ b/backends/advanced/src/advanced_omi_backend/workers/conversation_jobs.py @@ -1938,6 +1938,17 @@ async def generate_title_summary_job( logger.error(f"Conversation {conversation_id} not found") return {"success": False, "error": "Conversation not found"} + if conversation.memory_excluded: + logger.info( + f"Skipping title/summary generation for memory-excluded conversation {conversation_id[:8]}" + ) + return { + "success": True, + "skipped": True, + "reason": "memory_excluded", + "conversation_id": conversation_id, + } + set_span_attrs(user_id=str(conversation.user_id)) # Get transcript and segments (properties return data from active transcript version) @@ -2154,6 +2165,27 @@ async def dispatch_conversation_complete_event_job( if needs_save: await conversation.save() + if conversation.memory_excluded: + actual_end_reason = end_reason or "file_upload" + logger.info( + f"Skipping conversation.complete plugins for memory-excluded conversation {conversation_id[:8]}" + ) + publish_sse_event( + user_id, + "conversation.completed", + { + "conversation_id": conversation_id, + "end_reason": actual_end_reason, + }, + ) + return { + "success": True, + "skipped": True, + "reason": "memory_excluded", + "conversation_id": conversation_id, + "processing_time_seconds": time.time() - start_time, + } + # Get user email for event data user = await User.get(user_id) user_email = user.email if user else "" diff --git a/backends/advanced/src/advanced_omi_backend/workers/data_audit_jobs.py b/backends/advanced/src/advanced_omi_backend/workers/data_audit_jobs.py index e05124b2..bed4640f 100644 --- a/backends/advanced/src/advanced_omi_backend/workers/data_audit_jobs.py +++ b/backends/advanced/src/advanced_omi_backend/workers/data_audit_jobs.py @@ -668,6 +668,7 @@ async def export_annotation_dataset_job( exported = [s for s in conv_summaries if "skipped_reason" not in s] meta = { "export_id": export_id, + "schema_version": 1, "created_at": datetime.now(timezone.utc).isoformat(), "created_by": user_id, "params": { diff --git a/backends/advanced/src/advanced_omi_backend/workers/drift_jobs.py b/backends/advanced/src/advanced_omi_backend/workers/drift_jobs.py new file mode 100644 index 00000000..214c0513 --- /dev/null +++ b/backends/advanced/src/advanced_omi_backend/workers/drift_jobs.py @@ -0,0 +1,41 @@ +"""RQ jobs for speaker-label drift analysis support.""" + +from typing import Any, Dict + +from rq import get_current_job + +from advanced_omi_backend.controllers.drift_controller import ( + backfill_cluster_embeddings, +) +from advanced_omi_backend.models.job import async_job + + +@async_job(redis=False, beanie=True, timeout=7200) +async def cluster_embedding_backfill_job() -> Dict[str, Any]: + """Populate missing per-cluster embeddings used by the drift report.""" + job = get_current_job() + + def publish_progress( + processed: int, + total: int, + backfilled: int, + skipped: int, + failed: int, + ) -> None: + if not job: + return + job.meta["batch_progress"] = { + "percent": round(100 * processed / total) if total else 100, + "message": ( + f"Scanning {processed}/{total} · {backfilled} backfilled" + f" · {skipped} skipped · {failed} failed" + ), + "done": processed, + "total": total, + } + job.save_meta() + + return await backfill_cluster_embeddings( + only_missing=True, + progress_callback=publish_progress, + ) diff --git a/backends/advanced/src/advanced_omi_backend/workers/memory_jobs.py b/backends/advanced/src/advanced_omi_backend/workers/memory_jobs.py index 2422ecd9..0d119aec 100644 --- a/backends/advanced/src/advanced_omi_backend/workers/memory_jobs.py +++ b/backends/advanced/src/advanced_omi_backend/workers/memory_jobs.py @@ -13,6 +13,7 @@ """ import logging +import re import time from typing import Any, Dict, List @@ -45,6 +46,76 @@ logger = logging.getLogger(__name__) MIN_CONVERSATION_LENGTH = 10 +_OVERLAP_TOKEN = re.compile(r"[^\w']+") + + +def _normalise_overlap_token(token: str) -> str: + return _OVERLAP_TOKEN.sub("", token).casefold() + + +def _trim_repeated_prefix(previous: str, current: str) -> str: + """Remove a repeated word suffix/prefix caused by overlapping ASR windows.""" + previous_words = previous.split() + current_words = current.split() + limit = min(len(previous_words), len(current_words), 80) + for size in range(limit, 2, -1): + left = [_normalise_overlap_token(word) for word in previous_words[-size:]] + right = [_normalise_overlap_token(word) for word in current_words[:size]] + if left == right and all(left): + return " ".join(current_words[size:]).strip() + return current.strip() + + +def build_memory_transcript( + segments: list, raw_transcript: str | None +) -> tuple[str, set[str]]: + """Build bounded, speaker-labelled memory input from transcript segments. + + Provider window overlap is trimmed only for temporally adjacent speech segments. + If the resulting segment text is still much larger than the durable raw transcript, + the raw transcript wins so duplicated windows cannot multiply LLM input and facts. + """ + dialogue_lines: list[str] = [] + transcript_speakers: set[str] = set() + previous_speech_text = "" + previous_speech_end: float | None = None + + for segment in sorted(segments or [], key=lambda item: (item.start, item.end)): + text = segment.text.strip() + speaker = segment.speaker + seg_type = getattr(segment, "segment_type", "speech") + if not text: + continue + if seg_type == "event": + dialogue_lines.append(f"[{text}]" if not text.startswith("[") else text) + continue + if seg_type == "note": + dialogue_lines.append(f"[Note: {text}]") + continue + + if ( + previous_speech_end is not None + and segment.start <= previous_speech_end + 1.5 + ): + text = _trim_repeated_prefix(previous_speech_text, text) + if text: + dialogue_lines.append(f"{speaker}: {text}") + previous_speech_text = segment.text.strip() + previous_speech_end = segment.end + if speaker and not speaker.casefold().startswith("unknown"): + transcript_speakers.add(speaker.strip().lower()) + + assembled = "\n".join(dialogue_lines) + raw = raw_transcript.strip() if isinstance(raw_transcript, str) else "" + if raw and len(assembled) > max(int(len(raw) * 1.75), len(raw) + 1000): + logger.warning( + "Memory transcript segments amplified source text (%d vs %d chars); " + "using raw transcript", + len(assembled), + len(raw), + ) + return raw, transcript_speakers + return assembled, transcript_speakers def compute_speaker_diff( @@ -165,6 +236,20 @@ async def process_memory_job( logger.warning(f"No conversation found for {conversation_id}") return {"success": False, "error": "Conversation not found"} + # This is the final safety boundary. Scheduling paths should avoid enqueueing + # memory work for annotation datasets, but a stale or manual job must still be + # unable to mutate the user's vault. + if conversation_model.memory_excluded: + logger.info( + f"Skipping memory processing for excluded conversation {conversation_id}" + ) + return { + "success": True, + "skipped": True, + "reason": "memory_excluded", + "conversation_id": conversation_id, + } + # Get client_id, user_id, and user_email from conversation/user client_id = conversation_model.client_id user_id = conversation_model.user_id @@ -181,30 +266,10 @@ async def process_memory_job( f"🔄 Processing memory for conversation {conversation_id}, client={client_id}, user={user_id}" ) - # Extract conversation text and speakers from transcript segments in a single pass - dialogue_lines = [] - transcript_speakers = set() - segments = conversation_model.segments - if segments: - for segment in segments: - text = segment.text.strip() - speaker = segment.speaker - seg_type = getattr(segment, "segment_type", "speech") - if text: - if seg_type == "event": - # Non-speech event: include as context marker without speaker prefix - dialogue_lines.append( - f"[{text}]" if not text.startswith("[") else text - ) - elif seg_type == "note": - # User-inserted note: include as distinct context - dialogue_lines.append(f"[Note: {text}]") - else: - # Normal speech segment - dialogue_lines.append(f"{speaker}: {text}") - if speaker and speaker != "Unknown" and seg_type == "speech": - transcript_speakers.add(speaker.strip().lower()) - full_conversation = "\n".join(dialogue_lines) + full_conversation, transcript_speakers = build_memory_transcript( + conversation_model.segments, + conversation_model.transcript, + ) # Fallback: if segments have no text content but transcript exists, use transcript # This handles cases where speaker recognition fails/is disabled @@ -282,6 +347,13 @@ async def process_memory_job( user_id, user_email, allow_update=True, + source_date=conversation_model.created_at.isoformat(), + source_duration_minutes=( + conversation_model.audio_total_duration / 60 + if conversation_model.audio_total_duration is not None + else None + ), + source_title=conversation_model.title, ) except Exception as e: logger.error( @@ -537,6 +609,9 @@ def enqueue_memory_processing( *, cause: MemoryCause = MemoryCause.AUTO_EXTRACTION, strategy: UpdateStrategy = UpdateStrategy.FULL, + depends_on=None, + job_id: str | None = None, + job_timeout: int | None = None, ): """ Enqueue a memory processing job. @@ -560,12 +635,14 @@ def enqueue_memory_processing( # job_id uses [:12] to match the deterministic id the post-conversation chain and # _clear_post_conversation_chain use — so a standalone re-enqueue collides with # (replaces) the chain's memory job rather than creating an orphan twin. + resolved_job_id = job_id or f"memory_{conversation_id[:12]}" + resolved_timeout = job_timeout or timeout_mapping.get(priority, 1800) job = memory_queue.enqueue( process_memory_job, conversation_id, # Only argument needed - job fetches conversation data internally - job_timeout=timeout_mapping.get(priority, 1800), + job_timeout=resolved_timeout, result_ttl=JOB_RESULT_TTL, - job_id=f"memory_{conversation_id[:12]}", + job_id=resolved_job_id, description=f"Process memory for conversation {conversation_id[:8]}", **post_conv_enqueue_kwargs( "memory", @@ -574,6 +651,7 @@ def enqueue_memory_processing( "cause": cause.value, "strategy": strategy.value, }, + depends_on=depends_on, ), ) diff --git a/backends/advanced/src/advanced_omi_backend/workers/orchestrator/health_monitor.py b/backends/advanced/src/advanced_omi_backend/workers/orchestrator/health_monitor.py index c97812cc..f9f10613 100644 --- a/backends/advanced/src/advanced_omi_backend/workers/orchestrator/health_monitor.py +++ b/backends/advanced/src/advanced_omi_backend/workers/orchestrator/health_monitor.py @@ -6,6 +6,7 @@ """ import asyncio +import json import logging import time from typing import Optional @@ -13,6 +14,11 @@ from redis import Redis from rq import Worker +from advanced_omi_backend.heartbeat import ( + FLEET_HEALTH_KEY, + FLEET_HEARTBEAT_TTL_SECONDS, + is_rq_worker_fresh, +) from advanced_omi_backend.services.plugin_service import WORKER_RESTART_KEY from .config import OrchestratorConfig, WorkerType @@ -78,6 +84,11 @@ async def stop(self): except asyncio.CancelledError: pass + try: + self.redis.delete(FLEET_HEALTH_KEY) + except Exception as e: + logger.warning("Failed to clear worker fleet heartbeat: %s", e) + logger.info("Health monitor stopped") async def _monitor_loop(self): @@ -87,6 +98,7 @@ async def _monitor_loop(self): # Wait for startup grace period before starting checks elapsed = time.time() - self.start_time if elapsed < self.config.startup_grace_period: + self._publish_fleet_health("starting") remaining = self.config.startup_grace_period - elapsed logger.debug( f"In startup grace period - waiting {remaining:.0f}s before health checks" @@ -95,7 +107,11 @@ async def _monitor_loop(self): continue # Perform health checks - await self._check_health() + worker_health, rq_health, detail = await self._check_health() + self._publish_fleet_health( + "healthy" if worker_health and rq_health else "unhealthy", + detail=detail, + ) # Wait for next check await asyncio.sleep(self.config.check_interval) @@ -110,13 +126,13 @@ async def _monitor_loop(self): ) raise # Re-raise to ensure the monitor task fails properly - async def _check_health(self): + async def _check_health(self) -> tuple[bool, bool, str | None]: """Perform all health checks and restart failed workers""" try: # Check for plugin reload restart signal first if self._check_restart_signal(): # Workers are restarting — skip normal health checks this iteration - return + return True, True, "Workers restarting for plugin reload" # Check individual worker health worker_health = self._check_worker_health() @@ -137,8 +153,38 @@ async def _check_health(self): f"Health check: worker_health={worker_health}, rq_health={rq_health}" ) + detail = None + if not worker_health or not rq_health: + detail = ( + f"child_processes_healthy={worker_health}, " + f"rq_registration_healthy={rq_health}" + ) + return worker_health, rq_health, detail + except Exception as e: logger.error(f"Error during health check: {e}", exc_info=True) + return False, False, str(e) + + def _publish_fleet_health(self, status: str, detail: str | None = None) -> None: + """Publish the supervisor's view of the complete worker fleet.""" + try: + worker_status = self.process_manager.get_status() + payload = { + "status": status, + "timestamp": time.time(), + "workers_total": len(worker_status), + "workers_alive": sum( + 1 for worker in worker_status.values() if worker["is_alive"] + ), + "detail": detail, + } + self.redis.set( + FLEET_HEALTH_KEY, + json.dumps(payload), + ex=FLEET_HEARTBEAT_TTL_SECONDS, + ) + except Exception as e: + logger.warning("Failed to publish worker fleet heartbeat: %s", e) def _check_restart_signal(self) -> bool: """Check Redis for a plugin-reload restart signal and restart all workers if found. @@ -236,7 +282,11 @@ def _check_rq_worker_registration(self) -> bool: True if RQ worker count is sufficient """ try: - workers = Worker.all(connection=self.redis) + workers = [ + worker + for worker in Worker.all(connection=self.redis) + if is_rq_worker_fresh(worker) + ] worker_count = len(workers) if worker_count < self.config.min_rq_workers: @@ -418,7 +468,11 @@ def get_health_status(self) -> dict: # Check RQ worker registration try: - rq_workers = Worker.all(connection=self.redis) + rq_workers = [ + worker + for worker in Worker.all(connection=self.redis) + if is_rq_worker_fresh(worker) + ] rq_worker_count = len(rq_workers) except Exception: rq_worker_count = -1 # Error indicator diff --git a/backends/advanced/src/advanced_omi_backend/workers/rq_worker_entry.py b/backends/advanced/src/advanced_omi_backend/workers/rq_worker_entry.py index 3a9d7a62..d23cfae4 100755 --- a/backends/advanced/src/advanced_omi_backend/workers/rq_worker_entry.py +++ b/backends/advanced/src/advanced_omi_backend/workers/rq_worker_entry.py @@ -72,8 +72,12 @@ def main(): logger.info("✅ RQ worker ready") - # This blocks until worker is stopped - worker.work(logging_level="INFO") + # This blocks until worker is stopped. + # with_scheduler: required for Retry(interval=...) — retried jobs land in + # ScheduledJobRegistry and need a scheduler to promote them back onto the + # queue (RQ elects one scheduler per queue via a lock, so this is safe + # across multiple worker processes). + worker.work(logging_level="INFO", with_scheduler=True) if __name__ == "__main__": diff --git a/backends/advanced/src/advanced_omi_backend/workers/speaker_benchmark_jobs.py b/backends/advanced/src/advanced_omi_backend/workers/speaker_benchmark_jobs.py new file mode 100644 index 00000000..f18d56a0 --- /dev/null +++ b/backends/advanced/src/advanced_omi_backend/workers/speaker_benchmark_jobs.py @@ -0,0 +1,392 @@ +"""Conversation-grouped cross-validation for human-labeled speaker clips.""" + +import hashlib +import math +from datetime import datetime, timezone +from typing import Any, Dict, List + +import numpy as np +from rq import get_current_job + +from advanced_omi_backend.config import get_diarization_settings +from advanced_omi_backend.models.conversation import Conversation +from advanced_omi_backend.models.job import async_job +from advanced_omi_backend.speaker_recognition_client import SpeakerRecognitionClient +from advanced_omi_backend.utils.audio_chunk_utils import reconstruct_audio_segment + +FOLDS = 5 +FRACTIONS = (0.2, 0.4, 0.6, 0.8, 1.0) +MIN_SECONDS = 1.5 + + +def _unit(values: list) -> np.ndarray: + vector = np.asarray(values, dtype=np.float32).reshape(-1) + norm = np.linalg.norm(vector) + if not vector.size or not np.isfinite(vector).all() or norm == 0: + raise ValueError("invalid embedding") + return vector / norm + + +def _stable(value: str) -> str: + return hashlib.sha256(value.encode()).hexdigest() + + +def _active_segments(doc: dict) -> list: + versions = doc.get("transcript_versions") or [] + active_id = doc.get("active_transcript_version") + version = next( + (item for item in versions if item.get("version_id") == active_id), None + ) + return (version or (versions[-1] if versions else {})).get("segments") or [] + + +async def _labeled_clips(user_id: str) -> tuple[List[dict], Dict[str, int]]: + db = Conversation.get_pymongo_collection().database + clips: List[dict] = [] + exclusions = { + "bad_or_mixed": 0, + "missing_bounds": 0, + "too_short": 0, + "duplicate": 0, + } + + async for review in db["enrollment_reviews"].find( + {"reviewed_by": user_id}, + { + "conversation_id": 1, + "decision": 1, + "actual_speaker": 1, + "selected_start": 1, + "selected_end": 1, + "segment_start": 1, + "segment_end": 1, + }, + ): + if review.get("decision") not in ("accept", "another_speaker"): + exclusions["bad_or_mixed"] += 1 + continue + start = review.get("selected_start", review.get("segment_start")) + end = review.get("selected_end", review.get("segment_end")) + if start is None or end is None: + exclusions["missing_bounds"] += 1 + continue + clips.append( + { + "conversation_id": review["conversation_id"], + "start": float(start), + "end": float(end), + "speaker": review.get("actual_speaker"), + "source": "guided_enrollment", + } + ) + + async for annotation in db["annotations"].find( + { + "user_id": user_id, + "annotation_type": "diarization", + "status": "accepted", + "corrected_speaker": {"$nin": [None, "", "Noise", "Unknown Speaker"]}, + }, + { + "conversation_id": 1, + "segment_index": 1, + "segment_start_time": 1, + "corrected_speaker": 1, + }, + ): + doc = await Conversation.get_pymongo_collection().find_one( + {"conversation_id": annotation.get("conversation_id")}, + {"active_transcript_version": 1, "transcript_versions": 1}, + ) + segments = _active_segments(doc or {}) + index = annotation.get("segment_index") + segment = ( + segments[index] + if isinstance(index, int) and 0 <= index < len(segments) + else None + ) + if segment is None and annotation.get("segment_start_time") is not None: + target = float(annotation["segment_start_time"]) + segment = min( + segments, + key=lambda item: abs(float(item.get("start", 0)) - target), + default=None, + ) + if not segment: + exclusions["missing_bounds"] += 1 + continue + clips.append( + { + "conversation_id": annotation["conversation_id"], + "start": float(segment.get("start", 0)), + "end": float(segment.get("end", 0)), + "speaker": annotation["corrected_speaker"], + "source": "diarization_annotation", + } + ) + + unique = {} + for clip in clips: + if not clip.get("speaker") or clip["end"] - clip["start"] < MIN_SECONDS: + exclusions["too_short"] += 1 + continue + key = f'{clip["conversation_id"]}:{clip["start"]:.3f}:{clip["end"]:.3f}:{clip["speaker"]}' + if key in unique: + exclusions["duplicate"] += 1 + continue + clip["key"] = key + unique[key] = clip + return list(unique.values()), exclusions + + +async def _embed_clips( + clips: List[dict], user_id: str, embedding_model: str +) -> tuple[List[dict], int, List[dict]]: + db = Conversation.get_pymongo_collection().database + cache = db["speaker_evaluation_embeddings"] + client = SpeakerRecognitionClient() + embedded = [] + failures = [] + cache_hits = 0 + job = get_current_job() + for index, clip in enumerate(clips): + cached = await cache.find_one( + { + "user_id": user_id, + "clip_key": clip["key"], + "embedding_model": embedding_model, + } + ) + if cached: + result = cached + cache_hits += 1 + else: + try: + wav = await reconstruct_audio_segment( + clip["conversation_id"], clip["start"], clip["end"] + ) + result = await client.extract_speaker_embedding(wav) + if result.get("error"): + raise RuntimeError(str(result)) + await cache.update_one( + {"user_id": user_id, "clip_key": clip["key"]}, + { + "$set": { + **clip, + **result, + "user_id": user_id, + "created_at": datetime.now(timezone.utc), + } + }, + upsert=True, + ) + except Exception as exc: + failures.append({"clip_key": clip["key"], "error": str(exc)}) + continue + try: + embedded.append({**clip, "embedding": _unit(result["embedding"])}) + except Exception as exc: + failures.append({"clip_key": clip["key"], "error": str(exc)}) + if job and (index % 5 == 0 or index + 1 == len(clips)): + job.meta["batch_progress"] = { + "current": index + 1, + "total": len(clips), + "message": f"Embedding labeled clips {index + 1}/{len(clips)}", + } + job.save_meta() + return embedded, cache_hits, failures + + +def _centroids(samples: List[dict]) -> Dict[str, np.ndarray]: + grouped: Dict[str, list] = {} + for sample in samples: + grouped.setdefault(sample["speaker"], []).append(sample["embedding"]) + return { + speaker: _unit(np.mean(vectors, axis=0).tolist()) + for speaker, vectors in grouped.items() + } + + +def _metrics(train: List[dict], test: List[dict], threshold: float) -> dict: + centroids = _centroids(train) + evaluable = [sample for sample in test if sample["speaker"] in centroids] + if not evaluable or len(centroids) < 2: + return { + "test_clips": len(evaluable), + "top1_accuracy": None, + "macro_recall": None, + "false_accept_rate": None, + } + correct = 0 + by_speaker: Dict[str, list] = {} + confusion: Dict[str, Dict[str, int]] = {} + target_scores = [] + impostor_scores = [] + wrong_trials = wrong_accepts = 0 + for sample in evaluable: + scores = { + speaker: float(sample["embedding"] @ centroid) + for speaker, centroid in centroids.items() + } + predicted = max(scores, key=scores.get) + hit = predicted == sample["speaker"] + correct += int(hit) + by_speaker.setdefault(sample["speaker"], []).append(int(hit)) + confusion.setdefault(sample["speaker"], {})[predicted] = ( + confusion.setdefault(sample["speaker"], {}).get(predicted, 0) + 1 + ) + target_scores.append(scores[sample["speaker"]]) + for speaker, score in scores.items(): + if speaker != sample["speaker"]: + impostor_scores.append(score) + wrong_trials += 1 + wrong_accepts += int(score >= threshold) + recalls = [sum(values) / len(values) for values in by_speaker.values()] + thresholds = sorted(set(target_scores + impostor_scores)) + eer = None + if target_scores and impostor_scores and thresholds: + candidates = [] + for candidate in thresholds: + false_reject = sum(score < candidate for score in target_scores) / len( + target_scores + ) + false_accept = sum(score >= candidate for score in impostor_scores) / len( + impostor_scores + ) + candidates.append( + (abs(false_reject - false_accept), (false_reject + false_accept) / 2) + ) + eer = min(candidates, key=lambda item: item[0])[1] + return { + "test_clips": len(evaluable), + "speakers": len(centroids), + "top1_accuracy": round(correct / len(evaluable), 4), + "macro_recall": round(sum(recalls) / len(recalls), 4), + "false_accept_rate": ( + round(wrong_accepts / wrong_trials, 4) if wrong_trials else None + ), + "eer": round(eer, 4) if eer is not None else None, + "per_speaker_recall": { + speaker: round(sum(values) / len(values), 4) + for speaker, values in sorted(by_speaker.items()) + }, + "confusion": confusion, + } + + +def _evaluate(samples: List[dict], threshold: float) -> dict: + groups = sorted({sample["conversation_id"] for sample in samples}, key=_stable) + fold_by_group = {group: index % FOLDS for index, group in enumerate(groups)} + curves = [] + fold_results = [] + for fraction in FRACTIONS: + metrics_for_fraction = [] + for fold in range(FOLDS): + test = [ + sample + for sample in samples + if fold_by_group[sample["conversation_id"]] == fold + ] + available = [ + sample + for sample in samples + if fold_by_group[sample["conversation_id"]] != fold + ] + train = [] + for speaker in sorted({sample["speaker"] for sample in available}): + speaker_samples = sorted( + (sample for sample in available if sample["speaker"] == speaker), + key=lambda sample: _stable(sample["key"]), + ) + take = max(1, math.ceil(len(speaker_samples) * fraction)) + train.extend(speaker_samples[:take]) + metrics = _metrics(train, test, threshold) + metrics["fold"] = fold + 1 + metrics["train_clips"] = len(train) + metrics_for_fraction.append(metrics) + if fraction == 1.0: + fold_results.append(metrics) + valid = [ + item for item in metrics_for_fraction if item["top1_accuracy"] is not None + ] + valid_far = [item for item in valid if item["false_accept_rate"] is not None] + valid_eer = [item for item in valid if item.get("eer") is not None] + curves.append( + { + "fraction": fraction, + "train_clips_mean": round( + sum(item["train_clips"] for item in metrics_for_fraction) / FOLDS, 1 + ), + "top1_accuracy_mean": ( + round(sum(item["top1_accuracy"] for item in valid) / len(valid), 4) + if valid + else None + ), + "macro_recall_mean": ( + round(sum(item["macro_recall"] for item in valid) / len(valid), 4) + if valid + else None + ), + "false_accept_rate_mean": ( + round( + sum(item["false_accept_rate"] for item in valid_far) + / len(valid_far), + 4, + ) + if valid_far + else None + ), + "eer_mean": ( + round(sum(item["eer"] for item in valid_eer) / len(valid_eer), 4) + if valid_eer + else None + ), + } + ) + return { + "learning_curve": curves, + "folds": fold_results, + "conversation_groups": len(groups), + "fold_groups": { + str(fold + 1): sorted( + group for group, assigned in fold_by_group.items() if assigned == fold + ) + for fold in range(FOLDS) + }, + } + + +@async_job(redis=False, beanie=True, timeout=7200) +async def run_speaker_benchmark_job(user_id: str) -> Dict[str, Any]: + started_at = datetime.now(timezone.utc) + clips, exclusions = await _labeled_clips(user_id) + client = SpeakerRecognitionClient() + embedding_info = await client.get_embedding_info() + if embedding_info.get("error") or not embedding_info.get("embedding_model"): + raise RuntimeError(f"Cannot determine embedding model: {embedding_info}") + embedding_model = embedding_info["embedding_model"] + embedded, cache_hits, failures = await _embed_clips(clips, user_id, embedding_model) + threshold = float(get_diarization_settings().get("similarity_threshold", 0.5)) + report = { + "user_id": user_id, + "created_at": started_at, + "protocol": "5-fold conversation-grouped cross-validation", + "fractions": list(FRACTIONS), + "threshold": threshold, + "embedding_model": embedding_model, + "dataset": { + "labeled_clips": len(clips), + "embedded_clips": len(embedded), + "speakers": len({sample["speaker"] for sample in embedded}), + "cache_hits": cache_hits, + "embedding_failures": len(failures), + "exclusions": exclusions, + }, + **_evaluate(embedded, threshold), + "failures": failures[:20], + } + db = Conversation.get_pymongo_collection().database + await db["speaker_benchmark_runs"].insert_one(report) + report.pop("_id", None) + report["created_at"] = started_at.isoformat() + return report diff --git a/backends/advanced/src/advanced_omi_backend/workers/speaker_discovery_jobs.py b/backends/advanced/src/advanced_omi_backend/workers/speaker_discovery_jobs.py new file mode 100644 index 00000000..a2a96ab7 --- /dev/null +++ b/backends/advanced/src/advanced_omi_backend/workers/speaker_discovery_jobs.py @@ -0,0 +1,243 @@ +"""Reusable speech-segment embedding index for speaker corpus discovery.""" + +import asyncio +from datetime import datetime, timezone +from typing import Any + +from rq import get_current_job + +from advanced_omi_backend.models.conversation import Conversation +from advanced_omi_backend.models.job import async_job +from advanced_omi_backend.speaker_recognition_client import SpeakerRecognitionClient +from advanced_omi_backend.utils.audio_chunk_utils import reconstruct_audio_segment + +MIN_SECONDS = 3.0 +MAX_SECONDS = 30.0 +EMBED_CONCURRENCY = 4 +SCORE_BATCH_SIZE = 500 + + +def _active_segments(doc: dict) -> list: + versions = doc.get("transcript_versions") or [] + active_id = doc.get("active_transcript_version") + version = next( + (item for item in versions if item.get("version_id") == active_id), None + ) + return (version or (versions[-1] if versions else {})).get("segments") or [] + + +async def _speech_clips( + user_id: str, include_all_users: bool, include_deleted: bool = False +) -> list[dict]: + query: dict[str, Any] = { + "audio_archived": {"$ne": True}, + "audio_chunks_count": {"$gt": 0}, + } + if not include_deleted: + # Soft-deleted conversations keep their audio chunks; opt in to mine them. + query["deleted"] = {"$ne": True} + if not include_all_users: + query["user_id"] = user_id + database = Conversation.get_pymongo_collection().database + human_labels: dict[str, str] = {} + review_query = {} if include_all_users else {"reviewed_by": user_id} + async for review in database["enrollment_reviews"].find( + review_query, + {"conversation_id": 1, "segment_start": 1, "actual_speaker": 1, "enrolled": 1}, + ): + if review.get("enrolled") and review.get("actual_speaker"): + key = f'{review["conversation_id"]}:{round(float(review["segment_start"]), 2)}' + human_labels[key] = review["actual_speaker"] + annotation_query: dict[str, Any] = { + "annotation_type": "diarization", + "status": "accepted", + "corrected_speaker": {"$nin": [None, "", "Noise", "Unknown Speaker"]}, + "segment_start_time": {"$ne": None}, + } + if not include_all_users: + annotation_query["user_id"] = user_id + async for annotation in database["annotations"].find( + annotation_query, + {"conversation_id": 1, "segment_start_time": 1, "corrected_speaker": 1}, + ): + key = ( + f'{annotation["conversation_id"]}:' + f'{round(float(annotation["segment_start_time"]), 2)}' + ) + human_labels[key] = annotation["corrected_speaker"] + + projection = { + "conversation_id": 1, + "user_id": 1, + "title": 1, + "created_at": 1, + "audio_total_duration": 1, + "active_transcript_version": 1, + "transcript_versions": 1, + } + clips = [] + async for doc in Conversation.get_pymongo_collection().find(query, projection): + audio_duration = float(doc.get("audio_total_duration") or 0) + for index, segment in enumerate(_active_segments(doc)): + if segment.get("segment_type") not in (None, "speech"): + continue + start = float(segment.get("start") or 0) + end = min(float(segment.get("end") or 0), audio_duration or float("inf")) + end = min(end, start + MAX_SECONDS) + if end - start < MIN_SECONDS: + continue + clip_key = f'{doc["conversation_id"]}:{start:.3f}:{end:.3f}' + review_key = f'{doc["conversation_id"]}:{round(start, 2)}' + clips.append( + { + "clip_key": clip_key, + "review_key": review_key, + "conversation_id": doc["conversation_id"], + "corpus_user_id": doc.get("user_id"), + "conversation_title": doc.get("title"), + "conversation_date": doc.get("created_at"), + "conversation_duration": audio_duration, + "segment_index": index, + "start": start, + "end": end, + "duration": end - start, + "text": (segment.get("text") or "")[:300], + "current_label": segment.get("identified_as"), + "human_label": human_labels.get(review_key), + "stored_confidence": segment.get("confidence"), + } + ) + return clips + + +def _progress(current: int, total: int, message: str) -> None: + job = get_current_job() + if not job: + return + job.meta["batch_progress"] = { + "current": current, + "total": total, + "percent": round(current * 100 / total) if total else 100, + "message": message, + } + job.save_meta() + + +@async_job(redis=False, beanie=True, timeout=14400) +async def discover_speaker_candidates_job( + requested_by: str, + speaker_id: str, + speaker_name: str, + include_all_users: bool = False, + include_deleted: bool = False, +) -> dict: + db = Conversation.get_pymongo_collection().database + cache = db["speaker_corpus_embeddings"] + matches = db["speaker_corpus_matches"] + await cache.create_index([("clip_key", 1), ("embedding_model", 1)], unique=True) + await matches.create_index( + [("requested_by", 1), ("speaker_id", 1), ("review_key", 1)] + ) + client = SpeakerRecognitionClient() + info = await client.get_embedding_info() + model = info.get("embedding_model") + if not model: + raise RuntimeError(f"Could not resolve speaker embedding model: {info}") + + clips = await _speech_clips(requested_by, include_all_users, include_deleted) + sem = asyncio.Semaphore(EMBED_CONCURRENCY) + embedded: list[tuple[dict, list[float]]] = [] + failures = 0 + completed = 0 + progress_lock = asyncio.Lock() + + async def embed(index: int, clip: dict) -> None: + nonlocal completed, failures + async with sem: + cached = await cache.find_one( + {"clip_key": clip["clip_key"], "embedding_model": model}, + {"embedding": 1}, + ) + if cached and cached.get("embedding"): + embedded.append((clip, cached["embedding"])) + async with progress_lock: + completed += 1 + _progress( + completed, + len(clips), + f"Indexing corpus speech {completed}/{len(clips)}", + ) + return + try: + wav = await reconstruct_audio_segment( + clip["conversation_id"], clip["start"], clip["end"] + ) + result = await client.extract_speaker_embedding(wav) + if result.get("error") or not result.get("embedding"): + raise RuntimeError(str(result)) + await cache.update_one( + {"clip_key": clip["clip_key"], "embedding_model": model}, + { + "$set": { + **clip, + **result, + "indexed_at": datetime.now(timezone.utc), + } + }, + upsert=True, + ) + embedded.append((clip, result["embedding"])) + except Exception: + failures += 1 + finally: + async with progress_lock: + completed += 1 + _progress( + completed, + len(clips), + f"Indexing corpus speech {completed}/{len(clips)}", + ) + + await asyncio.gather(*(embed(index, clip) for index, clip in enumerate(clips))) + embedded.sort(key=lambda item: item[0]["clip_key"]) + + await matches.delete_many({"requested_by": requested_by, "speaker_id": speaker_id}) + scored_count = 0 + for offset in range(0, len(embedded), SCORE_BATCH_SIZE): + batch = embedded[offset : offset + SCORE_BATCH_SIZE] + response = await client.score_cached_embeddings( + speaker_id, [embedding for _clip, embedding in batch] + ) + scores = response.get("scores") + if not isinstance(scores, list) or len(scores) != len(batch): + raise RuntimeError(f"Speaker-service batch scoring failed: {response}") + now = datetime.now(timezone.utc) + for (clip, _embedding), score in zip(batch, scores): + if score.get("sim_centroid") is None: + continue + await matches.insert_one( + { + **clip, + "requested_by": requested_by, + "speaker_id": speaker_id, + "speaker_name": speaker_name, + "embedding_model": model, + "scores": score, + "scored_at": now, + } + ) + scored_count += 1 + _progress( + min(offset + len(batch), len(embedded)), + len(embedded), + f"Comparing {speaker_name} with indexed speech", + ) + + return { + "speaker_name": speaker_name, + "speech_clips": len(clips), + "embedded": len(embedded), + "scored": scored_count, + "failures": failures, + "embedding_model": model, + } diff --git a/backends/advanced/src/advanced_omi_backend/workers/speaker_mining_jobs.py b/backends/advanced/src/advanced_omi_backend/workers/speaker_mining_jobs.py new file mode 100644 index 00000000..3d78d773 --- /dev/null +++ b/backends/advanced/src/advanced_omi_backend/workers/speaker_mining_jobs.py @@ -0,0 +1,201 @@ +"""Speaker mining: ingest an unlabelled audio corpus and score it for one speaker. + +Backs the "mine more audio" flow on the speaker-enhancement page. Server-side +audio files (e.g. backup WAVs whose conversations were purged) are ingested +through the standard upload pipeline as annotation-only conversations — audio +chunks in Mongo, batch transcription, speaker identification, no memory +extraction — and a corpus-discovery job is chained behind the transcription +jobs so the new speech is embedded and scored against the target speaker's +gallery as soon as it has segments. Mined clips then surface in guided +enrollment like any other corpus match. +""" + +import json +import logging +from datetime import datetime, timezone +from pathlib import Path + +from fastapi.responses import JSONResponse +from rq import get_current_job +from rq.job import Dependency +from starlette.datastructures import UploadFile as StarletteUploadFile + +from advanced_omi_backend.controllers.queue_controller import ( + JOB_RESULT_TTL, + default_queue, +) +from advanced_omi_backend.models.conversation import Conversation +from advanced_omi_backend.models.job import async_job +from advanced_omi_backend.models.user import get_user_by_id +from advanced_omi_backend.workers.speaker_discovery_jobs import ( + discover_speaker_candidates_job, +) + +logger = logging.getLogger(__name__) + +# Server-side ingest only reads inside the mounted data dir. +ALLOWED_ROOT = Path("/app/data") +INGEST_BATCH_SIZE = 8 +MINING_DEVICE_NAME = "speaker-mining" + + +def _progress(current: int, total: int, message: str) -> None: + job = get_current_job() + if not job: + return + job.meta["batch_progress"] = { + "current": current, + "total": total, + "percent": round(current * 100 / total) if total else 100, + "message": message, + } + job.save_meta() + + +def _parse_upload_response(result) -> dict: + if isinstance(result, JSONResponse): + return json.loads(result.body) + return result + + +async def enqueue_discovery_after( + requested_by: str, + speaker_id: str, + speaker_name: str, + depends_on_job_ids: list[str], + include_all_users: bool, + include_deleted: bool = False, +) -> str: + """Queue corpus discovery behind the ingest's transcription jobs. + + ``allow_failure`` keeps discovery running even if some files fail to + transcribe — partial corpus coverage beats none. + """ + kwargs = dict( + requested_by=requested_by, + speaker_id=speaker_id, + speaker_name=speaker_name, + include_all_users=include_all_users, + include_deleted=include_deleted, + job_timeout=14400, + result_ttl=JOB_RESULT_TTL, + description=f"Speaker corpus discovery: {speaker_name} (after mining ingest)", + ) + if depends_on_job_ids: + kwargs["depends_on"] = Dependency(jobs=depends_on_job_ids, allow_failure=True) + job = default_queue.enqueue(discover_speaker_candidates_job, **kwargs) + # Record the run so the enhancement page reattaches to it on refresh. + database = Conversation.get_pymongo_collection().database + await database["speaker_discovery_runs"].update_one( + {"requested_by": requested_by, "speaker_id": speaker_id}, + { + "$set": { + "speaker_name": speaker_name, + "job_id": job.id, + "queued_at": datetime.now(timezone.utc), + } + }, + upsert=True, + ) + return job.id + + +@async_job(redis=False, beanie=True, timeout=14400) +async def mine_local_corpus_job( + requested_by: str, + speaker_id: str, + speaker_name: str, + paths: list[str], + include_all_users: bool = False, +) -> dict: + """Ingest server-side audio files and chain discovery for one speaker.""" + # Lazy import: audio_controller pulls in the transcription stack, which the + # guided-enrollment controller (importing this module's enqueue helper's + # sibling) must not load at import time. + from advanced_omi_backend.controllers.audio_controller import ( + upload_and_process_audio_files, + ) + + user = await get_user_by_id(requested_by) + if user is None: + raise RuntimeError(f"Unknown user {requested_by}") + + root = ALLOWED_ROOT.resolve() + valid: list[Path] = [] + skipped: list[dict] = [] + for raw in paths: + path = Path(raw).resolve() + if not str(path).startswith(str(root) + "/"): + skipped.append({"path": raw, "error": "outside the data directory"}) + continue + if not path.is_file(): + skipped.append({"path": raw, "error": "not a file"}) + continue + valid.append(path) + + transcript_job_ids: list[str] = [] + conversation_ids: list[str] = [] + failed: list[dict] = [] + processed = 0 + for offset in range(0, len(valid), INGEST_BATCH_SIZE): + batch = valid[offset : offset + INGEST_BATCH_SIZE] + handles = [path.open("rb") for path in batch] + try: + uploads = [ + StarletteUploadFile(file=handle, filename=path.name) + for handle, path in zip(handles, batch) + ] + result = _parse_upload_response( + await upload_and_process_audio_files( + user, + uploads, + device_name=MINING_DEVICE_NAME, + annotation_only=True, + ) + ) + finally: + for handle in handles: + handle.close() + for item in result.get("files", []): + if item.get("status") == "started": + conversation_ids.append(item.get("conversation_id")) + if item.get("transcript_job_id"): + transcript_job_ids.append(item["transcript_job_id"]) + else: + failed.append( + {"path": item.get("filename"), "error": item.get("error")} + ) + processed += len(batch) + _progress( + processed, + len(valid), + f"Ingested {processed}/{len(valid)} files for {speaker_name}", + ) + + discovery_job_id = None + if conversation_ids: + discovery_job_id = await enqueue_discovery_after( + requested_by, + speaker_id, + speaker_name, + transcript_job_ids, + include_all_users, + ) + + if not transcript_job_ids and conversation_ids: + logger.warning( + "Speaker mining: no transcription provider available; %d conversations " + "ingested without transcripts and will not be discoverable until " + "transcribed", + len(conversation_ids), + ) + + return { + "speaker_name": speaker_name, + "requested_files": len(paths), + "ingested": len(conversation_ids), + "transcription_jobs": len(transcript_job_ids), + "discovery_job_id": discovery_job_id, + "failed": failed[:20], + "skipped": skipped[:20], + } diff --git a/backends/advanced/src/advanced_omi_backend/workers/transcription_jobs.py b/backends/advanced/src/advanced_omi_backend/workers/transcription_jobs.py index b70cfaf5..8115ff94 100644 --- a/backends/advanced/src/advanced_omi_backend/workers/transcription_jobs.py +++ b/backends/advanced/src/advanced_omi_backend/workers/transcription_jobs.py @@ -69,6 +69,10 @@ ) from advanced_omi_backend.utils.job_utils import check_job_alive, update_job_meta from advanced_omi_backend.utils.segment_utils import classify_segment_text +from advanced_omi_backend.utils.silence_condense import ( + condense_silence, + remap_condensed_result, +) logger = logging.getLogger(__name__) @@ -268,6 +272,37 @@ async def transcribe_audio_range( if hasattr(provider, "get_capabilities_dict"): provider_capabilities = provider.get_capabilities_dict() + # Batch providers bill by audio duration, silence included — cut long + # silences with the local VAD before sending, and map timestamps back to + # the original timeline afterwards (see utils/silence_condense.py). + condense_map = None + condensed_pcm, mapping, speech_seconds = condense_silence( + pcm_data, actual_sample_rate, channels, sample_width + ) + if mapping is not None and not mapping: + logger.info( + f"🔇 No speech detected in [{start_time:.1f}s - {end_time or 'end'}] " + f"of {conversation_id[:8]} — skipping paid transcription entirely" + ) + return { + "text": "", + "segments": [], + "words": [], + "provider_name": provider.name, + "provider_capabilities": provider_capabilities, + "wav_size": 0, + "sample_rate": actual_sample_rate, + } + if mapping is not None: + condense_map = mapping + pcm_data = condensed_pcm + wav_data = _build_wav(pcm_data, actual_sample_rate, channels, sample_width) + duration = ( + len(pcm_data) / (actual_sample_rate * sample_width * channels) + if (actual_sample_rate * sample_width * channels) > 0 + else 0 + ) + if duration <= BATCH_CHUNK_SECONDS: # Single chunk — transcribe directly transcribe_kwargs: dict = { @@ -289,6 +324,8 @@ async def transcribe_audio_range( except Exception as e: raise RuntimeError(f"Transcription failed ({type(e).__name__}): {e}") + if condense_map: + result = remap_condensed_result(result, condense_map) return { "text": result.get("text", ""), "segments": result.get("segments", []), @@ -357,7 +394,7 @@ async def transcribe_audio_range( all_words.extend(result.get("words", [])) total_wav_size += len(chunk_wav) - return { + merged = { "text": " ".join(all_text), "segments": all_segments, "words": all_words, @@ -366,6 +403,10 @@ async def transcribe_audio_range( "wav_size": total_wav_size, "sample_rate": actual_sample_rate, } + # Chunk offsets above are condensed-timeline; map back to original time. + if condense_map: + remap_condensed_result(merged, condense_map) + return merged async def process_transcription_result( @@ -434,7 +475,7 @@ async def process_transcription_result( } # Trigger transcript-level plugins BEFORE speech validation - if transcript_text: + if transcript_text and not conversation.memory_excluded: try: await dispatch_plugin_event( event=PluginEvent.TRANSCRIPT_BATCH, @@ -453,6 +494,10 @@ async def process_transcription_result( logger.exception( f"⚠️ Error triggering transcript plugins in batch mode: {e}" ) + elif transcript_text: + logger.info( + f"Skipping transcript plugins for memory-excluded conversation {conversation_id[:8]}" + ) # Validate meaningful speech transcript_data = {"text": transcript_text, "words": words} @@ -1110,6 +1155,21 @@ async def _transcription_failure_context( chunk_count = "unknown" lines.append(f" Transcribed chunks: {chunk_count}") + try: + view = await store.read(session_id) + except Exception: # noqa: BLE001 + view = None + if view: + lines.append( + f" Provider connection: {view.transcription_provider_status or 'unknown'}" + ) + lines.append( + f" Last audio sent: {view.transcription_last_audio_sent_at or 'never'}" + ) + lines.append( + f" Last provider message: {view.transcription_last_message_at or 'never'}" + ) + lines.append(f" session={session_id} client={client_id}") return "\n".join(lines) @@ -1204,6 +1264,7 @@ async def stream_speech_detection_job( ) last_speech_analysis = None # Track last analysis for detailed logging max_runtime_reached = False # Distinguish the max-runtime exit from session close + no_activity_warning_logged = False # Main loop: Listen for speech while True: @@ -1222,7 +1283,12 @@ async def stream_speech_detection_job( # — in "off" mode nothing fills the aggregator by design, so skip it and let # the loop run until the session closes (full-audio batch fallback then runs). elapsed = time.time() - start_time - if expects_live_results and elapsed > 60 and not session_closed_at: + if ( + expects_live_results + and elapsed > 60 + and not session_closed_at + and not no_activity_warning_logged + ): watchdog_combined = await aggregator.get_combined_results(session_id) if not watchdog_combined.get("chunk_count", 0): # Only a real provider fault when audio is flowing in but the streaming @@ -1241,12 +1307,12 @@ async def stream_speech_detection_job( diag = await _transcription_failure_context( store, aggregator, session_id, client_id ) - logger.error( - f"❌ No transcription activity after {elapsed:.0f}s — " - f"check provider config (API key, connectivity, consumer running)\n" + logger.warning( + f"⚠️ No transcription activity after {elapsed:.0f}s — " + f"provider may be connected and awaiting recognizable speech\n" f"{diag}" ) - break + no_activity_warning_logged = True # Check if session has closed session_closed = await store.get_status(session_id) in ( @@ -1308,7 +1374,8 @@ async def stream_speech_detection_job( and grace_elapsed > 5 and not combined.get("chunk_count", 0) ): - # 5+ seconds with no transcription activity at all - likely API key issue + # A healthy provider may legitimately produce no messages for + # silence/noise, so lack of transcript alone is not a service error. diag = await _transcription_failure_context( store, aggregator, session_id, client_id ) @@ -1319,10 +1386,10 @@ async def stream_speech_detection_job( f"network drop, not a service fault\n{diag}" ) else: - logger.error( - f"❌ Session failed - check transcription service configuration\n" + logger.warning( + f"⚠️ Session ended without transcription\n" f" No transcription activity after {grace_elapsed:.1f}s " - f"(possible API key or connectivity issue)\n" + f"of finalization grace\n" f"{diag}" ) break @@ -1524,20 +1591,29 @@ async def stream_speech_detection_job( else: reason = "No transcription received" - # Distinguish between transcription failures (error) vs legitimate no speech (info) + # No transcript is not itself a provider failure: streaming APIs may emit no + # messages for silence/noise. Only an explicit transport/provider error is ERROR. if reason == "No transcription received": diag = await _transcription_failure_context( store, aggregator, session_id, client_id ) - if await _session_ended_by_disconnect(store, session_id): + provider_error = await store.get_transcription_error(session_id) + if provider_error: + logger.error( + f"❌ Session failed - transcription provider error\n" + f" Reason: {provider_error}\n" + f" Runtime: {time.time() - start_time:.1f}s\n" + f"{diag}" + ) + elif await _session_ended_by_disconnect(store, session_id): logger.warning( f"⚠️ Session ended by client disconnect with no transcription " f"(runtime {time.time() - start_time:.1f}s) — likely a network drop, " f"not a service fault\n{diag}" ) else: - logger.error( - f"❌ Session failed - transcription service did not respond\n" + logger.warning( + f"⚠️ Session ended without transcription\n" f" Reason: {reason}\n" f" Runtime: {time.time() - start_time:.1f}s\n" f"{diag}" diff --git a/backends/advanced/src/scripts/chronicle_data.py b/backends/advanced/src/scripts/chronicle_data.py new file mode 100644 index 00000000..a8b1d206 --- /dev/null +++ b/backends/advanced/src/scripts/chronicle_data.py @@ -0,0 +1,308 @@ +#!/usr/bin/env python3 +"""Administrative CLI for Chronicle data archives and memory reconstruction.""" + +from __future__ import annotations + +import argparse +import asyncio +import sys +from datetime import datetime, timezone +from pathlib import Path + +from rich.console import Console +from rich.panel import Panel +from rich.prompt import Confirm +from rich.table import Table + +from advanced_omi_backend.database import get_database +from advanced_omi_backend.services.data_archive import ( + ARCHIVE_SUFFIX, + ArchiveError, + create_data_archive, + import_data_archive, + verify_data_archive, +) +from advanced_omi_backend.services.memory.rebuild import ( + MemoryRebuildError, + RebuildStage, + build_rebuild_plan, + execute_memory_rebuild, +) + +console = Console() +DATA_DIR = Path("/app/data") + + +def _human_size(value: int) -> str: + size = float(value) + for unit in ("B", "KiB", "MiB", "GiB", "TiB"): + if size < 1024 or unit == "TiB": + return f"{size:.1f} {unit}" + size /= 1024 + raise AssertionError("unreachable") + + +def _default_archive_path() -> Path: + timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S") + return DATA_DIR / "backups" / f"chronicle_{timestamp}{ARCHIVE_SUFFIX}" + + +def _manifest_table(manifest: dict) -> Table: + table = Table(title="Archive contents") + table.add_column("Collection") + table.add_column("Documents", justify="right") + for name, metadata in manifest["collections"].items(): + table.add_row(name, str(metadata["documents"])) + data_files = sum( + metadata.get("kind") == "data_file" for metadata in manifest["files"].values() + ) + table.add_section() + table.add_row("Filesystem files", str(data_files)) + return table + + +def _require_confirmation(message: str, force: bool) -> None: + if force: + return + console.print(Panel(message, title="Destructive operation", border_style="red")) + if not Confirm.ask("Proceed?", default=False): + raise KeyboardInterrupt + + +async def _connect_database(): + database = get_database() + await database.command("ping") + return database + + +async def _run_export(args: argparse.Namespace) -> None: + database = await _connect_database() + output = args.output or _default_archive_path() + with console.status("Exporting Chronicle data..."): + summary = await create_data_archive( + database, + output, + data_dir=args.data_dir, + overwrite=args.overwrite, + ) + console.print( + Panel( + f"[green]Archive created[/green]\n" + f"Path: {summary.path}\n" + f"Collections: {summary.collections}\n" + f"Documents: {summary.documents}\n" + f"Filesystem files: {summary.files}\n" + f"Archive size: {_human_size(summary.bytes_written)}", + border_style="green", + ) + ) + + +async def _run_verify(args: argparse.Namespace) -> None: + with console.status("Verifying archive checksums..."): + manifest = verify_data_archive(args.archive) + console.print(_manifest_table(manifest)) + console.print("[green]Archive verification passed.[/green]") + + +async def _rebuild(database, args: argparse.Namespace): + from_stage = RebuildStage(args.rebuild_from) + plan = await build_rebuild_plan( + database, + args.user_id, + from_stage=from_stage, + ) + console.print( + f"Rebuild plan from {from_stage.value}: {plan.speaker_count} speaker inputs, " + f"{plan.memory_count} memory inputs across {len(plan.user_ids)} users." + ) + if from_stage is RebuildStage.SPEAKERS: + skipped_count = plan.count - plan.speaker_count + if skipped_count: + console.print( + f"[yellow]Speaker stage will skip {skipped_count} transcript-only " + "conversations with no stored audio.[/yellow]" + ) + if getattr(args, "dry_run", False): + return None + _require_confirmation( + "This deletes the selected users' current Markdown memory vaults and memory " + "audit history, then recreates them from active transcripts. Syncthing " + "pairing markers are retained.", + args.force, + ) + backup_dir = None if args.no_vault_backup else args.data_dir / "backups" + result = await execute_memory_rebuild( + database, + plan, + data_dir=args.data_dir, + backup_dir=backup_dir, + from_stage=from_stage, + ) + backup_text = str(result.vault_backup) if result.vault_backup else "not needed" + console.print( + Panel( + f"[green]Memory rebuild queued[/green]\n" + f"Run ID: {result.run_id}\n" + f"Speaker jobs: {len(result.speaker_jobs)}\n" + f"Speaker skipped (no audio): " + f"{len(result.skipped_speaker_conversations)}\n" + f"Memory jobs: {len(result.memory_jobs)}\n" + f"Users: {len(result.user_ids)}\n" + f"Deleted vault files: {result.deleted_vault_files}\n" + f"Deleted audit entries: {result.deleted_audit_entries}\n" + f"Previous vault backup: {backup_text}\n\n" + "Stages run chronologically within each user; different users may " + "rebuild in parallel.", + border_style="green", + ) + ) + return result + + +async def _run_import(args: argparse.Namespace) -> None: + if args.rebuild_from and args.user_id: + raise ArchiveError( + "--user-id cannot be combined with --rebuild-from import because the " + "archive contains all users. Import first, then run rebuild-memory " + "--user-id for a selective rebuild." + ) + with console.status("Verifying archive before import..."): + manifest = verify_data_archive(args.archive) + console.print(_manifest_table(manifest)) + destructive = args.replace or args.rebuild_from + if destructive: + _require_confirmation( + "Replace mode clears each archived Mongo collection before restore. Fresh " + "rebuild mode also deletes current derived vault and audit state.", + args.force, + ) + # One confirmation covers the complete import + derived-data rebuild. + if args.rebuild_from: + args.force = True + + database = await _connect_database() + restore_files = not args.database_only and not args.rebuild_from + with console.status("Importing verified Chronicle archive..."): + summary = await import_data_archive( + database, + args.archive, + data_dir=args.data_dir, + replace=args.replace, + restore_files=restore_files, + fresh_memory=bool(args.rebuild_from), + ) + console.print( + f"[green]Imported {summary.documents} documents from " + f"{summary.collections} collections and {summary.files} files.[/green]" + ) + if summary.skipped_collections: + console.print( + "Skipped derived collections: " + ", ".join(summary.skipped_collections) + ) + for warning in summary.duplicate_audio_warnings: + console.print( + "[yellow]Duplicate audio skipped:[/yellow] conversation " + f"{warning.skipped_conversation_id} matches {warning.kept_source} " + f"conversation {warning.kept_conversation_id}; kept the first copy." + ) + for warning in summary.duplicate_chunk_warnings: + console.print( + "[yellow]Duplicate audio chunk skipped:[/yellow] conversation " + f"{warning.conversation_id}, chunk {warning.chunk_index}; kept " + f"{warning.kept_source} chunk {warning.kept_chunk_id}." + ) + if args.rebuild_from: + await _rebuild(database, args) + + +async def _run_rebuild(args: argparse.Namespace) -> None: + database = await _connect_database() + await _rebuild(database, args) + + +def _add_common_data_dir(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "--data-dir", + type=Path, + default=DATA_DIR, + help="Chronicle data directory inside the container (default: /app/data)", + ) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Export, import, and reconstruct Chronicle's durable data" + ) + subparsers = parser.add_subparsers(dest="command", required=True) + + export_parser = subparsers.add_parser("export", help="Create a full data archive") + export_parser.add_argument("output", nargs="?", type=Path) + export_parser.add_argument("--overwrite", action="store_true") + _add_common_data_dir(export_parser) + export_parser.set_defaults(handler=_run_export) + + verify_parser = subparsers.add_parser( + "verify", help="Verify archive structure and checksums" + ) + verify_parser.add_argument("archive", type=Path) + verify_parser.set_defaults(handler=_run_verify) + + import_parser = subparsers.add_parser("import", help="Import a data archive") + import_parser.add_argument("archive", type=Path) + import_parser.add_argument( + "--replace", + action="store_true", + help="Clear each archived collection before restoring it", + ) + import_parser.add_argument( + "--database-only", + action="store_true", + help="Do not restore vault or legacy audio files", + ) + import_parser.add_argument( + "--rebuild-from", + choices=[stage.value for stage in RebuildStage], + help=( + "Skip archived derived memory and rebuild from this stage; " + "'speakers' runs speaker recognition before memory" + ), + ) + import_parser.add_argument("--user-id", action="append") + import_parser.add_argument("--no-vault-backup", action="store_true") + import_parser.add_argument("--force", action="store_true") + _add_common_data_dir(import_parser) + import_parser.set_defaults(handler=_run_import, dry_run=False) + + rebuild_parser = subparsers.add_parser( + "rebuild-memory", help="Recreate Markdown memories from active transcripts" + ) + rebuild_parser.add_argument("--user-id", action="append") + rebuild_parser.add_argument( + "--rebuild-from", + choices=[stage.value for stage in RebuildStage], + default=RebuildStage.MEMORY.value, + help="Earliest stage to rerun (default: memory)", + ) + rebuild_parser.add_argument("--dry-run", action="store_true") + rebuild_parser.add_argument("--no-vault-backup", action="store_true") + rebuild_parser.add_argument("--force", action="store_true") + _add_common_data_dir(rebuild_parser) + rebuild_parser.set_defaults(handler=_run_rebuild) + return parser + + +async def main() -> None: + args = build_parser().parse_args() + await args.handler(args) + + +if __name__ == "__main__": + try: + asyncio.run(main()) + except KeyboardInterrupt: + console.print("[yellow]Cancelled.[/yellow]") + sys.exit(130) + except (ArchiveError, MemoryRebuildError, OSError) as exc: + console.print(f"[bold red]Error:[/bold red] {exc}") + sys.exit(1) diff --git a/backends/advanced/tests/scripts/graph-validation/sample_data.py b/backends/advanced/tests/scripts/graph-validation/sample_data.py new file mode 100644 index 00000000..8157f8b0 --- /dev/null +++ b/backends/advanced/tests/scripts/graph-validation/sample_data.py @@ -0,0 +1,190 @@ +"""Insert sample conversation data into FalkorDB for validation. + +Reads sample .md files, chunks them by ### headers, generates real embeddings, +and stores as ConvDoc/ConvChunk/ConvEntity nodes with relationships. + +Usage: + uv run python tests/scripts/graph-validation/sample_data.py +""" + +import uuid +from pathlib import Path + +from utils import ( + generate_embeddings_sync, + get_graph, + parse_frontmatter, + parse_people_section, + print_fail, + print_header, + print_pass, + split_on_headers, +) + +SAMPLE_DIR = Path(__file__).parent / "sample_conversations" +TEST_USER_ID = "test_user_001" +OTHER_USER_ID = "other_user_002" # For user-scoping tests + + +def insert_conversation(graph, md_path: Path, user_id: str): + """Parse a conversation .md and insert into FalkorDB.""" + content = md_path.read_text(encoding="utf-8") + frontmatter = parse_frontmatter(content) + chunks = split_on_headers(content) + people = parse_people_section(content) + + conv_id = frontmatter.get("conversation_id", md_path.stem) + date = frontmatter.get("date", "2026-03-15T00:00:00") + speakers = frontmatter.get("speakers", "") + + # Extract title from first ## header + import re + + title_match = re.search(r"^##\s+(.+)$", content, re.MULTILINE) + title = title_match.group(1) if title_match else md_path.stem + + print(f"\n Processing: {title}") + print(f" Chunks: {len(chunks)}, People: {len(people)}") + + # Generate embeddings for all chunks + chunk_texts = [f"{c['section_title']}: {c['text']}" for c in chunks] + if not chunk_texts: + print_fail(f"No chunks extracted from {md_path.name}") + return + + print(f" Generating embeddings for {len(chunk_texts)} chunks...") + embeddings = generate_embeddings_sync(chunk_texts) + print(f" Embeddings generated ({len(embeddings[0])} dimensions)") + + # Create ConvDoc node + graph.query( + """ + MERGE (d:ConvDoc {conversation_id: $conv_id}) + SET d.title = $title, + d.date = $date, + d.user_id = $user_id, + d.speakers = $speakers, + d.file_path = $file_path + """, + params={ + "conv_id": conv_id, + "title": title, + "date": date, + "user_id": user_id, + "speakers": speakers, + "file_path": str(md_path.relative_to(SAMPLE_DIR.parent)), + }, + ) + + # Create ConvChunk nodes with embeddings and link to ConvDoc + for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)): + chunk_id = f"{conv_id}_chunk_{i:03d}" + graph.query( + """ + MATCH (d:ConvDoc {conversation_id: $conv_id}) + MERGE (c:ConvChunk {id: $chunk_id}) + SET c.text = $text, + c.section_title = $section_title, + c.embedding = vecf32($embedding), + c.user_id = $user_id, + c.created_at = $date, + c.chunk_index = $chunk_index + MERGE (d)-[:HAS_CHUNK]->(c) + """, + params={ + "conv_id": conv_id, + "chunk_id": chunk_id, + "text": chunk["text"], + "section_title": chunk["section_title"], + "embedding": embedding, + "user_id": user_id, + "date": date, + "chunk_index": i, + }, + ) + + # Create ConvEntity nodes from People section and link to chunks + for person in people: + entity_id = f"{user_id}_{person['name'].lower().replace(' ', '_')}" + graph.query( + """ + MERGE (e:ConvEntity {id: $entity_id}) + SET e.name = $name, + e.description = $description, + e.type = 'person', + e.user_id = $user_id + """, + params={ + "entity_id": entity_id, + "name": person["name"], + "description": person["description"], + "user_id": user_id, + }, + ) + + # Link entity to all chunks that mention their name + for i, chunk in enumerate(chunks): + if person["name"].lower() in chunk["text"].lower(): + chunk_id = f"{conv_id}_chunk_{i:03d}" + graph.query( + """ + MATCH (c:ConvChunk {id: $chunk_id}) + MATCH (e:ConvEntity {id: $entity_id}) + MERGE (c)-[:MENTIONS]->(e) + """, + params={ + "chunk_id": chunk_id, + "entity_id": entity_id, + }, + ) + + print_pass(f"Inserted: {title} ({len(chunks)} chunks, {len(people)} entities)") + + +def main(): + print_header("Inserting Sample Data") + + graph = get_graph() + + # Insert conversations for the primary test user + for md_file in sorted(SAMPLE_DIR.glob("*.md")): + insert_conversation(graph, md_file, TEST_USER_ID) + + # Insert one conversation for a different user (for scoping tests) + dentist_file = SAMPLE_DIR / "dentist_appointment.md" + if dentist_file.exists(): + print(f"\n Inserting duplicate for other user (scoping test)...") + # We need to modify the conv_id to avoid UNIQUE constraint clash + content = dentist_file.read_text(encoding="utf-8") + # Write a temp copy with different conv_id + temp_path = SAMPLE_DIR / "_other_user_dentist.md" + temp_path.write_text( + content.replace( + "conversation_id: conv_002", "conversation_id: conv_002_other" + ) + ) + insert_conversation(graph, temp_path, OTHER_USER_ID) + temp_path.unlink() + + # Verify counts + doc_result = graph.query("MATCH (d:ConvDoc) RETURN count(d) AS c") + doc_count = doc_result.result_set[0][0] + + chunk_result = graph.query("MATCH (c:ConvChunk) RETURN count(c) AS c") + chunk_count = chunk_result.result_set[0][0] + + entity_result = graph.query("MATCH (e:ConvEntity) RETURN count(e) AS c") + entity_count = entity_result.result_set[0][0] + + rel_result = graph.query("MATCH ()-[r:HAS_CHUNK|MENTIONS]->() RETURN count(r) AS c") + rel_count = rel_result.result_set[0][0] + + print_header("Summary") + print_pass(f"ConvDoc nodes: {doc_count}") + print_pass(f"ConvChunk nodes: {chunk_count}") + print_pass(f"ConvEntity nodes: {entity_count}") + print_pass(f"Relationships: {rel_count}") + + +if __name__ == "__main__": + main() diff --git a/backends/advanced/tests/scripts/graph-validation/setup_schema.py b/backends/advanced/tests/scripts/graph-validation/setup_schema.py new file mode 100644 index 00000000..7159baa1 --- /dev/null +++ b/backends/advanced/tests/scripts/graph-validation/setup_schema.py @@ -0,0 +1,139 @@ +"""Create FalkorDB schema for conversation memory validation. + +Creates constraints, vector index, and full-text index using ConvDoc/ConvChunk +labels to avoid conflicting with existing schema. + +Usage: + uv run python tests/scripts/graph-validation/setup_schema.py + uv run python tests/scripts/graph-validation/setup_schema.py --cleanup +""" + +import sys + +from utils import EMBEDDING_DIMENSIONS, get_graph, print_fail, print_header, print_pass + + +def setup(graph): + print_header("Creating Schema") + + # Constraints — FalkorDB doesn't support IF NOT EXISTS, use try/except + try: + graph.query( + "CREATE CONSTRAINT ON (d:ConvDoc) ASSERT d.conversation_id IS UNIQUE" + ) + print_pass("Constraint: ConvDoc.conversation_id UNIQUE") + except Exception as e: + if "already exists" in str(e).lower() or "already indexed" in str(e).lower(): + print_pass("Constraint: ConvDoc.conversation_id UNIQUE (already exists)") + else: + print_fail(f"Constraint ConvDoc.conversation_id: {e}") + + try: + graph.query("CREATE CONSTRAINT ON (c:ConvChunk) ASSERT c.id IS UNIQUE") + print_pass("Constraint: ConvChunk.id UNIQUE") + except Exception as e: + if "already exists" in str(e).lower() or "already indexed" in str(e).lower(): + print_pass("Constraint: ConvChunk.id UNIQUE (already exists)") + else: + print_fail(f"Constraint ConvChunk.id: {e}") + + try: + graph.query("CREATE CONSTRAINT ON (e:ConvEntity) ASSERT e.id IS UNIQUE") + print_pass("Constraint: ConvEntity.id UNIQUE") + except Exception as e: + if "already exists" in str(e).lower() or "already indexed" in str(e).lower(): + print_pass("Constraint: ConvEntity.id UNIQUE (already exists)") + else: + print_fail(f"Constraint ConvEntity.id: {e}") + + # Vector index + try: + graph.query( + f"CREATE VECTOR INDEX FOR (c:ConvChunk) ON (c.embedding) " + f"OPTIONS {{dimension: {EMBEDDING_DIMENSIONS}, similarityFunction: 'cosine'}}" + ) + print_pass( + f"Vector index: ConvChunk.embedding ({EMBEDDING_DIMENSIONS}d, cosine)" + ) + except Exception as e: + if "already exists" in str(e).lower() or "already indexed" in str(e).lower(): + print_pass( + f"Vector index: ConvChunk.embedding ({EMBEDDING_DIMENSIONS}d, cosine) (already exists)" + ) + else: + print_fail(f"Vector index: {e}") + + # Full-text index + try: + graph.query( + "CREATE FULLTEXT INDEX FOR (c:ConvChunk) ON (c.text, c.section_title)" + ) + print_pass("Full-text index: ConvChunk (text + section_title)") + except Exception as e: + if "already exists" in str(e).lower() or "already indexed" in str(e).lower(): + print_pass( + "Full-text index: ConvChunk (text + section_title) (already exists)" + ) + else: + print_fail(f"Full-text index: {e}") + + # Verify indexes exist + try: + result = graph.query("CALL db.indexes()") + idx_names = [] + for row in result.result_set: + # Each row is a list; index name/label varies by FalkorDB version + idx_names.append(str(row)) + print_pass(f"Indexes present: {len(result.result_set)} index(es) found") + except Exception as e: + print_fail(f"Could not verify indexes: {e}") + + print("\nSchema setup complete.") + + +def cleanup(graph): + print_header("Cleaning Up Test Data") + + # Delete all test nodes and relationships + result = graph.query( + "MATCH (n) WHERE n:ConvDoc OR n:ConvChunk OR n:ConvEntity " + "DETACH DELETE n RETURN count(n) AS deleted" + ) + if result.result_set: + deleted = result.result_set[0][0] + print_pass(f"Deleted {deleted} test nodes") + else: + print_pass("No test nodes to delete") + + # Drop indexes — FalkorDB doesn't support IF EXISTS for DROP, use try/except + for idx_query in [ + "DROP VECTOR INDEX ON :ConvChunk(embedding)", + "DROP FULLTEXT INDEX ON :ConvChunk(text, section_title)", + ]: + try: + graph.query(idx_query) + print_pass(f"Dropped index: {idx_query}") + except Exception as e: + print_fail(f"Failed to drop index ({idx_query}): {e}") + + # Drop constraints + for constraint_query in [ + "DROP CONSTRAINT ON (d:ConvDoc) ASSERT d.conversation_id IS UNIQUE", + "DROP CONSTRAINT ON (c:ConvChunk) ASSERT c.id IS UNIQUE", + "DROP CONSTRAINT ON (e:ConvEntity) ASSERT e.id IS UNIQUE", + ]: + try: + graph.query(constraint_query) + print_pass(f"Dropped constraint: {constraint_query}") + except Exception as e: + print_fail(f"Failed to drop constraint ({constraint_query}): {e}") + + print("\nCleanup complete.") + + +if __name__ == "__main__": + graph = get_graph() + if "--cleanup" in sys.argv: + cleanup(graph) + else: + setup(graph) diff --git a/backends/advanced/tests/scripts/graph-validation/test_conversation_doc.py b/backends/advanced/tests/scripts/graph-validation/test_conversation_doc.py new file mode 100644 index 00000000..7832e828 --- /dev/null +++ b/backends/advanced/tests/scripts/graph-validation/test_conversation_doc.py @@ -0,0 +1,236 @@ +"""Validate LLM conversation document generation and markdown parsing. + +Tests: +1. LLM produces a well-structured .md from a sample transcript +2. Parsed sections match expected structure (Summary, Key Facts, People, Action Items) +3. Entity extraction from People section is reliable +4. Header-based chunking produces correct splits + +This directly tests the review's concern about entity parsing fragility. + +Usage: + uv run python tests/scripts/graph-validation/test_conversation_doc.py +""" + +import asyncio +import json + +from openai import AsyncOpenAI +from utils import ( + OPENAI_API_KEY, + OPENAI_BASE_URL, + parse_action_items, + parse_frontmatter, + parse_people_section, + print_fail, + print_header, + print_pass, + split_on_headers, +) + +SAMPLE_TRANSCRIPT = """Speaker 0: Hey John, thanks for meeting. I wanted to go over the Q3 timeline. +John: Sure, so the main concern is the backend migration. It might block our September 15th deadline. +Speaker 0: Right. Sarah mentioned she wants to descope the auth rewrite. What do you think? +John: I agree with Sarah actually. The auth rewrite is nice to have but not critical for Q3. Let's push it to Q4. +Speaker 0: Makes sense. I'll schedule a meeting with Sarah to align on scope. Oh and John, do you prefer morning or afternoon standups? +John: Morning for sure. I'm much more productive before noon. +Speaker 0: Got it. One more thing - we need the budget review ready by Friday. Can you pull the numbers? +John: Yeah I'll have the draft ready by Thursday. Mike from finance wants to review it too. +Speaker 0: Perfect. I'll also check with the DevOps team about the migration timeline. Thanks John. +John: Sounds good, talk later.""" + +GENERATE_CONVERSATION_DOC_PROMPT = """\ +You are generating a structured conversation document from a transcript. + +Given a transcript with speaker labels, produce a markdown document with this EXACT structure: + +--- +conversation_id: {conversation_id} +date: {date} +speakers: [{speakers}] +duration_minutes: {duration} +--- + +## {Title - descriptive, 3-8 words} + +### Summary +{2-3 sentence summary of what was discussed} + +### Key Facts +{Bulleted list of specific facts, decisions, numbers, dates mentioned} + +### People +{Bulleted list in format: - Name (role/relationship, context)} +Include ALL named individuals — speakers, people mentioned, people referenced. +If a speaker is identified by name (e.g., "John" not "Speaker 0"), they MUST appear here. +Do not include unnamed roles or generic labels like "Speaker 0". + +### Action Items +{Bulleted list in format: - [ ] Action item description} +Use [x] for items already completed in the conversation. + +Rules: +- Every ### section MUST be present, even if empty (use "- None" for empty sections) +- People section: ONLY named individuals, format MUST be "- Name (description)" +- Key Facts: be specific — include dates, numbers, names +- Action Items: use checkbox format [ ] or [x] +- Do NOT add any sections beyond the four listed above +""" + + +async def generate_doc(transcript: str) -> str: + """Call LLM to generate a conversation document.""" + client = AsyncOpenAI(api_key=OPENAI_API_KEY, base_url=OPENAI_BASE_URL) + + response = await client.chat.completions.create( + model="gpt-4o-mini", + messages=[ + {"role": "system", "content": GENERATE_CONVERSATION_DOC_PROMPT}, + { + "role": "user", + "content": f"Transcript:\n{transcript}\n\nMetadata:\n- conversation_id: test_conv_llm\n- date: 2026-03-18T10:00:00\n- speakers: Speaker 0, John\n- duration: 8 minutes", + }, + ], + temperature=0.2, + ) + return response.choices[0].message.content.strip() + + +def test_structure(doc: str): + """Verify the generated doc has all required sections.""" + print_header("Test: Document Structure") + + required_sections = ["Summary", "Key Facts", "People", "Action Items"] + chunks = split_on_headers(doc) + found_sections = [c["section_title"] for c in chunks] + + print(f" Found sections: {found_sections}") + + for section in required_sections: + if section in found_sections: + print_pass(f"Section '{section}' present") + else: + print_fail(f"Section '{section}' MISSING") + + # Check frontmatter + fm = parse_frontmatter(doc) + if fm.get("conversation_id"): + print_pass(f"Frontmatter has conversation_id: {fm['conversation_id']}") + else: + print_fail("Frontmatter missing conversation_id") + + if fm.get("date"): + print_pass(f"Frontmatter has date: {fm['date']}") + else: + print_fail("Frontmatter missing date") + + return chunks + + +def test_entity_parsing(doc: str): + """Test that People section can be reliably parsed.""" + print_header("Test: Entity Parsing from People Section") + + people = parse_people_section(doc) + print(f" Parsed entities: {len(people)}") + for p in people: + print(f" - {p['name']} ({p['description']})") + + # Expected entities from the transcript + expected_names = {"John", "Sarah", "Mike"} + + found_names = {p["name"] for p in people} + for name in expected_names: + if name in found_names: + print_pass(f"Entity '{name}' extracted") + else: + print_fail(f"Entity '{name}' NOT extracted (found: {found_names})") + + # Check that descriptions are non-empty + has_descriptions = sum(1 for p in people if p["description"]) + if has_descriptions == len(people): + print_pass(f"All {len(people)} entities have descriptions") + elif has_descriptions > 0: + print_pass(f"{has_descriptions}/{len(people)} entities have descriptions") + else: + print_fail("No entities have descriptions") + + # Check for unnamed entries (these should NOT appear) + unnamed = [ + p for p in people if p["name"].lower().startswith(("speaker", "the ", "a ")) + ] + if not unnamed: + print_pass("No unnamed/generic entities (good)") + else: + print_fail(f"Found unnamed entities: {[p['name'] for p in unnamed]}") + + return people + + +def test_action_items(doc: str): + """Test Action Items parsing.""" + print_header("Test: Action Items Parsing") + + items = parse_action_items(doc) + print(f" Parsed items: {len(items)}") + for item in items: + status = "[x]" if item["done"] else "[ ]" + print(f" {status} {item['text']}") + + if items: + print_pass(f"Extracted {len(items)} action items") + else: + print_fail("No action items extracted") + + # Should have at least 2 action items from the transcript + if len(items) >= 2: + print_pass(f"At least 2 action items found (got {len(items)})") + else: + print_fail(f"Expected at least 2 action items, got {len(items)}") + + +def test_chunking(doc: str): + """Test header-based chunking produces correct splits.""" + print_header("Test: Header-Based Chunking") + + chunks = split_on_headers(doc) + print(f" Chunks: {len(chunks)}") + for c in chunks: + text_preview = c["text"][:80].replace("\n", " ") + print(f" [{c['section_title']}] {text_preview}...") + + # Should have 4 chunks (Summary, Key Facts, People, Action Items) + if len(chunks) >= 4: + print_pass(f"Got {len(chunks)} chunks (expected ≥4)") + else: + print_fail(f"Got {len(chunks)} chunks (expected ≥4)") + + # Each chunk should have non-empty text + empty_chunks = [c for c in chunks if not c["text"].strip()] + if not empty_chunks: + print_pass("All chunks have non-empty text") + else: + print_fail(f"{len(empty_chunks)} chunks have empty text") + + +async def main(): + print_header("Generating Conversation Document via LLM") + print(f" Transcript length: {len(SAMPLE_TRANSCRIPT)} chars") + print(f" Generating...") + + doc = await generate_doc(SAMPLE_TRANSCRIPT) + + print(f" Generated document: {len(doc)} chars") + print(f"\n{'─'*60}") + print(doc) + print(f"{'─'*60}\n") + + # Run all tests + test_structure(doc) + test_entity_parsing(doc) + test_action_items(doc) + test_chunking(doc) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/backends/advanced/tests/scripts/graph-validation/test_entity_graph.py b/backends/advanced/tests/scripts/graph-validation/test_entity_graph.py new file mode 100644 index 00000000..f80fedb3 --- /dev/null +++ b/backends/advanced/tests/scripts/graph-validation/test_entity_graph.py @@ -0,0 +1,161 @@ +"""Validate entity-based graph traversal search. + +Tests: +1. Find chunks mentioning a specific person via graph traversal +2. Cross-entity query (conversations mentioning both X and Y) +3. Entity listing per user + +Usage: + uv run python tests/scripts/graph-validation/test_entity_graph.py +""" + +from utils import get_graph, print_fail, print_header, print_pass, print_results_table + +TEST_USER = "test_user_001" + + +def _parse_results(result): + """Parse FalkorDB result_set into list of dicts using header names.""" + headers = [h[1] for h in result.header] + rows = [] + for row in result.result_set: + rows.append({headers[i]: row[i] for i in range(len(headers))}) + return rows + + +def test_entity_search(graph): + """Search for chunks mentioning 'John' via graph traversal.""" + print_header("Test: Entity Graph Traversal -- 'John'") + + result = graph.query( + """ + MATCH (e:ConvEntity {user_id: $user_id}) + WHERE toLower(e.name) CONTAINS toLower($entity_name) + WITH e + MATCH (chunk:ConvChunk)-[:MENTIONS]->(e) + MATCH (doc:ConvDoc)-[:HAS_CHUNK]->(chunk) + RETURN e.name AS entity, chunk.id AS chunk_id, + chunk.section_title AS section, chunk.text AS text, + doc.title AS title, doc.date AS date + ORDER BY doc.date DESC + """, + params={"user_id": TEST_USER, "entity_name": "John"}, + ) + results = _parse_results(result) + + print(f" Entity: 'John'") + print(f" Results: {len(results)}") + print_results_table(results, ["entity", "title", "section"]) + + if results: + print_pass(f"Found {len(results)} chunks mentioning John") + # All should be from Q3 meeting + titles = {r["title"] for r in results} + if all("Q3" in t for t in titles): + print_pass("All John mentions are in Q3 meeting (expected)") + else: + print_fail(f"Unexpected titles: {titles}") + else: + print_fail("No results for entity 'John'") + + +def test_entity_search_dr_chen(graph): + """Search for chunks mentioning 'Dr. Chen'.""" + print_header("Test: Entity Graph Traversal -- 'Dr. Chen'") + + result = graph.query( + """ + MATCH (e:ConvEntity {user_id: $user_id}) + WHERE toLower(e.name) CONTAINS toLower($entity_name) + WITH e + MATCH (chunk:ConvChunk)-[:MENTIONS]->(e) + MATCH (doc:ConvDoc)-[:HAS_CHUNK]->(chunk) + RETURN e.name AS entity, chunk.section_title AS section, + doc.title AS title, doc.date AS date + ORDER BY doc.date DESC + """, + params={"user_id": TEST_USER, "entity_name": "Dr. Chen"}, + ) + results = _parse_results(result) + + print(f" Entity: 'Dr. Chen'") + print(f" Results: {len(results)}") + print_results_table(results, ["entity", "title", "section"]) + + if results: + print_pass(f"Found {len(results)} chunks mentioning Dr. Chen") + if all("Dentist" in r.get("title", "") for r in results): + print_pass("All Dr. Chen mentions are in dentist conversation (expected)") + else: + print_fail("No results for entity 'Dr. Chen'") + + +def test_cross_entity_query(graph): + """Find conversations that mention BOTH John AND Sarah.""" + print_header("Test: Cross-Entity Query (John AND Sarah)") + + result = graph.query( + """ + MATCH (e1:ConvEntity {user_id: $user_id}) + WHERE toLower(e1.name) CONTAINS 'john' + MATCH (e2:ConvEntity {user_id: $user_id}) + WHERE toLower(e2.name) CONTAINS 'sarah' + MATCH (doc:ConvDoc)-[:HAS_CHUNK]->(c1:ConvChunk)-[:MENTIONS]->(e1) + MATCH (doc)-[:HAS_CHUNK]->(c2:ConvChunk)-[:MENTIONS]->(e2) + RETURN DISTINCT doc.conversation_id AS conv_id, doc.title AS title, + collect(DISTINCT e1.name) AS entity1, + collect(DISTINCT e2.name) AS entity2 + """, + params={"user_id": TEST_USER}, + ) + results = _parse_results(result) + + print(f" Cross-entity: John AND Sarah") + print(f" Results: {len(results)}") + print_results_table(results, ["title", "entity1", "entity2"]) + + if results: + print_pass(f"Found {len(results)} conversations mentioning both John and Sarah") + else: + print_fail("No conversations found with both John and Sarah") + + +def test_list_entities(graph): + """List all entities for a user with mention counts.""" + print_header("Test: List All Entities") + + result = graph.query( + """ + MATCH (e:ConvEntity {user_id: $user_id}) + OPTIONAL MATCH (chunk:ConvChunk)-[:MENTIONS]->(e) + RETURN e.name AS name, e.type AS type, e.description AS description, + count(chunk) AS mention_count + ORDER BY mention_count DESC + """, + params={"user_id": TEST_USER}, + ) + results = _parse_results(result) + + print(f" User: {TEST_USER}") + print(f" Entities: {len(results)}") + print_results_table(results, ["name", "type", "description", "mention_count"]) + + if results: + print_pass(f"Found {len(results)} entities for user") + # Should have at least John, Sarah, Dr. Chen, etc. + names = {r["name"] for r in results} + for expected in ["John", "Sarah", "Dr. Chen"]: + if expected in names: + print_pass(f"Entity '{expected}' found") + else: + print_fail(f"Entity '{expected}' not found in {names}") + else: + print_fail("No entities found") + + +if __name__ == "__main__": + graph = get_graph() + test_entity_search(graph) + test_entity_search_dr_chen(graph) + test_cross_entity_query(graph) + test_list_entities(graph) diff --git a/backends/advanced/tests/scripts/graph-validation/test_fulltext_search.py b/backends/advanced/tests/scripts/graph-validation/test_fulltext_search.py new file mode 100644 index 00000000..20d1ae86 --- /dev/null +++ b/backends/advanced/tests/scripts/graph-validation/test_fulltext_search.py @@ -0,0 +1,215 @@ +"""Validate FalkorDB full-text (BM25) search. + +Tests: +1. Single keyword search returns matching chunks +2. Multi-term AND search narrows results +3. BM25 scores are non-zero and ranked correctly +4. User scoping works with full-text results + +Usage: + uv run python tests/scripts/graph-validation/test_fulltext_search.py +""" + +from utils import get_graph, print_fail, print_header, print_pass, print_results_table + +TEST_USER = "test_user_001" + + +def _parse_results(result): + """Parse FalkorDB result_set into list of dicts using header names.""" + headers = [h[1] for h in result.header] + rows = [] + for row in result.result_set: + rows.append({headers[i]: row[i] for i in range(len(headers))}) + return rows + + +def test_single_keyword(graph): + """Search for 'deadline' -- should return Q3 meeting chunks.""" + print_header("Test: Single Keyword Full-Text Search") + + result = graph.query( + """ + CALL db.idx.fulltext.queryNodes('ConvChunk', $search_term) + YIELD node AS chunk, score + WHERE chunk.user_id = $user_id + MATCH (doc:ConvDoc)-[:HAS_CHUNK]->(chunk) + RETURN chunk.id AS chunk_id, chunk.section_title AS section, + chunk.text AS text, doc.title AS title, score + ORDER BY score DESC + """, + params={"search_term": "deadline", "user_id": TEST_USER}, + ) + results = _parse_results(result) + + print(f" Query: 'deadline'") + print(f" Results: {len(results)}") + print_results_table(results, ["score", "title", "section"]) + + if results: + print_pass(f"Full-text search returned {len(results)} results") + # Check that scores are positive + if all(r["score"] > 0 for r in results): + print_pass("All BM25 scores are positive") + else: + print_fail("Some scores are zero or negative") + else: + print_fail("No results for 'deadline'") + + return results + + +def test_multi_term_search(graph): + """Search for 'Q3 AND deadline' -- should narrow to Q3 meeting.""" + print_header("Test: Multi-Term AND Search") + + result = graph.query( + """ + CALL db.idx.fulltext.queryNodes('ConvChunk', $search_term) + YIELD node AS chunk, score + WHERE chunk.user_id = $user_id + MATCH (doc:ConvDoc)-[:HAS_CHUNK]->(chunk) + RETURN chunk.id AS chunk_id, chunk.section_title AS section, + doc.title AS title, score + ORDER BY score DESC + """, + params={"search_term": "Q3 AND deadline", "user_id": TEST_USER}, + ) + results = _parse_results(result) + + print(f" Query: 'Q3 AND deadline'") + print(f" Results: {len(results)}") + print_results_table(results, ["score", "title", "section"]) + + if results: + # All results should be from Q3 meeting + q3_results = [r for r in results if "Q3" in r.get("title", "")] + if len(q3_results) == len(results): + print_pass("All multi-term results are from Q3 meeting") + else: + non_q3 = [r["title"] for r in results if "Q3" not in r.get("title", "")] + print_fail(f"Some results are not Q3: {non_q3}") + else: + print_fail("No results for 'Q3 AND deadline'") + + +def test_dental_terms(graph): + """Search for dental-specific terms -- should return dentist chunks.""" + print_header("Test: Domain-Specific Terms") + + result = graph.query( + """ + CALL db.idx.fulltext.queryNodes('ConvChunk', $search_term) + YIELD node AS chunk, score + WHERE chunk.user_id = $user_id + MATCH (doc:ConvDoc)-[:HAS_CHUNK]->(chunk) + RETURN chunk.id AS chunk_id, chunk.section_title AS section, + doc.title AS title, score + ORDER BY score DESC + """, + params={"search_term": "crown molar", "user_id": TEST_USER}, + ) + results = _parse_results(result) + + print(f" Query: 'crown molar'") + print(f" Results: {len(results)}") + print_results_table(results, ["score", "title", "section"]) + + if results and "Dentist" in results[0].get("title", ""): + print_pass("Top result is dentist conversation") + elif results: + print_fail(f"Top result: {results[0].get('title')}") + else: + print_fail("No results for 'crown molar'") + + +def test_no_results(graph): + """Search for completely unrelated term -- should return empty.""" + print_header("Test: No Results for Unrelated Query") + + result = graph.query( + """ + CALL db.idx.fulltext.queryNodes('ConvChunk', $search_term) + YIELD node AS chunk, score + WHERE chunk.user_id = $user_id + RETURN chunk.id AS chunk_id, score + """, + params={"search_term": "quantum entanglement photon", "user_id": TEST_USER}, + ) + results = _parse_results(result) + + print(f" Query: 'quantum entanglement photon'") + print(f" Results: {len(results)}") + + if len(results) == 0: + print_pass("No results for unrelated query (expected)") + else: + print_fail(f"Got {len(results)} unexpected results") + + +def test_user_scoping_fulltext(graph): + """Verify full-text results are scoped to user.""" + print_header("Test: Full-Text User Scoping") + + # Search as test user -- should find all 3 conversations + result = graph.query( + """ + CALL db.idx.fulltext.queryNodes('ConvChunk', $search_term) + YIELD node AS chunk, score + WHERE chunk.user_id = $user_id + RETURN chunk.id AS chunk_id, chunk.user_id AS uid + """, + params={ + "search_term": "appointment OR deadline OR grocery", + "user_id": TEST_USER, + }, + ) + test_results = _parse_results(result) + + # Search as other user -- should find only their dentist copy + result = graph.query( + """ + CALL db.idx.fulltext.queryNodes('ConvChunk', $search_term) + YIELD node AS chunk, score + WHERE chunk.user_id = $user_id + RETURN chunk.id AS chunk_id, chunk.user_id AS uid + """, + params={ + "search_term": "appointment OR deadline OR grocery", + "user_id": "other_user_002", + }, + ) + other_results = _parse_results(result) + + print(f" TEST_USER results: {len(test_results)}") + print(f" OTHER_USER results: {len(other_results)}") + + if len(test_results) > len(other_results): + print_pass("TEST_USER has more results (expected -- more conversations)") + else: + print_pass( + f"Result counts: TEST={len(test_results)}, OTHER={len(other_results)}" + ) + + # Verify no cross-user contamination + test_uids = {r["uid"] for r in test_results} + other_uids = {r["uid"] for r in other_results} + + if test_uids <= {TEST_USER}: + print_pass("TEST_USER full-text results scoped correctly") + else: + print_fail(f"Leaked user_ids in TEST results: {test_uids}") + + if other_uids <= {"other_user_002"}: + print_pass("OTHER_USER full-text results scoped correctly") + else: + print_fail(f"Leaked user_ids in OTHER results: {other_uids}") + + +if __name__ == "__main__": + graph = get_graph() + test_single_keyword(graph) + test_multi_term_search(graph) + test_dental_terms(graph) + test_no_results(graph) + test_user_scoping_fulltext(graph) diff --git a/backends/advanced/tests/scripts/graph-validation/test_hybrid_search.py b/backends/advanced/tests/scripts/graph-validation/test_hybrid_search.py new file mode 100644 index 00000000..0105335e --- /dev/null +++ b/backends/advanced/tests/scripts/graph-validation/test_hybrid_search.py @@ -0,0 +1,199 @@ +"""Validate hybrid search: vector + full-text + recency bias. + +Tests: +1. Hybrid merge produces better ranking than either alone +2. Recency bias correctly boosts recent conversations +3. Score components are traceable + +Usage: + uv run python tests/scripts/graph-validation/test_hybrid_search.py +""" + +from utils import ( + compute_hybrid_scores, + generate_embeddings_sync, + get_graph, + print_fail, + print_header, + print_pass, + print_results_table, +) + +TEST_USER = "test_user_001" + + +def _parse_results(result): + """Parse FalkorDB result_set into list of dicts using header names.""" + headers = [h[1] for h in result.header] + rows = [] + for row in result.result_set: + rows.append({headers[i]: row[i] for i in range(len(headers))}) + return rows + + +def vector_search(graph, query: str, user_id: str, limit: int = 20): + """Run vector search and return results.""" + embedding = generate_embeddings_sync([query])[0] + + result = graph.query( + """ + CALL db.idx.vector.queryNodes('ConvChunk', 'embedding', $limit, vecf32($vector)) + YIELD node AS chunk, score + WHERE chunk.user_id = $user_id + MATCH (doc:ConvDoc)-[:HAS_CHUNK]->(chunk) + RETURN chunk.id AS chunk_id, chunk.section_title AS section, + doc.title AS title, doc.date AS date, score + ORDER BY score DESC + """, + params={"vector": embedding, "limit": limit, "user_id": user_id}, + ) + return _parse_results(result) + + +def fulltext_search(graph, query: str, user_id: str): + """Run full-text search and return results.""" + result = graph.query( + """ + CALL db.idx.fulltext.queryNodes('ConvChunk', $search_term) + YIELD node AS chunk, score + WHERE chunk.user_id = $user_id + MATCH (doc:ConvDoc)-[:HAS_CHUNK]->(chunk) + RETURN chunk.id AS chunk_id, chunk.section_title AS section, + doc.title AS title, doc.date AS date, score + ORDER BY score DESC + """, + params={"search_term": query, "user_id": user_id}, + ) + return _parse_results(result) + + +def test_hybrid_merge(graph): + """Test that hybrid merge combines vector + full-text results.""" + print_header("Test: Hybrid Merge (Vector + Full-Text)") + + query = "project deadline September" + + vec_results = vector_search(graph, query, TEST_USER) + ft_results = fulltext_search(graph, query, TEST_USER) + + print(f" Query: '{query}'") + print(f" Vector results: {len(vec_results)}") + print(f" Full-text results: {len(ft_results)}") + + # Merge + merged = compute_hybrid_scores(vec_results, ft_results) + + print(f"\n Hybrid results (top 5):") + print_results_table( + merged[:5], + ["final_score", "relevance_score", "recency_score", "title", "section"], + ) + + if merged: + print_pass(f"Hybrid merge produced {len(merged)} results") + + # Check that both sources contribute + has_vector = any(r["vector_score"] > 0 for r in merged) + has_text = any(r["text_score"] > 0 for r in merged) + if has_vector: + print_pass("Vector scores contributing to results") + else: + print_fail("No vector scores in merged results") + if has_text: + print_pass("Full-text scores contributing to results") + else: + print_fail( + "No full-text scores in merged results (might be expected if query terms not exact)" + ) + else: + print_fail("Hybrid merge returned no results") + + +def test_recency_bias(graph): + """Test that recency bias boosts recent conversations.""" + print_header("Test: Recency Bias") + + # Use a generic query that matches multiple conversations + query = "appointment schedule plans" + + vec_results = vector_search(graph, query, TEST_USER) + ft_results = fulltext_search(graph, query, TEST_USER) + + # Compare with and without recency + no_recency = compute_hybrid_scores( + vec_results, ft_results, recency_half_life_days=99999, recency_floor=1.0 + ) + with_recency = compute_hybrid_scores( + vec_results, ft_results, recency_half_life_days=30, recency_floor=0.5 + ) + + print(f" Query: '{query}'") + print(f"\n WITHOUT recency bias:") + print_results_table( + no_recency[:5], + ["final_score", "recency_score", "title", "date"], + ) + print(f"\n WITH recency bias (30-day half-life):") + print_results_table( + with_recency[:5], + ["final_score", "recency_score", "title", "date"], + ) + + if with_recency: + # Most recent conversation (grocery, 2026-03-16) should rank higher with recency + recency_scores = {r["title"]: r["recency_score"] for r in with_recency} + if recency_scores: + print_pass("Recency scores computed for all results") + + # Check that newer dates have higher recency scores + dates_and_recency = [(r["date"], r["recency_score"]) for r in with_recency] + print(f"\n Date -> Recency mapping:") + for d, rc in sorted(set(dates_and_recency)): + print(f" {d} -> {rc:.4f}") + else: + print_fail("No recency scores found") + else: + print_fail("No results to test recency on") + + +def test_exact_keyword_boost(graph): + """Test that exact keyword match (BM25) boosts relevance.""" + print_header("Test: Exact Keyword Boost") + + # Search for "Dr. Chen" -- should strongly prefer dentist via BM25 + query_semantic = "doctor dental health" + query_exact = "Dr. Chen" + + vec_results = vector_search(graph, query_semantic, TEST_USER) + ft_results_exact = fulltext_search(graph, query_exact, TEST_USER) + + # Hybrid with exact keyword + merged = compute_hybrid_scores(vec_results, ft_results_exact) + + print(f" Semantic query: '{query_semantic}'") + print(f" Exact keyword: '{query_exact}'") + print(f"\n Hybrid results:") + print_results_table( + merged[:5], ["final_score", "vector_score", "text_score", "title", "section"] + ) + + if merged: + # Chunks with exact "Dr. Chen" match should have text_score > 0 + exact_matches = [r for r in merged if r["text_score"] > 0] + if exact_matches: + print_pass(f"{len(exact_matches)} results boosted by exact keyword match") + # These should all be from the dentist conversation + dentist_exact = [ + r for r in exact_matches if "Dentist" in r.get("title", "") + ] + if dentist_exact: + print_pass("Exact keyword matches are from dentist conversation") + else: + print_fail("No results had text_score > 0 for exact keyword") + + +if __name__ == "__main__": + graph = get_graph() + test_hybrid_merge(graph) + test_recency_bias(graph) + test_exact_keyword_boost(graph) diff --git a/backends/advanced/tests/scripts/graph-validation/test_vector_search.py b/backends/advanced/tests/scripts/graph-validation/test_vector_search.py new file mode 100644 index 00000000..a2bec151 --- /dev/null +++ b/backends/advanced/tests/scripts/graph-validation/test_vector_search.py @@ -0,0 +1,174 @@ +"""Validate vector index search with user scoping. + +Tests: +1. Semantic search finds relevant chunks (project deadline -> Q3 meeting) +2. User scoping prevents cross-user data leakage +3. Score ordering is correct + +Usage: + uv run python tests/scripts/graph-validation/test_vector_search.py +""" + +from utils import ( + generate_embeddings_sync, + get_graph, + print_fail, + print_header, + print_pass, + print_results_table, +) + +TEST_USER = "test_user_001" +OTHER_USER = "other_user_002" + + +def _parse_vector_results(result, fields): + """Parse FalkorDB result_set into list of dicts using header names.""" + headers = [h[1] for h in result.header] + rows = [] + for row in result.result_set: + rows.append({headers[i]: row[i] for i in range(len(headers))}) + return rows + + +def test_basic_vector_search(graph): + """Search for project deadline -- should return Q3 meeting chunks.""" + print_header("Test: Basic Vector Search") + + query = "project deadline and timeline" + embedding = generate_embeddings_sync([query])[0] + + result = graph.query( + """ + CALL db.idx.vector.queryNodes('ConvChunk', 'embedding', $limit, vecf32($vector)) + YIELD node AS chunk, score + WHERE chunk.user_id = $user_id + MATCH (doc:ConvDoc)-[:HAS_CHUNK]->(chunk) + RETURN chunk.id AS chunk_id, chunk.section_title AS section, + chunk.text AS text, doc.title AS title, score + ORDER BY score DESC + """, + params={"vector": embedding, "limit": 10, "user_id": TEST_USER}, + ) + results = _parse_vector_results(result, []) + + print(f" Query: '{query}'") + print(f" Results: {len(results)}") + print_results_table(results, ["score", "title", "section"]) + + # Verify: top result should be from Q3 meeting + if results and "Q3" in results[0].get("title", ""): + print_pass("Top result is from Q3 meeting (expected)") + elif results: + print_fail(f"Top result is '{results[0].get('title')}', expected Q3 meeting") + else: + print_fail("No results returned") + + # Verify scores are in descending order + scores = [r["score"] for r in results] + if scores == sorted(scores, reverse=True): + print_pass("Scores in descending order") + else: + print_fail("Scores NOT in descending order") + + return results + + +def test_semantic_relevance(graph): + """Search for health/medical topic -- should return dentist chunks.""" + print_header("Test: Semantic Relevance") + + query = "dental health and medical procedures" + embedding = generate_embeddings_sync([query])[0] + + result = graph.query( + """ + CALL db.idx.vector.queryNodes('ConvChunk', 'embedding', $limit, vecf32($vector)) + YIELD node AS chunk, score + WHERE chunk.user_id = $user_id + MATCH (doc:ConvDoc)-[:HAS_CHUNK]->(chunk) + RETURN chunk.id AS chunk_id, chunk.section_title AS section, + doc.title AS title, score + ORDER BY score DESC + """, + params={"vector": embedding, "limit": 5, "user_id": TEST_USER}, + ) + results = _parse_vector_results(result, []) + + print(f" Query: '{query}'") + print_results_table(results, ["score", "title", "section"]) + + if results and "Dentist" in results[0].get("title", ""): + print_pass("Top result is dentist conversation (expected)") + elif results: + print_fail(f"Top result is '{results[0].get('title')}', expected Dentist") + else: + print_fail("No results returned") + + +def test_user_scoping(graph): + """Verify user_id filtering prevents cross-user results.""" + print_header("Test: User Scoping") + + query = "dentist crown procedure" + embedding = generate_embeddings_sync([query])[0] + + # Search as TEST_USER + result = graph.query( + """ + CALL db.idx.vector.queryNodes('ConvChunk', 'embedding', $limit, vecf32($vector)) + YIELD node AS chunk, score + WHERE chunk.user_id = $user_id + RETURN chunk.id AS chunk_id, chunk.user_id AS user_id, score + ORDER BY score DESC + """, + params={"vector": embedding, "limit": 20, "user_id": TEST_USER}, + ) + test_user_results = _parse_vector_results(result, []) + + # Search as OTHER_USER + result = graph.query( + """ + CALL db.idx.vector.queryNodes('ConvChunk', 'embedding', $limit, vecf32($vector)) + YIELD node AS chunk, score + WHERE chunk.user_id = $user_id + RETURN chunk.id AS chunk_id, chunk.user_id AS user_id, score + ORDER BY score DESC + """, + params={"vector": embedding, "limit": 20, "user_id": OTHER_USER}, + ) + other_user_results = _parse_vector_results(result, []) + + print(f" TEST_USER results: {len(test_user_results)}") + print(f" OTHER_USER results: {len(other_user_results)}") + + # Verify no cross-user leakage + test_user_ids = {r["user_id"] for r in test_user_results} + other_user_ids = {r["user_id"] for r in other_user_results} + + if test_user_ids <= {TEST_USER}: + print_pass("TEST_USER results contain only TEST_USER data") + else: + print_fail(f"TEST_USER results leaked: {test_user_ids}") + + if other_user_ids <= {OTHER_USER}: + print_pass("OTHER_USER results contain only OTHER_USER data") + else: + print_fail(f"OTHER_USER results leaked: {other_user_ids}") + + # Verify different result counts (other user has fewer conversations) + if len(test_user_results) > len(other_user_results): + print_pass( + f"TEST_USER has more results ({len(test_user_results)} vs {len(other_user_results)})" + ) + else: + print_pass( + f"Result counts: TEST={len(test_user_results)}, OTHER={len(other_user_results)}" + ) + + +if __name__ == "__main__": + graph = get_graph() + test_basic_vector_search(graph) + test_semantic_relevance(graph) + test_user_scoping(graph) diff --git a/backends/advanced/tests/scripts/graph-validation/utils.py b/backends/advanced/tests/scripts/graph-validation/utils.py new file mode 100644 index 00000000..0d5a8d1a --- /dev/null +++ b/backends/advanced/tests/scripts/graph-validation/utils.py @@ -0,0 +1,292 @@ +"""Shared utilities for FalkorDB graph memory validation scripts.""" + +import asyncio +import math +import os +import re +import sys +from datetime import datetime, timezone +from pathlib import Path +from typing import Dict, List, Optional, Tuple + +from dotenv import load_dotenv +from falkordb import FalkorDB + +# Load .env from backends/advanced/ +# __file__ is .../backends/advanced/tests/scripts/graph-validation/utils.py +# parents: [0]=graph-validation, [1]=scripts, [2]=tests, [3]=advanced +_env_path = Path(__file__).resolve().parents[3] / ".env" +load_dotenv(_env_path) + +# --- FalkorDB Connection --- + +_falkordb_host = os.getenv("FALKORDB_HOST", "localhost") +if _falkordb_host == "falkordb": + _falkordb_host = "localhost" +FALKORDB_PORT = int(os.getenv("FALKORDB_PORT", "6381")) +FALKORDB_GRAPH = os.getenv("FALKORDB_GRAPH", "chronicle") + +# --- OpenAI Embeddings --- + +OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "") +OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1") +EMBEDDING_MODEL = "text-embedding-3-small" +EMBEDDING_DIMENSIONS = 1536 + + +def get_graph(host=None, port=None, graph_name=None): + """Create a FalkorDB graph connected to the local instance.""" + host = host or _falkordb_host + port = port or FALKORDB_PORT + graph_name = graph_name or FALKORDB_GRAPH + db = FalkorDB(host=host, port=port) + return db.select_graph(graph_name) + + +async def generate_embeddings(texts: List[str]) -> List[List[float]]: + """Generate embeddings using OpenAI API.""" + from openai import AsyncOpenAI + + client = AsyncOpenAI(api_key=OPENAI_API_KEY, base_url=OPENAI_BASE_URL) + response = await client.embeddings.create(model=EMBEDDING_MODEL, input=texts) + return [data.embedding for data in response.data] + + +def generate_embeddings_sync(texts: List[str]) -> List[List[float]]: + """Synchronous wrapper for embedding generation.""" + return asyncio.run(generate_embeddings(texts)) + + +# --- Markdown Chunking --- + + +def split_on_headers(markdown: str) -> List[Dict[str, str]]: + """Split markdown into chunks by ### headers. + + Returns list of dicts with 'section_title' and 'text' keys. + Frontmatter (---...---) is stripped and returned as metadata. + """ + # Strip frontmatter + fm_match = re.match(r"^---\s*\n(.*?)\n---\s*\n", markdown, re.DOTALL) + if fm_match: + content = markdown[fm_match.end() :] + else: + content = markdown + + chunks = [] + # Split on ### headers (h3) + parts = re.split(r"^(###\s+.+)$", content, flags=re.MULTILINE) + + # parts alternates: [pre-header text, header, body, header, body, ...] + # First element is text before any header + i = 0 + if parts[0].strip(): + chunks.append({"section_title": "Introduction", "text": parts[0].strip()}) + i = 1 + else: + i = 1 + + while i < len(parts): + if parts[i].startswith("###"): + title = parts[i].replace("###", "").strip() + body = parts[i + 1].strip() if i + 1 < len(parts) else "" + if body: + chunks.append({"section_title": title, "text": body}) + i += 2 + else: + i += 1 + + return chunks + + +def parse_frontmatter(markdown: str) -> Dict[str, str]: + """Extract YAML frontmatter as a simple dict.""" + fm_match = re.match(r"^---\s*\n(.*?)\n---\s*\n", markdown, re.DOTALL) + if not fm_match: + return {} + result = {} + for line in fm_match.group(1).splitlines(): + if ":" in line: + key, _, value = line.partition(":") + result[key.strip()] = value.strip().strip('"').strip("'") + return result + + +# --- Entity Parsing --- + + +def parse_people_section(markdown: str) -> List[Dict[str, str]]: + """Parse the ### People section to extract entity names and descriptions. + + Expected format: + ### People + - John (coworker, project lead) + - Dr. Sarah Chen (dentist, referred by Jane) + - "The IT guy" (unnamed, fixed the printer) + + Returns list of dicts with 'name' and 'description'. + """ + # Find the People section + match = re.search( + r"^###\s+People\s*\n(.*?)(?=^###|\Z)", markdown, re.MULTILINE | re.DOTALL + ) + if not match: + return [] + + people = [] + for line in match.group(1).strip().splitlines(): + line = line.strip() + if not line.startswith("-"): + continue + line = line[1:].strip() + + # Try to parse "Name (description)" format + paren_match = re.match(r'^["""]?(.+?)["""]?\s*\((.+?)\)\s*$', line) + if paren_match: + people.append( + { + "name": paren_match.group(1).strip(), + "description": paren_match.group(2).strip(), + } + ) + else: + # Just a name with no description + people.append({"name": line.strip('"').strip("'"), "description": ""}) + + return people + + +def parse_action_items(markdown: str) -> List[Dict[str, str]]: + """Parse ### Action Items section. + + Expected format: + ### Action Items + - [ ] Send Q3 report to John by Friday + - [x] Book dentist follow-up + + Returns list of dicts with 'text' and 'done' (bool). + """ + match = re.search( + r"^###\s+Action Items\s*\n(.*?)(?=^###|\Z)", markdown, re.MULTILINE | re.DOTALL + ) + if not match: + return [] + + items = [] + for line in match.group(1).strip().splitlines(): + line = line.strip() + if not line.startswith("-"): + continue + line = line[1:].strip() + + done = False + if line.startswith("[x]") or line.startswith("[X]"): + done = True + line = line[3:].strip() + elif line.startswith("[ ]"): + line = line[3:].strip() + + if line: + items.append({"text": line, "done": done}) + + return items + + +# --- Hybrid Search Scoring --- + + +def compute_hybrid_scores( + vector_results: List[Dict], + fulltext_results: List[Dict], + vector_weight: float = 0.7, + text_weight: float = 0.3, + recency_half_life_days: float = 30.0, + recency_floor: float = 0.5, +) -> List[Dict]: + """Merge vector and full-text results with recency bias. + + Each result dict must have: 'chunk_id', 'score', 'date' (ISO string or datetime). + Additional fields are preserved. + """ + now = datetime.now(timezone.utc) + merged: Dict[str, Dict] = {} + + for r in vector_results: + cid = r["chunk_id"] + merged[cid] = {**r, "vector_score": r["score"], "text_score": 0.0} + + for r in fulltext_results: + cid = r["chunk_id"] + if cid in merged: + merged[cid]["text_score"] = r["score"] + else: + merged[cid] = {**r, "vector_score": 0.0, "text_score": r["score"]} + + results = [] + for entry in merged.values(): + # Parse date + d = entry.get("date") + if isinstance(d, str): + d = datetime.fromisoformat(d.replace("Z", "+00:00")) + if d is None: + d = now + + if d.tzinfo is None: + d = d.replace(tzinfo=timezone.utc) + + age_days = (now - d).total_seconds() / 86400.0 + + relevance = ( + vector_weight * entry["vector_score"] + text_weight * entry["text_score"] + ) + recency = max( + recency_floor, math.exp(-0.693 * age_days / recency_half_life_days) + ) + entry["relevance_score"] = relevance + entry["recency_score"] = recency + entry["final_score"] = relevance * recency + results.append(entry) + + results.sort(key=lambda x: x["final_score"], reverse=True) + return results + + +# --- Pretty Printing --- + +GREEN = "\033[92m" +RED = "\033[91m" +YELLOW = "\033[93m" +BOLD = "\033[1m" +RESET = "\033[0m" + + +def print_pass(msg: str): + print(f" {GREEN}PASS{RESET} {msg}") + + +def print_fail(msg: str): + print(f" {RED}FAIL{RESET} {msg}") + + +def print_header(msg: str): + print(f"\n{BOLD}{'='*60}{RESET}") + print(f"{BOLD}{msg}{RESET}") + print(f"{BOLD}{'='*60}{RESET}") + + +def print_results_table(results: List[Dict], fields: List[str], max_text_len: int = 60): + """Print results as a simple table.""" + if not results: + print(" (no results)") + return + for i, r in enumerate(results): + parts = [] + for f in fields: + val = r.get(f, "") + if isinstance(val, float): + parts.append(f"{f}={val:.4f}") + elif isinstance(val, str) and len(val) > max_text_len: + parts.append(f"{f}={val[:max_text_len]}...") + else: + parts.append(f"{f}={val}") + print(f" [{i+1}] {', '.join(parts)}") diff --git a/backends/advanced/tests/test_annotation_export.py b/backends/advanced/tests/test_annotation_export.py new file mode 100644 index 00000000..d2e1c43a --- /dev/null +++ b/backends/advanced/tests/test_annotation_export.py @@ -0,0 +1,78 @@ +"""Unit tests for annotation-export helpers (utils/annotation_export.py).""" + +from advanced_omi_backend.models.conversation import Conversation +from advanced_omi_backend.utils.annotation_export import ( + build_clip_record, + new_export_id, + validate_export_id, +) +from advanced_omi_backend.utils.transcript_slicing import slice_segments + + +def _segment(start: float, end: float, text: str, speaker: str = "speaker_0"): + return Conversation.SpeakerSegment(start=start, end=end, text=text, speaker=speaker) + + +class TestExportId: + def test_new_export_id_is_valid(self): + export_id = new_export_id() + assert validate_export_id(export_id) + + def test_rejects_path_traversal(self): + assert not validate_export_id("../etc") + assert not validate_export_id("annotation_20260611_120000_ab12/../x") + assert not validate_export_id("annotation_20260611_120000_AB12") # uppercase + assert not validate_export_id("other_20260611_120000_ab12") + + def test_ids_are_unique(self): + assert new_export_id() != new_export_id() + + +class TestBuildClipRecord: + def _record(self, segments, t0=100.0, t1=130.0, idx=2): + return build_clip_record( + conversation_id="conv-abc", + conversation_title="Morning chat", + client_id="user01-phone", + conversation_created_at="2026-06-11T08:00:00+00:00", + clip_index=idx, + region_start=t0, + region_end=t1, + sample_rate=16000, + segments=slice_segments(segments, t0, t1), + ) + + def test_basic_record(self): + segments = [ + _segment(90.0, 99.0, "before the clip"), + _segment(101.0, 105.0, "hello there"), + _segment(110.0, 120.0, "how are you", speaker="speaker_1"), + _segment(140.0, 150.0, "after the clip"), + ] + rec = self._record(segments) + + assert rec["clip_id"] == "conv-abc_002" + assert rec["audio_path"] == "audio/conv-abc_002.wav" + assert rec["source_start_seconds"] == 100.0 + assert rec["source_end_seconds"] == 130.0 + assert rec["duration_seconds"] == 30.0 + # Only the two in-clip segments survive, shifted to clip-relative time. + assert [s["text"] for s in rec["segments"]] == ["hello there", "how are you"] + assert rec["segments"][0]["start"] == 1.0 + assert rec["segments"][0]["end"] == 5.0 + assert rec["segments"][1]["speaker"] == "speaker_1" + assert rec["text"] == "hello there how are you" + # Annotation block present and empty (the annotator's contract). + assert rec["annotation"] == {"text": None, "segments": None, "notes": None} + + def test_clip_without_transcript(self): + rec = self._record([_segment(0.0, 5.0, "far away")]) + assert rec["text"] == "" + assert rec["segments"] == [] + + def test_segment_straddling_clip_start_is_clamped(self): + # Midpoint (101.5) inside [100, 130) → kept, clamped to clip start. + rec = self._record([_segment(98.0, 105.0, "straddler")]) + assert len(rec["segments"]) == 1 + assert rec["segments"][0]["start"] == 0.0 + assert rec["segments"][0]["end"] == 5.0 diff --git a/backends/advanced/tests/test_annotation_import.py b/backends/advanced/tests/test_annotation_import.py new file mode 100644 index 00000000..1a975172 --- /dev/null +++ b/backends/advanced/tests/test_annotation_import.py @@ -0,0 +1,107 @@ +"""Regression tests for annotation-dataset ZIP import.""" + +import io +import json +import zipfile + +import pytest + +from advanced_omi_backend.utils.annotation_import import ( + AnnotationDatasetError, + parse_annotation_dataset, +) + + +def _dataset_zip( + record: dict, *, export_id: str = "annotation_20260628_180119_adaf" +) -> bytes: + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, "w") as archive: + archive.writestr( + "export.json", + json.dumps({"export_id": export_id, "schema_version": 1}), + ) + archive.writestr("manifest.jsonl", json.dumps(record) + "\n") + archive.writestr(record["audio_path"], b"RIFF-test-wav") + return buffer.getvalue() + + +def _record() -> dict: + return { + "clip_id": "conv-abc_000", + "audio_path": "audio/conv-abc_000.wav", + "conversation_id": "conv-abc", + "conversation_title": "Morning chat", + "client_id": "user01-phone", + "duration_seconds": 12.5, + "sample_rate": 16000, + "text": "machine transcript", + "segments": [ + { + "start": 0.2, + "end": 3.0, + "speaker": "speaker_0", + "identified_as": None, + "text": "machine transcript", + } + ], + "annotation": { + "text": "human transcript", + "segments": [ + { + "start": 0.2, + "end": 3.0, + "speaker": "ankush", + "identified_as": "ankush", + "text": "human transcript", + } + ], + "notes": "checked", + }, + } + + +def test_parse_export_uses_human_annotation_as_active_transcript(): + dataset = parse_annotation_dataset(_dataset_zip(_record())) + + assert dataset.dataset_id == "annotation_20260628_180119_adaf" + assert len(dataset.clips) == 1 + clip = dataset.clips[0] + assert clip.transcript == "human transcript" + assert clip.transcript_source == "human_annotation" + assert clip.segments[0]["speaker"] == "ankush" + assert clip.audio_bytes == b"RIFF-test-wav" + + +def test_parse_existing_export_without_schema_version(): + record = _record() + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, "w") as archive: + archive.writestr( + "export.json", + json.dumps({"export_id": "annotation_20260628_180119_adaf"}), + ) + archive.writestr("manifest.jsonl", json.dumps(record)) + archive.writestr(record["audio_path"], b"RIFF-test-wav") + + dataset = parse_annotation_dataset(buffer.getvalue()) + + assert dataset.schema_version == 1 + + +def test_parse_rejects_audio_path_traversal(): + record = _record() + record["audio_path"] = "../outside.wav" + + with pytest.raises(AnnotationDatasetError, match="Unsafe"): + parse_annotation_dataset(_dataset_zip(record)) + + +def test_parse_rejects_manifest_without_audio(): + record = _record() + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, "w") as archive: + archive.writestr("manifest.jsonl", json.dumps(record)) + + with pytest.raises(AnnotationDatasetError, match="missing audio"): + parse_annotation_dataset(buffer.getvalue()) diff --git a/backends/advanced/tests/test_audio_persistence_singleflight.py b/backends/advanced/tests/test_audio_persistence_singleflight.py new file mode 100644 index 00000000..091dbf24 --- /dev/null +++ b/backends/advanced/tests/test_audio_persistence_singleflight.py @@ -0,0 +1,76 @@ +"""Single-flight guarantee for the per-session audio-persistence job. + +A WebSocket reconnect mid-session re-runs ``start_streaming_jobs``. The audio +persistence job MUST be single-flight per session: re-enqueuing while one is +already live (i.e. a reconnect) must reuse the live job, never start a second +consumer. + +Why this matters (the bug this guards against): two persistence jobs for one +session share the SAME Redis consumer name (``persistence-{session_id[:8]}``), +so the audio stream gets SPLIT between them — each new message goes to only one +consumer. The speech-detected conversations created after the reconnect then +find no audio chunks under their id and get deleted as ``audio_chunks_not_ready`` +while their transcripts are stranded on a different conversation document. +""" + +import pytest +from fakeredis import FakeStrictRedis +from rq import Queue + +pytestmark = pytest.mark.unit + + +@pytest.fixture +def qc(monkeypatch): + """queue_controller with its Redis + audio queue pointed at fakeredis.""" + fake = FakeStrictRedis() + from advanced_omi_backend.controllers import queue_controller as module + + monkeypatch.setattr(module, "redis_conn", fake) + monkeypatch.setattr(module, "audio_queue", Queue("audio", connection=fake)) + return module + + +def test_reconnect_reuses_live_persistence_job(qc): + """A second enqueue for the same session (a reconnect) reuses the live job.""" + first = qc.enqueue_audio_persistence( + "sess-1", "user-1", "sess-1", always_persist=True + ) + second = qc.enqueue_audio_persistence( + "sess-1", "user-1", "sess-1", always_persist=True + ) + + assert first == second, "reconnect must reuse the live persistence job id" + assert qc.audio_queue.count == 1, "must not enqueue a second persistence consumer" + + +def test_distinct_sessions_get_distinct_jobs(qc): + """Single-flight is per-session: different sessions each get their own job.""" + a = qc.enqueue_audio_persistence("sess-A", "user-1", "sess-A", always_persist=True) + b = qc.enqueue_audio_persistence("sess-B", "user-1", "sess-B", always_persist=True) + + assert a != b + assert qc.audio_queue.count == 2 + + +def test_ended_job_allows_a_fresh_enqueue(qc): + """After the persistence job terminates, a new session may enqueue again. + + Single-flight must gate on LIVENESS, not merely on the id ever having existed — + otherwise a clean session end would permanently block the next session. + """ + from rq.job import Job, JobStatus + + first = qc.enqueue_audio_persistence( + "sess-1", "user-1", "sess-1", always_persist=True + ) + # Simulate the job terminating (worker finished/abandoned it). + job = Job.fetch(first, connection=qc.redis_conn) + job.set_status(JobStatus.FINISHED) + + second = qc.enqueue_audio_persistence( + "sess-1", "user-1", "sess-1", always_persist=True + ) + assert qc._job_is_live( + second + ), "a fresh persistence job must be live after re-enqueue" diff --git a/backends/advanced/tests/test_client_lifecycle.py b/backends/advanced/tests/test_client_lifecycle.py new file mode 100644 index 00000000..e350d7f1 --- /dev/null +++ b/backends/advanced/tests/test_client_lifecycle.py @@ -0,0 +1,56 @@ +"""Unit tests for the in-memory client lifecycle. + +Covers the invariants the WebSocket evict-on-reconnect and idle-timeout/reaper +paths rely on: + - a freshly created client is present, connected, and freshly stamped + - touch() advances last_activity (drives the reaper + honest "connected") + - remove_client_with_cleanup() is the single removal path and flips connected + - create_client() rejects a duplicate id (the evict path must clean up first) + +These are pure in-memory — no Redis, Mongo, or API keys required. +""" + +import time + +import pytest + +from advanced_omi_backend.client import ClientState +from advanced_omi_backend.client_manager import ClientManager + + +def test_new_client_is_present_connected_and_fresh(): + mgr = ClientManager() + before = time.time() + state = mgr.create_client("u1-phone", "u1", "u1@example.com") + + assert mgr.has_client("u1-phone") + assert state.connected is True + assert state.last_activity >= before + + +def test_touch_advances_last_activity(): + state = ClientState("u1-phone", "u1", "u1@example.com") + original = state.last_activity = time.time() - 100 # pretend it's been idle + state.touch() + assert state.last_activity > original + + +@pytest.mark.asyncio +async def test_remove_with_cleanup_disconnects_and_removes(): + mgr = ClientManager() + state = mgr.create_client("u1-phone", "u1", "u1@example.com") + + removed = await mgr.remove_client_with_cleanup("u1-phone") + + assert removed is True + assert not mgr.has_client("u1-phone") + assert state.connected is False + # Removing a now-absent client is a graceful no-op (idempotent reaper/evict). + assert await mgr.remove_client_with_cleanup("u1-phone") is False + + +def test_create_duplicate_raises_so_evict_must_run_first(): + mgr = ClientManager() + mgr.create_client("u1-phone", "u1", "u1@example.com") + with pytest.raises(ValueError): + mgr.create_client("u1-phone", "u1", "u1@example.com") diff --git a/backends/advanced/tests/test_codex_executor.py b/backends/advanced/tests/test_codex_executor.py new file mode 100644 index 00000000..26eb8e9d --- /dev/null +++ b/backends/advanced/tests/test_codex_executor.py @@ -0,0 +1,174 @@ +"""Codex CLI memory-agent executor: selection, filesystem-diff auditing, failure paths.""" + +import contextlib +import subprocess +from types import SimpleNamespace + +import pytest + +from advanced_omi_backend.services.memory.agent import codex_agent, memory_agent +from advanced_omi_backend.services.memory.agent.codex_agent import CodexMemoryAgent +from advanced_omi_backend.services.memory.agent.memory_agent import ( + MemoryAgent, + MemoryAgentResult, +) +from advanced_omi_backend.services.memory.config import MemoryConfig +from advanced_omi_backend.services.memory.providers.chronicle import MemoryService + + +@contextlib.contextmanager +def _no_lock(_user_id, ttl_seconds=0): + yield + + +@pytest.fixture +def unlocked(monkeypatch): + monkeypatch.setattr( + "advanced_omi_backend.services.memory.vault_lock.vault_run_lock", _no_lock + ) + + +# --------------------------------------------------------------------------- +# Executor selection (chronicle._agent_class) +# --------------------------------------------------------------------------- + + +def test_agent_class_defaults_to_direct(): + service = MemoryService(MemoryConfig()) + assert service._agent_class() is MemoryAgent + + +def test_agent_class_uses_codex_when_available(monkeypatch): + monkeypatch.setattr( + codex_agent, "codex_executor_available", lambda: (True, "/usr/bin/codex") + ) + service = MemoryService(MemoryConfig(agent_executor="codex")) + assert service._agent_class() is CodexMemoryAgent + + +def test_agent_class_falls_back_when_codex_unavailable(monkeypatch): + monkeypatch.setattr( + codex_agent, "codex_executor_available", lambda: (False, "no binary") + ) + service = MemoryService(MemoryConfig(agent_executor="codex")) + assert service._agent_class() is MemoryAgent + + +# --------------------------------------------------------------------------- +# CodexMemoryAgent.run +# --------------------------------------------------------------------------- + + +def _fake_codex_run(vault_root, *, summary="Recorded the conversation.", returncode=0): + """A subprocess.run stand-in that mimics one codex exec editing the vault.""" + + def fake_run(cmd, **kwargs): + # Simulate the agent's edits: create the conversation note, update a + # person note, retire a topic note. + (vault_root / "Conversations").mkdir(exist_ok=True) + (vault_root / "Conversations" / "conv1.md").write_text("recorded") + (vault_root / "People" / "Old.md").write_text("updated content") + (vault_root / "Topics" / "Gone.md").unlink() + last_msg = cmd[cmd.index("--output-last-message") + 1] + with open(last_msg, "w") as f: + f.write(summary) + stdout = ( + '{"type":"item.completed","item":{"item_type":"command_execution"}}\n' + '{"type":"turn.completed"}\n' + ) + return SimpleNamespace(returncode=returncode, stdout=stdout, stderr="") + + return fake_run + + +def _seed_vault(tmp_path): + root = tmp_path / "user1" + (root / "People").mkdir(parents=True) + (root / "Topics").mkdir() + (root / "People" / "Old.md").write_text("original content") + (root / "Topics" / "Gone.md").write_text("doomed note") + return root + + +@pytest.mark.asyncio +async def test_run_derives_touched_and_removed_from_fs_diff( + tmp_path, monkeypatch, unlocked +): + root = _seed_vault(tmp_path) + monkeypatch.setattr( + codex_agent, "codex_executor_available", lambda: (True, "/usr/bin/codex") + ) + monkeypatch.setattr(subprocess, "run", _fake_codex_run(root)) + + result = await CodexMemoryAgent(root).run("a real transcript", "conv1") + + assert result.touched == ["Conversations/conv1.md", "People/Old.md"] + assert result.removed == [ + {"old_path": "Topics/Gone.md", "new_path": "", "before": "doomed note"} + ] + assert result.summary == "Recorded the conversation." + assert result.tool_calls == 1 + assert result.rounds == 1 + assert not result.truncated + assert result.errors == [] + + +@pytest.mark.asyncio +async def test_run_failure_is_truncated_with_errors(tmp_path, monkeypatch, unlocked): + root = _seed_vault(tmp_path) + monkeypatch.setattr( + codex_agent, "codex_executor_available", lambda: (True, "/usr/bin/codex") + ) + + def failing_run(cmd, **kwargs): + raise subprocess.TimeoutExpired(cmd, 5) + + monkeypatch.setattr(subprocess, "run", failing_run) + + result = await CodexMemoryAgent(root).run("a real transcript", "conv1") + + assert result.truncated + assert any("timed out" in e for e in result.errors) + assert result.touched == [] # nothing was written + + +@pytest.mark.asyncio +async def test_run_unavailable_executor_returns_truncated(tmp_path, monkeypatch): + root = _seed_vault(tmp_path) + monkeypatch.setattr( + codex_agent, "codex_executor_available", lambda: (False, "no auth") + ) + + result = await CodexMemoryAgent(root).run("a real transcript", "conv1") + + assert result.truncated + assert result.errors == ["codex executor unavailable: no auth"] + + +@pytest.mark.asyncio +async def test_force_fallback_delegates_to_direct_agent(tmp_path, monkeypatch): + root = _seed_vault(tmp_path) + seen = {} + + class FakeDirectAgent: + def __init__( + self, vault_root, operation="memory_agent", *, force_fallback=False + ): + seen["force_fallback"] = force_fallback + + async def run(self, transcript, conversation_id, **kwargs): + return MemoryAgentResult( + conversation_id=conversation_id, + rounds=1, + touched=["Conversations/conv1.md"], + summary="fallback ran", + ) + + monkeypatch.setattr(memory_agent, "MemoryAgent", FakeDirectAgent) + + result = await CodexMemoryAgent(root, force_fallback=True).run( + "a real transcript", "conv1" + ) + + assert seen["force_fallback"] is True + assert result.summary == "fallback ran" diff --git a/backends/advanced/tests/test_conversation_audio_salvage.py b/backends/advanced/tests/test_conversation_audio_salvage.py new file mode 100644 index 00000000..692b6965 --- /dev/null +++ b/backends/advanced/tests/test_conversation_audio_salvage.py @@ -0,0 +1,27 @@ +"""The no-data-loss guard for conversations whose audio never persisted. + +When a conversation's audio chunks are missing at finalize (e.g. a mid-session +reconnect routed the audio to the session's always_persist placeholder under a +different conversation_id), the conversation must NOT be discarded if it carries a +real transcript — losing a real transcript is worse than keeping an audio-less +conversation. Only a genuinely empty conversation (no meaningful transcript) is +discarded. +""" + +import pytest + +from advanced_omi_backend.workers.conversation_jobs import ( + should_discard_unbacked_conversation, +) + +pytestmark = pytest.mark.unit + + +def test_empty_conversation_is_discarded(): + # No transcript and no audio → nothing worth keeping. + assert should_discard_unbacked_conversation(has_meaningful_transcript=False) is True + + +def test_transcript_bearing_conversation_is_kept(): + # Real transcript but missing audio → keep it (don't lose the transcript). + assert should_discard_unbacked_conversation(has_meaningful_transcript=True) is False diff --git a/backends/advanced/tests/test_data_archive.py b/backends/advanced/tests/test_data_archive.py new file mode 100644 index 00000000..ba4ef0b9 --- /dev/null +++ b/backends/advanced/tests/test_data_archive.py @@ -0,0 +1,389 @@ +"""Tests for portable Chronicle data archives.""" + +import zipfile +from datetime import datetime, timezone +from pathlib import Path +from types import SimpleNamespace + +import pytest +from bson import ObjectId + +from advanced_omi_backend.services import data_archive +from advanced_omi_backend.services.data_archive import ( + ArchiveError, + clear_vault_contents, + create_data_archive, + import_data_archive, + verify_data_archive, +) + + +class AsyncCursor: + def __init__(self, documents): + self.documents = list(documents) + + def __aiter__(self): + self.iterator = iter(self.documents) + return self + + async def __anext__(self): + try: + return next(self.iterator) + except StopIteration as exc: + raise StopAsyncIteration from exc + + def sort(self, _fields): + return self + + +class FakeCollection: + def __init__(self, documents=()): + self.documents = {document["_id"]: document for document in documents} + self.delete_calls = 0 + + def find(self, _query, projection=None): + return AsyncCursor(self.documents.values()) + + async def delete_many(self, _query): + deleted = len(self.documents) + self.documents.clear() + self.delete_calls += 1 + return SimpleNamespace(deleted_count=deleted) + + async def bulk_write(self, operations, ordered): + assert ordered is False + for operation in operations: + self.documents[operation._filter["_id"]] = operation._doc + + +class FakeDatabase: + name = "chronicle_test" + + def __init__(self, collections): + self.collections = { + name: FakeCollection(documents) for name, documents in collections.items() + } + + async def list_collection_names(self): + return list(self.collections) + + def __getitem__(self, name): + self.collections.setdefault(name, FakeCollection()) + return self.collections[name] + + +@pytest.mark.asyncio +async def test_archive_round_trips_bson_audio_and_files(tmp_path: Path): + created_at = datetime(2026, 7, 15, 12, 30, tzinfo=timezone.utc) + conversation_id = "conversation-1" + source = FakeDatabase( + { + "conversations": [ + { + "_id": ObjectId(), + "conversation_id": conversation_id, + "created_at": created_at, + "transcript_versions": [ + {"version_id": "v1", "transcript": "Exact transcript"} + ], + } + ], + "audio_chunks": [ + { + "_id": ObjectId(), + "conversation_id": conversation_id, + "chunk_index": 0, + "audio_data": b"\x00opus\xffbytes", + "created_at": created_at, + } + ], + "memory_audit": [{"_id": ObjectId(), "user_id": "user-1"}], + } + ) + source_data = tmp_path / "source-data" + note = source_data / "conversation_docs" / "user-1" / "People" / "Ada.md" + note.parent.mkdir(parents=True) + note.write_text("# Ada\n", encoding="utf-8") + legacy_audio = source_data / "audio_chunks" / "legacy.wav" + legacy_audio.parent.mkdir(parents=True) + legacy_audio.write_bytes(b"RIFF-audio") + + archive_path = tmp_path / "backup.chronicle" + summary = await create_data_archive(source, archive_path, data_dir=source_data) + + assert summary.documents == 3 + assert summary.files == 2 + manifest = verify_data_archive(archive_path) + assert manifest["collections"]["audio_chunks"]["documents"] == 1 + + target = FakeDatabase({"conversations": [{"_id": ObjectId(), "old": True}]}) + target_data = tmp_path / "target-data" + stale_note = target_data / "conversation_docs" / "user-1" / "stale.md" + stale_note.parent.mkdir(parents=True) + stale_note.write_text("stale", encoding="utf-8") + imported = await import_data_archive( + target, + archive_path, + data_dir=target_data, + replace=True, + ) + + assert imported.documents == 3 + assert imported.files == 2 + assert len(target["conversations"].documents) == 1 + restored_audio = next(iter(target["audio_chunks"].documents.values())) + assert restored_audio["audio_data"] == b"\x00opus\xffbytes" + # BSON stores UTC milliseconds; the default Mongo codec returns a naive UTC value. + assert restored_audio["created_at"] == created_at.replace(tzinfo=None) + assert not stale_note.exists() + assert ( + target_data / "conversation_docs/user-1/People/Ada.md" + ).read_text() == "# Ada\n" + assert (target_data / "audio_chunks/legacy.wav").read_bytes() == b"RIFF-audio" + + +@pytest.mark.asyncio +async def test_fresh_memory_skips_derived_collection_and_files(tmp_path: Path): + source = FakeDatabase( + { + "conversations": [{"_id": ObjectId(), "conversation_id": "conv"}], + "memory_audit": [{"_id": ObjectId(), "user_id": "user"}], + } + ) + source_data = tmp_path / "source" + note = source_data / "conversation_docs/user/People/Old.md" + note.parent.mkdir(parents=True) + note.write_text("old", encoding="utf-8") + archive_path = tmp_path / "fresh.chronicle" + await create_data_archive(source, archive_path, data_dir=source_data) + + existing_audit = {"_id": ObjectId(), "user_id": "existing"} + target = FakeDatabase({"memory_audit": [existing_audit]}) + target_data = tmp_path / "target" + result = await import_data_archive( + target, + archive_path, + data_dir=target_data, + fresh_memory=True, + restore_files=False, + ) + + assert result.skipped_collections == ("memory_audit",) + assert list(target["memory_audit"].documents.values()) == [existing_audit] + assert not target_data.exists() + + +def test_verify_rejects_duplicate_or_tampered_members(tmp_path: Path): + archive_path = tmp_path / "invalid.chronicle" + with zipfile.ZipFile(archive_path, mode="w") as archive: + archive.writestr("manifest.json", "{}") + archive.writestr("manifest.json", "{}") + + with pytest.raises(ArchiveError, match="duplicate member"): + verify_data_archive(archive_path) + + +def test_clear_vault_preserves_syncthing_markers(tmp_path: Path): + user_root = tmp_path / "user" + (user_root / ".stfolder").mkdir(parents=True) + (user_root / ".stignore").write_text("ignore", encoding="utf-8") + (user_root / "People").mkdir() + (user_root / "People" / "Ada.md").write_text("memory", encoding="utf-8") + + deleted = clear_vault_contents(user_root) + + assert deleted == 1 + assert (user_root / ".stfolder").is_dir() + assert (user_root / ".stignore").is_file() + assert not (user_root / "People").exists() + + +@pytest.mark.asyncio +async def test_import_keeps_earliest_conversation_for_duplicate_audio( + tmp_path: Path, caplog, monkeypatch +): + async def decode_as_pcm(opus_data, _sample_rate, _channels): + return b"decoded-identical-pcm" + + monkeypatch.setattr(data_archive, "decode_opus_to_pcm", decode_as_pcm) + first_id = "first-conversation" + duplicate_id = "duplicate-conversation" + source = FakeDatabase( + { + "conversations": [ + { + "_id": ObjectId(), + "conversation_id": duplicate_id, + "created_at": datetime(2026, 7, 16, tzinfo=timezone.utc), + "transcript_versions": [{"transcript": "second version"}], + }, + { + "_id": ObjectId(), + "conversation_id": first_id, + "created_at": datetime(2026, 7, 15, tzinfo=timezone.utc), + "transcript_versions": [{"transcript": "first version"}], + }, + ], + "audio_chunks": [ + { + "_id": ObjectId(), + "conversation_id": duplicate_id, + "chunk_index": 0, + "audio_data": b"opus-encoding-b", + }, + { + "_id": ObjectId(), + "conversation_id": first_id, + "chunk_index": 0, + "audio_data": b"opus-encoding-a", + }, + ], + "annotations": [ + { + "_id": ObjectId(), + "conversation_id": duplicate_id, + "value": "skip me", + }, + { + "_id": ObjectId(), + "conversation_id": first_id, + "value": "keep me", + }, + ], + } + ) + archive_path = tmp_path / "duplicates.chronicle" + await create_data_archive(source, archive_path, data_dir=tmp_path / "source") + + target = FakeDatabase({}) + result = await import_data_archive( + target, + archive_path, + data_dir=tmp_path / "target", + replace=True, + ) + + conversations = list(target["conversations"].documents.values()) + assert [conversation["conversation_id"] for conversation in conversations] == [ + first_id + ] + assert conversations[0]["transcript_versions"] == [{"transcript": "first version"}] + assert len(target["audio_chunks"].documents) == 1 + assert [item["value"] for item in target["annotations"].documents.values()] == [ + "keep me" + ] + assert result.documents == 3 + assert len(result.duplicate_audio_warnings) == 1 + warning = result.duplicate_audio_warnings[0] + assert warning.kept_conversation_id == first_id + assert warning.skipped_conversation_id == duplicate_id + assert warning.kept_source == "archive" + assert "Duplicate audio skipped during import" in caplog.text + + +@pytest.mark.asyncio +async def test_merge_import_does_not_duplicate_audio_already_in_database( + tmp_path: Path, monkeypatch +): + async def decode_as_pcm(opus_data, _sample_rate, _channels): + return b"decoded-identical-pcm" + + monkeypatch.setattr(data_archive, "decode_opus_to_pcm", decode_as_pcm) + existing_id = "existing-conversation" + imported_id = "imported-conversation" + source = FakeDatabase( + { + "conversations": [ + { + "_id": ObjectId(), + "conversation_id": imported_id, + "created_at": datetime(2026, 7, 16, tzinfo=timezone.utc), + } + ], + "audio_chunks": [ + { + "_id": ObjectId(), + "conversation_id": imported_id, + "chunk_index": 0, + "audio_data": b"new-opus-encoding", + } + ], + } + ) + archive_path = tmp_path / "merge.chronicle" + await create_data_archive(source, archive_path, data_dir=tmp_path / "source") + target = FakeDatabase( + { + "conversations": [ + { + "_id": ObjectId(), + "conversation_id": existing_id, + "created_at": datetime(2026, 7, 14, tzinfo=timezone.utc), + } + ], + "audio_chunks": [ + { + "_id": ObjectId(), + "conversation_id": existing_id, + "chunk_index": 0, + "audio_data": b"existing-opus-encoding", + } + ], + } + ) + + result = await import_data_archive( + target, + archive_path, + data_dir=tmp_path / "target", + ) + + assert len(target["conversations"].documents) == 1 + assert len(target["audio_chunks"].documents) == 1 + warning = result.duplicate_audio_warnings[0] + assert warning.kept_conversation_id == existing_id + assert warning.skipped_conversation_id == imported_id + assert warning.kept_source == "existing_database" + + +@pytest.mark.asyncio +async def test_import_keeps_first_duplicate_chunk_and_warns(tmp_path: Path): + conversation_id = "conversation-1" + first_id = ObjectId() + duplicate_id = ObjectId() + source = FakeDatabase( + { + "conversations": [{"_id": ObjectId(), "conversation_id": conversation_id}], + "audio_chunks": [ + { + "_id": first_id, + "conversation_id": conversation_id, + "chunk_index": 0, + "audio_data": b"first", + "created_at": datetime(2026, 7, 15, tzinfo=timezone.utc), + }, + { + "_id": duplicate_id, + "conversation_id": conversation_id, + "chunk_index": 0, + "audio_data": b"second", + "created_at": datetime(2026, 7, 16, tzinfo=timezone.utc), + }, + ], + } + ) + archive_path = tmp_path / "duplicate-chunk.chronicle" + await create_data_archive(source, archive_path, data_dir=tmp_path / "source") + + target = FakeDatabase({}) + result = await import_data_archive( + target, archive_path, data_dir=tmp_path / "target", replace=True + ) + + chunks = list(target["audio_chunks"].documents.values()) + assert len(chunks) == 1 + assert chunks[0]["_id"] == first_id + assert len(result.duplicate_chunk_warnings) == 1 + warning = result.duplicate_chunk_warnings[0] + assert warning.kept_chunk_id == str(first_id) + assert warning.skipped_chunk_id == str(duplicate_id) diff --git a/backends/advanced/tests/test_data_audit_dataset_filter.py b/backends/advanced/tests/test_data_audit_dataset_filter.py new file mode 100644 index 00000000..9204c250 --- /dev/null +++ b/backends/advanced/tests/test_data_audit_dataset_filter.py @@ -0,0 +1,59 @@ +"""Dataset-scoping regressions for the Data Audit listing.""" + +from types import SimpleNamespace + +import pytest + +from advanced_omi_backend.controllers import data_audit_controller +from advanced_omi_backend.models.conversation import Conversation + + +class _Cursor: + def __init__(self, docs): + self.docs = docs + + def sort(self, *_args): + return self + + def limit(self, *_args): + return self + + async def to_list(self, *, length): + return self.docs[:length] + + +class _Collection: + def __init__(self): + self.queries = [] + + def find(self, query, projection): + self.queries.append((query, projection)) + call = len(self.queries) + if call == 2: + return _Cursor( + [ + {"external_source_id": "dataset-new:clip-2"}, + {"external_source_id": "dataset-new:clip-1"}, + {"external_source_id": "dataset-old:clip-1"}, + ] + ) + return _Cursor([]) + + +@pytest.mark.asyncio +async def test_list_for_audit_scopes_and_lists_annotation_datasets(monkeypatch): + collection = _Collection() + monkeypatch.setattr(Conversation, "get_pymongo_collection", lambda: collection) + user = SimpleNamespace(is_superuser=False, user_id="user-1") + + result = await data_audit_controller.list_for_audit( + user, + dataset_id="dataset.+(selected)", + ) + + listing_query = collection.queries[0][0] + assert listing_query["external_source_type"] == "annotation_dataset" + assert listing_query["external_source_id"] == { + "$regex": r"^dataset\.\+\(selected\):" + } + assert result["datasets"] == ["dataset-new", "dataset-old"] diff --git a/backends/advanced/tests/test_leading_silence_trim.py b/backends/advanced/tests/test_leading_silence_trim.py new file mode 100644 index 00000000..0717bcea --- /dev/null +++ b/backends/advanced/tests/test_leading_silence_trim.py @@ -0,0 +1,53 @@ +"""Computing where to trim leading silence off a conversation. + +With always_persist on, a session's placeholder accumulates audio from session +start — so a long pause before the user speaks ends up as leading silence on the +conversation. At finalize we split that leading silence off into a soft-deleted +remnant so the visible conversation begins at the first speech (the audio is kept +in Mongo on the remnant, just hidden). + +``leading_silence_trim_index`` is the pure decision at the heart of that: given the +chunk timeline and when speech first starts, which chunk_index should become the +first chunk of the trimmed conversation — or None when there isn't enough leading +silence to bother (so we never churn conversations over a few seconds of pre-roll). +""" + +import pytest + +from advanced_omi_backend.workers.conversation_jobs import leading_silence_trim_index + +pytestmark = pytest.mark.unit + + +def _ten_second_chunks(count): + return [ + {"chunk_index": i, "start_time": i * 10.0, "end_time": (i + 1) * 10.0} + for i in range(count) + ] + + +def test_long_leading_silence_returns_speech_boundary_chunk(): + # 1300s of audio in 10s chunks; speech first appears at 1200s (chunk 120). + chunks = _ten_second_chunks(130) + idx = leading_silence_trim_index( + chunks, speech_start_time=1200.0, min_trim_seconds=30.0 + ) + assert idx == 120 + + +def test_short_leading_silence_is_not_trimmed(): + # A few seconds of pre-roll is fine — don't churn the conversation for it. + chunks = _ten_second_chunks(10) + idx = leading_silence_trim_index( + chunks, speech_start_time=8.0, min_trim_seconds=30.0 + ) + assert idx is None + + +def test_speech_in_first_chunk_is_not_trimmed(): + # Even past the min threshold, if speech lands in chunk 0 there's nothing to trim. + chunks = _ten_second_chunks(10) + idx = leading_silence_trim_index( + chunks, speech_start_time=0.0, min_trim_seconds=30.0 + ) + assert idx is None diff --git a/backends/advanced/tests/test_leading_silence_trim_db.py b/backends/advanced/tests/test_leading_silence_trim_db.py new file mode 100644 index 00000000..098db494 --- /dev/null +++ b/backends/advanced/tests/test_leading_silence_trim_db.py @@ -0,0 +1,129 @@ +"""The leading-silence trim operation (real MongoDB). + +Verifies that ``trim_leading_silence`` moves the pre-speech chunks onto a +soft-deleted remnant (audio kept in Mongo, just hidden), re-bases the surviving +chunks in place so the visible conversation starts at the first speech, and loses +no audio in the process. Run against a real Mongo: + + MONGODB_URI=mongodb://localhost:27017 uv run pytest tests/test_leading_silence_trim_db.py +""" + +import os + +import pytest +import pytest_asyncio +from beanie import init_beanie +from motor.motor_asyncio import AsyncIOMotorClient + +from advanced_omi_backend.models.audio_chunk import AudioChunkDocument +from advanced_omi_backend.models.conversation import Conversation, create_conversation +from advanced_omi_backend.workers.conversation_jobs import trim_leading_silence + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +def _mongo_url(): + return os.getenv("MONGODB_URI", "mongodb://localhost:27018") + + +def _db_name(): + return os.getenv("TEST_DB_NAME", "test_silence_trim_db") + + +@pytest_asyncio.fixture(scope="session", loop_scope="session") +async def init_db(): + client = AsyncIOMotorClient(_mongo_url()) + await init_beanie( + database=client[_db_name()], + document_models=[AudioChunkDocument, Conversation], + ) + yield + await client.drop_database(_db_name()) + client.close() + + +@pytest_asyncio.fixture(loop_scope="session") +async def clean_db(init_db): + await AudioChunkDocument.delete_all() + await Conversation.delete_all() + yield + + +async def _make_conversation_with_chunks(n_chunks): + """A visible conversation with ``n_chunks`` 10s chunks (0..n*10s).""" + conv = create_conversation(user_id="u1", client_id="u1-phone", title="Recording...") + conv.audio_chunks_count = n_chunks + conv.audio_total_duration = n_chunks * 10.0 + await conv.insert() + for i in range(n_chunks): + await AudioChunkDocument( + conversation_id=conv.conversation_id, + chunk_index=i, + audio_data=b"x", + original_size=1, + compressed_size=1, + start_time=i * 10.0, + end_time=(i + 1) * 10.0, + duration=10.0, + ).insert() + return conv + + +async def test_leading_silence_is_moved_to_a_soft_deleted_remnant(clean_db): + # 1300s total: 1200s of leading silence (120 chunks) then 100s of speech (10 chunks). + conv = await _make_conversation_with_chunks(130) + + trimmed = await trim_leading_silence( + conv.conversation_id, speech_start_time=1200.0, min_trim_seconds=30.0 + ) + assert trimmed is True + + # Visible conversation now begins at the speech: 10 chunks, re-indexed from 0, + # times re-based so it starts at 0. + refreshed = await Conversation.find_one( + Conversation.conversation_id == conv.conversation_id + ) + assert refreshed.deleted is False + assert refreshed.audio_chunks_count == 10 + assert refreshed.audio_total_duration == 100.0 + + survivors = ( + await AudioChunkDocument.find( + AudioChunkDocument.conversation_id == conv.conversation_id + ) + .sort("+chunk_index") + .to_list() + ) + assert [c.chunk_index for c in survivors] == list(range(10)) + assert survivors[0].start_time == 0.0 + assert survivors[-1].end_time == 100.0 + + # The leading silence lives on a soft-deleted remnant — hidden, but its audio is kept. + remnant = await Conversation.find_one( + Conversation.deletion_reason == "leading_silence" + ) + assert remnant is not None + assert remnant.deleted is True + assert remnant.audio_chunks_count == 120 + remnant_chunks = await AudioChunkDocument.find( + AudioChunkDocument.conversation_id == remnant.conversation_id + ).to_list() + assert len(remnant_chunks) == 120 + + # No audio lost: every original chunk still exists somewhere. + assert await AudioChunkDocument.count() == 130 + + +async def test_short_leading_silence_is_left_untouched(clean_db): + conv = await _make_conversation_with_chunks(5) # 50s total + + trimmed = await trim_leading_silence( + conv.conversation_id, speech_start_time=8.0, min_trim_seconds=30.0 + ) + assert trimmed is False + + refreshed = await Conversation.find_one( + Conversation.conversation_id == conv.conversation_id + ) + assert refreshed.audio_chunks_count == 5 + assert await AudioChunkDocument.count() == 5 diff --git a/backends/advanced/tests/test_memory_agent_completion.py b/backends/advanced/tests/test_memory_agent_completion.py new file mode 100644 index 00000000..33f316b7 --- /dev/null +++ b/backends/advanced/tests/test_memory_agent_completion.py @@ -0,0 +1,257 @@ +"""Completion guarantees for the agentic memory provider.""" + +from types import SimpleNamespace + +import httpx +import openai +import pytest + +from advanced_omi_backend import llm_client +from advanced_omi_backend.services.memory.agent.memory_agent import MemoryAgentResult +from advanced_omi_backend.services.memory.conversation_note import ( + ConversationNoteError, + canonicalize_conversation_note, + write_source_fallback_conversation_note, +) +from advanced_omi_backend.services.memory.providers import chronicle + + +class _Completions: + def __init__(self, result=None, error=None): + self.result = result + self.error = error + self.calls = 0 + + async def create(self, **_kwargs): + self.calls += 1 + if self.error: + raise self.error + return self.result + + +class _Operation: + def __init__(self, completions, name): + self.model_name = name + self._client = SimpleNamespace(chat=SimpleNamespace(completions=completions)) + + def get_client(self, is_async=False): + assert is_async is True + return self._client + + def to_api_params(self): + return {"model": self.model_name} + + +@pytest.mark.asyncio +async def test_tool_chat_uses_configured_fallback_on_context_overflow(monkeypatch): + request = httpx.Request("POST", "http://local.test/v1/chat/completions") + response = httpx.Response(400, request=request) + error = openai.BadRequestError( + "context overflow", + response=response, + body={ + "error": { + "code": 400, + "type": "exceed_context_size_error", + "message": "request exceeds the available context size", + } + }, + ) + primary_calls = _Completions(error=error) + expected = SimpleNamespace(choices=[]) + fallback_calls = _Completions(result=expected) + primary = _Operation(primary_calls, "local") + fallback = _Operation(fallback_calls, "fallback") + registry = SimpleNamespace( + get_llm_operation=lambda _name: primary, + get_fallback_llm_operation=lambda _name, primary: fallback, + ) + monkeypatch.setattr(llm_client, "get_models_registry", lambda: registry) + + result = await llm_client.async_chat_with_tools( + [{"role": "user", "content": "long transcript"}], + operation="memory_agent", + ) + + assert result is expected + assert primary_calls.calls == 1 + assert fallback_calls.calls == 1 + + +@pytest.mark.asyncio +async def test_memory_provider_retries_fallback_when_conversation_note_is_missing( + monkeypatch, tmp_path +): + user_root = tmp_path / "user" + service = chronicle.MemoryService(SimpleNamespace()) + monkeypatch.setattr(service.vault, "user_root", lambda _user_id: user_root) + monkeypatch.setattr(chronicle, "seed_vault_scaffold", lambda _root: None) + + recorded = [] + + async def fake_record(*args, **kwargs): + recorded.append((args, kwargs)) + + monkeypatch.setattr(service, "_record_agent_touches", fake_record) + calls = [] + + class FakeMemoryAgent: + def __init__(self, vault_root, force_fallback=False): + self.vault_root = vault_root + self.force_fallback = force_fallback + + async def run(self, _transcript, conversation_id, **_kwargs): + calls.append(self.force_fallback) + touched = [] + if self.force_fallback: + note = self.vault_root / "Conversations" / f"{conversation_id}.md" + note.parent.mkdir(parents=True, exist_ok=True) + note.write_text( + """--- +categories: + - "[[Conversations]]" +conversation_id: conversation-1 +date: 2026-07-15T12:00:00+00:00 +people: [] +topics: [] +duration_minutes: 2.5 +--- +## A useful conversation + +### Summary +The speakers discussed a concrete plan for the project. + +### Key Facts +- The project plan was reviewed. + +### Action Items +- [ ] Follow up on the project plan. +""", + encoding="utf-8", + ) + touched.append(f"Conversations/{conversation_id}.md") + return MemoryAgentResult( + conversation_id=conversation_id, + rounds=1, + touched=touched, + summary="done", + ) + + import advanced_omi_backend.services.memory.agent as agent_module + + monkeypatch.setattr(agent_module, "MemoryAgent", FakeMemoryAgent) + + success, touched = await service._add_memory_agent( + "Speaker: enough transcript text", + "conversation-1", + "user-1", + source_date="2026-07-15T12:00:00+00:00", + source_duration_minutes=2.5, + source_title="A useful conversation", + ) + + assert success is True + assert calls == [False, True] + assert touched == ["Conversations/conversation-1.md"] + assert len(recorded) == 1 + + +def test_conversation_note_canonicalization_rejects_placeholder_content(tmp_path): + note = tmp_path / "conversation.md" + note.write_text( + """--- +categories: ["[[Conversations]]"] +conversation_id: wrong-id +date: 2026-07-16 +people: [] +topics: [] +duration_minutes: 0 +--- +## Untitled + +### Summary + +### Key Facts +- + +### Action Items +- [ ] +""", + encoding="utf-8", + ) + + with pytest.raises(ConversationNoteError): + canonicalize_conversation_note( + note, + conversation_id="conversation-1", + date="2026-07-15T12:00:00+00:00", + duration_minutes=2.5, + title="Source title", + ) + + +def test_conversation_note_canonicalization_uses_trusted_metadata(tmp_path): + note = tmp_path / "conversation.md" + note.write_text( + """## Model supplied title +--- +categories: ["[[Conversations]]"] +conversation_id: hallucinated +date: 2026-07-16 +people: ["[[Ankush]]", "[[Unknown Speaker 4]]", "[[Hermes]]"] +topics: ["[[Memory systems]]"] +duration_minutes: 999 +--- + +### Summary +The conversation covered a reliable memory rebuild process. + +### Key Facts +- The rebuild must preserve source metadata. +- The rebuild must preserve source metadata. + +### Action Items +- [ ] Run a canary before the full rebuild. +""", + encoding="utf-8", + ) + + canonicalize_conversation_note( + note, + conversation_id="conversation-1", + date="2026-07-15T12:00:00+00:00", + duration_minutes=2.5, + title="Trusted source title", + ) + + content = note.read_text(encoding="utf-8") + assert content.startswith("---\n") + assert 'conversation_id: "conversation-1"' in content + assert 'date: "2026-07-15T12:00:00+00:00"' in content + assert "duration_minutes: 2.5" in content + assert "## Model supplied title" in content + assert "Unknown Speaker" not in content + assert 'people:\n - "[[Ankush]]"' in content + assert ' - "[[Hermes]]"' in content + assert content.count("- The rebuild must preserve source metadata.") == 1 + + +def test_source_fallback_preserves_short_transcript(tmp_path): + note = tmp_path / "conversation.md" + write_source_fallback_conversation_note( + note, + transcript="Hey Hermes, why does it only work during the demo?", + conversation_id="conversation-1", + date="2026-07-15T12:00:00+00:00", + duration_minutes=0.2, + title="Hermes Discussion", + ) + + canonicalize_conversation_note( + note, + conversation_id="conversation-1", + date="2026-07-15T12:00:00+00:00", + duration_minutes=0.2, + title="Hermes Discussion", + ) + assert "why does it only work during the demo" in note.read_text(encoding="utf-8") diff --git a/backends/advanced/tests/test_memory_exclusion.py b/backends/advanced/tests/test_memory_exclusion.py new file mode 100644 index 00000000..696c525f --- /dev/null +++ b/backends/advanced/tests/test_memory_exclusion.py @@ -0,0 +1,32 @@ +"""Regression coverage for conversations excluded from user memory.""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest + +from advanced_omi_backend.workers import memory_jobs + + +@pytest.mark.asyncio +async def test_memory_worker_refuses_excluded_conversation(monkeypatch): + conversation = SimpleNamespace(memory_excluded=True) + find_one = AsyncMock(return_value=conversation) + get_memory_service = Mock() + conversation_model = SimpleNamespace( + conversation_id=object(), + find_one=find_one, + ) + monkeypatch.setattr(memory_jobs, "Conversation", conversation_model) + monkeypatch.setattr(memory_jobs, "get_memory_service", get_memory_service) + + undecorated_job = memory_jobs.process_memory_job.__wrapped__.__wrapped__ + result = await undecorated_job("excluded-conversation", redis_client=None) + + assert result == { + "success": True, + "skipped": True, + "reason": "memory_excluded", + "conversation_id": "excluded-conversation", + } + get_memory_service.assert_not_called() diff --git a/backends/advanced/tests/test_memory_rebuild.py b/backends/advanced/tests/test_memory_rebuild.py new file mode 100644 index 00000000..1142d875 --- /dev/null +++ b/backends/advanced/tests/test_memory_rebuild.py @@ -0,0 +1,345 @@ +"""Tests for clean, ordered Markdown-vault reconstruction.""" + +from datetime import datetime, timezone +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from advanced_omi_backend.services.memory import rebuild +from advanced_omi_backend.services.memory.rebuild import ( + RebuildConversation, + RebuildPlan, + RebuildStage, + build_rebuild_plan, + execute_memory_rebuild, +) + + +class AuditCollection: + def __init__(self, deleted_count): + self.deleted_count = deleted_count + self.query = None + + async def delete_many(self, query): + self.query = query + return SimpleNamespace(deleted_count=self.deleted_count) + + +class FakeDatabase: + def __init__(self, audit): + self.audit = audit + + def __getitem__(self, name): + assert name == "memory_audit" + return self.audit + + +class PlanCursor: + def __init__(self, documents): + self.documents = documents + + def sort(self, _fields): + return self + + def __aiter__(self): + self.iterator = iter(self.documents) + return self + + async def __anext__(self): + try: + return next(self.iterator) + except StopIteration as exc: + raise StopAsyncIteration from exc + + +class ConversationCollection: + def __init__(self, documents): + self.documents = documents + self.query = None + + def find(self, query, projection): + self.query = query + return PlanCursor(self.documents) + + +class AudioCollection: + def __init__(self, conversation_ids): + self.conversation_ids = conversation_ids + self.query = None + + async def distinct(self, field, query): + assert field == "conversation_id" + self.query = query + return self.conversation_ids + + +@pytest.mark.asyncio +async def test_build_rebuild_plan_collects_async_cursor(): + collection = ConversationCollection( + [ + { + "conversation_id": "conversation-1", + "user_id": "user-1", + "created_at": datetime(2026, 7, 15, tzinfo=timezone.utc), + "active_transcript_version": "version-1", + "transcript_versions": [], + } + ] + ) + database = {"conversations": collection} + + plan = await build_rebuild_plan(database) + + assert plan.count == 1 + assert plan.user_ids == ("user-1",) + assert plan.conversations[0].conversation_id == "conversation-1" + assert plan.conversations[0].transcript_version_id == "version-1" + + +@pytest.mark.asyncio +async def test_speaker_plan_includes_memory_excluded_conversations(): + collection = ConversationCollection( + [ + { + "conversation_id": "excluded-conversation", + "user_id": "user-1", + "created_at": datetime(2026, 7, 15, tzinfo=timezone.utc), + "active_transcript_version": "version-1", + "transcript_versions": [], + "memory_excluded": True, + } + ] + ) + + audio = AudioCollection([]) + plan = await build_rebuild_plan( + {"conversations": collection, "audio_chunks": audio}, + from_stage=RebuildStage.SPEAKERS, + ) + + assert plan.count == 1 + assert plan.speaker_count == 0 + assert plan.memory_count == 0 + assert plan.conversations[0].memory_excluded is True + assert plan.conversations[0].has_audio is False + assert "memory_excluded" not in collection.query + assert audio.query == { + "conversation_id": {"$in": ["excluded-conversation"]}, + "deleted": {"$ne": True}, + } + + +@pytest.mark.asyncio +async def test_speaker_plan_unwraps_previous_speaker_version(): + collection = ConversationCollection( + [ + { + "conversation_id": "conversation-1", + "user_id": "user-1", + "created_at": datetime(2026, 7, 15, tzinfo=timezone.utc), + "active_transcript_version": "speaker-version", + "transcript_versions": [ + {"version_id": "source-version", "metadata": {}}, + { + "version_id": "speaker-version", + "metadata": { + "reprocessing_type": "speaker_diarization", + "source_version_id": "source-version", + }, + }, + ], + } + ] + ) + + plan = await build_rebuild_plan( + { + "conversations": collection, + "audio_chunks": AudioCollection(["conversation-1"]), + }, + from_stage=RebuildStage.SPEAKERS, + ) + + item = plan.conversations[0] + assert item.transcript_version_id == "source-version" + assert item.active_transcript_version_id == "speaker-version" + + +@pytest.mark.asyncio +async def test_execute_rebuild_chains_each_user_chronologically( + monkeypatch, tmp_path: Path +): + created = datetime(2026, 7, 15, tzinfo=timezone.utc) + plan = RebuildPlan( + conversations=( + RebuildConversation("a-first", "user-a", created, "version-a1"), + RebuildConversation("a-second", "user-a", created, "version-a2"), + RebuildConversation("b-first", "user-b", created, "version-b1"), + ), + user_ids=("user-a", "user-b"), + ) + for user_id in plan.user_ids: + root = tmp_path / "conversation_docs" / user_id + root.mkdir(parents=True) + (root / "old.md").write_text("old", encoding="utf-8") + (root / ".stignore").write_text("sync", encoding="utf-8") + + monkeypatch.setattr(rebuild, "_active_rebuild_jobs", lambda _ids: []) + enqueued = [] + + def fake_enqueue(conversation_id, **kwargs): + job = SimpleNamespace(id=kwargs["job_id"]) + enqueued.append((conversation_id, kwargs, job)) + return job + + monkeypatch.setattr(rebuild, "enqueue_memory_processing", fake_enqueue) + audit = AuditCollection(deleted_count=7) + + result = await execute_memory_rebuild( + FakeDatabase(audit), + plan, + data_dir=tmp_path, + backup_dir=None, + ) + + assert [item[0] for item in enqueued] == ["a-first", "a-second", "b-first"] + assert enqueued[0][1]["depends_on"] is None + assert enqueued[1][1]["depends_on"] is enqueued[0][2] + assert enqueued[2][1]["depends_on"] is None + assert all(item[1]["cause"].value == "memory_rebuild" for item in enqueued) + assert all( + item[1]["job_timeout"] == rebuild.MEMORY_REBUILD_JOB_TIMEOUT + for item in enqueued + ) + assert result.deleted_vault_files == 2 + assert result.deleted_audit_entries == 7 + assert audit.query == {"user_id": {"$in": ["user-a", "user-b"]}} + for user_id in plan.user_ids: + assert (tmp_path / "conversation_docs" / user_id / ".stignore").exists() + assert not (tmp_path / "conversation_docs" / user_id / "old.md").exists() + + +@pytest.mark.asyncio +async def test_execute_rebuild_from_speakers_continues_after_failed_speaker( + monkeypatch, tmp_path: Path +): + created = datetime(2026, 7, 15, tzinfo=timezone.utc) + plan = RebuildPlan( + conversations=( + RebuildConversation("first", "user-a", created, "version-1"), + RebuildConversation("second", "user-a", created, "version-2"), + ), + user_ids=("user-a",), + ) + speaker_calls = [] + memory_calls = [] + + def fake_speaker(item, *, run_id, sequence, depends_on): + job = SimpleNamespace(id=f"speaker-{sequence}") + speaker_calls.append((item, run_id, sequence, depends_on, job)) + return job + + def fake_memory(conversation_id, **kwargs): + job = SimpleNamespace(id=f"memory-{conversation_id}") + memory_calls.append((conversation_id, kwargs, job)) + return job + + monkeypatch.setattr(rebuild, "_active_rebuild_jobs", lambda _ids: []) + monkeypatch.setattr(rebuild, "_enqueue_speaker_rebuild", fake_speaker) + monkeypatch.setattr(rebuild, "enqueue_memory_processing", fake_memory) + + result = await execute_memory_rebuild( + FakeDatabase(AuditCollection(deleted_count=0)), + plan, + data_dir=tmp_path, + backup_dir=None, + from_stage=RebuildStage.SPEAKERS, + ) + + assert [call[0].conversation_id for call in speaker_calls] == ["first", "second"] + assert speaker_calls[0][3] is None + assert speaker_calls[1][3] == "speaker-1" + assert [call[0] for call in memory_calls] == ["first", "second"] + first_dependency = memory_calls[0][1]["depends_on"] + assert first_dependency.dependencies == ["speaker-2"] + assert first_dependency.allow_failure is True + assert memory_calls[1][1]["depends_on"] is memory_calls[0][2] + assert result.from_stage is RebuildStage.SPEAKERS + assert result.speaker_jobs == ("speaker-1", "speaker-2") + + +def test_speaker_rebuild_dependency_allows_previous_failure(monkeypatch): + captured = {} + + def fake_enqueue(*args, **kwargs): + captured.update(kwargs) + return SimpleNamespace(id=kwargs["job_id"]) + + def fake_enqueue_kwargs(_stage, _meta, depends_on=None): + return {"depends_on": depends_on} + + monkeypatch.setattr( + rebuild, "transcription_queue", SimpleNamespace(enqueue=fake_enqueue) + ) + monkeypatch.setattr(rebuild, "post_conv_enqueue_kwargs", fake_enqueue_kwargs) + item = RebuildConversation( + "conversation-1", + "user-1", + datetime(2026, 7, 15, tzinfo=timezone.utc), + "version-1", + ) + + rebuild._enqueue_speaker_rebuild( + item, + run_id="run-id", + sequence=2, + depends_on="previous-speaker-job", + ) + + dependency = captured["depends_on"] + assert dependency.dependencies == ["previous-speaker-job"] + assert dependency.allow_failure is True + + +@pytest.mark.asyncio +async def test_speaker_rebuild_skips_conversation_without_audio_but_rebuilds_memory( + monkeypatch, tmp_path: Path +): + plan = RebuildPlan( + conversations=( + RebuildConversation( + "transcript-only", + "user-a", + datetime(2026, 7, 15, tzinfo=timezone.utc), + "version-1", + has_audio=False, + ), + ), + user_ids=("user-a",), + ) + memory_calls = [] + + monkeypatch.setattr(rebuild, "_active_rebuild_jobs", lambda _ids: []) + monkeypatch.setattr( + rebuild, + "_enqueue_speaker_rebuild", + lambda *args, **kwargs: pytest.fail("speaker job should not be enqueued"), + ) + + def fake_memory(conversation_id, **kwargs): + memory_calls.append(conversation_id) + return SimpleNamespace(id="memory-job") + + monkeypatch.setattr(rebuild, "enqueue_memory_processing", fake_memory) + + result = await execute_memory_rebuild( + FakeDatabase(AuditCollection(deleted_count=0)), + plan, + data_dir=tmp_path, + from_stage=RebuildStage.SPEAKERS, + ) + + assert memory_calls == ["transcript-only"] + assert result.speaker_jobs == () + assert result.skipped_speaker_conversations == ("transcript-only",) diff --git a/backends/advanced/tests/test_memory_transcript_input.py b/backends/advanced/tests/test_memory_transcript_input.py new file mode 100644 index 00000000..2bccaba7 --- /dev/null +++ b/backends/advanced/tests/test_memory_transcript_input.py @@ -0,0 +1,52 @@ +"""Memory extraction input must not amplify overlapping transcript windows.""" + +from types import SimpleNamespace + +from advanced_omi_backend.workers.memory_jobs import build_memory_transcript + + +def _segment(start, end, text, speaker="Speaker 0", segment_type="speech"): + return SimpleNamespace( + start=start, + end=end, + text=text, + speaker=speaker, + segment_type=segment_type, + ) + + +def test_build_memory_transcript_trims_overlapping_word_prefix(): + segments = [ + _segment(0, 30, "one two three four five"), + _segment(25, 55, "three four five six seven", "Speaker 1"), + ] + + transcript, speakers = build_memory_transcript(segments, raw_transcript=None) + + assert transcript == "Speaker 0: one two three four five\nSpeaker 1: six seven" + assert speakers == {"speaker 0", "speaker 1"} + + +def test_build_memory_transcript_falls_back_when_segments_amplify_raw_text(): + segments = [ + _segment(0, 100, "duplicated window text " * 100), + _segment(50, 150, "different duplicated text " * 100), + ] + raw = "This is the durable raw transcript and it should be used instead." + + transcript, speakers = build_memory_transcript(segments, raw_transcript=raw) + + assert transcript == raw + assert speakers == {"speaker 0"} + + +def test_build_memory_transcript_preserves_events_and_notes(): + segments = [ + _segment(0, 1, "music", segment_type="event"), + _segment(1, 2, "Remember this", segment_type="note"), + ] + + transcript, speakers = build_memory_transcript(segments, raw_transcript=None) + + assert transcript == "[music]\n[Note: Remember this]" + assert speakers == set() diff --git a/backends/advanced/tests/test_model_registry_reasoning.py b/backends/advanced/tests/test_model_registry_reasoning.py new file mode 100644 index 00000000..b0b5d53d --- /dev/null +++ b/backends/advanced/tests/test_model_registry_reasoning.py @@ -0,0 +1,26 @@ +"""Reasoning parameter compatibility for versioned GPT-5 models.""" + +from advanced_omi_backend.model_registry import ModelDef, ResolvedLLMOperation + + +def _operation(model_name: str) -> ResolvedLLMOperation: + return ResolvedLLMOperation( + model_def=ModelDef( + name="test", + model_name=model_name, + model_type="llm", + model_provider="openai", + model_url="https://api.openai.com/v1", + ), + temperature=0.2, + max_tokens=100, + reasoning_effort="none", + ) + + +def test_gpt_5_4_preserves_none_reasoning_effort(): + assert _operation("gpt-5.4-mini").to_api_params()["reasoning_effort"] == "none" + + +def test_unversioned_gpt_5_uses_minimal_reasoning_effort(): + assert _operation("gpt-5-mini").to_api_params()["reasoning_effort"] == "minimal" diff --git a/backends/advanced/tests/test_openai_compat_routes.py b/backends/advanced/tests/test_openai_compat_routes.py new file mode 100644 index 00000000..76418b6f --- /dev/null +++ b/backends/advanced/tests/test_openai_compat_routes.py @@ -0,0 +1,36 @@ +import pytest + +from advanced_omi_backend.model_registry import ModelDef +from advanced_omi_backend.routers.modules import openai_compat_routes as routes + + +@pytest.fixture(autouse=True) +def clear_unavailable_models(): + routes._unavailable_models.clear() + yield + routes._unavailable_models.clear() + + +def test_unavailable_model_is_bypassed_until_cooldown_expires(monkeypatch): + now = 1_000.0 + monkeypatch.setattr(routes.time, "monotonic", lambda: now) + + routes._mark_model_unavailable("llamacpp-llm") + + assert routes._model_is_unavailable("llamacpp-llm") + + now += routes._UPSTREAM_FAILURE_COOLDOWN_SECONDS + 0.1 + assert not routes._model_is_unavailable("llamacpp-llm") + + +@pytest.mark.asyncio +async def test_proxy_skips_a_model_with_an_open_circuit(): + model = ModelDef( + name="llamacpp-llm", + model_type="llm", + model_url="http://unreachable.example/v1", + ) + routes._mark_model_unavailable(model.name) + + with pytest.raises(routes._UpstreamTransportError, match="cooldown active"): + await routes._proxy_chat_completion(model, {"messages": []}, stream=False) diff --git a/backends/advanced/tests/test_session_store.py b/backends/advanced/tests/test_session_store.py new file mode 100644 index 00000000..1420f1e1 --- /dev/null +++ b/backends/advanced/tests/test_session_store.py @@ -0,0 +1,479 @@ +"""Unit tests for the SessionStore facade and SessionView read-model. + +Covers the two things the facade exists to guarantee: +1. Decoding is identical whether the redis client returns bytes or str. +2. Lifecycle writes are atomic (single hset) and the signal publish follows the + hash write. +""" + +import asyncio +import json + +import pytest +from fakeredis import aioredis as fake_aioredis + +from advanced_omi_backend.services.audio_stream.session_store import ( + SessionStatus, + SessionStore, + SessionView, + SpeakerCheckStatus, +) + +pytestmark = pytest.mark.unit + + +def _fake_redis(decode_responses=False): + return fake_aioredis.FakeRedis(decode_responses=decode_responses) + + +# --------------------------------------------------------------------------- # +# SessionView.from_hash +# --------------------------------------------------------------------------- # + + +def _sample_str_hash(): + return { + "user_id": "507f1f77bcf86cd799439011", + "client_id": "a39011-phone", + "status": "finalizing", + "websocket_connected": "true", + "chunks_published": "42", + "started_at": "1704067200.5", + "last_chunk_at": "1704067260.0", + "finalized_at": "1704067300.0", + "speaker_check_status": "enrolled", + "audio_format": json.dumps({"rate": 48000, "channels": 2, "width": 2}), + "markers": json.dumps([{"type": "button_event"}]), + "identified_speakers": "Alice,Bob", + "speech_detected_at": "2026-01-01T00:00:00+00:00", + "completion_reason": "user_stopped", + } + + +def test_from_hash_bytes_and_str_are_identical(): + str_hash = _sample_str_hash() + bytes_hash = {k.encode(): v.encode() for k, v in str_hash.items()} + + view_from_str = SessionView.from_hash("sess-1", str_hash) + view_from_bytes = SessionView.from_hash("sess-1", bytes_hash) + + assert view_from_str == view_from_bytes + + +def test_from_hash_coercions(): + view = SessionView.from_hash("sess-1", _sample_str_hash()) + + assert view.status is SessionStatus.FINALIZING + assert view.speaker_check_status is SpeakerCheckStatus.ENROLLED + assert view.websocket_connected is True + assert view.chunks_published == 42 + assert view.started_at == 1704067200.5 + assert view.finalized_at == 1704067300.0 + assert view.completed_at is None # missing -> None (not 0.0) + assert view.audio_format == {"rate": 48000, "channels": 2, "width": 2} + assert view.audio_format_tuple == (48000, 2, 2) + assert view.markers == [{"type": "button_event"}] + assert view.identified_speakers == ["Alice", "Bob"] + assert view.speech_detected_at == "2026-01-01T00:00:00+00:00" + + +def test_from_hash_unknown_enum_and_empty_defaults(): + view = SessionView.from_hash( + "sess-1", + {"status": "bogus", "chunks_published": "", "websocket_connected": "false"}, + ) + assert view.status is None # unknown enum value -> None, not a raise + assert view.chunks_published == 0 + assert view.websocket_connected is False + assert view.started_at == 0.0 + assert view.audio_format is None + assert view.markers == [] + assert view.identified_speakers == [] + + +def test_from_hash_malformed_audio_format_falls_back(): + view = SessionView.from_hash("sess-1", {"audio_format": "not-json"}) + assert view.audio_format is None + assert view.audio_format_tuple == (16000, 1, 2) + + +# --------------------------------------------------------------------------- # +# SessionStore lifecycle writes +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize("decode_responses", [False, True]) +async def test_init_session_writes_expected_mapping(decode_responses): + client = _fake_redis(decode_responses) + store = SessionStore(client) + + await store.init_session( + "sess-1", + user_id="u1", + client_id="c1", + stream_name="audio:stream:c1", + user_email="e@x.com", + mode="streaming", + provider="deepgram", + ) + + view = await store.read("sess-1") + assert view is not None + assert view.user_id == "u1" + assert view.client_id == "c1" + assert view.stream_name == "audio:stream:c1" + assert view.user_email == "e@x.com" + assert view.provider == "deepgram" + assert view.mode == "streaming" + assert view.status is SessionStatus.ACTIVE + assert view.websocket_connected is True + assert view.chunks_published == 0 + assert view.started_at > 0 + + +async def test_mark_finalizing_sets_status_and_publishes_signal(): + client = _fake_redis() + store = SessionStore(client) + await store.init_session( + "sess-1", user_id="u1", client_id="c1", stream_name="audio:stream:c1" + ) + + pubsub = client.pubsub() + await pubsub.subscribe("session:signal:sess-1") + await pubsub.get_message(timeout=1) # consume subscribe ack + + await store.mark_finalizing("sess-1", "websocket_disconnect") + + view = await store.read("sess-1") + assert view.status is SessionStatus.FINALIZING + assert view.finalized_at is not None + assert view.completion_reason == "websocket_disconnect" + assert view.websocket_connected is False # set on websocket_disconnect + + msg = await pubsub.get_message(ignore_subscribe_messages=True, timeout=1) + assert msg is not None + payload = json.loads(msg["data"]) + assert payload == {"type": "finalize", "reason": "websocket_disconnect"} + + +async def test_mark_complete_atomic_hset_then_publish(): + """status + completed_at + completion_reason in one hset, publish after.""" + + class Recorder: + def __init__(self, inner): + self._inner = inner + self.calls = [] + + def __getattr__(self, name): + attr = getattr(self._inner, name) + if name in ("hset", "publish"): + + async def wrapper(*a, **k): + self.calls.append(name) + return await attr(*a, **k) + + return wrapper + return attr + + rec = Recorder(_fake_redis()) + store = SessionStore(rec) + await store.init_session( + "sess-1", user_id="u1", client_id="c1", stream_name="audio:stream:c1" + ) + rec.calls.clear() + + await store.mark_complete("sess-1", "user_stopped") + + # exactly one hset, then the publish + assert rec.calls == ["hset", "publish"] + view = await store.read("sess-1") + assert view.status is SessionStatus.FINISHED + assert view.completion_reason == "user_stopped" + assert view.completed_at is not None + + +async def test_take_close_request_returns_then_clears(): + client = _fake_redis() + store = SessionStore(client) + await store.init_session( + "sess-1", user_id="u1", client_id="c1", stream_name="audio:stream:c1" + ) + assert await store.request_close("sess-1", "user_requested") is True + + assert await store.take_close_request("sess-1") == "user_requested" + assert await store.take_close_request("sess-1") is None # consumed + + +async def test_request_close_returns_false_when_missing(): + store = SessionStore(_fake_redis()) + assert await store.request_close("nope", "user_requested") is False + + +async def test_get_audio_format_fallback_on_missing_and_garbage(): + client = _fake_redis() + store = SessionStore(client) + await store.init_session( + "sess-1", user_id="u1", client_id="c1", stream_name="audio:stream:c1" + ) + + assert await store.get_audio_format("sess-1") == (16000, 1, 2) # not set yet + + await client.hset("audio:session:sess-1", "audio_format", "not-json") + assert await store.get_audio_format("sess-1") == (16000, 1, 2) # garbage -> default + + await store.set_audio_format("sess-1", {"rate": 8000, "channels": 1, "width": 2}) + assert await store.get_audio_format("sess-1") == (8000, 1, 2) + + +async def test_bump_chunk_count_increments(): + store = SessionStore(_fake_redis()) + await store.init_session( + "sess-1", user_id="u1", client_id="c1", stream_name="audio:stream:c1" + ) + await store.bump_chunk_count("sess-1") + await store.bump_chunk_count("sess-1") + view = await store.read("sess-1") + assert view.chunks_published == 2 + assert view.last_chunk_at > 0 + + +async def test_transcription_provider_health_lifecycle(): + store = SessionStore(_fake_redis()) + await store.init_session( + "sess-1", user_id="u1", client_id="c1", stream_name="audio:stream:c1" + ) + + initial = await store.read("sess-1") + assert initial.transcription_provider_status == "disconnected" + + await store.mark_transcription_provider_connected("sess-1") + await store.mark_transcription_audio_sent("sess-1") + await store.mark_transcription_provider_message("sess-1") + + connected = await store.read("sess-1") + assert connected.transcription_provider_status == "connected" + assert connected.transcription_provider_connected_at > 0 + assert connected.transcription_last_audio_sent_at > 0 + assert connected.transcription_last_message_at > 0 + + await store.set_transcription_error("sess-1", "socket closed") + failed = await store.read("sess-1") + assert failed.transcription_provider_status == "error" + assert failed.transcription_error == "socket closed" + + await store.mark_transcription_provider_disconnected("sess-1") + still_failed = await store.read("sess-1") + assert still_failed.transcription_provider_status == "error" + + await store.mark_transcription_provider_connected("sess-1") + await store.mark_transcription_provider_disconnected("sess-1") + disconnected = await store.read("sess-1") + assert disconnected.transcription_provider_status == "disconnected" + + +async def test_get_status_ws_reason_batched(): + store = SessionStore(_fake_redis()) + await store.init_session( + "sess-1", user_id="u1", client_id="c1", stream_name="audio:stream:c1" + ) + await store.mark_complete("sess-1", "inactivity_timeout") + status, ws, reason = await store.get_status_ws_reason("sess-1") + assert status is SessionStatus.FINISHED + # inactivity_timeout does NOT clear websocket_connected (only websocket_disconnect does) + assert ws is True + assert reason == "inactivity_timeout" + + +async def test_iter_views_strips_prefix_and_skips_empty(): + client = _fake_redis() + store = SessionStore(client) + await store.init_session( + "alpha", user_id="u1", client_id="c1", stream_name="audio:stream:c1" + ) + await store.init_session( + "beta", user_id="u2", client_id="c2", stream_name="audio:stream:c2" + ) + # an empty hash key (e.g. fully hdel'd) should be skipped by iter_views + await client.hset("audio:session:ghost", "x", "1") + await client.hdel("audio:session:ghost", "x") + + ids = sorted([sid async for sid in store.scan_session_ids()]) + assert "alpha" in ids and "beta" in ids + + views = {v.session_id: v async for v in store.iter_views()} + assert set(views) == {"alpha", "beta"} # ghost skipped + assert views["alpha"].user_id == "u1" + + +async def test_increment_conversation_count_sets_ttl(): + client = _fake_redis() + store = SessionStore(client) + assert await store.get_conversation_count("sess-1") == 0 + assert await store.increment_conversation_count("sess-1") == 1 + assert await store.increment_conversation_count("sess-1") == 2 + assert await store.get_conversation_count("sess-1") == 2 + ttl = await client.ttl("session:conversation_count:sess-1") + assert 0 < ttl <= 3600 + + +async def test_record_event_and_speaker_check(): + store = SessionStore(_fake_redis()) + await store.init_session( + "sess-1", user_id="u1", client_id="c1", stream_name="audio:stream:c1" + ) + await store.record_event("sess-1", "speech_detected") + await store.set_speaker_check("sess-1", SpeakerCheckStatus.ENROLLED) + await store.set_identified_speakers("sess-1", ["Alice", "Bob"]) + + view = await store.read("sess-1") + assert view.last_event.startswith("speech_detected:") + assert view.speaker_check_status is SpeakerCheckStatus.ENROLLED + assert view.identified_speakers == ["Alice", "Bob"] + + +# --------------------------------------------------------------------------- # +# Conversation pointer (conversation:current) +# +# The pointer has per-caller TTLs that used to be passed inline at ~11 sites +# (86400 rotation / none for the always_persist placeholder / 3600 disconnect +# backstop). These tests pin that behaviour now that it lives on the facade. +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize("decode_responses", [False, True]) +async def test_set_current_conversation_default_ttl(decode_responses): + client = _fake_redis(decode_responses) + store = SessionStore(client) + + await store.set_current_conversation("sess-1", "conv-A") + + assert await store.get_current_conversation_id("sess-1") == "conv-A" + ttl = await client.ttl("conversation:current:sess-1") + assert 0 < ttl <= 86400 # rotation pointer carries the 24h TTL + + +async def test_set_current_conversation_no_ttl_for_placeholder(): + client = _fake_redis() + store = SessionStore(client) + + await store.set_current_conversation("sess-1", "conv-A", ttl=None) + + assert await store.get_current_conversation_id("sess-1") == "conv-A" + assert await client.ttl("conversation:current:sess-1") == -1 # persistent + + +async def test_set_current_conversation_overwrites_on_rotation(): + store = SessionStore(_fake_redis()) + await store.set_current_conversation("sess-1", "conv-A") + await store.set_current_conversation("sess-1", "conv-B") + assert await store.get_current_conversation_id("sess-1") == "conv-B" + + +async def test_get_current_conversation_none_when_absent(): + store = SessionStore(_fake_redis()) + assert await store.get_current_conversation_id("sess-1") is None + + +async def test_clear_current_conversation_is_noop_when_absent(): + store = SessionStore(_fake_redis()) + await store.clear_current_conversation("sess-1") # must not raise + await store.set_current_conversation("sess-1", "conv-A") + await store.clear_current_conversation("sess-1") + assert await store.get_current_conversation_id("sess-1") is None + + +async def test_expire_current_conversation_only_when_present(): + client = _fake_redis() + store = SessionStore(client) + + assert await store.expire_current_conversation("sess-1", 3600) is False + + await store.set_current_conversation("sess-1", "conv-A", ttl=None) + assert await store.expire_current_conversation("sess-1", 3600) is True + ttl = await client.ttl("conversation:current:sess-1") + assert 0 < ttl <= 3600 + + +# --------------------------------------------------------------------------- # +# conversation_create_lock — the dual-creation guarantee +# +# This lock is the entire reason one streaming session can't produce two +# conversations (the persistence job and open_conversation_job both run a +# get→create→set on the pointer). These are its permanent regression tests. +# --------------------------------------------------------------------------- # + + +async def test_conversation_create_lock_is_mutually_exclusive(): + """Two concurrent holders of the same session lock never overlap.""" + store = SessionStore(_fake_redis()) + inside = 0 + max_concurrent = 0 + order = [] + + async def worker(tag): + nonlocal inside, max_concurrent + async with store.conversation_create_lock("sess-1") as acquired: + assert acquired is True + inside += 1 + max_concurrent = max(max_concurrent, inside) + order.append(f"enter-{tag}") + await asyncio.sleep(0.05) # hold long enough for the other to contend + order.append(f"exit-{tag}") + inside -= 1 + + await asyncio.gather(worker("a"), worker("b")) + + assert max_concurrent == 1 # serialized — never both inside the section + # whichever wins fully completes before the other enters (no interleave) + assert order in ( + ["enter-a", "exit-a", "enter-b", "exit-b"], + ["enter-b", "exit-b", "enter-a", "exit-a"], + ) + + +async def test_conversation_create_lock_different_sessions_run_concurrently(): + store = SessionStore(_fake_redis()) + inside = 0 + max_concurrent = 0 + + async def worker(sid): + nonlocal inside, max_concurrent + async with store.conversation_create_lock(sid): + inside += 1 + max_concurrent = max(max_concurrent, inside) + await asyncio.sleep(0.05) + inside -= 1 + + await asyncio.gather(worker("sess-1"), worker("sess-2")) + + assert max_concurrent == 2 # independent sessions don't block each other + + +async def test_conversation_create_lock_fails_open_on_timeout(): + """If the lock can't be acquired in time, the body still runs (yields False).""" + client = _fake_redis() + store = SessionStore(client) + # Pre-hold the lock so the contender can never acquire within wait_timeout. + await client.set("conversation:create_lock:sess-1", "1", ex=30) + + ran = False + async with store.conversation_create_lock( + "sess-1", wait_timeout=0.2, poll=0.02 + ) as acquired: + ran = True + assert acquired is False # degraded to unlocked, not deadlocked + + assert ran is True + # we never acquired, so the pre-existing lock is left untouched + assert await client.get("conversation:create_lock:sess-1") is not None + + +async def test_conversation_create_lock_releases_on_exit(): + client = _fake_redis() + store = SessionStore(client) + + async with store.conversation_create_lock("sess-1") as acquired: + assert acquired is True + assert await client.get("conversation:create_lock:sess-1") is not None + + assert await client.get("conversation:create_lock:sess-1") is None # released diff --git a/backends/advanced/tests/test_smallest_diarization_config.py b/backends/advanced/tests/test_smallest_diarization_config.py new file mode 100644 index 00000000..06057a4d --- /dev/null +++ b/backends/advanced/tests/test_smallest_diarization_config.py @@ -0,0 +1,26 @@ +from pathlib import Path + +import yaml + +ROOT = Path(__file__).resolve().parents[3] + + +def _smallest_batch_model(config_path: Path) -> dict: + config = yaml.safe_load(config_path.read_text()) + return next(model for model in config["models"] if model["name"] == "stt-smallest") + + +def test_smallest_batch_enables_advertised_diarization(): + model = _smallest_batch_model(ROOT / "config" / "defaults.yml") + + assert "diarization" in model["capabilities"] + query = model["operations"]["stt_transcribe"]["query"] + assert query["diarize"] == "true" + + +def test_smallest_template_enables_advertised_diarization(): + model = _smallest_batch_model(ROOT / "config" / "config.yml.template") + + assert "diarization" in model["capabilities"] + query = model["operations"]["stt_transcribe"]["query"] + assert query["diarize"] == "true" diff --git a/backends/advanced/tests/test_speaker_benchmark.py b/backends/advanced/tests/test_speaker_benchmark.py new file mode 100644 index 00000000..6c64a1f6 --- /dev/null +++ b/backends/advanced/tests/test_speaker_benchmark.py @@ -0,0 +1,30 @@ +import numpy as np + +from advanced_omi_backend.workers.speaker_benchmark_jobs import FRACTIONS, _evaluate + + +def test_learning_curve_uses_disjoint_conversation_folds(): + samples = [] + for speaker, base in ( + ("ankush", np.array([1.0, 0.0])), + ("janhavi", np.array([0.0, 1.0])), + ): + for conversation in range(10): + vector = base + np.array([conversation * 0.001, -conversation * 0.001]) + vector = vector / np.linalg.norm(vector) + samples.append( + { + "key": f"{speaker}-{conversation}", + "speaker": speaker, + "conversation_id": f"conversation-{conversation}", + "embedding": vector, + } + ) + + report = _evaluate(samples, threshold=0.5) + + assert [point["fraction"] for point in report["learning_curve"]] == list(FRACTIONS) + assert report["learning_curve"][-1]["top1_accuracy_mean"] == 1.0 + assigned = [group for groups in report["fold_groups"].values() for group in groups] + assert sorted(assigned) == sorted({sample["conversation_id"] for sample in samples}) + assert len(assigned) == len(set(assigned)) diff --git a/backends/advanced/tests/test_speaker_discovery.py b/backends/advanced/tests/test_speaker_discovery.py new file mode 100644 index 00000000..79139a7b --- /dev/null +++ b/backends/advanced/tests/test_speaker_discovery.py @@ -0,0 +1,61 @@ +from advanced_omi_backend.controllers.guided_enrollment_controller import ( + _information_score, +) +from advanced_omi_backend.workers.speaker_discovery_jobs import ( + _active_segments, + discover_speaker_candidates_job, +) + + +def test_discovery_job_is_rq_importable(): + assert discover_speaker_candidates_job.__module__ == ( + "advanced_omi_backend.workers.speaker_discovery_jobs" + ) + assert discover_speaker_candidates_job.__name__ == "discover_speaker_candidates_job" + + +def test_active_segments_prefers_active_version(): + document = { + "active_transcript_version": "active", + "transcript_versions": [ + {"version_id": "old", "segments": [{"text": "old"}]}, + {"version_id": "active", "segments": [{"text": "current"}]}, + ], + } + + assert _active_segments(document) == [{"text": "current"}] + + +def test_automatic_label_disagreement_remains_reviewable(): + candidate = { + "conversation_id": "conversation-1", + "start": 10.0, + "duration": 8.0, + "current_label": "anushpa", + "speaker_name": "Janhavi", + "scores": { + "sim_centroid": 0.52, + "max_clip_sim": 0.35, + "best_other": {"name": "anushpa", "score": 0.49}, + }, + } + + scored = _information_score(candidate, threshold=0.5) + + assert scored is not None + assert "currently labeled anushpa — possible mismatch" in scored["reasons"] + + +def test_clear_other_speaker_is_gated_out(): + candidate = { + "duration": 8.0, + "current_label": "anushpa", + "speaker_name": "Janhavi", + "scores": { + "sim_centroid": 0.40, + "max_clip_sim": 0.35, + "best_other": {"name": "anushpa", "score": 0.55}, + }, + } + + assert _information_score(candidate, threshold=0.5) is None diff --git a/backends/advanced/tests/test_speaker_identification_policy.py b/backends/advanced/tests/test_speaker_identification_policy.py new file mode 100644 index 00000000..9ac3ff6d --- /dev/null +++ b/backends/advanced/tests/test_speaker_identification_policy.py @@ -0,0 +1,31 @@ +"""Conservative label-level speaker identification policy.""" + +import pytest + +from advanced_omi_backend.speaker_recognition_client import _select_label_mappings + + +def test_single_low_confidence_clip_does_not_name_speaker(): + mappings = _select_label_mappings( + {"Speaker 0": [("Ankush", 0.55)]}, similarity_threshold=0.5 + ) + assert mappings == {} + + +def test_two_agreeing_clips_name_speaker(): + mappings = _select_label_mappings( + {"Speaker 0": [("Ankush", 0.56), ("Ankush", 0.61)]}, + similarity_threshold=0.5, + ) + assert mappings["Speaker 0"][0] == "Ankush" + + +def test_one_identity_cannot_be_assigned_to_two_labels(): + mappings = _select_label_mappings( + { + "Speaker 0": [("Ankush", 0.75), ("Ankush", 0.72)], + "Speaker 1": [("Ankush", 0.61), ("Ankush", 0.60)], + }, + similarity_threshold=0.5, + ) + assert mappings == {"Speaker 0": ("Ankush", pytest.approx(0.735))} diff --git a/backends/advanced/tests/test_streaming_clock_offset.py b/backends/advanced/tests/test_streaming_clock_offset.py new file mode 100644 index 00000000..8eb86f54 --- /dev/null +++ b/backends/advanced/tests/test_streaming_clock_offset.py @@ -0,0 +1,180 @@ +"""Regression tests for the streaming-transcription session-relative clock. + +Streaming providers stamp word timestamps relative to their own WebSocket +session. When the provider connection drops mid-session and a new stream is +opened, the provider clock restarts at 0. Without offsetting, late audio gets +timestamps that collide with the start of the conversation and downstream +segment-building interleaves the two timelines (observed in production: +conversation e3c1dabd, 2026-06-12 — smallest.ai Pulse dropped the WS at +~20.5 min; everything transcribed after reconnect was stamped from 0 and got +bucketed into the first ~10 minutes of diarized segments). + +These tests cover the consumer-level seam: chunks flow through +process_audio_chunk, the connection dies, _reconnect_session re-opens the +provider stream, and all stored results must stay on one monotonic timeline. +""" + +import json + +import pytest +from fakeredis import aioredis as fake_aioredis + +import advanced_omi_backend.services.transcription.streaming_consumer as sc_module +from advanced_omi_backend.services.transcription.streaming_consumer import ( + StreamingTranscriptionConsumer, + _apply_time_offset, +) + +pytestmark = pytest.mark.unit + +SAMPLE_RATE = 16000 +BYTES_PER_SECOND = SAMPLE_RATE * 2 # PCM 16-bit mono + + +# --------------------------------------------------------------------------- # +# _apply_time_offset +# --------------------------------------------------------------------------- # + + +def test_apply_time_offset_shifts_words_and_segments(): + result = { + "words": [{"word": "hi", "start": 0.5, "end": 1.0}], + "segments": [ + { + "start": 0.5, + "end": 1.0, + "text": "hi", + "words": [{"word": "hi", "start": 0.5, "end": 1.0}], + } + ], + } + _apply_time_offset(result, 100.0) + + assert result["words"][0]["start"] == pytest.approx(100.5) + assert result["words"][0]["end"] == pytest.approx(101.0) + assert result["segments"][0]["start"] == pytest.approx(100.5) + assert result["segments"][0]["end"] == pytest.approx(101.0) + assert result["segments"][0]["words"][0]["start"] == pytest.approx(100.5) + + +def test_apply_time_offset_zero_is_noop_and_tolerates_missing_fields(): + result = {"words": [{"word": "hi", "start": 0.5, "end": 1.0}, {"word": "x"}]} + _apply_time_offset(result, 0.0) + assert result["words"][0]["start"] == pytest.approx(0.5) + + # Words without timestamps must not raise + _apply_time_offset(result, 10.0) + assert result["words"][0]["start"] == pytest.approx(10.5) + assert "start" not in result["words"][1] + + +# --------------------------------------------------------------------------- # +# Consumer seam: provider clock reset across reconnect +# --------------------------------------------------------------------------- # + + +class FakeStreamingProvider: + """Minimal provider whose word clock is relative to each start_stream call, + mirroring real streaming providers (Pulse, Deepgram).""" + + capabilities: set = set() + + def __init__(self): + self.clock = 0.0 + self.sessions_started = 0 + self.dead = False + + async def start_stream(self, client_id, sample_rate=16000, diarize=False): + self.sessions_started += 1 + self.clock = 0.0 # the bug trigger: every new WS session restarts at 0 + self.dead = False + + async def process_audio_chunk(self, client_id, audio_chunk): + if self.dead: + raise ConnectionError("provider socket closed") + secs = len(audio_chunk) / BYTES_PER_SECOND + start = self.clock + self.clock += secs + return { + "text": "word", + "words": [ + {"word": "word", "start": start, "end": self.clock, "confidence": 1.0} + ], + "segments": [], + "is_final": True, + "confidence": 1.0, + } + + async def end_stream(self, client_id): + if self.dead: + raise ConnectionError("provider socket closed") + return {"text": "", "words": [], "segments": []} + + +@pytest.fixture +def consumer(monkeypatch): + redis = fake_aioredis.FakeRedis() + provider = FakeStreamingProvider() + monkeypatch.setattr(sc_module, "get_transcription_provider", lambda mode: provider) + c = StreamingTranscriptionConsumer(redis_client=redis) + return c, provider, redis + + +async def _stored_word_starts(redis, session_id): + entries = await redis.xrange(f"transcription:results:{session_id}") + words = [] + for _, fields in entries: + words.extend(json.loads(fields[b"words"])) + return [w["start"] for w in words] + + +async def test_timestamps_stay_monotonic_across_reconnect(consumer): + c, provider, redis = consumer + session_id = "a421c9-havpe-test" + one_second_chunk = b"\x00" * BYTES_PER_SECOND + + await c.start_session_stream(session_id, sample_rate=SAMPLE_RATE) + for i in range(3): + await c.process_audio_chunk(session_id, one_second_chunk, f"pre-{i}") + + # Provider kills the connection (e.g. max WS session duration) + provider.dead = True + with pytest.raises(ConnectionError): + await c.process_audio_chunk(session_id, one_second_chunk, "dying") + + assert await c._reconnect_session(session_id) + assert provider.sessions_started == 2 # fresh provider session, clock back at 0 + + for i in range(2): + await c.process_audio_chunk(session_id, one_second_chunk, f"post-{i}") + + starts = await _stored_word_starts(redis, session_id) + assert starts == sorted(starts), "word timeline must stay monotonic" + # 3s sent before the drop (the dying chunk never reached the provider), + # so the post-reconnect words must continue at +3s — not restart at 0. + assert starts[3] == pytest.approx(3.0) + assert starts[4] == pytest.approx(4.0) + + +async def test_clock_survives_consumer_restart_via_session_store(consumer, monkeypatch): + """The offset also resumes from the session hash when the consumer process + itself restarts (in-memory counter lost).""" + c, provider, redis = consumer + session_id = "a421c9-havpe-test2" + one_second_chunk = b"\x00" * BYTES_PER_SECOND + + await c.start_session_stream(session_id, sample_rate=SAMPLE_RATE) + for i in range(3): + await c.process_audio_chunk(session_id, one_second_chunk, f"pre-{i}") + # end_session_stream persists the clock to the session hash + await c.end_session_stream(session_id) + assert session_id not in c._session_audio_seconds + + # New consumer instance (same Redis) — e.g. workers container restart + monkeypatch.setattr(sc_module, "get_transcription_provider", lambda mode: provider) + c2 = StreamingTranscriptionConsumer(redis_client=redis) + await c2.start_session_stream(session_id, sample_rate=SAMPLE_RATE) + await c2.process_audio_chunk(session_id, one_second_chunk, "post-0") + + starts = await _stored_word_starts(redis, session_id) + assert starts[-1] == pytest.approx(3.0) diff --git a/backends/advanced/tests/test_streaming_provider_health.py b/backends/advanced/tests/test_streaming_provider_health.py new file mode 100644 index 00000000..21278a01 --- /dev/null +++ b/backends/advanced/tests/test_streaming_provider_health.py @@ -0,0 +1,37 @@ +"""Transport-health regression tests for registry streaming transcription.""" + +from types import SimpleNamespace + +import pytest + +from advanced_omi_backend.services.transcription import ( + RegistryStreamingTranscriptionProvider, +) + +pytestmark = pytest.mark.unit + + +class FailingReceiveWebSocket: + async def send(self, _payload): + return None + + async def recv(self): + raise OSError("provider transport closed") + + +async def test_process_audio_chunk_propagates_transport_failure(): + provider = RegistryStreamingTranscriptionProvider.__new__( + RegistryStreamingTranscriptionProvider + ) + provider.model = SimpleNamespace(operations={"expect": {}}) + provider._streams = { + "session-1": { + "ws": FailingReceiveWebSocket(), + "sample_rate": 16000, + "final": None, + "interim": [], + } + } + + with pytest.raises(OSError, match="provider transport closed"): + await provider.process_audio_chunk("session-1", b"\x00\x00") diff --git a/backends/advanced/tests/test_transcript_slicing.py b/backends/advanced/tests/test_transcript_slicing.py new file mode 100644 index 00000000..2c77d621 --- /dev/null +++ b/backends/advanced/tests/test_transcript_slicing.py @@ -0,0 +1,112 @@ +"""Unit tests for transcript slicing/shifting used by conversation split/merge.""" + +from advanced_omi_backend.models.conversation import Conversation +from advanced_omi_backend.utils.transcript_slicing import ( + build_transcript_text, + shift_segments, + shift_words, + slice_segments, + slice_words, +) + + +def word(text, start, end): + return Conversation.Word(word=text, start=start, end=end) + + +def segment(text, start, end, speaker="Speaker 1", segment_type="speech", words=None): + return Conversation.SpeakerSegment( + start=start, + end=end, + text=text, + speaker=speaker, + segment_type=segment_type, + words=words or [], + ) + + +class TestSliceWords: + def test_membership_and_shift(self): + words = [word("a", 5.0, 5.5), word("b", 10.0, 10.5), word("c", 20.0, 20.5)] + sliced = slice_words(words, 10.0, 20.0) + assert [w.word for w in sliced] == ["b"] + assert sliced[0].start == 0.0 + assert sliced[0].end == 0.5 + + def test_end_clamped_to_range(self): + words = [word("a", 19.5, 21.0)] + sliced = slice_words(words, 10.0, 20.0) + assert sliced[0].start == 9.5 + assert sliced[0].end == 10.0 + + +class TestSliceSegments: + def test_midpoint_membership(self): + segments = [ + segment("before", 0.0, 8.0), + segment("inside", 12.0, 18.0), + segment("after", 22.0, 30.0), + ] + sliced = slice_segments(segments, 10.0, 20.0) + assert [s.text for s in sliced] == ["inside"] + assert sliced[0].start == 2.0 + assert sliced[0].end == 8.0 + + def test_clamps_overhanging_segment(self): + # Midpoint at 15 → included; start before range start gets clamped + segments = [segment("overhang", 8.0, 22.0)] + sliced = slice_segments(segments, 10.0, 20.0) + assert sliced[0].start == 0.0 + assert sliced[0].end == 10.0 + + def test_nested_words_sliced_and_shifted(self): + words = [word("in", 12.0, 12.5), word("out", 25.0, 25.5)] + segments = [segment("seg", 11.0, 19.0, words=words)] + sliced = slice_segments(segments, 10.0, 20.0) + assert len(sliced[0].words) == 1 + assert sliced[0].words[0].word == "in" + assert sliced[0].words[0].start == 2.0 + + def test_speaker_labels_preserved(self): + segments = [segment("hello", 12.0, 14.0, speaker="Speaker 2")] + sliced = slice_segments(segments, 10.0, 20.0) + assert sliced[0].speaker == "Speaker 2" + + +class TestShift: + def test_shift_segments_with_words(self): + segments = [segment("s", 0.0, 5.0, words=[word("w", 1.0, 1.5)])] + shifted = shift_segments(segments, 100.0) + assert shifted[0].start == 100.0 + assert shifted[0].end == 105.0 + assert shifted[0].words[0].start == 101.0 + + def test_shift_words(self): + shifted = shift_words([word("w", 1.0, 1.5)], 10.0) + assert shifted[0].start == 11.0 + assert shifted[0].end == 11.5 + + def test_originals_not_mutated(self): + seg = segment("s", 0.0, 5.0) + shift_segments([seg], 100.0) + assert seg.start == 0.0 + + +class TestBuildTranscriptText: + def test_joins_speech_skips_notes_and_events(self): + segments = [ + segment("Hello there.", 0.0, 2.0), + segment( + "[merged: 20 min gap elided]", + 2.0, + 2.0, + speaker="system", + segment_type="note", + ), + segment("[laughter]", 3.0, 4.0, segment_type="event"), + segment("Bye.", 5.0, 6.0), + ] + assert build_transcript_text(segments) == "Hello there. Bye." + + def test_empty(self): + assert build_transcript_text([]) == "" diff --git a/backends/advanced/tests/test_vault_rename_person.py b/backends/advanced/tests/test_vault_rename_person.py new file mode 100644 index 00000000..f63c19eb --- /dev/null +++ b/backends/advanced/tests/test_vault_rename_person.py @@ -0,0 +1,115 @@ +"""Unit tests for ``VaultTools.rename_person``. + +Regression cover for the audit blind spot and the lossy merge that made a renamed +note vanish from the ledger (it re-appeared later as an unexplained ``create``): + +- a rename/merge must be *recorded* in ``VaultTools.removed`` so the provider can + emit a ``rename`` audit-ledger entry, and +- a merge must migrate the retiring note's facts into the target *before* deleting + it — never rely on a follow-up ``edit_note`` that may never come. +""" + +import pytest + +from advanced_omi_backend.services.memory.agent.vault_tools import ( + VaultToolError, + VaultTools, +) + +PERSON_TEMPLATE = """--- +categories: + - "[[People]]" +created: 2026-06-13 +updated: 2026-06-15 +--- +## About +{about} + +## Conversations +![[Conversations.base#Person]] + +## Mentions +{mentions} +""" + + +def _person(about: str, mentions: str) -> str: + return PERSON_TEMPLATE.format(about=about, mentions=mentions) + + +def test_unknown_speaker_person_note_is_rejected(tmp_path): + tools = VaultTools(tmp_path) + with pytest.raises(VaultToolError, match="diarization placeholders"): + tools.write_note("People/Unknown Speaker 4.md", _person("- x", "- y")) + + +def test_hermes_person_note_is_rejected(tmp_path): + tools = VaultTools(tmp_path) + with pytest.raises(VaultToolError, match="Hermes assistant"): + tools.write_note("People/Hermes.md", _person("- x", "- y")) + + +class TestRenamePersonMerge: + def test_merge_migrates_facts_before_deleting_old(self, tmp_path): + tools = VaultTools(tmp_path) + (tmp_path / "People").mkdir(parents=True, exist_ok=True) + (tmp_path / "People" / "kenneth.md").write_text( + _person( + about="- Talked about API usage.", + mentions="- 2026-06-13 — commented on non-technical people.", + ), + encoding="utf-8", + ) + (tmp_path / "People" / "Naren.md").write_text( + _person(about="- Interested in agents.", mentions="- 2026-06-29 — spoke."), + encoding="utf-8", + ) + + msg = tools.rename_person("kenneth", "Naren") + + # Old note is gone; target absorbed the retiring note's facts. + assert not (tmp_path / "People" / "kenneth.md").exists() + naren = (tmp_path / "People" / "Naren.md").read_text(encoding="utf-8") + assert "API usage" in naren + assert "non-technical people" in naren + # No facts were dropped and no orphan section was needed (both headings exist). + assert "## Merged from" not in naren + assert "migrated 2 fact bullet(s)" in msg + + def test_merge_records_removal_for_audit(self, tmp_path): + tools = VaultTools(tmp_path) + (tmp_path / "People").mkdir(parents=True, exist_ok=True) + (tmp_path / "People" / "kenneth.md").write_text( + _person(about="- x.", mentions="- y."), encoding="utf-8" + ) + (tmp_path / "People" / "Naren.md").write_text( + _person(about="- a.", mentions="- b."), encoding="utf-8" + ) + + tools.rename_person("kenneth", "Naren") + + assert len(tools.removed) == 1 + rec = tools.removed[0] + assert rec["old_path"] == "People/kenneth.md" + assert rec["new_path"] == "People/Naren.md" + assert "- x." in rec["before"] # pre-removal content captured for the ledger + # The vanished path must not linger in `touched` (nothing to re-read there). + assert "People/kenneth.md" not in tools.touched + assert "People/Naren.md" in tools.touched + + +class TestRenamePersonMove: + def test_plain_rename_records_removal(self, tmp_path): + tools = VaultTools(tmp_path) + (tmp_path / "People").mkdir(parents=True, exist_ok=True) + (tmp_path / "People" / "kenneth.md").write_text( + _person(about="- fact.", mentions="- m."), encoding="utf-8" + ) + + tools.rename_person("kenneth", "Kenneth Long") + + assert not (tmp_path / "People" / "kenneth.md").exists() + assert (tmp_path / "People" / "Kenneth Long.md").exists() + assert [r["old_path"] for r in tools.removed] == ["People/kenneth.md"] + assert tools.removed[0]["new_path"] == "People/Kenneth Long.md" + assert "People/kenneth.md" not in tools.touched diff --git a/backends/advanced/tests/test_vault_scaffold.py b/backends/advanced/tests/test_vault_scaffold.py new file mode 100644 index 00000000..876533d2 --- /dev/null +++ b/backends/advanced/tests/test_vault_scaffold.py @@ -0,0 +1,97 @@ +"""Unit tests for the Kepano-style vault scaffold. + +Covers the spine layout (templates under Templates/, bases under Templates/Bases/, hubs at +root), the enumeration guard that keeps scaffolding out of captured memories, and organic +category creation. +""" + +from advanced_omi_backend.services.memory.vault_scaffold import ( + BASES_DIR, + TEMPLATES_DIR, + build_category_files, + is_scaffold_note, + seed_vault_scaffold, + write_category, +) +from advanced_omi_backend.services.memory.vault_templates import SPINE_TEMPLATES + + +class TestSeedVaultScaffold: + def test_spine_layout(self, tmp_path): + created = set(seed_vault_scaffold(tmp_path)) + # Templates live under Templates/, bases under Templates/Bases/, hubs at root. + for name in SPINE_TEMPLATES: + assert f"{TEMPLATES_DIR}/{name}" in created + for base in ("People.base", "Conversations.base", "Topics.base"): + assert f"{BASES_DIR}/{base}" in created + assert (tmp_path / BASES_DIR / base).exists() + for hub in ("People.md", "Conversations.md", "Topics.md"): + assert hub in created + assert (tmp_path / hub).exists() + + def test_idempotent(self, tmp_path): + seed_vault_scaffold(tmp_path) + # A user edit must not be clobbered on re-seed. + hub = tmp_path / "People.md" + hub.write_text("my own notes", encoding="utf-8") + assert seed_vault_scaffold(tmp_path) == [] + assert hub.read_text(encoding="utf-8") == "my own notes" + + def test_person_template_embeds_base_view(self, tmp_path): + seed_vault_scaffold(tmp_path) + person = (tmp_path / TEMPLATES_DIR / "Person Template.md").read_text( + encoding="utf-8" + ) + assert "![[Conversations.base#Person]]" in person + assert 'categories:\n - "[[People]]"' in person + + +class TestIsScaffoldNote: + def test_excludes_templates_and_hubs_keeps_real_notes(self, tmp_path): + seed_vault_scaffold(tmp_path) + (tmp_path / "People").mkdir(exist_ok=True) + real = tmp_path / "People" / "Alice.md" + real.write_text("x", encoding="utf-8") + + assert is_scaffold_note(tmp_path / "People.md", tmp_path) is True + assert ( + is_scaffold_note(tmp_path / TEMPLATES_DIR / "Person Template.md", tmp_path) + is True + ) + assert is_scaffold_note(real, tmp_path) is False + + enumerated = [ + p.relative_to(tmp_path).as_posix() + for p in tmp_path.rglob("*.md") + if not is_scaffold_note(p, tmp_path) + ] + assert enumerated == ["People/Alice.md"] + + +class TestOrganicCategory: + def test_build_category_files_shape(self): + files = build_category_files("Places", ["location", "type"]) + assert set(files) == { + f"{TEMPLATES_DIR}/Places Template.md", + f"{BASES_DIR}/Places.base", + "Places.md", + } + template = files[f"{TEMPLATES_DIR}/Places Template.md"] + assert 'categories:\n - "[[Places]]"' in template + assert "location:" in template and "type:" in template + assert ( + 'categories.contains(link("Places"))' in files[f"{BASES_DIR}/Places.base"] + ) + + def test_invalid_property_names_dropped(self): + # Only lowercase identifier-style keys are emitted (no spaces/uppercase/wikilinks). + template = build_category_files("Books", ["author", "Bad Key", "[[x]]"])[ + f"{TEMPLATES_DIR}/Books Template.md" + ] + assert "author:" in template + assert "Bad Key" not in template and "[[x]]" not in template + + def test_write_category_idempotent(self, tmp_path): + first = write_category(tmp_path, "Projects", ["status"]) + assert f"{TEMPLATES_DIR}/Projects Template.md" in first + assert write_category(tmp_path, "Projects", ["status"]) == [] diff --git a/backends/advanced/tests/test_worker_fleet_health.py b/backends/advanced/tests/test_worker_fleet_health.py new file mode 100644 index 00000000..3cbfa2aa --- /dev/null +++ b/backends/advanced/tests/test_worker_fleet_health.py @@ -0,0 +1,175 @@ +import json +from datetime import datetime, timezone +from types import SimpleNamespace + +import pytest + +from advanced_omi_backend.heartbeat import ( + FLEET_HEALTH_KEY, + evaluate_fleet_health, + is_rq_worker_fresh, +) +from advanced_omi_backend.routers.modules import health_routes +from advanced_omi_backend.services.observability import health_poller +from advanced_omi_backend.workers.orchestrator.health_monitor import HealthMonitor + + +class FakeAsyncRedis: + def __init__(self, values=None): + self.values = dict(values or {}) + self.hashes = {} + + async def get(self, key): + return self.values.get(key) + + async def hget(self, key, field): + return self.hashes.get(key, {}).get(field) + + async def hset(self, key, field, value): + self.hashes.setdefault(key, {})[field] = value + + +class FakeMongoAdmin: + async def command(self, _command): + return {"ok": 1} + + +class FakeMongoClient: + admin = FakeMongoAdmin() + + +class FakeSyncRedis: + def __init__(self, heartbeat=None): + self.heartbeat = heartbeat + + def ping(self): + return True + + def get(self, key): + assert key == FLEET_HEALTH_KEY + return self.heartbeat + + +class FakeFleetRedis: + def __init__(self): + self.writes = [] + + def set(self, key, value, ex): + self.writes.append((key, json.loads(value), ex)) + + +class FakeProcessManager: + def get_status(self): + return { + "rq-1": {"is_alive": True}, + "streaming-stt": {"is_alive": True}, + "audio-persistence": {"is_alive": False}, + } + + +def test_evaluate_worker_fleet_rejects_missing_and_stale_heartbeats(): + missing = evaluate_fleet_health(None, now=1_000) + assert missing["healthy"] is False + assert missing["status"] == "missing" + + stale = evaluate_fleet_health( + json.dumps({"status": "healthy", "timestamp": 900, "workers_total": 12}), + now=1_000, + max_age_seconds=30, + ) + assert stale["healthy"] is False + assert stale["status"] == "stale" + assert stale["age_seconds"] == pytest.approx(100) + + +def test_evaluate_worker_fleet_accepts_fresh_healthy_heartbeat(): + result = evaluate_fleet_health( + json.dumps({"status": "healthy", "timestamp": 995, "workers_total": 12}), + now=1_000, + max_age_seconds=30, + ) + assert result["healthy"] is True + assert result["status"] == "healthy" + assert result["workers_total"] == 12 + + +def test_rq_worker_freshness_rejects_stale_redis_registrations(): + fresh = SimpleNamespace( + last_heartbeat=datetime.fromtimestamp(990, tz=timezone.utc), worker_ttl=60 + ) + stale = SimpleNamespace( + last_heartbeat=datetime.fromtimestamp(900, tz=timezone.utc), worker_ttl=60 + ) + + assert is_rq_worker_fresh(fresh, now=1_000) is True + assert is_rq_worker_fresh(stale, now=1_000) is False + + +def test_health_monitor_publishes_actual_child_process_counts(monkeypatch): + redis = FakeFleetRedis() + monitor = HealthMonitor(FakeProcessManager(), SimpleNamespace(), redis) + monkeypatch.setattr( + "advanced_omi_backend.workers.orchestrator.health_monitor.time.time", + lambda: 1_000, + ) + + monitor._publish_fleet_health("unhealthy", "one child is down") + + key, payload, ttl = redis.writes[0] + assert key == FLEET_HEALTH_KEY + assert ttl > 30 + assert payload == { + "status": "unhealthy", + "timestamp": 1_000, + "workers_total": 3, + "workers_alive": 2, + "detail": "one child is down", + } + + +@pytest.mark.asyncio +async def test_worker_fleet_poller_records_outage_once_then_recovery(monkeypatch): + redis = FakeAsyncRedis() + recorded = [] + + async def capture_event(**event): + recorded.append(event) + + monkeypatch.setattr(health_poller, "record_event", capture_event) + + await health_poller._poll_worker_fleet(redis, now=1_000) + await health_poller._poll_worker_fleet(redis, now=1_015) + + assert len(recorded) == 1 + assert recorded[0]["severity"] == "critical" + assert recorded[0]["category"] == "service" + assert recorded[0]["source"] == "workers" + assert recorded[0]["metadata"]["health"] == "missing" + + redis.values[FLEET_HEALTH_KEY] = json.dumps( + {"status": "healthy", "timestamp": 1_020, "workers_total": 12} + ) + await health_poller._poll_worker_fleet(redis, now=1_025) + + assert len(recorded) == 2 + assert recorded[1]["severity"] == "info" + assert "recovered" in recorded[1]["title"].lower() + + +@pytest.mark.asyncio +async def test_readiness_requires_fresh_worker_fleet_heartbeat(monkeypatch): + monkeypatch.setattr(health_routes, "mongo_client", FakeMongoClient()) + + monkeypatch.setattr(health_routes, "redis_conn", FakeSyncRedis()) + unavailable = await health_routes.readiness_check() + assert unavailable.status_code == 503 + assert json.loads(unavailable.body)["status"] == "not_ready" + + heartbeat = json.dumps( + {"status": "healthy", "timestamp": 1_000, "workers_total": 12} + ) + monkeypatch.setattr(health_routes.time, "time", lambda: 1_005) + monkeypatch.setattr(health_routes, "redis_conn", FakeSyncRedis(heartbeat)) + ready = await health_routes.readiness_check() + assert ready.status_code == 200 + assert json.loads(ready.body)["status"] == "ready" diff --git a/backends/advanced/webui/src/App.tsx b/backends/advanced/webui/src/App.tsx index 5bc0272c..d4759e58 100644 --- a/backends/advanced/webui/src/App.tsx +++ b/backends/advanced/webui/src/App.tsx @@ -23,6 +23,7 @@ const Plugins = lazy(() => import('./pages/Plugins')) const Finetuning = lazy(() => import('./pages/Finetuning')) const Network = lazy(() => import('./pages/Network')) const DataAudit = lazy(() => import('./pages/DataAudit')) +const SpeakerEnrollment = lazy(() => import('./pages/SpeakerEnrollment')) const WakeWordLab = lazy(() => import('./pages/WakeWordLab')) const MemoryLedger = lazy(() => import('./pages/MemoryLedger')) const SystemEvents = lazy(() => import('./pages/SystemEvents')) @@ -177,6 +178,13 @@ function App() { </Suspense> </PageErrorBoundary> } /> + <Route path="speaker-enrollment" element={ + <PageErrorBoundary> + <Suspense fallback={<PageSkeleton />}> + <SpeakerEnrollment /> + </Suspense> + </PageErrorBoundary> + } /> <Route path="wakeword-lab" element={ <PageErrorBoundary> <Suspense fallback={<PageSkeleton />}> diff --git a/backends/advanced/webui/src/components/MemoryAuditCard.tsx b/backends/advanced/webui/src/components/MemoryAuditCard.tsx index b312fe24..96fa31f0 100644 --- a/backends/advanced/webui/src/components/MemoryAuditCard.tsx +++ b/backends/advanced/webui/src/components/MemoryAuditCard.tsx @@ -1,5 +1,5 @@ import { useEffect, useState } from 'react' -import { Loader2 } from 'lucide-react' +import { ArrowUp, Loader2 } from 'lucide-react' import { conversationsApi } from '../services/api' interface MemoryAuditEntry { @@ -55,10 +55,15 @@ export default function MemoryAuditCard({ conversationId }: { conversationId: st if (!loading && !error && entries.length === 0) return null return ( - <div className="bg-gray-50 dark:bg-gray-800/50 border border-gray-200 dark:border-gray-700 rounded-lg p-4"> - <h3 className="text-sm font-semibold text-gray-500 dark:text-gray-400 uppercase mb-3"> - Memory History - </h3> + <div id="memory-history" className="bg-gray-50 dark:bg-gray-800/50 border border-gray-200 dark:border-gray-700 rounded-lg p-4 scroll-mt-6"> + <div className="mb-3 flex items-center justify-between gap-3"> + <h3 className="text-sm font-semibold text-gray-500 dark:text-gray-400 uppercase"> + Memory History + </h3> + <a href="#transcript" className="inline-flex items-center gap-1 text-xs font-medium text-blue-600 hover:underline dark:text-blue-400"> + <ArrowUp className="h-3.5 w-3.5" /> Back to transcript + </a> + </div> {loading && ( <div className="flex items-center space-x-2 text-sm text-gray-500 dark:text-gray-400"> @@ -70,7 +75,7 @@ export default function MemoryAuditCard({ conversationId }: { conversationId: st {error && <div className="text-sm text-red-600 dark:text-red-400">{error}</div>} {!loading && !error && ( - <ul className="space-y-2 text-sm"> + <ul className="grid grid-cols-1 xl:grid-cols-2 gap-x-8 gap-y-3 text-sm"> {entries.map((entry) => ( <li key={entry.id} className="flex items-start justify-between gap-3"> <div className="min-w-0"> @@ -83,7 +88,7 @@ export default function MemoryAuditCard({ conversationId }: { conversationId: st > {entry.operation} </span> - <span className="truncate text-gray-900 dark:text-gray-100"> + <span className="text-gray-900 dark:text-gray-100 break-all"> {entry.note_path || '(whole vault)'} </span> </div> diff --git a/backends/advanced/webui/src/components/SpeakerInlineInput.tsx b/backends/advanced/webui/src/components/SpeakerInlineInput.tsx index f7fe6389..71dc3541 100644 --- a/backends/advanced/webui/src/components/SpeakerInlineInput.tsx +++ b/backends/advanced/webui/src/components/SpeakerInlineInput.tsx @@ -1,6 +1,7 @@ import { useState, useRef, useEffect } from 'react' -import { Plus } from 'lucide-react' +import { Plus, UserX } from 'lucide-react' import { useSortedSpeakers } from '../hooks/useSortedSpeakers' +import { nextUnknownSpeakerLabel } from '../utils/speakerLabels' interface SpeakerInlineInputProps { value: string @@ -9,6 +10,8 @@ interface SpeakerInlineInputProps { enrolledSpeakers: Array<{ speaker_id: string; name: string }> recentSpeakers?: string[] placeholder?: string + allowCreate?: boolean + usedSpeakerNames?: string[] } /** @@ -22,14 +25,21 @@ export default function SpeakerInlineInput({ enrolledSpeakers, recentSpeakers = [], placeholder = 'Type speaker name...', + allowCreate = true, + usedSpeakerNames = [], }: SpeakerInlineInputProps) { const [isFocused, setIsFocused] = useState(false) const containerRef = useRef<HTMLDivElement>(null) const sortedSpeakers = useSortedSpeakers(enrolledSpeakers, value, recentSpeakers) const hasResults = sortedSpeakers.recent.length > 0 || sortedSpeakers.rest.length > 0 - const showDropdown = isFocused && (hasResults || value.trim()) - const canCreate = value.trim() && !enrolledSpeakers.some(s => s.name.toLowerCase() === value.trim().toLowerCase()) + const unknownSpeakers = [nextUnknownSpeakerLabel(usedSpeakerNames)] + const matchingUnknowns = unknownSpeakers.filter((name) => + !value.trim() || name.toLowerCase().includes(value.trim().toLowerCase()) + ) + const showDropdown = isFocused && (hasResults || matchingUnknowns.length > 0 || value.trim()) + const isKnownUnknown = unknownSpeakers.some((name) => name.toLowerCase() === value.trim().toLowerCase()) + const canCreate = allowCreate && value.trim() && !isKnownUnknown && !enrolledSpeakers.some(s => s.name.toLowerCase() === value.trim().toLowerCase()) useEffect(() => { const handleClickOutside = (event: MouseEvent) => { @@ -68,8 +78,11 @@ export default function SpeakerInlineInput({ onKeyDown={(e) => { if (e.key === 'Enter' && value.trim()) { // If there's an exact match or results, select first; otherwise create + const exactUnknown = unknownSpeakers.find((name) => name.toLowerCase() === value.trim().toLowerCase()) const first = sortedSpeakers.recent[0] || sortedSpeakers.rest[0] - handleSelect(first ? first.name : value.trim()) + if (exactUnknown) handleSelect(exactUnknown) + else if (first) handleSelect(first.name) + else if (allowCreate) handleSelect(value.trim()) } }} placeholder={placeholder} @@ -78,6 +91,16 @@ export default function SpeakerInlineInput({ {showDropdown && ( <div className="absolute top-full left-0 mt-1 w-full min-w-[180px] bg-white dark:bg-gray-800 rounded-lg shadow-lg border border-gray-200 dark:border-gray-700 z-50 max-h-40 overflow-y-auto"> + {matchingUnknowns.map((name) => ( + <button + key={name} + onMouseDown={(e) => { e.preventDefault(); handleSelect(name) }} + className="w-full text-left px-3 py-1.5 text-xs hover:bg-gray-100 dark:hover:bg-gray-700 text-gray-500 dark:text-gray-400 flex items-center gap-1.5" + > + <UserX className="h-3 w-3" /> {name} + </button> + ))} + {matchingUnknowns.length > 0 && hasResults && <div className="border-t border-gray-100 dark:border-gray-700" />} {sortedSpeakers.recent.length > 0 && ( <> <div className="px-3 py-0.5 text-[10px] text-gray-400 dark:text-gray-500 uppercase tracking-wider bg-gray-50 dark:bg-gray-800/50"> diff --git a/backends/advanced/webui/src/components/SpeakerNameDropdown.tsx b/backends/advanced/webui/src/components/SpeakerNameDropdown.tsx index 9380531c..7b46a159 100644 --- a/backends/advanced/webui/src/components/SpeakerNameDropdown.tsx +++ b/backends/advanced/webui/src/components/SpeakerNameDropdown.tsx @@ -1,8 +1,7 @@ import { useState, useRef, useEffect } from 'react' import { Check, Plus, UserX } from 'lucide-react' import { useSortedSpeakers } from '../hooks/useSortedSpeakers' - -const UNKNOWN_SPEAKER = 'Unknown Speaker' +import { nextUnknownSpeakerLabel } from '../utils/speakerLabels' interface SpeakerNameDropdownProps { currentSpeaker: string @@ -13,6 +12,7 @@ interface SpeakerNameDropdownProps { annotated?: boolean speakerColor?: string recentSpeakers?: string[] + usedSpeakerNames?: string[] } export default function SpeakerNameDropdown({ @@ -22,6 +22,7 @@ export default function SpeakerNameDropdown({ annotated = false, speakerColor = 'text-blue-700 dark:text-blue-300', recentSpeakers = [], + usedSpeakerNames = [], }: SpeakerNameDropdownProps) { const [isOpen, setIsOpen] = useState(false) const [searchQuery, setSearchQuery] = useState('') @@ -29,6 +30,10 @@ export default function SpeakerNameDropdown({ const sortedSpeakers = useSortedSpeakers(enrolledSpeakers, searchQuery, recentSpeakers) const hasResults = sortedSpeakers.recent.length > 0 || sortedSpeakers.rest.length > 0 + const unknownSpeakers = [nextUnknownSpeakerLabel(usedSpeakerNames)] + const matchingUnknowns = unknownSpeakers.filter((name) => + !searchQuery || name.toLowerCase().includes(searchQuery.toLowerCase()) + ) useEffect(() => { const handleClickOutside = (event: MouseEvent) => { @@ -74,7 +79,9 @@ export default function SpeakerNameDropdown({ </button> ) - const canCreate = searchQuery && !enrolledSpeakers.some(s => s.name.toLowerCase() === searchQuery.toLowerCase()) + const canCreate = searchQuery && + !unknownSpeakers.some((name) => name.toLowerCase() === searchQuery.trim().toLowerCase()) && + !enrolledSpeakers.some(s => s.name.toLowerCase() === searchQuery.toLowerCase()) return ( <div className="relative inline-block" ref={dropdownRef}> @@ -103,7 +110,10 @@ export default function SpeakerNameDropdown({ autoFocus onKeyDown={(e) => { if (e.key === 'Enter') { - if (hasResults) { + const exactUnknown = unknownSpeakers.find((name) => name.toLowerCase() === searchQuery.trim().toLowerCase()) + if (exactUnknown) { + handleSpeakerSelect(exactUnknown) + } else if (hasResults) { const first = sortedSpeakers.recent[0] || sortedSpeakers.rest[0] if (first) handleSpeakerSelect(first.name) } else if (searchQuery) { @@ -115,20 +125,21 @@ export default function SpeakerNameDropdown({ </div> {/* Unknown Speaker option */} - {(!searchQuery || UNKNOWN_SPEAKER.toLowerCase().includes(searchQuery.toLowerCase())) && ( + {matchingUnknowns.length > 0 && ( <div className="border-b border-gray-200 dark:border-gray-700"> - <button - onClick={() => handleSpeakerSelect(UNKNOWN_SPEAKER)} - className="w-full text-left px-4 py-2 text-sm hover:bg-gray-100 dark:hover:bg-gray-700 flex items-center justify-between" - > - <span className="flex items-center space-x-2 text-gray-500 dark:text-gray-400"> - <UserX className="h-4 w-4" /> - <span>{UNKNOWN_SPEAKER}</span> - </span> - {currentSpeaker === UNKNOWN_SPEAKER && ( - <Check className="h-4 w-4 text-green-600" /> - )} - </button> + {matchingUnknowns.map((name) => ( + <button + key={name} + onClick={() => handleSpeakerSelect(name)} + className="w-full text-left px-4 py-2 text-sm hover:bg-gray-100 dark:hover:bg-gray-700 flex items-center justify-between" + > + <span className="flex items-center space-x-2 text-gray-500 dark:text-gray-400"> + <UserX className="h-4 w-4" /> + <span>{name}</span> + </span> + {currentSpeaker === name && <Check className="h-4 w-4 text-green-600" />} + </button> + ))} </div> )} diff --git a/backends/advanced/webui/src/components/audio/PlayheadWaveform.tsx b/backends/advanced/webui/src/components/audio/PlayheadWaveform.tsx index 53738641..a9fef954 100644 --- a/backends/advanced/webui/src/components/audio/PlayheadWaveform.tsx +++ b/backends/advanced/webui/src/components/audio/PlayheadWaveform.tsx @@ -23,6 +23,8 @@ interface PlayheadWaveformProps { // belongs to this conversation. segmentMarker?: PlayerSegmentMarker | null hoverMarker?: { start: number; end: number } | null + coloredSegments?: { start: number; end: number; color: string; segmentIndex?: number; label?: string }[] + onSegmentClick?: (index: number) => void } export const PlayheadWaveform: React.FC<PlayheadWaveformProps> = ({ @@ -33,6 +35,8 @@ export const PlayheadWaveform: React.FC<PlayheadWaveformProps> = ({ segments, segmentMarker, hoverMarker, + coloredSegments, + onSegmentClick, }) => { const currentTime = usePlayheadTime(cid) @@ -58,6 +62,8 @@ export const PlayheadWaveform: React.FC<PlayheadWaveformProps> = ({ height={height} segments={segments} segmentMarkers={markers} + coloredSegments={coloredSegments} + onSegmentClick={onSegmentClick} /> ) } diff --git a/backends/advanced/webui/src/components/audio/WaveformDisplay.tsx b/backends/advanced/webui/src/components/audio/WaveformDisplay.tsx index f597aa45..17172420 100644 --- a/backends/advanced/webui/src/components/audio/WaveformDisplay.tsx +++ b/backends/advanced/webui/src/components/audio/WaveformDisplay.tsx @@ -1,4 +1,4 @@ -import React, { useEffect, useRef } from 'react'; +import React, { useEffect, useRef, useState } from 'react'; import { useWaveformData } from './useWaveformData'; // A transcript-segment band rendered on the waveform. @@ -19,6 +19,8 @@ interface WaveformDisplayProps { height?: number; // Canvas height in pixels (default: 100) segments?: { start: number; end: number }[]; // All transcript segments — drawn as faint base bands when >1 segmentMarkers?: SegmentMarker[]; // Transcript segment bands (playing/anchor/hover) + coloredSegments?: { start: number; end: number; color: string; segmentIndex?: number; label?: string }[]; + onSegmentClick?: (index: number) => void; } export const WaveformDisplay: React.FC<WaveformDisplayProps> = ({ @@ -29,9 +31,36 @@ export const WaveformDisplay: React.FC<WaveformDisplayProps> = ({ height = 100, segments, segmentMarkers, + coloredSegments, + onSegmentClick, }) => { const { data: waveformData, loading, error } = useWaveformData(conversationId); const canvasRef = useRef<HTMLCanvasElement>(null); + const [hoveredSegment, setHoveredSegment] = useState<number | null>(null); + + const segmentAtPointer = (clientX: number) => { + if (!canvasRef.current || !coloredSegments?.length || duration <= 0) return null; + const rect = canvasRef.current.getBoundingClientRect(); + const x = Math.max(0, Math.min(clientX - rect.left, rect.width)); + const time = (x / rect.width) * duration; + const hitSlopSeconds = (10 / rect.width) * duration; + const candidates = coloredSegments.map((segment, index) => { + const distance = time < segment.start + ? segment.start - time + : time > segment.end + ? time - segment.end + : 0; + return { segment, index, distance }; + }).filter(({ distance }) => distance <= hitSlopSeconds); + + candidates.sort((a, b) => { + if (a.distance !== b.distance) return a.distance - b.distance; + const aCenterDistance = Math.abs(time - ((a.segment.start + a.segment.end) / 2)); + const bCenterDistance = Math.abs(time - ((b.segment.start + b.segment.end) / 2)); + return aCenterDistance - bCenterDistance; + }); + return candidates[0] ?? null; + }; // Draw waveform when data changes useEffect(() => { @@ -53,6 +82,34 @@ export const WaveformDisplay: React.FC<WaveformDisplayProps> = ({ // Draw waveform bars drawWaveform(ctx, waveformData.samples, rect.width, height); + // Speaker annotation mode: color the waveform by speaker while keeping the + // waveform itself visible. These are deliberately rendered before selection + // markers and the playhead. + if (coloredSegments && duration > 0) { + coloredSegments.forEach((seg, index) => { + const x1 = (seg.start / duration) * rect.width; + const x2 = Math.max((seg.end / duration) * rect.width, x1 + 3); + ctx.fillStyle = `${seg.color}38`; + ctx.fillRect(x1, 0, x2 - x1, height); + ctx.strokeStyle = seg.color; + ctx.lineWidth = 1.5; + ctx.beginPath(); + ctx.moveTo(x1, 0); + ctx.lineTo(x1, height); + ctx.moveTo(x2, 0); + ctx.lineTo(x2, height); + ctx.stroke(); + + if (index === hoveredSegment) { + ctx.fillStyle = `${seg.color}55`; + ctx.fillRect(x1, 0, x2 - x1, height); + ctx.strokeStyle = '#111827'; + ctx.lineWidth = 2.5; + ctx.strokeRect(Math.max(1, x1), 1, Math.max(3, x2 - x1), height - 2); + } + }); + } + // Draw faint base bands for every transcript segment (only meaningful when there's >1) if (segments && segments.length > 1 && duration > 0) { segments.forEach(seg => @@ -70,7 +127,7 @@ export const WaveformDisplay: React.FC<WaveformDisplayProps> = ({ drawPlaybackIndicator(ctx, currentTime, duration, rect.width, height); } // eslint-disable-next-line react-hooks/exhaustive-deps - }, [waveformData, currentTime, duration, height, segments, JSON.stringify(segmentMarkers)]); + }, [waveformData, currentTime, duration, height, segments, hoveredSegment, JSON.stringify(segmentMarkers), JSON.stringify(coloredSegments)]); const drawWaveform = ( ctx: CanvasRenderingContext2D, @@ -163,6 +220,14 @@ export const WaveformDisplay: React.FC<WaveformDisplayProps> = ({ const seekProgress = x / rect.width; const seekTime = seekProgress * duration; + if (onSegmentClick && coloredSegments?.length) { + const hit = segmentAtPointer(e.clientX); + if (hit) { + onSegmentClick(hit.segment.segmentIndex ?? hit.index); + return; + } + } + console.log(`🎵 Waveform seek: clicked at ${x}px (${(seekProgress * 100).toFixed(1)}%) → ${seekTime.toFixed(2)}s`); onSeek(seekTime); @@ -197,9 +262,15 @@ export const WaveformDisplay: React.FC<WaveformDisplayProps> = ({ <canvas ref={canvasRef} onClick={handleClick} - className="w-full cursor-pointer hover:opacity-80 transition-opacity rounded" + onMouseMove={(event) => setHoveredSegment(segmentAtPointer(event.clientX)?.index ?? null)} + onMouseLeave={() => setHoveredSegment(null)} + className="w-full cursor-crosshair rounded focus:outline-none focus:ring-2 focus:ring-blue-500" style={{ height: `${height}px` }} - title="Click to seek to position" + title={onSegmentClick && hoveredSegment !== null + ? `${coloredSegments?.[hoveredSegment]?.label || 'Speaker segment'} · click to select` + : onSegmentClick + ? 'Hover a speaker segment, then click to select it' + : 'Click to seek to position'} /> ); }; diff --git a/backends/advanced/webui/src/components/audio/WaveformRegionEditor.tsx b/backends/advanced/webui/src/components/audio/WaveformRegionEditor.tsx index f21fa42b..fd8f18b7 100644 --- a/backends/advanced/webui/src/components/audio/WaveformRegionEditor.tsx +++ b/backends/advanced/webui/src/components/audio/WaveformRegionEditor.tsx @@ -27,6 +27,8 @@ interface WaveformRegionEditorProps { onCancel: () => void /** Audition the current region. */ onPlay?: (region: Region) => void + /** Place the playhead at a clicked time and start playback. Drag gestures do not fire this. */ + onSeekPlay?: (time: number, region: Region | null) => void /** * Picker mode: no own commit buttons — just a region selector. Fires `onChange` * whenever the region changes; an external form (the insert menu) does the commit. @@ -41,6 +43,8 @@ interface WaveformRegionEditorProps { */ commitMode?: 'self' | 'linked' height?: number + /** Absolute conversation time supplied by an external audio player. */ + playheadTime?: number | null } const EDGE_PX = 7 // hit zone for grabbing a region edge @@ -75,16 +79,19 @@ export const WaveformRegionEditor: React.FC<WaveformRegionEditorProps> = ({ addLabel = '+ New', onCancel, onPlay, + onSeekPlay, pickerMode = false, onChange, commitMode = 'self', height = 96, + playheadTime, }) => { const ownCommit = !pickerMode && commitMode === 'self' const { data, loading, error } = useWaveformData(conversationId) const canvasRef = useRef<HTMLCanvasElement>(null) // Live playback position (so the playhead moves while playing, even when zoomed in). - const currentTime = usePlayheadTime(conversationId) + const sharedCurrentTime = usePlayheadTime(conversationId) + const currentTime = playheadTime !== undefined ? playheadTime : sharedCurrentTime const playheadRef = useRef<number | null>(currentTime ?? null) // view = visible [t0,t1] window; region = current selection (null = free-select mode). @@ -114,8 +121,8 @@ export const WaveformRegionEditor: React.FC<WaveformRegionEditorProps> = ({ onChange?.(r) } - const dragRef = useRef<{ mode: DragMode; anchorT: number; offset: number; panT0: number; panX: number }>( - { mode: null, anchorT: 0, offset: 0, panT0: 0, panX: 0 } + const dragRef = useRef<{ mode: DragMode; anchorT: number; offset: number; panT0: number; panX: number; downX: number; moved: boolean }>( + { mode: null, anchorT: 0, offset: 0, panT0: 0, panX: 0, downX: 0, moved: false } ) // Smooth auto-zoom from the full clip down to the segment (or to the focus time, in @@ -279,6 +286,8 @@ export const WaveformRegionEditor: React.FC<WaveformRegionEditorProps> = ({ const t = xToTime(x, rect.width) const r = regionRef.current const d = dragRef.current + d.downX = x + d.moved = false if (r) { const xs = timeToX(r.start, rect.width) @@ -306,6 +315,7 @@ export const WaveformRegionEditor: React.FC<WaveformRegionEditorProps> = ({ const el = canvasRef.current! const rect = el.getBoundingClientRect() const x = clamp(e.clientX - rect.left, 0, rect.width) + if (Math.abs(x - d.downX) > 4) d.moved = true const t = clamp(xToTime(x, rect.width), 0, duration) const r = regionRef.current @@ -333,11 +343,16 @@ export const WaveformRegionEditor: React.FC<WaveformRegionEditorProps> = ({ const endDrag = (e: React.PointerEvent<HTMLCanvasElement>) => { const d = dragRef.current + const wasClick = !d.moved + const el = canvasRef.current + const rect = el?.getBoundingClientRect() + const clickTime = rect ? clamp(xToTime(clamp(e.clientX - rect.left, 0, rect.width), rect.width), 0, duration) : null if (d.mode === 'draw') { const r = regionRef.current if (r && r.end - r.start < MIN_REGION) setRegion(null) // treat as a stray click } d.mode = null + if (wasClick && clickTime !== null && onSeekPlay) onSeekPlay(clickTime, regionRef.current) try { canvasRef.current?.releasePointerCapture(e.pointerId) } catch { @@ -366,7 +381,7 @@ export const WaveformRegionEditor: React.FC<WaveformRegionEditorProps> = ({ onPointerCancel={endDrag} className="w-full rounded bg-white/60 dark:bg-gray-900/40 touch-none select-none" style={{ height, cursor }} - title="Scroll to zoom · drag the band to move · drag edges to resize · drag outside to pan" + title={`${onSeekPlay ? 'Click to play from that point · ' : ''}scroll to zoom · drag the band to move · drag edges to resize · drag outside to pan`} /> )} diff --git a/backends/advanced/webui/src/components/dataAudit/AuditFilterBar.tsx b/backends/advanced/webui/src/components/dataAudit/AuditFilterBar.tsx index 05fe962a..26bc6591 100644 --- a/backends/advanced/webui/src/components/dataAudit/AuditFilterBar.tsx +++ b/backends/advanced/webui/src/components/dataAudit/AuditFilterBar.tsx @@ -95,7 +95,12 @@ export default function AuditFilterBar({ title={`Edit ${def.label.toLowerCase()} filter`} > <Icon className="h-3.5 w-3.5" /> - <span>{def.chipLabel(value)}</span> + <span + className="max-w-48 truncate sm:max-w-72" + title={def.chipLabel(value)} + > + {def.chipLabel(value)} + </span> </button> <button onClick={() => { diff --git a/backends/advanced/webui/src/components/dataAudit/DriftPanel.tsx b/backends/advanced/webui/src/components/dataAudit/DriftPanel.tsx index c8e9c71d..26c7d3c7 100644 --- a/backends/advanced/webui/src/components/dataAudit/DriftPanel.tsx +++ b/backends/advanced/webui/src/components/dataAudit/DriftPanel.tsx @@ -1,6 +1,6 @@ import { useCallback, useState } from 'react' -import { ChevronDown, ChevronRight, Waypoints, RefreshCw } from 'lucide-react' -import { conversationsApi } from '../../services/api' +import { ChevronDown, ChevronRight, DatabaseZap, RefreshCw, Waypoints } from 'lucide-react' +import { conversationsApi, dataAuditApi } from '../../services/api' interface Transition { from: string | null; to: string | null; count: number } interface DriftConversation { @@ -18,6 +18,7 @@ interface DriftReport { no_centroid_data: number similarity_threshold: number | null } +interface BackfillResult { backfilled: number; skipped: number; failed: number } const lbl = (s: string | null) => s || 'Unknown' @@ -36,6 +37,9 @@ export default function DriftPanel() { const [error, setError] = useState<string | null>(null) const [busy, setBusy] = useState<string | null>(null) const [reprocessed, setReprocessed] = useState<Set<string>>(new Set()) + const [backfilling, setBackfilling] = useState(false) + const [backfillProgress, setBackfillProgress] = useState<string | null>(null) + const [backfillResult, setBackfillResult] = useState<BackfillResult | null>(null) const load = useCallback(async () => { setLoading(true) @@ -68,6 +72,39 @@ export default function DriftPanel() { } } + const backfill = async () => { + setBackfilling(true) + setBackfillProgress('Queueing cluster-embedding backfill…') + setBackfillResult(null) + setError(null) + try { + const queued = await conversationsApi.backfillDriftClusterEmbeddings() + const jobId = queued.data.job_id + while (true) { + const statusResponse = await dataAuditApi.getJobStatus(jobId) + const { status, batch_progress: progress } = statusResponse.data + setBackfillProgress( + progress?.message || (status === 'queued' ? 'Waiting for a worker…' : 'Backfilling…') + ) + if (status === 'finished') { + const resultResponse = await dataAuditApi.getJobResult<BackfillResult>(jobId) + setBackfillResult(resultResponse.data.result) + await load() + break + } + if (status === 'failed' || status === 'canceled' || status === 'stopped') { + throw new Error(`Backfill job ${status}`) + } + await new Promise((resolve) => window.setTimeout(resolve, 1500)) + } + } catch { + setError('Cluster-embedding backfill failed') + } finally { + setBackfilling(false) + setBackfillProgress(null) + } + } + return ( <div className="border border-gray-200 dark:border-gray-700 rounded-lg"> <button @@ -115,18 +152,39 @@ export default function DriftPanel() { {error && <p className="text-sm text-red-600">{error}</p>} {data && data.no_centroid_data > 0 && ( - <p className="text-[11px] text-amber-600 dark:text-amber-400"> - {data.no_centroid_data} conversations have no stored cluster embeddings yet and - can't be analyzed (run the cluster-embedding backfill to cover them). - </p> + <div className="flex flex-wrap items-center gap-2 text-[11px] text-amber-600 dark:text-amber-400"> + <span> + {data.no_centroid_data} conversations have no stored cluster embeddings yet and + can't be analyzed. + </span> + <button + onClick={backfill} + disabled={backfilling} + className="inline-flex items-center gap-1.5 rounded bg-amber-100 px-2.5 py-1 font-medium text-amber-800 hover:bg-amber-200 disabled:opacity-50 dark:bg-amber-900/40 dark:text-amber-200 dark:hover:bg-amber-900/60" + > + <DatabaseZap className={`h-3 w-3 ${backfilling ? 'animate-pulse' : ''}`} /> + {backfilling ? 'Backfilling…' : `Backfill ${data.no_centroid_data}`} + </button> + {backfillProgress && <span>{backfillProgress}</span>} + </div> )} - {data && data.total_drifted === 0 && !loading && ( - <p className="text-sm text-green-600 dark:text-green-400"> - No drift — every analyzable conversation still matches the current gallery. + {backfillResult && ( + <p className={`text-[11px] ${backfillResult.failed ? 'text-amber-600 dark:text-amber-400' : 'text-green-600 dark:text-green-400'}`}> + Backfill complete: {backfillResult.backfilled} added, {backfillResult.skipped} skipped, + {' '}{backfillResult.failed} failed. Drift analysis refreshed. </p> )} + {data && + data.total_drifted === 0 && + data.conversations_scanned > data.no_centroid_data && + !loading && ( + <p className="text-sm text-green-600 dark:text-green-400"> + No drift — every analyzable conversation still matches the current gallery. + </p> + )} + {data && data.total_drifted > 0 && ( <div className="overflow-x-auto"> <table className="min-w-full text-xs"> diff --git a/backends/advanced/webui/src/components/dataAudit/GuidedEnrollment.tsx b/backends/advanced/webui/src/components/dataAudit/GuidedEnrollment.tsx new file mode 100644 index 00000000..82b4774c --- /dev/null +++ b/backends/advanced/webui/src/components/dataAudit/GuidedEnrollment.tsx @@ -0,0 +1,1163 @@ +import { useCallback, useEffect, useRef, useState } from 'react' +import { + Check, + AudioLines, + Eraser, + Loader2, + Mic, + Users, + Pause, + Play, + RefreshCw, + Search, + Scissors, + SkipForward, + Trash2, + X, +} from 'lucide-react' +import { + BACKEND_URL, + dataAuditApi, + GuidedEnrollmentClip, + GuidedEnrollmentGalleryResponse, + GuidedEnrollmentSuggestResponse, + GuidedEnrollmentSession, + speakerApi, + SpeakerBenchmarkReport, + SpeakerGalleryBaseline, +} from '../../services/api' +import { getStorageKey } from '../../utils/storage' +import { formatClock } from './format' +import SpeakerInlineInput from '../SpeakerInlineInput' +import { Region, WaveformRegionEditor } from '../audio/WaveformRegionEditor' +import { useJobPolling } from '../../hooks/useJobPolling' + +type EnrolledSpeaker = { speaker_id: string; name: string } + +type Decision = { + kind: 'accept' | 'reject' | 'skip' | 'bad_clip' | 'multiple_speakers' | 'another_speaker' + actualSpeaker?: string +} + +function submissionErrorMessage(error: any): string { + const status = error?.response?.status + const data = error?.response?.data + const prefix = status ? `Submission failed (HTTP ${status})` : 'Submission failed' + + if (Array.isArray(data?.detail)) { + const details = data.detail + .slice(0, 3) + .map((item: any) => { + const location = Array.isArray(item.loc) ? item.loc.slice(1).join('.') : '' + return `${location ? `${location}: ` : ''}${item.msg || 'Invalid value'}` + }) + const remaining = data.detail.length - details.length + return `${prefix}: ${details.join('; ')}${remaining > 0 ? `; and ${remaining} more` : ''}` + } + if (typeof data?.detail === 'string') return `${prefix}: ${data.detail}` + if (typeof data?.error === 'string') return `${prefix}: ${data.error}` + if (typeof error?.message === 'string') return `${prefix}: ${error.message}` + return prefix +} + +function qualityDelta(before: any, after: any, novelty: number | null): string { + if (!before || !after) return '' + const cohesion = + before.median_self != null && after.median_self != null + ? `cohesion ${before.median_self.toFixed(3)} → ${after.median_self.toFixed(3)}` + : 'cohesion unavailable' + const outliers = `outliers ${before.n_flagged}/${before.n_clips} → ${after.n_flagged}/${after.n_clips}` + const coverage = novelty != null ? `accepted novelty ${(novelty * 100).toFixed(0)}%` : null + return ` · ${cohesion} · ${outliers}${coverage ? ` · ${coverage}` : ''}` +} + +function EnrollmentTrend({ sessions }: { sessions: GuidedEnrollmentSession[] }) { + const points = [...sessions] + .reverse() + .filter((session) => session.health_after?.median_self != null) + .map((session) => ({ + clips: session.health_after!.n_clips, + cohesion: session.health_after!.median_self!, + })) + if (points.length < 2) return null + + const width = 640 + const height = 150 + const pad = 28 + const minClips = Math.min(...points.map((point) => point.clips)) + const maxClips = Math.max(...points.map((point) => point.clips)) + const rawMin = Math.min(...points.map((point) => point.cohesion)) + const rawMax = Math.max(...points.map((point) => point.cohesion)) + const minCohesion = Math.max(0, rawMin - 0.015) + const maxCohesion = Math.min(1, rawMax + 0.015) + const x = (value: number) => + pad + ((value - minClips) / Math.max(1, maxClips - minClips)) * (width - pad * 2) + const y = (value: number) => + height - pad - ((value - minCohesion) / Math.max(0.001, maxCohesion - minCohesion)) * (height - pad * 2) + const line = points.map((point) => `${x(point.clips)},${y(point.cohesion)}`).join(' ') + + return ( + <div className="mb-4"> + <div className="flex items-baseline justify-between gap-3 mb-1"> + <h3 className="text-xs font-medium text-gray-700 dark:text-gray-200">Observed gallery cohesion</h3> + <span className="text-[11px] text-gray-500 dark:text-gray-400">higher is more internally consistent</span> + </div> + <svg viewBox={`0 0 ${width} ${height}`} className="w-full h-[150px] border border-gray-200 dark:border-gray-700 rounded bg-gray-50 dark:bg-gray-900/30" role="img" aria-label="Gallery cohesion by clip count"> + <line x1={pad} y1={height - pad} x2={width - pad} y2={height - pad} stroke="currentColor" className="text-gray-300 dark:text-gray-600" /> + <line x1={pad} y1={pad} x2={pad} y2={height - pad} stroke="currentColor" className="text-gray-300 dark:text-gray-600" /> + <polyline points={line} fill="none" stroke="#2563eb" strokeWidth="2.5" /> + {points.map((point) => ( + <g key={`${point.clips}:${point.cohesion}`}> + <circle cx={x(point.clips)} cy={y(point.cohesion)} r="4" fill="#2563eb" /> + <text x={x(point.clips)} y={y(point.cohesion) - 8} textAnchor="middle" fontSize="10" fill="currentColor">{point.cohesion.toFixed(3)}</text> + </g> + ))} + <text x={width / 2} y={height - 7} textAnchor="middle" fontSize="10" fill="currentColor">gallery clips</text> + <text x={pad} y={height - 12} textAnchor="middle" fontSize="10" fill="currentColor">{minClips}</text> + <text x={width - pad} y={height - 12} textAnchor="middle" fontSize="10" fill="currentColor">{maxClips}</text> + </svg> + </div> + ) +} + +function BenchmarkPanel({ speakerName }: { speakerName: string }) { + const [report, setReport] = useState<SpeakerBenchmarkReport | null>(null) + const [running, setRunning] = useState(false) + const [progress, setProgress] = useState('') + const [benchmarkError, setBenchmarkError] = useState<string | null>(null) + const [baseline, setBaseline] = useState<SpeakerGalleryBaseline | null>(null) + const { pollJob } = useJobPolling() + + const loadLatest = useCallback(async () => { + const response = await dataAuditApi.getLatestSpeakerBenchmark() + setReport(response.data.report) + }, []) + + useEffect(() => { + loadLatest().catch(() => setBenchmarkError('Failed to load the latest benchmark')) + dataAuditApi.getSpeakerGalleryBaseline().then((response) => setBaseline(response.data)).catch(() => undefined) + }, [loadLatest]) + + const run = async () => { + setRunning(true) + setBenchmarkError(null) + setProgress('Queued') + try { + const response = await dataAuditApi.runSpeakerBenchmark() + const status = await pollJob(response.data.job_id, (_status, batch) => { + setProgress(batch?.message || _status) + }) + if (status !== 'finished') throw new Error('Benchmark job failed; inspect Queue & Events for details') + await loadLatest() + setProgress('') + } catch (error: any) { + setBenchmarkError(submissionErrorMessage(error)) + } finally { + setRunning(false) + } + } + + const latest = report?.learning_curve[report.learning_curve.length - 1] + const speakerBaseline = baseline?.speakers.find( + (item) => item.name.toLocaleLowerCase() === speakerName.toLocaleLowerCase(), + ) + + return ( + <section className="space-y-3 border-t border-gray-200 dark:border-gray-700 pt-4"> + <div> + <h2 className="text-sm font-semibold text-gray-900 dark:text-gray-100">{speakerName} gallery baseline</h2> + {speakerBaseline ? ( + <div className="grid grid-cols-3 gap-3 mt-2 text-xs"> + <div><span className="text-gray-500">Clips</span><div className="font-semibold text-gray-900 dark:text-gray-100">{speakerBaseline.baseline.n_clips} → {speakerBaseline.current?.n_clips ?? '—'}</div></div> + <div><span className="text-gray-500">Cohesion</span><div className="font-semibold text-gray-900 dark:text-gray-100">{speakerBaseline.baseline.median_self?.toFixed(3) ?? '—'} → {speakerBaseline.current?.median_self?.toFixed(3) ?? '—'}</div></div> + <div><span className="text-gray-500">Outliers</span><div className="font-semibold text-gray-900 dark:text-gray-100">{speakerBaseline.baseline.n_flagged} → {speakerBaseline.current?.n_flagged ?? '—'}</div></div> + </div> + ) : ( + <p className="mt-1 text-xs text-gray-500 dark:text-gray-400">No pre-enhancement baseline is available for this speaker.</p> + )} + {speakerBaseline && baseline?.cutoff && ( + <p className="mt-1 text-[11px] text-gray-500 dark:text-gray-400">Baseline reconstructed from enrolled clips at {new Date(baseline.cutoff).toLocaleString()}</p> + )} + </div> + + <details className="border-t border-gray-200 dark:border-gray-700 pt-3"> + <summary className="cursor-pointer text-xs font-medium text-gray-700 dark:text-gray-200"> + Overall recognition benchmark{latest?.top1_accuracy_mean != null ? ` · Top-1 ${Math.round(latest.top1_accuracy_mean * 100)}%` : ''}{latest?.eer_mean != null ? ` · EER ${(latest.eer_mean * 100).toFixed(1)}%` : ''} + </summary> + <div className="space-y-3 pt-3"> + <div className="flex flex-wrap items-center justify-between gap-3"> + <p className="text-xs text-gray-500 dark:text-gray-400">Five folds grouped by conversation; cached embeddings; live galleries unchanged.</p> + <button onClick={run} disabled={running} className="inline-flex items-center gap-1.5 px-3 py-1.5 rounded bg-blue-600 text-white text-sm disabled:opacity-50"> + {running ? <Loader2 className="h-4 w-4 animate-spin" /> : <RefreshCw className="h-4 w-4" />} + {running ? 'Benchmarking' : report ? 'Run again' : 'Run benchmark'} + </button> + </div> + {progress && <p className="text-xs text-blue-700 dark:text-blue-300">{progress}</p>} + {benchmarkError && <p className="text-xs text-red-600 dark:text-red-400">{benchmarkError}</p>} + {report && ( + <> + <div className="grid grid-cols-2 sm:grid-cols-5 gap-2 text-xs"> + <div><span className="text-gray-500">Top-1</span><div className="font-semibold text-gray-900 dark:text-gray-100">{latest?.top1_accuracy_mean != null ? `${Math.round(latest.top1_accuracy_mean * 100)}%` : 'Insufficient data'}</div></div> + <div><span className="text-gray-500">Macro recall</span><div className="font-semibold text-gray-900 dark:text-gray-100">{latest?.macro_recall_mean != null ? `${Math.round(latest.macro_recall_mean * 100)}%` : '—'}</div></div> + <div><span className="text-gray-500">False accepts</span><div className="font-semibold text-gray-900 dark:text-gray-100">{latest?.false_accept_rate_mean != null ? `${(latest.false_accept_rate_mean * 100).toFixed(1)}%` : '—'}</div></div> + <div><span className="text-gray-500">EER</span><div className="font-semibold text-gray-900 dark:text-gray-100">{latest?.eer_mean != null ? `${(latest.eer_mean * 100).toFixed(1)}%` : '—'}</div></div> + <div><span className="text-gray-500">Dataset</span><div className="font-semibold text-gray-900 dark:text-gray-100">{report.dataset.embedded_clips} clips · {report.dataset.speakers} speakers</div></div> + </div> + <div className="space-y-1.5"> + {report.learning_curve.map((point) => ( + <div key={point.fraction} className="grid grid-cols-[3rem_1fr_3.5rem] items-center gap-2 text-xs"> + <span className="text-gray-500">{Math.round(point.fraction * 100)}%</span> + <div className="h-2 bg-gray-100 dark:bg-gray-700 rounded overflow-hidden"><div className="h-full bg-blue-600" style={{ width: `${(point.top1_accuracy_mean || 0) * 100}%` }} /></div> + <span className="text-right font-mono text-gray-700 dark:text-gray-200">{point.top1_accuracy_mean != null ? `${Math.round(point.top1_accuracy_mean * 100)}%` : '—'}</span> + </div> + ))} + </div> + <p className="text-[11px] text-gray-500 dark:text-gray-400">{report.protocol} · {report.conversation_groups} conversations · threshold {report.threshold.toFixed(2)} · {report.embedding_model} · {new Date(report.created_at).toLocaleString()}</p> + </> + )} + </div> + </details> + </section> + ) +} + +function flagBadge(clip: GuidedEnrollmentGalleryResponse['clips'][number]) { + if (clip.flags.includes('mislabel')) + return ( + <span className="text-[11px] px-1.5 py-0.5 rounded bg-red-100 dark:bg-red-900/40 text-red-700 dark:text-red-300"> + sounds like {clip.suggested?.name || 'another speaker'} + </span> + ) + if (clip.flags.includes('junk')) + return ( + <span className="text-[11px] px-1.5 py-0.5 rounded bg-red-100 dark:bg-red-900/40 text-red-700 dark:text-red-300"> + junk + </span> + ) + if (clip.flags.includes('weak')) + return ( + <span className="text-[11px] px-1.5 py-0.5 rounded bg-amber-100 dark:bg-amber-900/40 text-amber-700 dark:text-amber-300"> + weak match + </span> + ) + return null +} + +/** + * The clips currently backing a speaker's voiceprint. Play each by ear, remove + * bad ones (quarantined server-side, centroid recomputed), or clean the whole + * profile: forget every review decision, session snapshot, and corpus match + * recorded under the name — so after a mislabeled run the speaker can be + * re-enrolled from scratch and previously reviewed clips are suggested again. + */ +function GallerySection({ + speakerName, + refreshKey, + onReset, +}: { + speakerName: string + refreshKey: number + onReset: (galleryPurged: boolean, message: string) => void +}) { + const [gallery, setGallery] = useState<GuidedEnrollmentGalleryResponse | null>(null) + const [loading, setLoading] = useState(false) + const [error, setError] = useState<string | null>(null) + const [pendingDelete, setPendingDelete] = useState<number | null>(null) + const [deleting, setDeleting] = useState<number | null>(null) + const [playingId, setPlayingId] = useState<number | null>(null) + const audioRef = useRef<HTMLAudioElement | null>(null) + const [confirmingReset, setConfirmingReset] = useState(false) + const [purgeGallery, setPurgeGallery] = useState(false) + const [resetting, setResetting] = useState(false) + + const token = () => localStorage.getItem(getStorageKey('token')) || '' + + const load = useCallback(async () => { + setLoading(true) + setError(null) + setPendingDelete(null) + try { + const res = await dataAuditApi.getEnrollmentGallery(speakerName) + setGallery(res.data) + } catch (e: any) { + setGallery(null) + setError(e?.response?.data?.error || 'Failed to load enrolled clips') + } finally { + setLoading(false) + } + }, [speakerName]) + + useEffect(() => { + load() + }, [load, refreshKey]) + + const stop = useCallback(() => { + audioRef.current?.pause() + audioRef.current = null + setPlayingId(null) + }, []) + + useEffect(() => stop, [stop]) + + const play = (segmentId: number) => { + if (playingId === segmentId) { + stop() + return + } + stop() + const audio = new Audio( + `${BACKEND_URL}/api/data-audit/enrollment/guided/gallery/segments/${segmentId}/audio?token=${token()}` + ) + audioRef.current = audio + setPlayingId(segmentId) + audio.addEventListener('ended', () => stop()) + audio.addEventListener('error', () => stop()) + audio.play().catch(() => stop()) + } + + const removeClip = async (segmentId: number) => { + setDeleting(segmentId) + setError(null) + try { + if (playingId === segmentId) stop() + await dataAuditApi.deleteEnrollmentGalleryClip(speakerName, segmentId) + setPendingDelete(null) + await load() + } catch (e: any) { + setError(submissionErrorMessage(e)) + } finally { + setDeleting(null) + } + } + + const reset = async () => { + setResetting(true) + setError(null) + try { + const res = await dataAuditApi.resetGuidedEnrollment(speakerName, purgeGallery) + const d = res.data.deleted + setConfirmingReset(false) + setPurgeGallery(false) + onReset( + res.data.gallery_deleted, + `Cleaned ${speakerName}: forgot ${d.reviews} review decisions, ${d.sessions} sessions, ` + + `${d.discovery_matches} corpus matches` + + (res.data.gallery_deleted ? ' — voiceprint deleted from the speaker service' : '') + ) + if (!res.data.gallery_deleted) await load() + } catch (e: any) { + setError(submissionErrorMessage(e)) + } finally { + setResetting(false) + } + } + + return ( + <section className="space-y-3 border-t border-gray-200 dark:border-gray-700 pt-4"> + <div className="flex flex-wrap items-center justify-between gap-2"> + <h2 className="text-sm font-semibold text-gray-900 dark:text-gray-100"> + Enrolled clips + {gallery && ( + <span className="ml-2 font-normal text-xs text-gray-500 dark:text-gray-400"> + {gallery.clips.length} clips + {gallery.median_self != null && ` · cohesion ${gallery.median_self.toFixed(3)}`} + {gallery.verdict && ` · ${gallery.verdict}`} + </span> + )} + </h2> + <button + onClick={() => load()} + disabled={loading} + className="inline-flex items-center gap-1 text-xs px-2 py-1 rounded border border-gray-300 dark:border-gray-600 text-gray-600 dark:text-gray-300 disabled:opacity-50" + > + {loading ? <Loader2 className="h-3.5 w-3.5 animate-spin" /> : <RefreshCw className="h-3.5 w-3.5" />} + Reload + </button> + </div> + + {error && <p className="text-sm text-red-600">{error}</p>} + + {gallery && gallery.clips.length === 0 && !loading && ( + <p className="text-sm text-gray-500 dark:text-gray-400"> + No per-clip records for this voiceprint. Older enrollments may need the + segment backfill on the speaker service before they can be managed here. + </p> + )} + + {gallery && gallery.clips.length > 0 && ( + <ul className="space-y-1.5"> + {gallery.clips.map((clip) => ( + <li + key={clip.segment_id} + className="flex flex-wrap items-center gap-2 rounded border border-gray-200 dark:border-gray-700 px-3 py-2" + > + <button + onClick={() => play(clip.segment_id)} + className="shrink-0 p-1.5 rounded bg-gray-100 dark:bg-gray-700 text-gray-700 dark:text-gray-200" + title="Play clip" + > + {playingId === clip.segment_id ? ( + <Pause className="h-4 w-4" /> + ) : ( + <Play className="h-4 w-4" /> + )} + </button> + <span className="text-xs text-gray-700 dark:text-gray-200 truncate max-w-[16rem]" title={clip.filename}> + {clip.filename} + </span> + <span className="text-xs text-gray-500 dark:text-gray-400">{clip.duration.toFixed(1)}s</span> + {clip.self_score != null && ( + <span + className="text-xs font-mono text-gray-600 dark:text-gray-300" + title="Similarity to the rest of this speaker's gallery" + > + self {clip.self_score.toFixed(3)} + </span> + )} + {flagBadge(clip)} + <div className="ml-auto flex items-center gap-1.5"> + {pendingDelete === clip.segment_id ? ( + <> + <button + onClick={() => removeClip(clip.segment_id)} + disabled={deleting === clip.segment_id} + className="inline-flex items-center gap-1 text-xs px-2 py-1 rounded bg-red-600 text-white disabled:opacity-50" + > + {deleting === clip.segment_id && <Loader2 className="h-3.5 w-3.5 animate-spin" />} + Remove from voiceprint + </button> + <button + onClick={() => setPendingDelete(null)} + className="text-xs px-2 py-1 rounded border border-gray-300 dark:border-gray-600 text-gray-600 dark:text-gray-300" + > + Cancel + </button> + </> + ) : ( + <button + onClick={() => setPendingDelete(clip.segment_id)} + className="p-1.5 rounded bg-gray-100 dark:bg-gray-700 text-gray-500 dark:text-gray-400 hover:text-red-600 dark:hover:text-red-400" + title="Remove this clip from the voiceprint (recoverable — audio is quarantined)" + > + <Trash2 className="h-4 w-4" /> + </button> + )} + </div> + </li> + ))} + </ul> + )} + + <div className="border-t border-gray-200 dark:border-gray-700 pt-3"> + {!confirmingReset ? ( + <button + onClick={() => setConfirmingReset(true)} + className="inline-flex items-center gap-1.5 text-xs px-2.5 py-1.5 rounded border border-red-300 dark:border-red-800 text-red-700 dark:text-red-300" + > + <Eraser className="h-3.5 w-3.5" /> + Clean profile… + </button> + ) : ( + <div className="space-y-2 rounded border border-red-300 dark:border-red-800 bg-red-50 dark:bg-red-900/15 p-3"> + <p className="text-xs text-gray-700 dark:text-gray-200"> + Forget everything recorded for “{speakerName}”: review decisions, enrollment + session history, and corpus-discovery matches. Previously reviewed clips become + suggestible again — use this after a mislabeled run to start the profile over. + </p> + <label className="flex items-center gap-2 text-xs text-gray-700 dark:text-gray-200"> + <input + type="checkbox" + checked={purgeGallery} + onChange={(e) => setPurgeGallery(e.target.checked)} + /> + Also delete the voiceprint and its enrollment audio from the speaker service + </label> + <div className="flex items-center gap-2"> + <button + onClick={reset} + disabled={resetting} + className="inline-flex items-center gap-1.5 text-xs px-2.5 py-1.5 rounded bg-red-600 text-white disabled:opacity-50" + > + {resetting && <Loader2 className="h-3.5 w-3.5 animate-spin" />} + Clean profile + </button> + <button + onClick={() => { + setConfirmingReset(false) + setPurgeGallery(false) + }} + className="text-xs px-2.5 py-1.5 rounded border border-gray-300 dark:border-gray-600 text-gray-600 dark:text-gray-300" + > + Cancel + </button> + </div> + </div> + )} + </div> + </section> + ) +} + +/** + * Guided speaker enrollment. Pick an enrolled speaker and the backend scans the + * corpus for the clips whose confirmation would improve that speaker's + * voiceprint the most (new acoustic conditions for the gallery, matches near + * the decision boundary, longer clips — max 2 per conversation so enrollment + * spans sessions). Review each clip by ear, accept/reject, submit, repeat. + * Decisions are recorded server-side, so a decided clip is never re-suggested. + */ +export default function GuidedEnrollment() { + const [speakers, setSpeakers] = useState<EnrolledSpeaker[]>([]) + const [speaker, setSpeaker] = useState('') + const [speakerSearch, setSpeakerSearch] = useState('') + + const [suggestion, setSuggestion] = useState<GuidedEnrollmentSuggestResponse | null>(null) + const [decisions, setDecisions] = useState<Record<string, Decision>>({}) + const [loading, setLoading] = useState(false) + const [submitting, setSubmitting] = useState(false) + const [discovering, setDiscovering] = useState(false) + const [discoveryProgress, setDiscoveryProgress] = useState('') + const [error, setError] = useState<string | null>(null) + const [lastResult, setLastResult] = useState<string | null>(null) + const [history, setHistory] = useState<GuidedEnrollmentSession[]>([]) + const [galleryVersion, setGalleryVersion] = useState(0) + const [editingKey, setEditingKey] = useState<string | null>(null) + const [trimmedRegions, setTrimmedRegions] = useState<Record<string, Region>>({}) + + const [playingKey, setPlayingKey] = useState<string | null>(null) + const [playheadTime, setPlayheadTime] = useState<number | null>(null) + const [includeDeleted, setIncludeDeleted] = useState(false) + const [sortOrder, setSortOrder] = useState<'informative' | 'confidence'>('informative') + const [mining, setMining] = useState(false) + const audioRef = useRef<HTMLAudioElement | null>(null) + const playheadFrameRef = useRef<number | null>(null) + const miningInputRef = useRef<HTMLInputElement | null>(null) + const selectedSpeakerRef = useRef('') + const { pollJob } = useJobPolling() + + const token = () => localStorage.getItem(getStorageKey('token')) || '' + const clipKey = (c: GuidedEnrollmentClip) => `${c.conversation_id}:${c.start}` + + useEffect(() => { + if (speakers.length) return + speakerApi + .getEnrolledSpeakers() + .then((res) => { + const enrolled = (res.data.speakers || []) + .filter((s: any) => s.name) + .map((s: any) => ({ speaker_id: s.id || s.speaker_id || s.name, name: s.name })) + .sort((a: EnrolledSpeaker, b: EnrolledSpeaker) => a.name.localeCompare(b.name)) + setSpeakers(enrolled) + }) + .catch(() => setError('Failed to load enrolled speakers')) + }, [speakers.length]) + + const stopAudio = useCallback(() => { + if (playheadFrameRef.current != null) cancelAnimationFrame(playheadFrameRef.current) + playheadFrameRef.current = null + audioRef.current?.pause() + audioRef.current = null + setPlayingKey(null) + setPlayheadTime(null) + }, []) + + const trackPlayhead = (audio: HTMLAudioElement, absoluteStart: number) => { + const update = () => { + if (audioRef.current !== audio || audio.paused || audio.ended) return + setPlayheadTime(absoluteStart + audio.currentTime) + playheadFrameRef.current = requestAnimationFrame(update) + } + setPlayheadTime(absoluteStart) + playheadFrameRef.current = requestAnimationFrame(update) + } + + useEffect(() => stopAudio, [stopAudio]) + + const loadHistory = useCallback(async (name: string) => { + try { + const res = await dataAuditApi.guidedEnrollmentHistory(name) + setHistory(res.data.sessions) + } catch { + setHistory([]) + } + }, []) + + const playClip = (clip: GuidedEnrollmentClip) => { + const key = clipKey(clip) + if (playingKey === key) { + stopAudio() + return + } + stopAudio() + const region = trimmedRegions[key] || { start: clip.start, end: clip.end } + const url = + `${BACKEND_URL}/api/audio/chunks/${clip.conversation_id}` + + `?start_time=${region.start.toFixed(2)}&end_time=${region.end.toFixed(2)}&format=wav&token=${token()}` + const audio = new Audio(url) + audioRef.current = audio + setPlayingKey(key) + audio.addEventListener('ended', () => stopAudio()) + audio.addEventListener('error', () => stopAudio()) + audio.play().then(() => trackPlayhead(audio, region.start)).catch(() => stopAudio()) + } + + const playRegion = (clip: GuidedEnrollmentClip, region: Region) => { + const key = clipKey(clip) + setTrimmedRegions((current) => ({ ...current, [key]: region })) + stopAudio() + const url = + `${BACKEND_URL}/api/audio/chunks/${clip.conversation_id}` + + `?start_time=${region.start.toFixed(2)}&end_time=${region.end.toFixed(2)}&format=wav&token=${token()}` + const audio = new Audio(url) + audioRef.current = audio + setPlayingKey(key) + audio.addEventListener('ended', () => stopAudio()) + audio.addEventListener('error', () => stopAudio()) + audio.play().then(() => trackPlayhead(audio, region.start)).catch(() => stopAudio()) + } + + const suggest = useCallback( + async (name: string, order?: 'informative' | 'confidence') => { + stopAudio() + setLoading(true) + setError(null) + setSuggestion(null) + setDecisions({}) + setTrimmedRegions({}) + setEditingKey(null) + try { + const res = await dataAuditApi.guidedEnrollmentSuggest(name, order ?? sortOrder) + setSuggestion(res.data) + if (!res.data.batch.length) { + setLastResult( + res.data.discovery_indexed + ? 'The indexed corpus has no unreviewed clips that plausibly match this speaker.' + : res.data.pool_remaining > 0 + ? 'No clips passed the information gate this round — the remaining pool is likely other speakers or already covered.' + : 'No label-derived clips found. Search corpus speech to discover missed or incorrectly labeled clips.' + ) + } + } catch (e: any) { + setError(e?.response?.data?.error || 'Failed to fetch suggestions') + } finally { + setLoading(false) + } + }, + [stopAudio, sortOrder] + ) + + const submit = async () => { + if (!suggestion) return + const decided = suggestion.batch.filter((c) => decisions[clipKey(c)]) + if (!decided.length) return + setSubmitting(true) + setError(null) + try { + const res = await dataAuditApi.guidedEnrollmentDecide( + suggestion.speaker.speaker_name, + decided.map((c) => ({ + conversation_id: c.conversation_id, + start: (trimmedRegions[clipKey(c)] || c).start, + end: (trimmedRegions[clipKey(c)] || c).end, + original_start: c.start, + original_end: c.end, + decision: decisions[clipKey(c)].kind, + actual_speaker: decisions[clipKey(c)].actualSpeaker, + scores: c.scores, + })) + ) + const g = res.data.speaker + setLastResult( + `Enrolled ${res.data.enrolled}, rejected ${res.data.rejected}, bad clips ${res.data.bad_clips}, multiple speakers ${res.data.multiple_speakers}, skipped ${res.data.skipped}` + + (res.data.errors.length + ? `, ${res.data.errors.length} failed: ${res.data.errors + .slice(0, 2) + .map((item) => item.error) + .join('; ')}` + : '') + + (res.data.health_before && res.data.health_after + ? ` — gallery ${res.data.health_before.n_clips} → ${res.data.health_after.n_clips} clips` + : g?.n_clips != null + ? ` — gallery now ${g.n_clips} clips` + : '') + + qualityDelta( + res.data.health_before, + res.data.health_after, + res.data.coverage.accepted_novelty_mean + ) + + (res.data.benchmark_job_id ? ' · cross-validation queued' : '') + ) + // Every decided clip is now excluded server-side; fetch the next batch. + setGalleryVersion((v) => v + 1) + await suggest(suggestion.speaker.speaker_name) + await loadHistory(suggestion.speaker.speaker_name) + if (res.data.discovery_job_id) { + setDiscoveryProgress('Updating corpus matches for the new gallery') + void followDiscoveryJob( + suggestion.speaker.speaker_name, + res.data.discovery_job_id, + ) + } + } catch (e: any) { + setError(submissionErrorMessage(e)) + setSubmitting(false) + return + } + setSubmitting(false) + } + + const freshBatch = async () => { + if (!suggestion) return + setSubmitting(true) + setError(null) + try { + await dataAuditApi.guidedEnrollmentDecide( + suggestion.speaker.speaker_name, + suggestion.batch.map((c) => ({ + conversation_id: c.conversation_id, + start: (trimmedRegions[clipKey(c)] || c).start, + end: (trimmedRegions[clipKey(c)] || c).end, + original_start: c.start, + original_end: c.end, + decision: decisions[clipKey(c)]?.kind || 'skip', + actual_speaker: decisions[clipKey(c)]?.actualSpeaker, + scores: c.scores, + })) + ) + await suggest(suggestion.speaker.speaker_name) + } catch (e: any) { + setError(submissionErrorMessage(e)) + } finally { + setSubmitting(false) + } + } + + const followDiscoveryJob = async (name: string, jobId: string) => { + setDiscovering(true) + setError(null) + try { + const status = await pollJob(jobId, (_status, batch) => { + setDiscoveryProgress(batch?.message || _status) + }) + if (status !== 'finished') throw new Error('Corpus discovery failed; inspect Queue & Events for details') + setDiscoveryProgress('') + if (selectedSpeakerRef.current === name) await suggest(name) + } catch (e: any) { + setError(submissionErrorMessage(e)) + } finally { + setDiscovering(false) + } + } + + const discoverCorpus = async (name: string) => { + setDiscoveryProgress('Queued') + try { + const response = await dataAuditApi.discoverSpeakerCorpus(name, includeDeleted) + await followDiscoveryJob(name, response.data.job_id) + } catch (e: any) { + setError(submissionErrorMessage(e)) + setDiscovering(false) + } + } + + const mineFiles = async (name: string, files: File[]) => { + if (!files.length) return + setMining(true) + setError(null) + setDiscoveryProgress(`Uploading ${files.length} file(s) for mining…`) + try { + const res = await dataAuditApi.mineCorpusAudio(name, files) + const d = res.data + const failNote = d.failed.length ? `, ${d.failed.length} failed` : '' + if (!d.transcription_available) { + setError( + `Ingested ${d.ingested} file(s)${failNote}, but no batch transcription provider is ` + + 'configured — mined audio cannot be segmented until transcription runs.' + ) + setDiscoveryProgress('') + return + } + setLastResult( + `Mining ${d.ingested} file(s) for ${name}${failNote} — transcribing, then scanning for matches.` + ) + if (d.discovery_job_id) { + setDiscoveryProgress('Waiting for transcription, then scanning corpus') + void followDiscoveryJob(name, d.discovery_job_id) + } else { + setDiscoveryProgress('') + } + } catch (e: any) { + setError(submissionErrorMessage(e)) + setDiscoveryProgress('') + } finally { + setMining(false) + if (miningInputRef.current) miningInputRef.current.value = '' + } + } + + const attachOrCreateDiscovery = async (name: string) => { + try { + const response = await dataAuditApi.getSpeakerCorpusDiscovery(name) + const { job_id: jobId, status, matched_segments: matchedSegments } = response.data + if (jobId && ['queued', 'started', 'deferred', 'scheduled'].includes(status || '')) { + setDiscoveryProgress('Reattaching to corpus search') + await followDiscoveryJob(name, jobId) + } else if (matchedSegments === 0) { + await discoverCorpus(name) + } + } catch (e: any) { + setError(submissionErrorMessage(e)) + } + } + + const decidedCount = suggestion + ? suggestion.batch.filter((c) => decisions[clipKey(c)]).length + : 0 + + return ( + <div className="space-y-5"> + <section className="space-y-3"> + <div className="flex items-center gap-2"> + <Mic className="h-5 w-5 text-blue-600" /> + <h1 className="text-xl font-semibold text-gray-900 dark:text-gray-100">Speaker enrollment</h1> + </div> + <div className="max-w-md"> + <SpeakerInlineInput + value={speakerSearch} + onChange={setSpeakerSearch} + onSelect={(name) => { + selectedSpeakerRef.current = name + setSpeaker(name) + setSpeakerSearch(name) + suggest(name) + loadHistory(name) + attachOrCreateDiscovery(name) + }} + enrolledSpeakers={speakers} + placeholder="Search for a speaker to enhance..." + allowCreate={false} + /> + </div> + {!speaker && lastResult && ( + <p className="text-sm text-green-700 dark:text-green-400">{lastResult}</p> + )} + </section> + + {speaker && ( + <section className="space-y-3 border-t border-gray-200 dark:border-gray-700 pt-4"> + <div className="flex flex-wrap items-center gap-2"> + <button + onClick={() => speaker && suggest(speaker)} + disabled={!speaker || loading || submitting} + className="inline-flex items-center gap-1.5 text-sm px-3 py-1.5 rounded bg-blue-600 text-white disabled:opacity-50" + > + {loading ? ( + <Loader2 className="h-4 w-4 animate-spin" /> + ) : ( + <RefreshCw className="h-4 w-4" /> + )} + Fresh batch + </button> + <select + value={sortOrder} + onChange={(e) => { + const next = e.target.value as 'informative' | 'confidence' + setSortOrder(next) + if (speaker) suggest(speaker, next) + }} + disabled={loading || submitting} + className="text-sm px-2 py-1.5 rounded border border-gray-300 dark:border-gray-600 bg-white dark:bg-gray-800 text-gray-700 dark:text-gray-200 disabled:opacity-50" + title="Most informative surfaces clips the model learns most from (often uncertain, lower-similarity ones); best match first surfaces the highest-similarity clips" + > + <option value="informative">Sort: most informative</option> + <option value="confidence">Sort: best match first</option> + </select> + <button + onClick={() => discoverCorpus(speaker)} + disabled={discovering} + className="inline-flex items-center gap-1.5 text-sm px-3 py-1.5 rounded border border-gray-300 dark:border-gray-600 text-gray-700 dark:text-gray-200 disabled:opacity-50" + title="Refresh the reusable speech-embedding index and rescore it against this gallery" + > + {discovering ? <Loader2 className="h-4 w-4 animate-spin" /> : <Search className="h-4 w-4" />} + {discovering ? 'Searching corpus' : 'Refresh corpus'} + </button> + <label + className="inline-flex items-center gap-1.5 text-xs text-gray-600 dark:text-gray-300" + title="Also scan speech from soft-deleted conversations whose audio is still stored" + > + <input + type="checkbox" + checked={includeDeleted} + onChange={(e) => setIncludeDeleted(e.target.checked)} + /> + include deleted + </label> + <input + ref={miningInputRef} + type="file" + accept="audio/*,video/*,.wav,.mp3,.m4a,.ogg,.opus,.flac" + multiple + className="hidden" + onChange={(e) => mineFiles(speaker, Array.from(e.target.files || []))} + /> + <button + onClick={() => miningInputRef.current?.click()} + disabled={mining || discovering} + className="inline-flex items-center gap-1.5 text-sm px-3 py-1.5 rounded border border-gray-300 dark:border-gray-600 text-gray-700 dark:text-gray-200 disabled:opacity-50" + title="Upload unlabelled audio (recordings, exports) and mine it for this speaker's voice. Files are kept out of memory processing." + > + {mining ? <Loader2 className="h-4 w-4 animate-spin" /> : <AudioLines className="h-4 w-4" />} + {mining ? 'Uploading…' : 'Mine audio files'} + </button> + {suggestion && ( + <span className="text-xs text-gray-500 dark:text-gray-400"> + gallery: {suggestion.speaker.n_clips ?? '?'} clips + {suggestion.speaker.total_duration_s != null && + ` / ${Math.round(suggestion.speaker.total_duration_s)}s`} + {' · '} + {suggestion.reviewed_total} reviewed · ~{suggestion.pool_remaining} in pool + {suggestion.discovery_indexed && ` · ${suggestion.discovery_candidates} indexed segments`} + </span> + )} + </div> + + {error && <p className="text-sm text-red-600">{error}</p>} + {lastResult && !loading && ( + <p className="text-sm text-green-700 dark:text-green-400">{lastResult}</p> + )} + {loading && ( + <p className="text-sm text-gray-500"> + Scanning the corpus and scoring candidate clips… + </p> + )} + {discoveryProgress && ( + <p className="text-xs text-blue-700 dark:text-blue-300">{discoveryProgress}</p> + )} + + {suggestion && suggestion.batch.length > 0 && ( + <> + <ul className="space-y-2"> + {suggestion.batch.map((clip) => { + const key = clipKey(clip) + const decision = decisions[key] + const region = trimmedRegions[key] || { start: clip.start, end: clip.end } + return ( + <li + key={key} + className={`rounded border p-3 space-y-1.5 ${ + decision + ? decision.kind === 'accept' + ? 'border-green-400 bg-green-50 dark:bg-green-900/20' + : decision.kind === 'another_speaker' + ? 'border-purple-400 bg-purple-50 dark:bg-purple-900/20' + : decision.kind === 'reject' + ? 'border-red-300 bg-red-50 dark:bg-red-900/20' + : 'border-gray-300 bg-gray-50 dark:border-gray-600 dark:bg-gray-800' + : 'border-gray-200 dark:border-gray-700' + }`} + > + <div className="flex flex-col gap-1.5 sm:flex-row sm:items-center sm:gap-2"> + <div className="flex min-w-0 items-center gap-2"> + <button + onClick={() => playClip(clip)} + className="shrink-0 p-1.5 rounded bg-gray-100 dark:bg-gray-700 text-gray-700 dark:text-gray-200" + title="Play clip" + > + {playingKey === key ? ( + <Pause className="h-4 w-4" /> + ) : ( + <Play className="h-4 w-4" /> + )} + </button> + <span className="shrink-0 text-xs text-gray-500 dark:text-gray-400"> + {formatClock(region.start)}–{formatClock(region.end)} ·{' '} + {(region.end - region.start).toFixed(1)}s + {trimmedRegions[key] && ' · trimmed'} + </span> + <span className="truncate text-xs text-gray-500 dark:text-gray-400"> + {clip.conversation_title || clip.conversation_id.slice(0, 8)} + {clip.conversation_date && ` · ${clip.conversation_date.slice(0, 10)}`} + </span> + </div> + <div className="flex items-center justify-end gap-2 sm:ml-auto"> + <span className="text-xs font-mono text-gray-600 dark:text-gray-300"> + match {clip.scores.sim_centroid.toFixed(2)} + </span> + <div className="flex gap-1"> + <button + onClick={() => setEditingKey(editingKey === key ? null : key)} + className={`p-1.5 rounded ${editingKey === key ? 'bg-emerald-600 text-white' : 'bg-gray-100 dark:bg-gray-700 text-gray-600 dark:text-gray-300'}`} + title="Trim this enrollment clip" + > + <Scissors className="h-4 w-4" /> + </button> + <button + onClick={() => + setDecisions((d) => ({ ...d, [key]: { kind: 'accept' } })) + } + className={`p-1.5 rounded ${ + decision?.kind === 'accept' + ? 'bg-green-600 text-white' + : 'bg-gray-100 dark:bg-gray-700 text-gray-600 dark:text-gray-300' + }`} + title={`Yes — this is ${suggestion.speaker.speaker_name}`} + > + <Check className="h-4 w-4" /> + </button> + <button + onClick={() => + setDecisions((d) => ({ ...d, [key]: { kind: 'reject' } })) + } + className={`p-1.5 rounded ${ + decision?.kind === 'reject' + ? 'bg-red-600 text-white' + : 'bg-gray-100 dark:bg-gray-700 text-gray-600 dark:text-gray-300' + }`} + title="No — not this speaker" + > + <X className="h-4 w-4" /> + </button> + </div> + </div> + </div> + {editingKey === key && clip.conversation_duration > 0 && ( + <WaveformRegionEditor + conversationId={clip.conversation_id} + duration={clip.conversation_duration} + initialRegion={region} + pickerMode + onChange={(next) => { + if (next) setTrimmedRegions((current) => ({ ...current, [key]: next })) + }} + onCancel={() => setEditingKey(null)} + onPlay={(next) => playRegion(clip, next)} + playheadTime={playingKey === key ? playheadTime : null} + height={88} + /> + )} + <div className="flex flex-wrap items-center gap-2 sm:pl-9"> + <span className="w-full text-xs text-gray-500 dark:text-gray-400 sm:w-auto">Actually:</span> + <div className="w-full sm:w-52"> + <SpeakerInlineInput + value={decision?.actualSpeaker || ''} + onChange={(name) => setDecisions((d) => ({ ...d, [key]: { kind: 'another_speaker', actualSpeaker: name } }))} + onSelect={(name) => setDecisions((d) => ({ ...d, [key]: { kind: 'another_speaker', actualSpeaker: name } }))} + enrolledSpeakers={speakers.filter((item) => item.name !== speaker)} + placeholder="Search speakers..." + allowCreate={false} + /> + </div> + {decision?.kind === 'another_speaker' && decision.actualSpeaker && ( + <span className="text-xs px-2 py-0.5 rounded bg-purple-600 text-white"> + relabel → {decision.actualSpeaker} + </span> + )} + <button onClick={() => setDecisions((d) => ({ ...d, [key]: { kind: 'skip' } }))} className={`inline-flex items-center gap-1 text-xs px-2 py-1 rounded ${decision?.kind === 'skip' ? 'bg-gray-600 text-white' : 'border border-gray-300 dark:border-gray-600 text-gray-600 dark:text-gray-300'}`}><SkipForward className="h-3.5 w-3.5" />Skip</button> + <button onClick={() => setDecisions((d) => ({ ...d, [key]: { kind: 'multiple_speakers' } }))} className={`inline-flex items-center gap-1 text-xs px-2 py-1 rounded ${decision?.kind === 'multiple_speakers' ? 'bg-amber-600 text-white' : 'border border-gray-300 dark:border-gray-600 text-gray-600 dark:text-gray-300'}`} title="Several people are talking in this clip; do not enroll or label it"><Users className="h-3.5 w-3.5" />Multiple speakers</button> + <button onClick={() => setDecisions((d) => ({ ...d, [key]: { kind: 'bad_clip' } }))} className={`inline-flex items-center gap-1 text-xs px-2 py-1 rounded ${decision?.kind === 'bad_clip' ? 'bg-amber-600 text-white' : 'border border-gray-300 dark:border-gray-600 text-gray-600 dark:text-gray-300'}`} title="Poor enrollment audio, such as mostly noise or badly bounded speech"><AudioLines className="h-3.5 w-3.5" />Bad clip</button> + </div> + {clip.text && ( + <p className="text-sm text-gray-800 dark:text-gray-200 line-clamp-2"> + “{clip.text}” + </p> + )} + <p className="text-xs text-gray-500 dark:text-gray-400"> + {clip.current_label + ? `currently labeled ${clip.current_label}` + : 'currently unlabeled'} + {clip.reasons.length > 0 && ` · ${clip.reasons.join(' · ')}`} + </p> + </li> + ) + })} + </ul> + <button + onClick={submit} + disabled={submitting || decidedCount === 0} + className="inline-flex items-center gap-1.5 text-sm px-3 py-1.5 rounded bg-blue-600 text-white disabled:opacity-50" + > + {submitting && <Loader2 className="h-4 w-4 animate-spin" />} + Submit {decidedCount}/{suggestion.batch.length} & next batch + </button> + <button onClick={freshBatch} disabled={submitting} className="ml-2 inline-flex items-center gap-1.5 text-sm px-3 py-1.5 rounded border border-gray-300 dark:border-gray-600 text-gray-700 dark:text-gray-200 disabled:opacity-50"><RefreshCw className="h-4 w-4" />Skip remaining & fresh batch</button> + </> + )} + + <GallerySection + speakerName={speaker} + refreshKey={galleryVersion} + onReset={(galleryPurged, message) => { + setLastResult(message) + setError(null) + if (galleryPurged) { + stopAudio() + selectedSpeakerRef.current = '' + setSpeaker('') + setSpeakerSearch('') + setSuggestion(null) + setDecisions({}) + setHistory([]) + setSpeakers([]) // triggers a refetch of the enrolled-speaker list + } else { + suggest(speaker) + loadHistory(speaker) + } + }} + /> + + {suggestion && <BenchmarkPanel speakerName={speaker} />} + + {history.length > 0 && ( + <section className="border-t border-gray-200 dark:border-gray-700 pt-4"> + <h2 className="text-sm font-semibold text-gray-900 dark:text-gray-100 mb-2">Enrollment history</h2> + <div className="grid sm:grid-cols-3 gap-2 mb-4 text-xs"> + <div className="border border-gray-200 dark:border-gray-700 rounded p-2"> + <div className="text-gray-500 dark:text-gray-400">Clean speech quantity</div> + <div className="font-medium text-gray-900 dark:text-gray-100"> + {(suggestion?.speaker.total_duration_s ?? 0) >= 60 ? 'Sufficient' : 'Still building'} + {suggestion?.speaker.total_duration_s != null && ` · ${Math.round(suggestion.speaker.total_duration_s)}s`} + </div> + <div className="text-gray-500 dark:text-gray-400">Research gains usually flatten after 30–60s</div> + </div> + <div className="border border-gray-200 dark:border-gray-700 rounded p-2"> + <div className="text-gray-500 dark:text-gray-400">Gallery cleanliness</div> + <div className="font-medium text-gray-900 dark:text-gray-100">{history[0].health_after?.n_flagged ?? '—'} outliers</div> + <div className="text-gray-500 dark:text-gray-400">Cohesion {history[0].health_after?.median_self?.toFixed(3) ?? 'unavailable'}</div> + </div> + <div className="border border-gray-200 dark:border-gray-700 rounded p-2"> + <div className="text-gray-500 dark:text-gray-400">Recognition performance</div> + <div className="font-medium text-amber-700 dark:text-amber-300">Not benchmarked</div> + <div className="text-gray-500 dark:text-gray-400">Requires held-out identification clips</div> + </div> + </div> + <EnrollmentTrend sessions={history} /> + <div className="overflow-x-auto"> + <table className="w-full text-xs text-left"> + <thead className="text-gray-500 dark:text-gray-400 border-b border-gray-200 dark:border-gray-700"> + <tr><th className="py-2 pr-3">Date</th><th className="py-2 pr-3">Added</th><th className="py-2 pr-3">Gallery</th><th className="py-2 pr-3">Cohesion</th><th className="py-2 pr-3">Outliers</th><th className="py-2">Novelty</th></tr> + </thead> + <tbody> + {history.map((session, index) => { + const before = session.health_before + const after = session.health_after + return ( + <tr key={`${session.created_at}:${index}`} className="border-b border-gray-100 dark:border-gray-700/60 text-gray-700 dark:text-gray-200"> + <td className="py-2 pr-3 whitespace-nowrap">{new Date(session.created_at).toLocaleString()}</td> + <td className="py-2 pr-3">{session.decisions.enrolled}</td> + <td className="py-2 pr-3 whitespace-nowrap">{before && after ? `${before.n_clips} → ${after.n_clips}` : after?.n_clips ?? '—'}</td> + <td className="py-2 pr-3 whitespace-nowrap">{before?.median_self != null && after?.median_self != null ? `${before.median_self.toFixed(3)} → ${after.median_self.toFixed(3)}` : '—'}</td> + <td className="py-2 pr-3 whitespace-nowrap">{before && after ? `${before.n_flagged}/${before.n_clips} → ${after.n_flagged}/${after.n_clips}` : '—'}</td> + <td className="py-2">{session.coverage.accepted_novelty_mean != null ? `${Math.round(session.coverage.accepted_novelty_mean * 100)}%` : '—'}</td> + </tr> + ) + })} + </tbody> + </table> + </div> + </section> + )} + </section> + )} + </div> + ) +} diff --git a/backends/advanced/webui/src/components/dataAudit/filters.tsx b/backends/advanced/webui/src/components/dataAudit/filters.tsx index 78f3df00..2bd84ed7 100644 --- a/backends/advanced/webui/src/components/dataAudit/filters.tsx +++ b/backends/advanced/webui/src/components/dataAudit/filters.tsx @@ -12,6 +12,7 @@ import { CalendarRange, CheckCheck, Clock, + FileArchive, LucideIcon, Mic, Search, @@ -23,6 +24,7 @@ export type SpeakerFilterState = 'include' | 'exclude' export interface FilterContext { speakers: string[] + datasets: string[] } export interface EditorProps<V> { @@ -372,6 +374,40 @@ const dateFilter: FilterDef<DateValue> = { ), } +// --------------------------------------------------------------------------- +// Imported annotation dataset +// --------------------------------------------------------------------------- + +const datasetFilter: FilterDef<string> = { + key: 'dataset', + label: 'Dataset', + icon: FileArchive, + defaultValue: '', + isActive: (v) => v.trim() !== '', + chipLabel: (v) => `Dataset: ${v}`, + toParams: (v) => ({ dataset_id: v.trim() || undefined }), + Editor: ({ value, onChange, ctx }) => ( + <div className="w-[min(20rem,calc(100vw-3rem))] space-y-2"> + <label className="block text-xs font-medium text-gray-600 dark:text-gray-300"> + Annotation dataset + </label> + <select + value={value} + onChange={(event) => onChange(event.target.value)} + className="w-full rounded border border-gray-300 bg-white px-2 py-1.5 text-sm text-gray-700 dark:border-gray-600 dark:bg-gray-800 dark:text-gray-200" + > + <option value="">All datasets</option> + {!ctx.datasets.includes(value) && value && <option value={value}>{value}</option>} + {ctx.datasets.map((dataset) => ( + <option key={dataset} value={dataset}> + {dataset} + </option> + ))} + </select> + </div> + ), +} + // --------------------------------------------------------------------------- // Hide failed (processing_status == 'failed') // --------------------------------------------------------------------------- @@ -407,6 +443,7 @@ export const AUDIT_FILTERS: FilterDef[] = [ durationFilter, speakersFilter, dateFilter, + datasetFilter, hideFailedFilter, hideReviewedFilter, ] diff --git a/backends/advanced/webui/src/components/transcript/InsertSegmentForm.tsx b/backends/advanced/webui/src/components/transcript/InsertSegmentForm.tsx index af08f41c..f4989195 100644 --- a/backends/advanced/webui/src/components/transcript/InsertSegmentForm.tsx +++ b/backends/advanced/webui/src/components/transcript/InsertSegmentForm.tsx @@ -7,6 +7,7 @@ interface InsertSegmentFormProps { afterIndex: number // -1 = before first segment allSpeakers: { speaker_id: string; name: string }[] recentSpeakers: string[] + usedSpeakerNames?: string[] onSpeakerUsed?: (speaker: string) => void /** Optional waveform-drawn span for the new segment (else a zero-duration boundary marker). */ region?: { start: number; end: number } | null @@ -26,6 +27,7 @@ export default function InsertSegmentForm({ afterIndex, allSpeakers, recentSpeakers, + usedSpeakerNames = [], onSpeakerUsed, region, onDone, @@ -88,6 +90,7 @@ export default function InsertSegmentForm({ }} enrolledSpeakers={allSpeakers} recentSpeakers={recentSpeakers} + usedSpeakerNames={usedSpeakerNames} placeholder="Type or select speaker..." /> </div> diff --git a/backends/advanced/webui/src/components/transcript/TranscriptEditor.tsx b/backends/advanced/webui/src/components/transcript/TranscriptEditor.tsx index db6a1f18..39213106 100644 --- a/backends/advanced/webui/src/components/transcript/TranscriptEditor.tsx +++ b/backends/advanced/webui/src/components/transcript/TranscriptEditor.tsx @@ -1,8 +1,9 @@ import { useCallback, useEffect, useMemo, useState } from 'react' -import { Play, Pause, X, Check, RefreshCw, Trash2, Eye, EyeOff, Plus } from 'lucide-react' +import { Play, Pause, X, Check, RefreshCw, Trash2, Eye, EyeOff, Plus, Users, AlignLeft, Infinity, Scissors, ChevronLeft, ChevronRight } from 'lucide-react' import { annotationsApi } from '../../services/api' import { useGaplessPlayer } from '../../hooks/useGaplessPlayer' import SpeakerNameDropdown from '../SpeakerNameDropdown' +import SpeakerInlineInput from '../SpeakerInlineInput' import { PlayheadWaveform, PlayheadTimeLabel } from '../audio/PlayheadWaveform' import { WaveformRegionEditor, Region } from '../audio/WaveformRegionEditor' import InsertSegmentForm from './InsertSegmentForm' @@ -43,6 +44,8 @@ const SPEAKER_COLOR_PALETTE = [ 'text-indigo-700 dark:text-indigo-300', ] +const SPEAKER_HEX_PALETTE = ['#2563eb', '#059669', '#9333ea', '#ea580c', '#db2777', '#0891b2', '#d97706', '#4f46e5'] + const formatDuration = (seconds: number) => { const m = Math.floor(seconds / 60) const s = Math.floor(seconds % 60) @@ -99,6 +102,15 @@ export default function TranscriptEditor({ const [preview, setPreview] = useState(false) const [applying, setApplying] = useState(false) const [clearing, setClearing] = useState(false) + const [annotationMode, setAnnotationMode] = useState<'transcript' | 'speakers'>('transcript') + const [selectedSpeakerSegment, setSelectedSpeakerSegment] = useState<number | null>(null) + const [continuePastSegment, setContinuePastSegment] = useState(false) + const [autoPlayOnClick, setAutoPlayOnClick] = useState(false) + const [speakerCreationMode, setSpeakerCreationMode] = useState<'snip' | 'draw' | null>(null) + const [newSpeaker, setNewSpeaker] = useState('') + const [newSpeakerRegion, setNewSpeakerRegion] = useState<Region | null>(null) + const [speakerSnipTime, setSpeakerSnipTime] = useState<number | null>(null) + const [speakerFilters, setSpeakerFilters] = useState<Record<string, 'include' | 'exclude'>>({}) // While inserting with the waveform open, the region drawn on it for the new segment. const [insertRegion, setInsertRegion] = useState<Region | null>(null) // Whether the insert menu drives the top waveform (draw the new segment's span). @@ -142,6 +154,15 @@ export default function TranscriptEditor({ return list }, [enrolledSpeakers, diar]) + const usedSpeakerNames = useMemo(() => { + const names = new Set(segments.map((segment) => segment.speaker).filter((name): name is string => !!name)) + pendingDiar.forEach((annotation) => names.add(annotation.corrected_speaker)) + pendingInsert.forEach((annotation) => { + if (annotation.insert_speaker) names.add(annotation.insert_speaker) + }) + return [...names] + }, [segments, pendingDiar, pendingInsert]) + const speakerColorMap = useMemo(() => { const map: Record<string, string> = {} let i = 0 @@ -155,6 +176,89 @@ export default function TranscriptEditor({ return map }, [segments]) + const speakerHexMap = useMemo(() => { + const map: Record<string, string> = {} + let i = 0 + segments.forEach((seg, idx) => { + const corrected = pendingDiar.find((a) => a.segment_index === idx)?.corrected_speaker + const sp = corrected || seg.speaker || 'Unknown' + if (!map[sp]) map[sp] = SPEAKER_HEX_PALETTE[i++ % SPEAKER_HEX_PALETTE.length] + }) + return map + }, [segments, pendingDiar]) + + const displaySpeakerForSegment = useCallback((segment: Segment, idx: number) => ( + pendingDiar.find((annotation) => annotation.segment_index === idx)?.corrected_speaker + || segment.speaker + || 'Unknown' + ), [pendingDiar]) + + const transcriptSpeakers = useMemo(() => { + const speakers: string[] = [] + const seen = new Set<string>() + segments.forEach((segment, idx) => { + if (segment.segment_type === 'event' || segment.segment_type === 'note') return + const speaker = displaySpeakerForSegment(segment, idx) + if (!seen.has(speaker)) { + seen.add(speaker) + speakers.push(speaker) + } + }) + return speakers + }, [segments, displaySpeakerForSegment]) + + useEffect(() => { + setSpeakerFilters((current) => { + const available = new Set(transcriptSpeakers) + const next = Object.fromEntries(Object.entries(current).filter(([speaker]) => available.has(speaker))) as Record<string, 'include' | 'exclude'> + return Object.keys(next).length === Object.keys(current).length ? current : next + }) + }, [transcriptSpeakers]) + + const cycleSpeakerFilter = (speaker: string) => { + setSpeakerFilters((current) => { + const next = { ...current } + if (!current[speaker]) next[speaker] = 'include' + else if (current[speaker] === 'include') next[speaker] = 'exclude' + else delete next[speaker] + return next + }) + } + + const speakerIsVisible = (speaker: string) => { + const includes = Object.entries(speakerFilters).filter(([, state]) => state === 'include').map(([name]) => name) + if (speakerFilters[speaker] === 'exclude') return false + return includes.length === 0 || includes.includes(speaker) + } + + useEffect(() => { + if (selectedSpeakerSegment === null || Object.keys(speakerFilters).length === 0) return + const selected = segments[selectedSpeakerSegment] + if (!selected || !speakerIsVisible(displaySpeakerForSegment(selected, selectedSpeakerSegment))) { + setSelectedSpeakerSegment(null) + } + }, [speakerFilters, selectedSpeakerSegment, segments, displaySpeakerForSegment]) + + const speakerTimelineSegments = useMemo( + () => segments.flatMap((seg, idx) => { + if (seg.segment_type === 'event' || seg.segment_type === 'note') return [] + const region = (() => { + const pending = pendingTiming.find((a) => a.segment_index === idx) + return pending ? { start: pending.new_start, end: pending.new_end } : seg + })() + const speaker = displaySpeakerForSegment(seg, idx) + if (!speakerIsVisible(speaker)) return [] + return [{ + start: region.start, + end: region.end, + color: speakerHexMap[speaker] || SPEAKER_HEX_PALETTE[0], + segmentIndex: idx, + label: `${speaker} · ${formatDuration(region.start)}–${formatDuration(region.end)}`, + }] + }), + [segments, pendingTiming, speakerHexMap, displaySpeakerForSegment, speakerFilters] + ) + // Auto-open the timing editor on the main waveform when editing a segment's text // (unless the user disabled waveform zoom). Closing the text edit closes it too. useEffect(() => { @@ -375,6 +479,7 @@ export default function TranscriptEditor({ afterIndex={afterIndex} allSpeakers={allSpeakers} recentSpeakers={recentSpeakers} + usedSpeakerNames={usedSpeakerNames} onSpeakerUsed={noteRecent} onDone={async () => { closeInsert() @@ -403,8 +508,137 @@ export default function TranscriptEditor({ // transcript (no scrolling back up to the player). const editorActive = timingEditSegment !== null || insertOpen !== null + const playFromSpeakerPoint = (time: number, region: Region | null, segmentId: string) => { + setSpeakerSnipTime(time) + if (!autoPlayOnClick) return + if (continuePastSegment) { + player.play(conversationId, time, { totalDuration: duration! }) + return + } + if (region && time >= region.start && time < region.end) { + player.playSegment(conversationId, segmentId, time, region.end) + return + } + // With bounded playback selected, a click outside a span is positioning only. + // Stop any prior continuous program so it cannot appear to ignore the toggle. + player.pause() + } + + const selectSpeakerSegment = (idx: number) => { + const segment = segments[idx] + if (!segment) return + closeSpeakerCreation() + setSpeakerSnipTime(null) + setSelectedSpeakerSegment(idx) + const region = regionForSegment(idx) + playFromSpeakerPoint(region.start, region, `${conversationId}-${idx}`) + } + + const closeSpeakerCreation = () => { + setSpeakerCreationMode(null) + setNewSpeaker('') + setNewSpeakerRegion(null) + } + + const createSpeakerSpan = async () => { + if (selectedSpeakerSegment === null || !newSpeaker.trim()) return + const idx = selectedSpeakerSegment + const selectedRegion = regionForSegment(idx) + let region: Region | null = newSpeakerRegion + + if (speakerCreationMode === 'snip') { + const splitAt = speakerSnipTime + if (splitAt == null || splitAt <= selectedRegion.start + 0.05 || splitAt >= selectedRegion.end - 0.05) { + setRegionError('Place the red playhead inside the selected span before snipping.') + return + } + region = { start: splitAt, end: selectedRegion.end } + await handleSaveTiming(idx, { start: selectedRegion.start, end: splitAt }) + } + + if (!region || region.end - region.start < 0.05) { + setRegionError('Drag a speaker span on the waveform first.') + return + } + + await annotationsApi.createInsertAnnotation({ + conversation_id: conversationId, + insert_after_index: idx, + insert_text: '', + insert_segment_type: 'speech', + insert_speaker: newSpeaker.trim(), + insert_start: region.start, + insert_end: region.end, + }) + noteRecent(newSpeaker.trim()) + closeSpeakerCreation() + await reload() + } + return ( <div className="space-y-3"> + {showAudio && ( + <div className="flex items-center justify-between gap-3"> + <div className="inline-flex rounded-lg bg-gray-100 dark:bg-gray-700 p-1" aria-label="Annotation mode"> + <button + onClick={() => setAnnotationMode('transcript')} + className={`inline-flex items-center gap-1.5 rounded-md px-3 py-1.5 text-xs font-medium ${annotationMode === 'transcript' ? 'bg-white dark:bg-gray-800 shadow text-gray-900 dark:text-white' : 'text-gray-500 dark:text-gray-300'}`} + > + <AlignLeft className="h-3.5 w-3.5" /> Transcript + </button> + <button + onClick={() => { + setAnnotationMode('speakers') + setEditingSegment(null) + setInsertOpen(null) + }} + className={`inline-flex items-center gap-1.5 rounded-md px-3 py-1.5 text-xs font-medium ${annotationMode === 'speakers' ? 'bg-white dark:bg-gray-800 shadow text-gray-900 dark:text-white' : 'text-gray-500 dark:text-gray-300'}`} + > + <Users className="h-3.5 w-3.5" /> Edit speakers & timing + </button> + </div> + {annotationMode === 'speakers' && ( + <span className="text-xs text-gray-500 dark:text-gray-400">Hover a colored span to preview it, then click to edit</span> + )} + </div> + )} + + {transcriptSpeakers.length > 1 && ( + <div className="flex flex-wrap items-center gap-1.5" aria-label="Filter by speaker"> + <span className="mr-1 text-xs text-gray-500 dark:text-gray-400">Speakers:</span> + {transcriptSpeakers.map((speaker) => { + const state = speakerFilters[speaker] + return ( + <button + key={speaker} + type="button" + onClick={() => cycleSpeakerFilter(speaker)} + className={`inline-flex items-center gap-1.5 rounded-full border px-2.5 py-1 text-xs transition-colors ${ + state === 'include' + ? 'bg-blue-100 border-blue-400 text-blue-700 dark:bg-blue-900 dark:text-blue-100 dark:border-blue-600' + : state === 'exclude' + ? 'bg-red-100 border-red-400 text-red-700 line-through dark:bg-red-900/40 dark:text-red-200 dark:border-red-600' + : 'border-gray-300 bg-white text-gray-600 hover:bg-gray-100 dark:border-gray-600 dark:bg-gray-800 dark:text-gray-300 dark:hover:bg-gray-700' + }`} + title={`${speaker}: ${state || 'off'} — click to cycle`} + > + <span className="h-2.5 w-2.5 rounded-sm" style={{ backgroundColor: speakerHexMap[speaker] || SPEAKER_HEX_PALETTE[0] }} /> + {speaker} + </button> + ) + })} + {Object.keys(speakerFilters).length > 0 && ( + <button + type="button" + onClick={() => setSpeakerFilters({})} + className="px-2 py-1 text-xs text-gray-400 hover:text-gray-600 dark:hover:text-gray-200" + > + Clear + </button> + )} + </div> + )} + {/* Waveform — doubles as the timing editor while editing a segment */} {showAudio && ( <div @@ -414,7 +648,221 @@ export default function TranscriptEditor({ : undefined } > - {timingEditSegment !== null ? ( + {annotationMode === 'speakers' ? ( + <div className="space-y-3"> + <PlayheadWaveform + cid={conversationId} + duration={duration!} + onSeek={(t) => { + playFromSpeakerPoint(t, null, `${conversationId}-overview`) + }} + height={104} + coloredSegments={speakerTimelineSegments} + onSegmentClick={selectSpeakerSegment} + segmentMarker={player.segmentMarker} + hoverMarker={selectedSpeakerSegment === null ? null : regionForSegment(selectedSpeakerSegment)} + /> + + {selectedSpeakerSegment !== null ? (() => { + const idx = selectedSpeakerSegment + const segment = segments[idx] + const originalSpeaker = segment.speaker || 'Unknown' + const diarA = pendingDiar.find((a) => a.segment_index === idx) + const displaySpeaker = diarA?.corrected_speaker || originalSpeaker + return ( + <div className="rounded-lg border border-gray-200 dark:border-gray-700 p-3"> + <div className="flex items-center justify-between gap-3 mb-2"> + <div className="flex items-center gap-2"> + <span className="h-3 w-3 rounded-sm" style={{ backgroundColor: speakerHexMap[displaySpeaker] }} /> + <SpeakerNameDropdown + currentSpeaker={displaySpeaker} + enrolledSpeakers={allSpeakers} + onSpeakerChange={(speaker) => handleSpeakerChange(idx, originalSpeaker, speaker, segment.start)} + segmentIndex={idx} + conversationId={conversationId} + annotated={!!diarA} + speakerColor="text-gray-900 dark:text-white" + recentSpeakers={recentSpeakers} + usedSpeakerNames={usedSpeakerNames} + /> + <span className="text-xs text-gray-400">Segment {idx + 1} of {segments.length}</span> + <button + type="button" + disabled={idx <= 0} + onClick={() => selectSpeakerSegment(idx - 1)} + className="p-1 text-gray-500 hover:bg-gray-100 dark:hover:bg-gray-700 rounded disabled:opacity-30" + title="Previous segment" + aria-label="Previous segment" + > + <ChevronLeft className="h-4 w-4" /> + </button> + <button + type="button" + disabled={idx >= segments.length - 1} + onClick={() => selectSpeakerSegment(idx + 1)} + className="p-1 text-gray-500 hover:bg-gray-100 dark:hover:bg-gray-700 rounded disabled:opacity-30" + title="Next segment" + aria-label="Next segment" + > + <ChevronRight className="h-4 w-4" /> + </button> + </div> + <div className="flex items-center gap-1"> + <button + type="button" + aria-pressed={autoPlayOnClick} + onClick={() => setAutoPlayOnClick((value) => !value)} + className={`inline-flex items-center gap-1 rounded px-1.5 py-1 text-xs ${autoPlayOnClick ? 'bg-blue-100 text-blue-700 dark:bg-blue-900/40 dark:text-blue-300' : 'text-gray-400 hover:bg-gray-100 dark:hover:bg-gray-700'}`} + title={autoPlayOnClick ? 'Auto-play is on: waveform clicks start playback' : 'Auto-play is off: waveform clicks only position the snip cursor'} + > + <Play className="h-3.5 w-3.5" /> + <span className="hidden sm:inline">Auto-play</span> + </button> + <button + type="button" + onClick={() => { + setRegionError(null) + setSpeakerCreationMode('snip') + setNewSpeaker('') + setNewSpeakerRegion(null) + }} + className={`inline-flex items-center gap-1 rounded px-1.5 py-1 text-xs ${speakerCreationMode === 'snip' ? 'bg-purple-100 text-purple-700 dark:bg-purple-900/40 dark:text-purple-300' : 'text-gray-500 hover:bg-gray-100 dark:hover:bg-gray-700'}`} + title="Split this speaker span at the red playhead" + > + <Scissors className="h-3.5 w-3.5" /> Snip + </button> + <button + type="button" + onClick={() => { + setRegionError(null) + setSpeakerCreationMode('draw') + setNewSpeaker('') + setNewSpeakerRegion(null) + }} + className={`inline-flex items-center gap-1 rounded px-1.5 py-1 text-xs ${speakerCreationMode === 'draw' ? 'bg-purple-100 text-purple-700 dark:bg-purple-900/40 dark:text-purple-300' : 'text-gray-500 hover:bg-gray-100 dark:hover:bg-gray-700'}`} + title="Drag a new independent or overlapping speaker span" + > + <Plus className="h-3.5 w-3.5" /> New span + </button> + <button + type="button" + aria-pressed={continuePastSegment} + onClick={() => setContinuePastSegment((value) => { + if (value) player.stop() + return !value + })} + className={`inline-flex items-center gap-1 rounded px-1.5 py-1 text-xs ${continuePastSegment ? 'bg-blue-100 text-blue-700 dark:bg-blue-900/40 dark:text-blue-300' : 'text-gray-400 hover:bg-gray-100 dark:hover:bg-gray-700'}`} + title={continuePastSegment ? 'Continuous playback enabled: clicks play beyond this segment' : 'Stop at segment end: click to enable continuous playback'} + > + <Infinity className="h-3.5 w-3.5" /> + <span className="hidden sm:inline">Continue</span> + </button> + <button onClick={() => setSelectedSpeakerSegment(null)} className="p-1 text-gray-400 hover:text-gray-700" title="Close selection"><X className="h-4 w-4" /></button> + </div> + </div> + <WaveformRegionEditor + key={`speaker-boundary-${idx}-${regionForSegment(idx).start}-${regionForSegment(idx).end}`} + conversationId={conversationId} + duration={duration!} + initialRegion={regionForSegment(idx)} + onSaveTiming={async (region) => handleSaveTiming(idx, region)} + onCancel={() => setSelectedSpeakerSegment(null)} + onPlay={handlePlayRegion} + onSeekPlay={(time, region) => { + playFromSpeakerPoint(time, region, `${conversationId}-${idx}`) + }} + playheadTime={autoPlayOnClick ? undefined : speakerSnipTime} + height={96} + /> + <p className="mt-2 text-xs text-gray-500 dark:text-gray-400 line-clamp-1" title={segment.text}> + Words are reference only: {segment.text || '(no transcript text)'} + </p> + + {speakerCreationMode && ( + <div className="mt-3 rounded-lg border border-purple-200 dark:border-purple-800 bg-purple-50/60 dark:bg-purple-900/10 p-3 space-y-2"> + <div className="flex items-center justify-between gap-2"> + <p className="text-xs text-purple-700 dark:text-purple-300"> + {speakerCreationMode === 'snip' + ? `The new speaker starts at the last point clicked in the zoomed waveform (${speakerSnipTime == null ? 'click the waveform first' : `${speakerSnipTime.toFixed(2)}s`}) and takes the remainder of this span.` + : 'Drag the exact new speaker span below. It may overlap an existing speaker.'} + </p> + <button onClick={closeSpeakerCreation} className="p-1 text-gray-400 hover:text-gray-700" title="Cancel"><X className="h-3.5 w-3.5" /></button> + </div> + + {speakerCreationMode === 'draw' && ( + <WaveformRegionEditor + key={`new-speaker-span-${idx}`} + conversationId={conversationId} + duration={duration!} + initialRegion={null} + focusTime={speakerSnipTime ?? regionForSegment(idx).end} + pickerMode + onChange={setNewSpeakerRegion} + onCancel={closeSpeakerCreation} + onPlay={handlePlayRegion} + onSeekPlay={(time, region) => { + playFromSpeakerPoint(time, region, `${conversationId}-new-speaker-span`) + }} + playheadTime={autoPlayOnClick ? undefined : speakerSnipTime} + height={88} + /> + )} + + <div className="flex items-center gap-2"> + <span className="text-xs text-gray-500 whitespace-nowrap">New speaker:</span> + <div className="flex-1 min-w-0"> + <SpeakerInlineInput + value={newSpeaker} + onChange={setNewSpeaker} + onSelect={(speaker) => { + setNewSpeaker(speaker) + noteRecent(speaker) + }} + enrolledSpeakers={allSpeakers} + recentSpeakers={recentSpeakers} + usedSpeakerNames={usedSpeakerNames} + placeholder="Select who starts here…" + /> + </div> + <button + onClick={createSpeakerSpan} + disabled={!newSpeaker.trim() || (speakerCreationMode === 'draw' && !newSpeakerRegion)} + className="inline-flex items-center gap-1 rounded bg-purple-600 px-2.5 py-1.5 text-xs font-medium text-white hover:bg-purple-700 disabled:opacity-40" + > + {speakerCreationMode === 'snip' ? <Scissors className="h-3.5 w-3.5" /> : <Plus className="h-3.5 w-3.5" />} + {speakerCreationMode === 'snip' ? 'Create split' : 'Create span'} + </button> + </div> + {regionError && <p className="text-xs text-red-500">{regionError}</p>} + </div> + )} + </div> + ) + })() : ( + <div className="rounded-lg border border-dashed border-gray-300 dark:border-gray-600 px-3 py-5 text-center text-sm text-gray-500"> + Hover a colored span to see its speaker and time, then click to edit it. + </div> + )} + + <div className="flex items-center gap-3"> + <button + onClick={() => { + if (player.isActive(conversationId) && player.isPlaying) { + setAutoPlayOnClick(false) + player.pause() + } else { + player.togglePlay(conversationId, duration!) + } + }} + className="p-2 rounded-full bg-blue-600 hover:bg-blue-700 text-white" + title={player.isActive(conversationId) && player.isPlaying ? 'Pause and disable auto-play' : 'Play'} + > + {player.isActive(conversationId) && player.isPlaying ? <Pause className="w-4 h-4" /> : <Play className="w-4 h-4" />} + </button> + <PlayheadTimeLabel cid={conversationId} total={duration} className="text-sm text-gray-600 dark:text-gray-400 font-mono" /> + </div> + </div> + ) : timingEditSegment !== null ? ( <> <WaveformRegionEditor key={`timing-${timingEditSegment}`} @@ -454,6 +902,7 @@ export default function TranscriptEditor({ afterIndex={insertOpen} allSpeakers={allSpeakers} recentSpeakers={recentSpeakers} + usedSpeakerNames={usedSpeakerNames} onSpeakerUsed={noteRecent} region={insertRegion} onDone={async () => { @@ -534,14 +983,18 @@ export default function TranscriptEditor({ )} {/* Segments */} - {segments.length > 0 ? ( + {annotationMode === 'transcript' && (segments.length > 0 ? ( <div className="space-y-0.5"> <InsertDivider afterIndex={-1} /> {segments.map((segment, idx) => { const speaker = segment.speaker || 'Unknown' const isEvent = segment.segment_type === 'event' const isNote = segment.segment_type === 'note' - if (hideUnknownSpeakers && !isNote && isUnknownSpeakerLabel(speaker)) return <InsertDivider key={`d${idx}`} afterIndex={idx} /> + const displaySpeaker = displaySpeakerForSegment(segment, idx) + if (Object.keys(speakerFilters).length > 0 && (isEvent || isNote || !speakerIsVisible(displaySpeaker))) { + return null + } + if (hideUnknownSpeakers && !isNote && isUnknownSpeakerLabel(speaker)) return null const speakerColor = speakerColorMap[speaker] || SPEAKER_COLOR_PALETTE[0] const isEditing = editingSegment === idx @@ -549,15 +1002,14 @@ export default function TranscriptEditor({ const textA = pendingText.find((a) => a.segment_index === idx) const timingA = pendingTiming.find((a) => a.segment_index === idx) const delA = pendingDeletion.find((a) => a.segment_index === idx) - const displaySpeaker = diarA ? diarA.corrected_speaker : speaker const displayText = textA ? textA.corrected_text : segment.text if (isEvent || isNote) { return ( <div key={idx}> <div - className={`group flex items-center gap-2 py-1 px-3 rounded ${ - isEvent ? 'bg-yellow-50 dark:bg-yellow-900/20 border-l-2 border-yellow-400' : 'bg-green-50 dark:bg-green-900/20 border-l-2 border-green-400' + className={`group flex items-center gap-2 py-1.5 px-3 rounded ${ + isEvent ? 'bg-amber-50 dark:bg-amber-900/25 border-l-2 border-amber-500' : 'bg-green-50 dark:bg-green-900/25 border-l-2 border-green-500' }`} onMouseEnter={isEvent ? () => setHoverMarker({ start: segment.start, end: segment.end }) : undefined} onMouseLeave={isEvent ? () => setHoverMarker(null) : undefined} @@ -567,8 +1019,8 @@ export default function TranscriptEditor({ {player.playingSegmentId === `${conversationId}-${idx}` ? <Pause className="h-3 w-3 text-yellow-600" /> : <Play className="h-3 w-3 text-yellow-600" />} </button> )} - <span className="text-xs font-medium text-gray-500 uppercase mr-2">{isEvent ? 'event' : 'note'}</span> - <span className="text-sm text-gray-700 dark:text-gray-300 italic">{displayText}</span> + <span className="text-xs font-semibold text-gray-700 dark:text-gray-300 uppercase mr-2">{isEvent ? 'event' : 'note'}</span> + <span className="text-sm text-gray-800 dark:text-gray-200 italic">{displayText}</span> </div> <InsertDivider afterIndex={idx} /> </div> @@ -579,7 +1031,7 @@ export default function TranscriptEditor({ <div key={idx}> <div className={`group flex items-start gap-2 py-1 rounded hover:bg-gray-50 dark:hover:bg-gray-700/50 ${ - delA ? 'bg-red-50 dark:bg-red-900/10 opacity-60' : textA ? 'bg-yellow-50 dark:bg-yellow-900/10' : '' + delA ? 'bg-red-50 dark:bg-red-900/10 opacity-60' : (!preview && textA) ? 'bg-yellow-50 dark:bg-yellow-900/10' : '' }`} onMouseEnter={() => setHoverMarker({ start: segment.start, end: segment.end })} onMouseLeave={() => setHoverMarker(null)} @@ -609,6 +1061,7 @@ export default function TranscriptEditor({ annotated={!!diarA} speakerColor={speakerColor} recentSpeakers={recentSpeakers} + usedSpeakerNames={usedSpeakerNames} /> </> )} @@ -646,11 +1099,22 @@ export default function TranscriptEditor({ </div> ) : ( <p - className={`text-sm text-gray-700 dark:text-gray-300 px-1 rounded ${delA ? 'line-through text-red-500 dark:text-red-400' : ''} ${preview ? '' : 'cursor-pointer hover:bg-yellow-50 dark:hover:bg-yellow-900/10'}`} + className={`min-h-5 text-sm text-gray-700 dark:text-gray-300 px-1 rounded ${delA ? 'line-through text-red-500 dark:text-red-400' : ''} ${preview ? '' : 'cursor-pointer hover:bg-yellow-50 dark:hover:bg-yellow-900/10'}`} onClick={preview ? undefined : () => handleStartEdit(idx, segment.text)} - title={preview ? undefined : 'Click to edit'} + title={preview ? undefined : (textA ? `Suggested: ${textA.corrected_text}` : (displayText ? 'Click to edit' : 'Click to add transcript text'))} > - {displayText} + {/* Review mode shows the pending change as a diff (original struck + + correction in green); preview shows the clean applied result. */} + {textA && !preview && !delA ? ( + <> + <span className="line-through text-gray-400 dark:text-gray-500">{segment.text}</span>{' '} + <span className="text-green-700 dark:text-green-400">{textA.corrected_text}</span> + </> + ) : ( + displayText || (!preview && ( + <span className="italic text-gray-400 dark:text-gray-500">Click to add transcript text</span> + )) + )} </p> )} </div> @@ -688,7 +1152,7 @@ export default function TranscriptEditor({ </div> ) : ( <p className="text-sm text-gray-500 dark:text-gray-400 italic">{isLive ? 'Waiting for speech...' : 'No transcript segments available'}</p> - )} + ))} </div> ) } diff --git a/backends/advanced/webui/src/pages/ConversationDetail.tsx b/backends/advanced/webui/src/pages/ConversationDetail.tsx index 304a57d1..05305958 100644 --- a/backends/advanced/webui/src/pages/ConversationDetail.tsx +++ b/backends/advanced/webui/src/pages/ConversationDetail.tsx @@ -355,7 +355,7 @@ export default function ConversationDetail() { } return ( - <div className="max-w-5xl mx-auto p-6 space-y-6"> + <div className="max-w-7xl mx-auto p-6 space-y-6"> {/* Header */} <div className="flex items-center justify-between"> <button @@ -492,11 +492,11 @@ export default function ConversationDetail() { /> {/* Main Content Grid */} - <div className="grid grid-cols-1 lg:grid-cols-3 gap-6"> + <div className="grid grid-cols-1 lg:grid-cols-4 gap-6"> {/* Left Column - Main Content */} - <div className="lg:col-span-2 space-y-6"> + <div className="lg:col-span-3 space-y-6"> {/* Title */} - <div className="bg-white dark:bg-gray-800 border border-gray-200 dark:border-gray-700 rounded-lg p-6"> + <div id="transcript" className="bg-white dark:bg-gray-800 border border-gray-200 dark:border-gray-700 rounded-lg p-6 scroll-mt-6"> {editingTitle ? ( <div className="space-y-2"> <div className="flex items-center space-x-2"> @@ -599,7 +599,7 @@ export default function ConversationDetail() { </div> {/* Right Column - Sidebar */} - <div className="space-y-6"> + <div className="space-y-6 lg:sticky lg:top-6 lg:self-start"> {/* Metadata Card */} <div className="bg-gray-50 dark:bg-gray-800/50 border border-gray-200 dark:border-gray-700 rounded-lg p-4"> <h3 className="text-sm font-semibold text-gray-500 dark:text-gray-400 uppercase mb-3"> @@ -663,28 +663,21 @@ export default function ConversationDetail() { </dl> </div> - {/* Version Info Card */} - {(conversation.transcript_version_count || 0) > 0 && ( - <div className="bg-gray-50 dark:bg-gray-800/50 border border-gray-200 dark:border-gray-700 rounded-lg p-4"> - <h3 className="text-sm font-semibold text-gray-500 dark:text-gray-400 uppercase mb-3"> - Versions - </h3> - <dl className="space-y-3 text-sm"> - <div className="flex justify-between items-start"> - <dt className="text-gray-600 dark:text-gray-400">Transcript</dt> - <dd className="text-gray-900 dark:text-gray-100"> - v{conversation.active_transcript_version_number || 1} of {conversation.transcript_version_count} - </dd> - </div> - </dl> - </div> - )} + <a + href="#memory-history" + className="flex w-full items-center justify-between rounded-lg border border-blue-200 bg-blue-50 px-4 py-3 text-sm font-medium text-blue-700 transition-colors hover:bg-blue-100 dark:border-blue-800 dark:bg-blue-900/20 dark:text-blue-300 dark:hover:bg-blue-900/35" + > + <span>Memory history</span> + <span aria-hidden="true">↓</span> + </a> - {/* Memory change history (audit ledger) */} - <MemoryAuditCard conversationId={conversation.conversation_id} /> </div> </div> + {/* Memory change history is intentionally full-width: paths, summaries, and + timestamps become unreadable in the narrow metadata rail. */} + <MemoryAuditCard conversationId={conversation.conversation_id} /> + {/* Split modal — on success the conversation is soft-deleted, so leave */} {showSplitModal && ( <SplitConversationModal diff --git a/backends/advanced/webui/src/pages/DataAudit.tsx b/backends/advanced/webui/src/pages/DataAudit.tsx index 696687f4..0da9f6ca 100644 --- a/backends/advanced/webui/src/pages/DataAudit.tsx +++ b/backends/advanced/webui/src/pages/DataAudit.tsx @@ -1,5 +1,6 @@ import { useCallback, useEffect, useState } from 'react' -import { Sparkles, Archive as ArchiveIcon, AlertTriangle } from 'lucide-react' +import { Sparkles, Archive as ArchiveIcon, AlertTriangle, Mic, ArrowRight } from 'lucide-react' +import { Link, useSearchParams } from 'react-router-dom' import { dataAuditApi, AuditConversation } from '../services/api' import { useJobPolling } from '../hooks/useJobPolling' import AuditFilterBar from '../components/dataAudit/AuditFilterBar' @@ -40,8 +41,9 @@ function loadSelection(): Set<string> { return new Set() } -function loadFilters(): Record<string, unknown> { +function loadFilters(datasetId: string | null): Record<string, unknown> { const defaults = defaultFilterValues() + if (datasetId) return { ...defaults, dataset: datasetId } try { const raw = sessionStorage.getItem(FILTERS_STORAGE_KEY) if (raw) return { ...defaults, ...JSON.parse(raw) } @@ -51,7 +53,8 @@ function loadFilters(): Record<string, unknown> { return defaults } -function loadArchivedView(): boolean { +function loadArchivedView(datasetId: string | null): boolean { + if (datasetId) return false try { return sessionStorage.getItem(ARCHIVED_VIEW_STORAGE_KEY) === 'true' } catch { @@ -66,14 +69,21 @@ const REASON_LABELS: Record<ArchiveReason, string> = { } export default function DataAudit() { + const [searchParams, setSearchParams] = useSearchParams() + const initialDatasetId = searchParams.get('dataset') // Filter values keyed by AUDIT_FILTERS registry keys — initialized from // sessionStorage (lazy) so they survive navigating to a conversation detail // page and back. - const [filters, setFilters] = useState<Record<string, unknown>>(() => loadFilters()) - const [archivedOnly, setArchivedOnly] = useState(() => loadArchivedView()) + const [filters, setFilters] = useState<Record<string, unknown>>(() => + loadFilters(initialDatasetId) + ) + const [archivedOnly, setArchivedOnly] = useState(() => + loadArchivedView(initialDatasetId) + ) // Data const [speakers, setSpeakers] = useState<string[]>([]) + const [datasets, setDatasets] = useState<string[]>([]) const [rows, setRows] = useState<AuditConversation[]>([]) const [total, setTotal] = useState(0) const [scanCapped, setScanCapped] = useState(false) @@ -131,6 +141,7 @@ export default function DataAudit() { // Speaker filter options come from the same scan, so the list only // contains speakers actually present in the current view. setSpeakers(res.data.speakers || []) + setDatasets(res.data.datasets || []) // Prune (not clear) the selection: keep selected rows that are still // visible so navigation/refresh doesn't lose a curation in progress. const visible = new Set(res.data.conversations.map((c) => c.conversation_id)) @@ -142,6 +153,15 @@ export default function DataAudit() { } }, [filters, archivedOnly]) + useEffect(() => { + const dataset = typeof filters.dataset === 'string' ? filters.dataset.trim() : '' + if ((searchParams.get('dataset') || '') === dataset) return + const next = new URLSearchParams(searchParams) + if (dataset) next.set('dataset', dataset) + else next.delete('dataset') + setSearchParams(next, { replace: true }) + }, [filters.dataset, searchParams, setSearchParams]) + // Persist selection whenever it changes. useEffect(() => { try { @@ -347,15 +367,25 @@ export default function DataAudit() { return ( <div className="space-y-6"> {/* Header */} - <div className="flex items-center space-x-3"> + <div className="flex flex-wrap items-center gap-3"> <Sparkles className="h-6 w-6 text-blue-600" /> - <div> + <div className="flex-1 min-w-[240px]"> <h2 className="text-xl font-semibold text-gray-900 dark:text-gray-100">Data Audit</h2> <p className="text-sm text-gray-500 dark:text-gray-400"> Inspect recordings: find speech-free or mis-attributed audio, split long recordings at silence gaps, merge adjacent conversations, and archive audio. </p> </div> + {!archivedOnly && ( + <Link + to="/speaker-enrollment" + className="inline-flex items-center gap-2 px-3 py-2 rounded border border-gray-300 dark:border-gray-600 text-sm font-medium text-gray-700 dark:text-gray-200 hover:bg-gray-50 dark:hover:bg-gray-700" + > + <Mic className="h-4 w-4" /> + Speaker enrollment + <ArrowRight className="h-4 w-4" /> + </Link> + )} </div> {/* View toggle */} @@ -405,7 +435,7 @@ export default function DataAudit() { loadConversations(next) }} onApply={() => loadConversations()} - ctx={{ speakers }} + ctx={{ speakers, datasets }} loading={loading} /> )} @@ -432,6 +462,8 @@ export default function DataAudit() { {/* Speaker confidence overview (per-speaker baselines + noise magnets) */} {!archivedOnly && <SpeakerConfidencePanel />} + {/* Guided enrollment: confirm high-information clips to improve a voiceprint */} + {/* Drift: conversations whose speaker labels would change under the current gallery */} {!archivedOnly && <DriftPanel />} diff --git a/backends/advanced/webui/src/pages/SpeakerEnrollment.tsx b/backends/advanced/webui/src/pages/SpeakerEnrollment.tsx new file mode 100644 index 00000000..6c3971c9 --- /dev/null +++ b/backends/advanced/webui/src/pages/SpeakerEnrollment.tsx @@ -0,0 +1,5 @@ +import GuidedEnrollment from '../components/dataAudit/GuidedEnrollment' + +export default function SpeakerEnrollment() { + return <GuidedEnrollment /> +} diff --git a/backends/advanced/webui/src/pages/Upload.tsx b/backends/advanced/webui/src/pages/Upload.tsx index 92a360bd..38cb4ba7 100644 --- a/backends/advanced/webui/src/pages/Upload.tsx +++ b/backends/advanced/webui/src/pages/Upload.tsx @@ -1,10 +1,23 @@ -import React, { useState, useCallback } from 'react' -import { Upload as UploadIcon, File, X, CheckCircle, AlertCircle, RefreshCw } from 'lucide-react' -import { uploadApi } from '../services/api' +import React, { useCallback, useRef, useState } from 'react' +import { + AlertCircle, + ArrowRight, + Brain, + CheckCircle, + File, + FileArchive, + PenLine, + RefreshCw, + Upload as UploadIcon, + X, +} from 'lucide-react' +import { useNavigate } from 'react-router-dom' +import { dataAuditApi, uploadApi } from '../services/api' import { useAuth } from '../contexts/AuthContext' const SUPPORTED_EXTENSIONS = ['.wav', '.mp3', '.m4a', '.flac', '.ogg', '.mp4', '.webm'] const VIDEO_EXTENSIONS = ['.mp4', '.webm'] +type UploadMode = 'memory' | 'annotation' interface UploadFile { file: File @@ -14,12 +27,22 @@ interface UploadFile { } export default function Upload() { + const navigate = useNavigate() const [files, setFiles] = useState<UploadFile[]>([]) const [isUploading, setIsUploading] = useState(false) const [dragActive, setDragActive] = useState(false) const [uploadProgress, setUploadProgress] = useState(0) const [gdriveFolderId, setGdriveFolderId] = useState('') const [videoWarning, setVideoWarning] = useState(false) + const [uploadMode, setUploadMode] = useState<UploadMode>('memory') + const [uploadSummary, setUploadSummary] = useState('') + const [importedDataset, setImportedDataset] = useState<{ + id: string + clipCount: number + } | null>(null) + const fileInputRef = useRef<HTMLInputElement>(null) + + const annotationOnly = uploadMode === 'annotation' const { isAdmin } = useAuth() @@ -44,6 +67,7 @@ export default function Upload() { await uploadApi.uploadFromGDriveFolder({ gdrive_folder_id: gdriveFolderId, device_name: 'upload', + annotation_only: annotationOnly, }) setGdriveUploadStatus({ @@ -67,6 +91,7 @@ export default function Upload() { const acceptedFiles = Array.from(selectedFiles).filter((file) => { const ext = '.' + file.name.split('.').pop()?.toLowerCase() + if (annotationOnly && ext === '.zip') return true return ( file.type.startsWith('audio/') || file.type.startsWith('video/') || @@ -74,6 +99,20 @@ export default function Upload() { ) }) + const datasetFiles = acceptedFiles.filter((file) => file.name.toLowerCase().endsWith('.zip')) + if (datasetFiles.length > 0) { + const datasetFile = datasetFiles[0] + setFiles([{ + file: datasetFile, + id: generateId(), + status: 'pending', + }]) + setVideoWarning(false) + setUploadSummary('') + setImportedDataset(null) + return + } + const hasVideo = acceptedFiles.some((file) => { const ext = '.' + file.name.split('.').pop()?.toLowerCase() return file.type.startsWith('video/') || VIDEO_EXTENSIONS.includes(ext) @@ -87,6 +126,8 @@ export default function Upload() { })) setFiles((prev) => [...prev, ...newFiles]) + setUploadSummary('') + setImportedDataset(null) } const removeFile = (id: string) => { @@ -105,41 +146,79 @@ export default function Upload() { e.stopPropagation() setDragActive(false) handleFileSelect(e.dataTransfer.files) - }, []) + }, [annotationOnly]) + + const selectMode = (mode: UploadMode) => { + if (mode === uploadMode) return + setUploadMode(mode) + setFiles([]) + setVideoWarning(false) + setUploadProgress(0) + setUploadSummary('') + setImportedDataset(null) + if (fileInputRef.current) fileInputRef.current.value = '' + } const uploadFiles = async () => { - if (files.length === 0) return + const pendingFiles = files.filter((file) => file.status === 'pending') + if (pendingFiles.length === 0) return setIsUploading(true) setUploadProgress(0) try { + const pendingIds = new Set(pendingFiles.map((file) => file.id)) + const isDatasetImport = + annotationOnly && + pendingFiles.length === 1 && + pendingFiles[0].file.name.toLowerCase().endsWith('.zip') const formData = new FormData() - files.forEach(({ file }) => { - formData.append('files', file) - }) + if (isDatasetImport) { + formData.append('dataset', pendingFiles[0].file) + } else { + pendingFiles.forEach(({ file }) => formData.append('files', file)) + } setFiles((prev) => - prev.map((f) => ({ ...f, status: 'uploading' })) + prev.map((f) => pendingIds.has(f.id) ? { ...f, status: 'uploading' } : f) ) - await uploadApi.uploadAudioFiles(formData, (progress) => { - setUploadProgress(progress) - }) + const response = isDatasetImport + ? await dataAuditApi.importDataset(formData, setUploadProgress) + : await uploadApi.uploadAudioFiles(formData, setUploadProgress, { annotationOnly }) setFiles((prev) => - prev.map((f) => ({ ...f, status: 'success' })) + prev.map((f) => pendingIds.has(f.id) ? { ...f, status: 'success' } : f) + ) + setUploadSummary( + isDatasetImport + ? `${response.data.summary.imported} clips imported for review${response.data.summary.skipped ? `, ${response.data.summary.skipped} already present` : ''}.` + : annotationOnly + ? `${pendingFiles.length} file${pendingFiles.length === 1 ? '' : 's'} sent to the annotation workspace.` + : `${pendingFiles.length} file${pendingFiles.length === 1 ? '' : 's'} sent for processing.` ) + if (isDatasetImport) { + setImportedDataset({ + id: response.data.dataset_id, + clipCount: response.data.summary.imported + response.data.summary.skipped, + }) + } } catch (err: any) { console.error('Upload failed:', err) setFiles((prev) => - prev.map((f) => ({ - ...f, - status: 'error', - error: err.message || 'Upload failed', - })) + prev.map((f) => + f.status === 'uploading' + ? { + ...f, + status: 'error', + error: err?.response?.data?.error || err.message || 'Upload failed', + } + : f + ) ) + setUploadSummary('') + setImportedDataset(null) } finally { setIsUploading(false) setUploadProgress(100) @@ -194,31 +273,67 @@ export default function Upload() { <div className="flex items-center space-x-2 mb-6"> <UploadIcon className="h-6 w-6 text-blue-600" /> <h1 className="text-2xl font-bold text-gray-900 dark:text-gray-100"> - Upload Audio Files + Upload Audio </h1> </div> + <div className="mb-6"> + <div className="inline-grid grid-cols-2 gap-1 rounded-lg bg-gray-100 p-1 dark:bg-gray-800"> + <button + type="button" + aria-pressed={uploadMode === 'memory'} + onClick={() => selectMode('memory')} + className={`flex min-h-10 items-center justify-center gap-2 rounded-md px-4 py-2 text-sm font-medium transition-colors ${ + uploadMode === 'memory' + ? 'bg-white text-gray-900 shadow-sm dark:bg-gray-700 dark:text-gray-100' + : 'text-gray-600 hover:text-gray-900 dark:text-gray-400 dark:hover:text-gray-100' + }`} + > + <Brain className="h-4 w-4" /> + Process memories + </button> + <button + type="button" + aria-pressed={uploadMode === 'annotation'} + onClick={() => selectMode('annotation')} + className={`flex min-h-10 items-center justify-center gap-2 rounded-md px-4 py-2 text-sm font-medium transition-colors ${ + uploadMode === 'annotation' + ? 'bg-white text-gray-900 shadow-sm dark:bg-gray-700 dark:text-gray-100' + : 'text-gray-600 hover:text-gray-900 dark:text-gray-400 dark:hover:text-gray-100' + }`} + > + <PenLine className="h-4 w-4" /> + Annotation workspace + </button> + </div> + <p className="mt-2 text-sm text-gray-600 dark:text-gray-400"> + {annotationOnly + ? 'Import a Chronicle dataset with transcripts, or transcribe new audio without changing memory.' + : 'Transcribe audio and run the normal memory pipeline.'} + </p> + </div> + {/* Google Drive Folder Upload */} <div className="mb-6 p-4 bg-gray-50 dark:bg-gray-700 rounded-lg border border-gray-200 dark:border-gray-600"> <label className="block mb-2 font-medium text-gray-900 dark:text-gray-100"> - Paste Google Drive Folder ID: + Google Drive folder ID </label> - <div className="flex space-x-2"> + <div className="flex flex-col gap-2 sm:flex-row"> <input type="text" value={gdriveFolderId} onChange={(e) => setGdriveFolderId(e.target.value)} placeholder="1AbCdEfGhIjKlMnOpQrStUvWxYz123456" - className="flex-1 px-3 py-2 border rounded-lg dark:bg-gray-800 dark:text-gray-100" + className="min-w-0 flex-1 px-3 py-2 border rounded-lg dark:bg-gray-800 dark:text-gray-100" /> <button onClick={handleGDriveSubmit} disabled={isUploading || !gdriveFolderId} - className="px-4 py-2 bg-blue-600 text-white rounded-lg hover:bg-blue-700 disabled:opacity-50" + className="w-full whitespace-nowrap px-4 py-2 bg-blue-600 text-white rounded-lg hover:bg-blue-700 disabled:opacity-50 sm:w-auto" > - {isUploading ? 'Submitting...' : 'Submit Folder'} + {isUploading ? 'Submitting...' : annotationOnly ? 'Import for review' : 'Process folder'} </button> </div> @@ -249,25 +364,29 @@ export default function Upload() { > <UploadIcon className="h-12 w-12 mx-auto mb-4 text-gray-400" /> <h3 className="text-lg font-medium text-gray-900 dark:text-gray-100 mb-2"> - Drop audio files here or click to browse + {annotationOnly ? 'Drop audio or a Chronicle dataset ZIP' : 'Drop audio files here'} </h3> <p className="text-sm text-gray-600 dark:text-gray-400 mb-4"> - Supported formats: WAV, MP3, M4A, FLAC, OGG, MP4, WebM + {annotationOnly + ? 'ZIP datasets keep their existing transcripts; audio files are transcribed.' + : 'WAV, MP3, M4A, FLAC, OGG, MP4, or WebM'} </p> <input + ref={fileInputRef} type="file" multiple - accept="audio/*,video/mp4,video/webm,.wav,.mp3,.m4a,.flac,.ogg,.mp4,.webm" + accept={`${annotationOnly ? '.zip,' : ''}audio/*,video/mp4,video/webm,.wav,.mp3,.m4a,.flac,.ogg,.mp4,.webm`} onChange={(e) => handleFileSelect(e.target.files)} className="absolute inset-0 w-full h-full opacity-0 cursor-pointer" /> <button - onClick={() => (document.querySelector('input[type="file"]') as HTMLInputElement)?.click()} + type="button" + onClick={() => fileInputRef.current?.click()} className="px-4 py-2 bg-blue-600 text-white rounded-lg hover:bg-blue-700" > - Select Files + Select files </button> </div> @@ -300,7 +419,13 @@ export default function Upload() { disabled={isUploading || files.every((f) => f.status !== 'pending')} className="px-4 py-2 bg-blue-600 text-white rounded-lg hover:bg-blue-700 disabled:opacity-50" > - {isUploading ? 'Uploading...' : 'Upload All'} + {isUploading + ? 'Uploading...' + : annotationOnly + ? files.some((f) => f.file.name.toLowerCase().endsWith('.zip')) + ? 'Import dataset' + : 'Add to annotation workspace' + : 'Process files'} </button> </div> </div> @@ -358,6 +483,27 @@ export default function Upload() { </div> )} + {uploadSummary && ( + <div className="mt-6 flex flex-col gap-3 rounded-lg border border-green-200 bg-green-50 p-4 text-sm text-green-800 dark:border-green-800 dark:bg-green-900/20 dark:text-green-300 sm:flex-row sm:items-center sm:justify-between"> + <div className="flex items-start gap-2"> + <CheckCircle className="mt-0.5 h-4 w-4 flex-shrink-0" /> + <span>{uploadSummary}</span> + </div> + {importedDataset && importedDataset.clipCount > 0 && ( + <button + type="button" + onClick={() => + navigate(`/data-audit?dataset=${encodeURIComponent(importedDataset.id)}`) + } + className="flex min-h-10 w-full items-center justify-center gap-2 rounded-md bg-green-700 px-4 py-2 font-medium text-white hover:bg-green-800 sm:w-auto" + > + Review {importedDataset.clipCount} clip{importedDataset.clipCount === 1 ? '' : 's'} + <ArrowRight className="h-4 w-4" /> + </button> + )} + </div> + )} + {/* Upload Progress */} {isUploading && ( <div className="mt-6 p-4 bg-blue-50 dark:bg-blue-900/20 border border-blue-200 dark:border-blue-800 rounded-lg"> @@ -381,18 +527,18 @@ export default function Upload() { </div> )} - {/* Instructions */} - <div className="mt-8 bg-yellow-50 dark:bg-yellow-900/20 border border-yellow-200 dark:border-yellow-800 rounded-lg p-4"> - <h3 className="font-medium text-yellow-800 dark:text-yellow-200 mb-2"> - 📝 Upload Instructions - </h3> - <ul className="text-sm text-yellow-700 dark:text-yellow-300 space-y-1"> - <li>• Audio files will be processed sequentially for transcription and memory extraction</li> - <li>• Processing time varies based on audio length (roughly 3× duration + 60s)</li> - <li>• Large files or multiple files may cause timeout errors</li> - <li>• Check the Conversations tab for processed results</li> - <li>• Supported formats: WAV, MP3, M4A, FLAC, OGG, MP4, WebM</li> - </ul> + <div className="mt-8 border-t border-gray-200 pt-4 text-sm text-gray-600 dark:border-gray-700 dark:text-gray-400"> + {annotationOnly ? ( + <div className="flex items-center gap-2"> + <FileArchive className="h-4 w-4" /> + Imported clips remain editable conversations but are permanently excluded from memory. + </div> + ) : ( + <div className="flex items-center gap-2"> + <Brain className="h-4 w-4" /> + Uploaded audio follows the full transcription and memory pipeline. + </div> + )} </div> </div> diff --git a/backends/advanced/webui/src/services/api.ts b/backends/advanced/webui/src/services/api.ts index aae93687..bf7e04d3 100644 --- a/backends/advanced/webui/src/services/api.ts +++ b/backends/advanced/webui/src/services/api.ts @@ -148,6 +148,10 @@ export const conversationsApi = { // Conversations whose speaker labels would change under the current gallery (admin). getDrift: () => api.get('/api/conversations/drift'), + backfillDriftClusterEmbeddings: () => + api.post<{ job_id: string; status: string }>( + '/api/conversations/drift/backfill-cluster-embeddings' + ), // Version management (transcript only — memory is no longer versioned) activateTranscriptVersion: (conversationId: string, versionId: string) => api.post(`/api/conversations/${conversationId}/activate-transcript/${versionId}`), @@ -641,8 +645,12 @@ export const queueApi = { } export const uploadApi = { - uploadAudioFiles: (files: FormData, onProgress?: (progress: number) => void) => - api.post('/api/audio/upload', files, { + uploadAudioFiles: ( + files: FormData, + onProgress?: (progress: number) => void, + options?: { annotationOnly?: boolean } + ) => + api.post(options?.annotationOnly ? '/api/audio/upload/annotation' : '/api/audio/upload', files, { headers: { 'Content-Type': 'multipart/form-data' }, timeout: 300000, // 5 minutes onUploadProgress: (progressEvent) => { @@ -653,8 +661,8 @@ export const uploadApi = { }, }), - uploadFromGDriveFolder: (payload: { gdrive_folder_id: string; device_name?: string }) => - api.post('/api/audio/upload_audio_from_gdrive', null, { + uploadFromGDriveFolder: (payload: { gdrive_folder_id: string; device_name?: string; annotation_only?: boolean }) => + api.post(payload.annotation_only ? '/api/audio/upload_audio_from_gdrive/annotation' : '/api/audio/upload_audio_from_gdrive', null, { params: { gdrive_folder_id: payload.gdrive_folder_id, device_name: payload.device_name, @@ -810,6 +818,26 @@ export interface AuditListResponse { // Distinct speaker labels present in the scanned working set, so the // speaker filter offers only labels that exist in the current view. speakers: string[] + // Annotation dataset IDs available to the dataset filter. + datasets: string[] +} + +export interface AnnotationImportResponse { + dataset_id: string + schema_version: number + message: string + results: Array<{ + clip_id: string + status: 'imported' | 'skipped' | 'error' + conversation_id?: string + error?: string + }> + summary: { + total: number + imported: number + skipped: number + failed: number + } } export interface SpeakerConfidenceRow { @@ -968,6 +996,171 @@ export interface ScreenResult { totals: { conversation_count: number; flagged_segments: number } } +export interface GuidedEnrollmentClip { + conversation_id: string + conversation_title: string | null + conversation_date: string + conversation_duration: number + segment_index: number + start: number + end: number + duration: number + text: string + current_label: string | null + stored_confidence: number | null + scores: { + duration: number | null + sim_centroid: number + max_clip_sim: number | null + n_gallery_clips: number + best_other: { speaker_id: string; name: string; score: number } | null + } + info_score: number + novelty: number + uncertainty: number + reasons: string[] +} + +export interface GuidedEnrollmentSpeaker { + speaker_id: string + speaker_name: string + n_clips: number | null + total_duration_s: number | null +} + +export interface GuidedEnrollmentSuggestResponse { + speaker: GuidedEnrollmentSpeaker + threshold: number + batch: GuidedEnrollmentClip[] + scanned: number + gated_out: number + pool_remaining: number + reviewed_total: number + discovery_indexed: boolean + discovery_candidates: number +} + +export interface GuidedEnrollmentDecideResponse { + speaker: GuidedEnrollmentSpeaker | null + enrolled: number + reassigned: number + rejected: number + skipped: number + multiple_speakers: number + bad_clips: number + health_before: GuidedEnrollmentHealth | null + health_after: GuidedEnrollmentHealth | null + coverage: { accepted_novelty_mean: number | null } + benchmark_job_id: string | null + discovery_job_id: string | null + errors: { clip: unknown; error: string }[] + status: string +} + +export interface GuidedEnrollmentHealth { + n_clips: number + median_self: number | null + n_flagged: number + flagged_rate: number + verdict: string +} + +export interface GuidedEnrollmentSession { + created_at: string + health_before: GuidedEnrollmentHealth | null + health_after: GuidedEnrollmentHealth | null + coverage: { accepted_novelty_mean: number | null } + decisions: { + enrolled: number + reassigned: number + rejected: number + skipped: number + multiple_speakers: number + bad_clips: number + } +} + +export interface SpeakerBenchmarkReport { + created_at: string + protocol: string + threshold: number + embedding_model: string + conversation_groups: number + dataset: { + labeled_clips: number + embedded_clips: number + speakers: number + cache_hits: number + embedding_failures: number + exclusions: Record<string, number> + } + learning_curve: Array<{ + fraction: number + train_clips_mean: number + top1_accuracy_mean: number | null + macro_recall_mean: number | null + false_accept_rate_mean: number | null + eer_mean: number | null + }> + folds: Array<{ + fold: number + train_clips: number + test_clips: number + top1_accuracy: number | null + macro_recall: number | null + false_accept_rate: number | null + eer: number | null + per_speaker_recall: Record<string, number> + confusion: Record<string, Record<string, number>> + }> +} + +export interface SpeakerGalleryBaseline { + cutoff: string | null + status: string + limitations?: string + speakers: Array<{ + speaker_id: string + name: string + baseline: { n_clips: number; median_self: number | null; n_flagged: number; verdict: string } + current: { n_clips: number; median_self: number | null; n_flagged: number; verdict: string } | null + }> +} + +export interface GuidedEnrollmentGalleryClip { + segment_id: number + filename: string + duration: number + self_score: number | null + best_other: { speaker_id: string; name: string; score: number } | null + flags: string[] + suggested: { speaker_id: string; name: string; score: number } | null +} + +export interface GuidedEnrollmentGalleryResponse { + speaker: { + speaker_id: string + speaker_name: string + n_clips: number | null + total_duration_s: number | null + } + verdict: string | null + median_self: number | null + clips: GuidedEnrollmentGalleryClip[] +} + +export interface GuidedEnrollmentResetResponse { + speaker_name: string + deleted: { + reviews: number + sessions: number + discovery_matches: number + discovery_runs: number + } + gallery_deleted: boolean + status: string +} + export const dataAuditApi = { // Enqueue batch VAD analysis. Returns { job_id, status }. analyze: (conversationIds?: string[], force: boolean = false) => @@ -1054,6 +1247,108 @@ export const dataAuditApi = { { start, end } ), + // Guided enrollment: next batch of highest-information clips for a speaker. + guidedEnrollmentSuggest: ( + speakerName: string, + order: 'informative' | 'confidence' = 'informative' + ) => + api.post<GuidedEnrollmentSuggestResponse>( + '/api/data-audit/enrollment/guided/suggest', + { speaker_name: speakerName, batch_size: 5, order } + ), + + // Guided enrollment: record accept/reject decisions; accepted clips enroll. + guidedEnrollmentDecide: ( + speakerName: string, + decisions: { + conversation_id: string + start: number + end: number + original_start: number + original_end: number + decision: 'accept' | 'reject' | 'skip' | 'bad_clip' | 'multiple_speakers' | 'another_speaker' + actual_speaker?: string + scores?: GuidedEnrollmentClip['scores'] + }[] + ) => + api.post<GuidedEnrollmentDecideResponse>( + '/api/data-audit/enrollment/guided/decide', + { speaker_name: speakerName, decisions } + ), + + guidedEnrollmentHistory: (speakerName: string) => + api.get<{ speaker_name: string; sessions: GuidedEnrollmentSession[] }>( + '/api/data-audit/enrollment/guided/history', + { params: { speaker_name: speakerName } } + ), + + discoverSpeakerCorpus: (speakerName: string, includeDeleted = false) => + api.post<{ job_id: string; status: string; reused: boolean }>( + '/api/data-audit/enrollment/guided/discover', + { speaker_name: speakerName, batch_size: 5, include_deleted: includeDeleted } + ), + + // Upload an unlabelled audio corpus and mine it for one speaker: files become + // annotation-only conversations; discovery is chained after transcription. + mineCorpusAudio: (speakerName: string, files: File[]) => { + const form = new FormData() + form.append('speaker_name', speakerName) + files.forEach((f) => form.append('files', f)) + return api.post<{ + speaker_name: string + ingested: number + failed: { filename: string; error: string }[] + transcription_jobs: number + transcription_available: boolean + discovery_job_id: string | null + }>('/api/data-audit/enrollment/guided/mine', form, { + headers: { 'Content-Type': 'multipart/form-data' }, + }) + }, + + getSpeakerCorpusDiscovery: (speakerName: string) => + api.get<{ + speaker_name: string + job_id: string | null + status: string | null + matched_segments: number + }>('/api/data-audit/enrollment/guided/discover', { + params: { speaker_name: speakerName }, + }), + + runSpeakerBenchmark: () => + api.post<{ job_id: string; status: string }>('/api/data-audit/enrollment/benchmark'), + + getLatestSpeakerBenchmark: () => + api.get<{ report: SpeakerBenchmarkReport | null }>( + '/api/data-audit/enrollment/benchmark/latest' + ), + + getSpeakerGalleryBaseline: () => + api.get<SpeakerGalleryBaseline>('/api/data-audit/enrollment/baseline'), + + // Enrolled clips for one speaker with per-clip contamination flags. + getEnrollmentGallery: (speakerName: string) => + api.get<GuidedEnrollmentGalleryResponse>( + '/api/data-audit/enrollment/guided/gallery', + { params: { speaker_name: speakerName } } + ), + + // Remove one enrolled clip from the voiceprint (quarantined by default). + deleteEnrollmentGalleryClip: (speakerName: string, segmentId: number, hard = false) => + api.post( + `/api/data-audit/enrollment/guided/gallery/segments/${segmentId}/delete`, + { speaker_name: speakerName, hard } + ), + + // Forget all guided-enrollment state for a speaker name (reviews, history, + // discovery); optionally also purge the voiceprint gallery. + resetGuidedEnrollment: (speakerName: string, purgeGallery = false) => + api.post<GuidedEnrollmentResetResponse>( + '/api/data-audit/enrollment/guided/reset', + { speaker_name: speakerName, purge_gallery: purgeGallery } + ), + // Count of unapplied triage decisions and conversations they span. getTriagePending: () => api.get<{ pending_count: number; conversation_count: number }>( @@ -1108,6 +1403,19 @@ export const dataAuditApi = { ...params, }), + // Import an export-compatible ZIP as memory-excluded conversations with the + // manifest transcripts already active in the editor. + importDataset: (dataset: FormData, onProgress?: (progress: number) => void) => + api.post<AnnotationImportResponse>('/api/data-audit/import', dataset, { + headers: { 'Content-Type': 'multipart/form-data' }, + timeout: 300000, + onUploadProgress: (progressEvent) => { + if (onProgress && progressEvent.total) { + onProgress(Math.round((progressEvent.loaded * 100) / progressEvent.total)) + } + }, + }), + // List completed annotation exports listExports: () => api.get<{ exports: ExportRecord[] }>('/api/data-audit/exports'), diff --git a/backends/advanced/webui/src/utils/speakerLabels.ts b/backends/advanced/webui/src/utils/speakerLabels.ts new file mode 100644 index 00000000..f86ceda0 --- /dev/null +++ b/backends/advanced/webui/src/utils/speakerLabels.ts @@ -0,0 +1,10 @@ +const NUMBERED_UNKNOWN_SPEAKER = /^unknown speaker\s+(\d+)$/i + +export function nextUnknownSpeakerLabel(speakerNames: string[]): string { + const highestNumber = speakerNames.reduce((highest, name) => { + const match = NUMBERED_UNKNOWN_SPEAKER.exec(name.trim()) + return match ? Math.max(highest, Number(match[1])) : highest + }, 0) + + return `Unknown Speaker ${highestNumber + 1}` +} diff --git a/backends/advanced/worker_healthcheck.py b/backends/advanced/worker_healthcheck.py index 6ceb6f60..46001654 100644 --- a/backends/advanced/worker_healthcheck.py +++ b/backends/advanced/worker_healthcheck.py @@ -7,7 +7,7 @@ a stream-consumer that's alive-but-wedged. This probe makes those visible by failing (exit 1) when: - 1. fewer than ``MIN_RQ_WORKERS`` RQ workers are registered in Redis, or + 1. fewer than ``MIN_RQ_WORKERS`` fresh RQ workers are registered in Redis, or 2. any stream-consumer heartbeat (``worker:heartbeat:*``) is stale. Used as the ``workers`` service healthcheck in docker-compose.yml. @@ -19,6 +19,8 @@ from redis import Redis +from advanced_omi_backend.heartbeat import is_rq_worker_fresh + REDIS_URL = os.getenv("REDIS_URL", "redis://redis:6379/0") MIN_RQ_WORKERS = int(os.getenv("MIN_RQ_WORKERS", "6")) # A stream-consumer beats once per ~1s main-loop iteration; 90s of silence means @@ -35,11 +37,14 @@ def main() -> int: print(f"unhealthy: redis unreachable: {e}") return 1 - # 1. RQ worker registration count (a dead/wedged RQ worker drops out of this). + # 1. Fresh RQ worker registrations. Redis can retain registrations from dead + # containers, so Worker.all() alone is not a liveness signal. try: from rq import Worker - rq_count = len(Worker.all(connection=redis_client)) + rq_count = sum( + is_rq_worker_fresh(worker) for worker in Worker.all(connection=redis_client) + ) except Exception as e: # noqa: BLE001 print(f"unhealthy: could not count RQ workers: {e}") return 1 From 667f57162c1ac4943a0c4a289aff1d800eba5a80 Mon Sep 17 00:00:00 2001 From: Ankush Malaker <43288948+AnkushMalaker@users.noreply.github.com> Date: Fri, 17 Jul 2026 14:10:37 +0000 Subject: [PATCH 02/16] feat(asr): add fine-tuning and evaluation tooling --- .../gemma4/UNIFIED_ENDPOINT_DESIGN.md | 178 +++++++++ .../providers/gemma4/finetune/README.md | 128 +++++++ .../providers/gemma4/finetune/ab_empties.py | 97 +++++ .../providers/gemma4/finetune/align_spike.py | 79 ++++ .../gemma4/finetune/bench_coshe_12b.py | 342 ++++++++++++++++++ .../gemma4/finetune/build_windowed_dataset.py | 130 +++++++ .../providers/gemma4/finetune/data.py | 251 +++++++++++++ .../gemma4/finetune/data_interleave.py | 132 +++++++ .../gemma4/finetune/data_windowed.py | 57 +++ .../providers/gemma4/finetune/data_wispr.py | 15 + .../gemma4/finetune/debug_gen_nocache.py | 66 ++++ .../gemma4/finetune/debug_inprocess.py | 137 +++++++ .../gemma4/finetune/debug_overfit.py | 186 ++++++++++ .../gemma4/finetune/debug_reload_loss.py | 51 +++ .../gemma4/finetune/eval_interleave.py | 107 ++++++ .../gemma4/finetune/eval_test_split.py | 152 ++++++++ .../providers/gemma4/finetune/eval_wer.py | 176 +++++++++ .../gemma4/finetune/eval_windowed_stitch.py | 319 ++++++++++++++++ .../gemma4/finetune/eval_wispr_stitch.py | 148 ++++++++ .../gemma4/finetune/gemini_transcribe.py | 69 ++++ .../gemma4/finetune/gemma4_cache_patch.py | 64 ++++ .../providers/gemma4/finetune/infer.py | 169 +++++++++ .../providers/gemma4/finetune/introspect.py | 162 +++++++++ .../gemma4/finetune/poll_until_target.sh | 32 ++ .../gemma4/finetune/run_split_pipeline.sh | 49 +++ .../gemma4/finetune/smoke_interleave.py | 93 +++++ .../providers/gemma4/finetune/split_status.py | 36 ++ .../gemma4/finetune/stream_client_test.py | 71 ++++ .../providers/gemma4/finetune/sweep_12b_v2.py | 254 +++++++++++++ .../providers/gemma4/finetune/test_cache57.py | 90 +++++ .../gemma4/finetune/test_longaudio.py | 64 ++++ .../gemma4/finetune/test_patch_loss.py | 131 +++++++ .../providers/gemma4/finetune/train.py | 203 +++++++++++ .../gemma4/finetune/train_interleave.py | 150 ++++++++ .../gemma4/finetune/train_lora_split.py | 199 ++++++++++ .../gemma4/finetune/train_lora_windowed.py | 160 ++++++++ .../gemma4/finetune/train_until_wer.py | 302 ++++++++++++++++ .../providers/gemma4/finetune/train_wispr.py | 168 +++++++++ .../providers/gemma4/finetune/verify_fix.py | 53 +++ .../providers/gemma4/finetune/vm_exec.py | 90 +++++ .../providers/gemma4/finetune/window_coshe.py | 129 +++++++ .../gemma4/finetune/window_target.py | 25 ++ .../providers/nemo/finetune/cut_segments.py | 94 +++++ .../providers/nemo/finetune/eval_full.py | 120 ++++++ .../providers/nemo/finetune/fa_align.py | 114 ++++++ .../providers/nemo/finetune/make_manifest.py | 167 +++++++++ .../providers/nemo/finetune/run_deepgram.py | 125 +++++++ .../nemo/finetune/segment_by_alignment.py | 206 +++++++++++ .../providers/nemo/finetune/train_full.py | 154 ++++++++ .../providers/nemo/finetune/train_overfit.py | 204 +++++++++++ .../providers/nemo/finetune/train_split.py | 148 ++++++++ .../providers/qwen3_asr/finetune/README.md | 94 +++++ .../qwen3_asr/finetune/coshe_infer.py | 168 +++++++++ .../qwen3_asr/finetune/introspect.py | 62 ++++ .../qwen3_asr/finetune/make_manifest.py | 80 ++++ .../qwen3_asr/finetune/make_manifest_full.py | 70 ++++ .../qwen3_asr/finetune/qwen3_asr_sft.py | 334 +++++++++++++++++ .../finetune/train_qwen3_until_wer.py | 314 ++++++++++++++++ .../providers/qwen3_asr/finetune/verify.py | 76 ++++ extras/asr-services/pyproject.toml | 29 ++ .../asr-services/scripts/validate_nemotron.py | 143 ++++++++ .../tests/capture_vibevoice_ground_truth.py | 168 +++++++++ extras/asr-services/uv.lock | 153 +++++++- 63 files changed, 8523 insertions(+), 14 deletions(-) create mode 100644 extras/asr-services/providers/gemma4/UNIFIED_ENDPOINT_DESIGN.md create mode 100644 extras/asr-services/providers/gemma4/finetune/README.md create mode 100644 extras/asr-services/providers/gemma4/finetune/ab_empties.py create mode 100644 extras/asr-services/providers/gemma4/finetune/align_spike.py create mode 100644 extras/asr-services/providers/gemma4/finetune/bench_coshe_12b.py create mode 100644 extras/asr-services/providers/gemma4/finetune/build_windowed_dataset.py create mode 100644 extras/asr-services/providers/gemma4/finetune/data.py create mode 100644 extras/asr-services/providers/gemma4/finetune/data_interleave.py create mode 100644 extras/asr-services/providers/gemma4/finetune/data_windowed.py create mode 100644 extras/asr-services/providers/gemma4/finetune/data_wispr.py create mode 100644 extras/asr-services/providers/gemma4/finetune/debug_gen_nocache.py create mode 100644 extras/asr-services/providers/gemma4/finetune/debug_inprocess.py create mode 100644 extras/asr-services/providers/gemma4/finetune/debug_overfit.py create mode 100644 extras/asr-services/providers/gemma4/finetune/debug_reload_loss.py create mode 100644 extras/asr-services/providers/gemma4/finetune/eval_interleave.py create mode 100644 extras/asr-services/providers/gemma4/finetune/eval_test_split.py create mode 100644 extras/asr-services/providers/gemma4/finetune/eval_wer.py create mode 100644 extras/asr-services/providers/gemma4/finetune/eval_windowed_stitch.py create mode 100644 extras/asr-services/providers/gemma4/finetune/eval_wispr_stitch.py create mode 100644 extras/asr-services/providers/gemma4/finetune/gemini_transcribe.py create mode 100644 extras/asr-services/providers/gemma4/finetune/gemma4_cache_patch.py create mode 100644 extras/asr-services/providers/gemma4/finetune/infer.py create mode 100644 extras/asr-services/providers/gemma4/finetune/introspect.py create mode 100755 extras/asr-services/providers/gemma4/finetune/poll_until_target.sh create mode 100644 extras/asr-services/providers/gemma4/finetune/run_split_pipeline.sh create mode 100644 extras/asr-services/providers/gemma4/finetune/smoke_interleave.py create mode 100644 extras/asr-services/providers/gemma4/finetune/split_status.py create mode 100644 extras/asr-services/providers/gemma4/finetune/stream_client_test.py create mode 100644 extras/asr-services/providers/gemma4/finetune/sweep_12b_v2.py create mode 100644 extras/asr-services/providers/gemma4/finetune/test_cache57.py create mode 100644 extras/asr-services/providers/gemma4/finetune/test_longaudio.py create mode 100644 extras/asr-services/providers/gemma4/finetune/test_patch_loss.py create mode 100644 extras/asr-services/providers/gemma4/finetune/train.py create mode 100644 extras/asr-services/providers/gemma4/finetune/train_interleave.py create mode 100644 extras/asr-services/providers/gemma4/finetune/train_lora_split.py create mode 100644 extras/asr-services/providers/gemma4/finetune/train_lora_windowed.py create mode 100644 extras/asr-services/providers/gemma4/finetune/train_until_wer.py create mode 100644 extras/asr-services/providers/gemma4/finetune/train_wispr.py create mode 100644 extras/asr-services/providers/gemma4/finetune/verify_fix.py create mode 100644 extras/asr-services/providers/gemma4/finetune/vm_exec.py create mode 100644 extras/asr-services/providers/gemma4/finetune/window_coshe.py create mode 100644 extras/asr-services/providers/gemma4/finetune/window_target.py create mode 100644 extras/asr-services/providers/nemo/finetune/cut_segments.py create mode 100644 extras/asr-services/providers/nemo/finetune/eval_full.py create mode 100644 extras/asr-services/providers/nemo/finetune/fa_align.py create mode 100644 extras/asr-services/providers/nemo/finetune/make_manifest.py create mode 100644 extras/asr-services/providers/nemo/finetune/run_deepgram.py create mode 100644 extras/asr-services/providers/nemo/finetune/segment_by_alignment.py create mode 100644 extras/asr-services/providers/nemo/finetune/train_full.py create mode 100644 extras/asr-services/providers/nemo/finetune/train_overfit.py create mode 100644 extras/asr-services/providers/nemo/finetune/train_split.py create mode 100644 extras/asr-services/providers/qwen3_asr/finetune/README.md create mode 100644 extras/asr-services/providers/qwen3_asr/finetune/coshe_infer.py create mode 100644 extras/asr-services/providers/qwen3_asr/finetune/introspect.py create mode 100644 extras/asr-services/providers/qwen3_asr/finetune/make_manifest.py create mode 100644 extras/asr-services/providers/qwen3_asr/finetune/make_manifest_full.py create mode 100644 extras/asr-services/providers/qwen3_asr/finetune/qwen3_asr_sft.py create mode 100644 extras/asr-services/providers/qwen3_asr/finetune/train_qwen3_until_wer.py create mode 100644 extras/asr-services/providers/qwen3_asr/finetune/verify.py create mode 100644 extras/asr-services/scripts/validate_nemotron.py create mode 100644 extras/asr-services/tests/capture_vibevoice_ground_truth.py diff --git a/extras/asr-services/providers/gemma4/UNIFIED_ENDPOINT_DESIGN.md b/extras/asr-services/providers/gemma4/UNIFIED_ENDPOINT_DESIGN.md new file mode 100644 index 00000000..4d773b42 --- /dev/null +++ b/extras/asr-services/providers/gemma4/UNIFIED_ENDPOINT_DESIGN.md @@ -0,0 +1,178 @@ +# Gemma 4 unified multimodal endpoint — design + +## Problem + +Gemma 4 E2B is **one combined multimodal model**: audio understanding and text generation +are the *same* forward pass over interleaved text+audio tokens. But the service exposes it +as a set of task-specific wrappers, and none of them accept **more than one audio clip per +request**: + +| endpoint | audio in | limitation | +|----------|----------|-----------| +| `/transcribe` | 1 file (multipart) | single audio; ASR-only framing | +| `/v1/chat/completions` | none (text-only) | `ChatMessage.content: str`; no audio at all | +| `/judge` | 1 file | single audio + transcript | +| `/stream` (ws) | PCM frames | streaming ASR only | + +Consequences: +1. **Multi-audio prompting is impossible.** Few-shot wake-word classification (show K example + clips + a candidate in one prompt) can't be expressed → today it requires loading a *second* + copy of the model in a one-off container, which OOMs the 24 GB GPU, so we **stop the live ASR + service, run, and restart** every time. Redundant model load + ASR downtime + GPU contention. +2. **Redundant encode/decode round-trips.** The OpenAI `input_audio` transport is base64. For a + *local* caller that already has a wav on disk, the path is `file → base64 → (wire) → base64-decode + → temp file → processor re-reads+decodes`. Two of those steps are pure overhead. + +## Core insight + +Everything the model does is one operation: + +``` +generate(messages) -> text # messages = interleaved [text | audio] parts +``` + +Transcribe, classify, judge, and chat are all **presets** of this. The clean shape is: + +- a single **canonical multimodal generate** path on the loaded model, +- the **OpenAI-compatible chat/completions** endpoint as the *general* interface (this is the + industry-standard way to serve a combined model — same schema vLLM/Gemini/OpenRouter use for + audio: `content: [{type:"text"}, {type:"input_audio", input_audio:{data,format}}]`), +- the ASR-specific wrappers kept thin on top. + +"Use the combined model as a general endpoint for both chat and ASR" = make `/v1/chat/completions` +multimodal; transcription becomes "chat with a transcribe prompt + one audio". + +## Design — three layers + +``` +HTTP adapters (thin; format + presets) + POST /v1/chat/completions GENERAL multimodal chat (text+audio messages) ← chat, classify, judge, 1-shot ASR + POST /transcribe ASR preset: silence-gate + >30s window/stitch + diarization-parse + NDJSON + POST /classify (optional) convenience: example clips+labels+candidate -> label + POST /judge, WS /stream existing presets + │ normalize_audio() ── decode EVERY transport to np.float32@16k mono, ONCE + ▼ +Transcriber.generate(messages_with_array_audio) -> (text, in_tok, out_tok) ← single canonical path, loaded model +``` + +### Layer 1 — canonical generate (transcriber.py) +`generate(messages, *, max_tokens, temperature=0, top_p, top_k) -> (text, in_tok, out_tok)` +- `messages`: `[{role, content:[parts]}]`; parts are `{"type":"text",...}` or + `{"type":"audio","audio": <np.ndarray>}` (mono float32 @ 16 kHz). +- The ONLY place that calls `apply_chat_template` + `model.generate` + `_decode_response`. +- No silence gate, no batching, no temp files — pure model access on the in-memory weights. +- `generate_chat()` (today text-only) becomes a thin caller → full back-compat for text chat. + +### Layer 2 — audio normalization at the boundary (the round-trip fix) +One helper `normalize_audio_parts(content) -> content` that accepts audio in any transport and +decodes **exactly once to an array** (`soundfile.read(BytesIO)` / `load_audio_file`), never to disk: + +| transport | when | cost | +|-----------|------|------| +| `input_audio` base64 | remote / cloud-parity | b64→BytesIO→array (no temp file) | +| multipart raw bytes | local, avoid +33% base64 tax | bytes→array | +| `audio_path` ref (shared RO volume) | trusted local batch | `load_audio_file(path)` → array, **zero re-encode** | + +All converge to `{"type":"audio","audio": <array>}`. The `file→b64→file→decode` chain collapses to +a single decode regardless of caller. (Validation note: confirm the installed processor accepts an +ndarray under the `audio` key; today `_transcribe_single` passes a path. If a given transformers +version wants a path/URL, normalize to the processor's accepted form once — still one decode, in a +single place.) + +### Layer 3 — HTTP endpoints (thin adapters) +- **`/v1/chat/completions`** — widen `ChatMessage.content` from `str` to `str | list[part]`; parts may + include `input_audio`. Pass through `normalize_audio_parts` → `generate`. This is the general + endpoint; classification/judge/one-shot-ASR are just prompts over it. +- **`/transcribe`** — unchanged contract; still owns the ASR-only concerns the general endpoint + shouldn't carry (silence gate, >30 s windowing+stitch, diarization parse, NDJSON progress). + Internally routed through `normalize_audio_parts` + `generate`. +- **`/classify`** (optional sugar) — multipart `examples[]` + `labels[]` + `candidate` → builds the + few-shot messages server-side → returns `{label, raw}`. Saves callers from hand-building messages. +- `/judge`, `/stream` — refactor to call `generate`; behavior unchanged. + +## Consistency: OpenAI ↔ Gemma + +- The OpenAI `input_audio` part maps **deterministically** to Gemma's `{"type":"audio"}` part — a + faithful 1:1 adapter, *if normalized at the boundary*. Same payload works against OpenRouter/Gemini + **and** the local service → identical client code, just swap `base_url`. (This is what lets + `threeway_probe.py` target cloud or local unchanged.) +- One Gemma-specific nuance: it prefers the audio **after** its text label in a turn. We preserve + client-provided part order; presets emit the correct order. So consistency holds. +- The chat path deliberately **skips the silence gate** (that's an ASR-output concern) — correct for + classification, where we want a verdict even on near-silent clips. + +## Transport guidance (answering "just random round trips") + +- **Remote/cloud:** base64 `input_audio` is unavoidable (no shared FS) — but we decode once to an + array, never to a temp file. +- **Local batch (triage over pending clips):** prefer **`audio_path` refs** with the wakeword + `data/samples` volume mounted read-only into the asr container → the service reads each wav + straight to an array. No base64, no temp files, no copies. Multipart raw is the middle option. + +## Impact / migration +- **Back-compat:** `/transcribe` and text-only `/v1/chat` behave exactly as today. `ChatMessage.content` + widened to `str | list` (additive). One image rebuild. +- **Deletes the stop/restart workflow:** multi-audio few-shot classification is served by the *live* + model — no second load, no ASR downtime, no extra idle VRAM (same weights). +- **Less code, one behavior:** transcribe/classify/judge/chat unified on one generate path. +- **Concurrency:** GPU generate is serial; a batch labeling job shares the card with live ASR/wake + transcription. Fine for background triage; if it matters, route batch on the existing ASR "normal" + lane (priority lane keeps wake clips ahead). See `asr-priority-lane`. + +## Status — IMPLEMENTED + validated (2026-06-14) + +Done (minimal, backward-compatible): +- `common/audio_utils.py`: added `load_audio_bytes(bytes)` — in-memory WAV→array (no temp file), + reuses `convert_audio_to_numpy`/`convert_to_mono`/`resample_audio`. +- `providers/gemma4/transcriber.py`: added `_normalize_chat_content()` and wired it into + `generate_chat()`. `input_audio` base64 → decode once to ndarray; **path refs + ndarrays + text + pass straight through** (processor loads paths natively). No new endpoint, no schema change — + `/v1/chat/completions` already forwards `messages: list[dict]`, so audio parts just flow through. +- Verified the processor accepts an **ndarray** under `audio` (not only a path), so no temp files. + +Validated on the LIVE service (no stop/restart, no second model load): +- Multi-audio few-shot through `/v1/chat/completions` works on the running model. +- 3-way on the 27 held-out clips via the endpoint = **96.3% (26/27)** vs **100% in-container**. + The single flip is a borderline `hermes`: in-container passed **file paths** (transformers' own + loader), the endpoint sent **base64** decoded by our `wave` loader (`/32767` scaling) — slightly + different decode flips one boundary case (both greedy). + +Decode-path takeaway: for **local** callers, prefer **path refs** (mount `data/samples` RO into the +asr container, send `{"type":"audio","audio":"<container path>"}`) → exact `/transcribe` parity, no +base64, no decode. base64 stays for remote/cloud where there's no shared FS (accept the 1-in-27 +boundary wobble, or align the loaders for bit-parity later). + +NOTE: the running image was built before the final "path refs pass through" tweak; rebuild to +activate path-ref parity. base64 flow is already live and correct. + +## Conclusion (where we stopped, 2026-06-14) + +**Validated + serving primitive built; the wakeword-service integration is deferred.** +- Approach proven: 3-way few-shot ("hey hermes" / "hermes" / "nothing"), gemma4 E2B best + (98% labeled, 100%/96% on the user's 27 held-out hand-labels), local + free. +- `/v1/chat/completions` now serves multi-audio on the live model — no stop/restart. +- NOT built: any loop inside `extras/wakeword-service/` (no trigger, no pending triage, no + suggestion write/surface). Probe scripts under `ml-experiments/.../wakeword_gemma_classify/` + are eval harnesses only. + +**Eventual shape (per user):** mostly a **button** in the review UI ("suggest labels") that calls the +classifier over `pending/` and pre-fills suggestions for confirm/correct — NOT an autonomous labeler. + +**Key open problem before that ships: a reliable "unsure"/abstain.** The classifier currently makes a +forced 3-way choice; for a triage button it must *defer* hard cases to the human instead of guessing. +Candidate signals (none built yet): +- explicit 4th label "unsure" in the prompt (cheap, but models rarely self-abstain reliably); +- self-consistency — run N times with shuffled few-shot order / different ref sets, abstain on + disagreement (needs sampling or perturbation since greedy is deterministic); +- two-stage classify→verify, abstain when they disagree; +- token-logprob confidence if we expose it from `generate` (top-token prob as a calibrated score); +- cheap energy/VAD pre-gate to auto-route near-silence to "nothing" (note: a real wake existed at + rms 0.0023, so a global floor is unsafe — use only as a soft signal). +The abstain mechanism, not the classifier, is the remaining design work. + +## Build order +1. `transcriber.generate()` canonical path + refactor `generate_chat` to delegate. +2. `normalize_audio_parts()` (base64 + multipart + path), decode-once-to-array. +3. Widen `/v1/chat/completions` schema + wire normalization. +4. (optional) `/classify` convenience preset. +5. Rebuild image. Point `threeway_probe.py` / a new `classify_pending.py` at `localhost:8767` — zero restarts. diff --git a/extras/asr-services/providers/gemma4/finetune/README.md b/extras/asr-services/providers/gemma4/finetune/README.md new file mode 100644 index 00000000..457b258e --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/README.md @@ -0,0 +1,128 @@ +# Gemma 4 audio QLoRA fine-tuning + +QLoRA (4-bit NF4) LoRA fine-tuning for `google/gemma-4-E2B-it` / `E4B-it` on an +audio→text ASR task. Built to validate the training+inference loop and the +single-4090 hardware setup by **overfitting** the 7-clip CoSHE-Eval `sample7` +subset (Hinglish, ~55s clips truncated to the model's 30s audio window). + +## Approach (grounded in the official recipes) + +Mirrors HuggingFace's official `fine_tune_gemma3n_on_audio.ipynb` data pipeline, +adapted to the **gemma4** family (this repo's model is `Gemma4ForConditionalGeneration`, +loaded via `AutoModelForMultimodalLM`, transformers >= 5.5): + +- **Chat-message training examples**: `{user: [audio, prompt]}` + `{assistant: target}`, + rendered with `apply_chat_template(tokenize=False)` then `processor(text=, audio=, padding=True)`. +- **Loss only on the transcription**: labels = `input_ids` with the prompt prefix + (user turn incl. expanded audio soft tokens + `<|turn>model\n`), pad, and all + multimodal special tokens masked to `-100`. (The official notebook leaves the + prompt unmasked; we mask it — the cleaner ASR objective the research flagged.) +- **QLoRA**: NF4 + double-quant + bf16 compute. LoRA (r=16, α=32) on the **text + decoder only** (`language_model.*.{q,k,v,o,gate,up,down}_proj`, regex-scoped so + it doesn't hit the identically-named audio-tower linears). Audio/vision towers + and multimodal embedders are **frozen and kept in bf16** (not quantized). + +### Two gemma4-specific gotchas (vs. the gemma3n recipes) + +1. **Don't 4-bit-quantize the audio/vision towers.** Their `Gemma4ClippableLinear` + calls `torch.finfo(weight.dtype)` for gradient clipping, which throws on + uint8-stored 4-bit weights. We pass `llm_int8_skip_modules=["model.audio_tower", + "model.vision_tower", "model.embed_audio", "model.embed_vision", "lm_head"]`. + Note the `model.` prefix is required — transformers' `should_convert_module` + matches skip patterns with `re.match` (anchored at start), so bare `audio_tower` + never matches `model.audio_tower....`. +2. **Don't use `prepare_model_for_kbit_training`.** It upcasts every non-4bit + param to fp32, making the text embeddings fp32 while the (bf16) audio tower + stays bf16 → Gemma4's multimodal `masked_scatter` merge errors on the dtype + clash. We do manual prep (freeze base + `gradient_checkpointing_enable(use_reentrant=False)` + + `enable_input_require_grads`), keeping everything bf16. + +## Environment + +Runs inside the `chronicle-asr-gemma4` image (transformers 5.5, torch cu126) with +`peft` + `bitsandbytes` added. A persistent container is the fastest iteration loop: + +```bash +cd extras/asr-services +docker run -d --name gemma4-train --gpus all \ + -v "$PWD/model_cache:/models" \ + -v "$PWD/providers/gemma4/finetune:/train" \ + -v "$PWD/../ml-experiments/data/coshe-eval:/data/coshe-eval" \ + -e HF_HOME=/models \ + chronicle-asr-gemma4:latest sleep infinity +docker exec gemma4-train uv pip install --python /app/.venv/bin/python \ + "peft>=0.17.0" "bitsandbytes>=0.46.1" +``` + +## Run + +```bash +# train (E2B overfit) +docker exec gemma4-train bash -c 'export PATH=/app/.venv/bin:$PATH; cd /train && \ + python train.py --model google/gemma-4-E2B-it --output_dir /train/out/e2b-overfit \ + --epochs 60 --lr 2e-4' + +# verify: transcribe the training clips with the adapter, compare to ground truth +docker exec gemma4-train bash -c 'export PATH=/app/.venv/bin:$PATH; cd /train && \ + python infer.py --model google/gemma-4-E2B-it --adapter /train/out/e2b-overfit' +``` + +Switch `--model google/gemma-4-E4B-it` for the larger model (also fits 24GB in 4-bit). + +## Files + +- `data.py` — dataset (manifest → 16k mono, ≤30s) + multimodal collator with label masking +- `train.py` — 4-bit load, LoRA, HF `Trainer` overfit loop +- `infer.py` — adapter inference + char-similarity vs. ground truth +- `introspect.py` — dumps processor keys / token ids / module names for a model + +## Results (E2B overfit smoke test) — IT OVERFITS + +The overfit works: with `r=16` attn+MLP LoRA on the text decoder, **5 of 7 clips +reproduce the ground-truth transcript exactly** (char-sim 1.000) and mean char-sim +is **0.869** (the two misses, 0.66 / 0.43, are perfect for the first ~150 chars +then drift in the back half — the part of the 540-char target that runs past the +audible 30s, i.e. correct behaviour). Final teacher-forced loss ~0.006–0.12. + +### The bug that hid it: gemma4 produces wrong logits with the KV cache on + +Until we found this, generation looked broken (char-sim ≈ base 0.07). Root cause: + +**`use_cache=True` makes gemma4's multimodal forward produce wrong audio-conditioned +logits (transformers 5.5.0).** Measured on a fully-overfit adapter, same weights, +same input: + +| forward | loss on a memorized clip | +|---------|--------------------------| +| `use_cache=False` (training path) | **0.10** | +| `use_cache=True` (generation default) | **4.49** (≈ base 4.84) | + +Training runs with `use_cache=False`, so loss collapsed correctly. But +`model.generate()` defaults to the cache, so every generation silently used the +broken path and reverted to near-base output. **Fix: generate with +`use_cache=False`** (`infer.py` sets `model.config.use_cache=False` and passes +`use_cache=False` to `generate`). Slower (no KV cache → full recompute per token), +but correct. The proper fix is upstream in the gemma4 modeling code. + +This was NOT exposure bias and NOT a capacity problem — the earlier v1/v2 "fixes" +(full transcript, lm_head+embed at r64) were all measured through the broken cached +path and so looked like failures/degeneration. + +## What this validated + +- **Overfit demonstrated end-to-end**: load 4-bit → LoRA → train → save → reload → + generate reproduces the training transcripts (mean char-sim 0.869, 5/7 exact). +- 4-bit QLoRA on a single 4090 works for gemma4 audio. +- VRAM: lean r16 ≈ 18–19 GB; r64 + lm_head/embed_tokens ≈ 23.7 GB (fits, barely). + bf16 audio/vision towers stay resident; activations gradient-checkpointed over + ~750 audio + text tokens. +- Speed: ~7 min / 60 epochs (r16). E4B (~5 GB in 4-bit) also fits with room. +- **Debugging note**: `debug_inprocess.py` (train+eval, no save/reload), + `debug_reload_loss.py` (use_cache toggle), `debug_gen_nocache.py` (no-cache + generation) are the scripts that localized the cache bug — keep for reference. + +## For a real fine-tune (next step) + +Need many ≤30s clips with transcripts aligned to the audible window (not full-clip +transcripts truncated by char count). Generation must use `use_cache=False` until +the upstream cache bug is fixed; budget for the slower decode. diff --git a/extras/asr-services/providers/gemma4/finetune/ab_empties.py b/extras/asr-services/providers/gemma4/finetune/ab_empties.py new file mode 100644 index 00000000..53b69530 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/ab_empties.py @@ -0,0 +1,97 @@ +"""A/B diagnostic for the 12B empty-output / bail behavior on CoSHE. + +Runs the first N deterministic clips under several decode configs, prints +GT vs HYP, and dumps the RAW (pre-strip) output for empties so we can see +exactly what the model emits (immediate EOS? thinking-only? refusal?). +""" + +import glob +import os +import re +import sys + +sys.path.insert(0, "/home/gemma4ft") +import pyarrow.parquet as pq +import torch +from bench_coshe_12b import SR, VERBATIM_PROMPT, decode_audio, select_names +from transformers import AutoModelForMultimodalLM, AutoProcessor + +MODEL = "google/gemma-4-12B-it" +DATA = "/home/coshe-data/data" +N = 25 + +proc = AutoProcessor.from_pretrained(MODEL) +model = AutoModelForMultimodalLM.from_pretrained( + MODEL, dtype=torch.bfloat16, device_map="auto" +).eval() +print("model loaded\n", flush=True) + +sel, shards = select_names(DATA, 500, "coshe-eval") +clips = [] +for s in shards: + t = pq.read_table(s, columns=["audio_file_name", "transcription", "audio"]) + for nm, tr, au in zip( + t.column(0).to_pylist(), t.column(1).to_pylist(), t.column(2).to_pylist() + ): + if nm in sel and au and au.get("bytes"): + clips.append((nm, decode_audio(au["bytes"]), tr)) + if len(clips) >= N: + break + if len(clips) >= N: + break + + +def strip_ch(raw): + parts = [] + for part in raw.split("<channel|>"): + parts.append(part.split("<|channel>")[0] if "<|channel>" in part else part) + t = "".join(parts) + t = re.sub(r"<\|?[a-z_]+\|?>", " ", t) + return re.sub(r"\s+", " ", t).strip() + + +def gen(audio, enable_thinking, min_new): + msgs = [ + { + "role": "user", + "content": [ + {"type": "text", "text": VERBATIM_PROMPT}, + {"type": "audio", "audio": audio}, + ], + } + ] + inp = proc.apply_chat_template( + msgs, + tokenize=True, + return_dict=True, + return_tensors="pt", + add_generation_prompt=True, + enable_thinking=enable_thinking, + ).to(model.device) + il = inp["input_ids"].shape[-1] + kw = dict(max_new_tokens=1024, do_sample=False, no_repeat_ngram_size=3) + if min_new: + kw["min_new_tokens"] = min_new + with torch.inference_mode(): + out = model.generate(**inp, **kw) + raw = proc.decode(out[0][il:], skip_special_tokens=False) + return raw, strip_ch(raw) + + +CONFIGS = [ + ("A: think=False min0 (current)", False, 0), + ("B: think=False min32 (force gen)", False, 32), + ("C: think=True min0 (real think)", True, 0), +] +for label, et, mn in CONFIGS: + empt = 0 + print(f"\n########## CONFIG {label} ##########", flush=True) + for nm, audio, gt in clips: + raw, hyp = gen(audio, et, mn) + if not hyp: + empt += 1 + print(f"[{nm} {len(audio)/SR:.0f}s] EMPTY RAW={raw[:140]!r}", flush=True) + else: + print(f"[{nm} {len(audio)/SR:.0f}s] GT : {gt[:80]}", flush=True) + print(f"{' '*len(nm)} HYP: {hyp[:80]}", flush=True) + print(f"==> {label}: empty={empt}/{len(clips)}", flush=True) diff --git a/extras/asr-services/providers/gemma4/finetune/align_spike.py b/extras/asr-services/providers/gemma4/finetune/align_spike.py new file mode 100644 index 00000000..f36ebdd0 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/align_spike.py @@ -0,0 +1,79 @@ +"""Spike: forced-align a handful of CoSHE clips (Devanagari+Latin Hinglish) with the +MMS-based ctc-forced-aligner to check timestamp quality before building the full +windowed dataset. Prints per-word timestamps + sanity stats.""" + +import glob +import io +import sys +import tempfile +import wave + +import numpy as np +import pyarrow.parquet as pq +import soundfile as sf +import torch +from ctc_forced_aligner import ( + generate_emissions, + get_alignments, + get_spans, + load_alignment_model, + load_audio, + postprocess_results, + preprocess_text, +) + +N = int(sys.argv[1]) if len(sys.argv) > 1 else 5 +device = "cuda" if torch.cuda.is_available() else "cpu" +print(f"device={device}", flush=True) + +# grab N clips from first shard's first row group +f = sorted(glob.glob("/mnt/d/datasets/CoSHE-Eval/data/eval-*.parquet"))[0] +t = pq.ParquetFile(f).read_row_group( + 0, columns=["audio_file_name", "transcription", "audio"] +) +clips = [] +for i in range(min(N, t.num_rows)): + clips.append( + ( + t["audio_file_name"][i].as_py(), + t["transcription"][i].as_py(), + t["audio"][i].as_py()["bytes"], + ) + ) + +model, tokenizer = load_alignment_model( + device, dtype=torch.float16 if device == "cuda" else torch.float32 +) + +for name, text, ab in clips: + wav, sr = sf.read(io.BytesIO(ab)) + if wav.ndim > 1: + wav = wav.mean(1) + dur = len(wav) / sr + # write 16k mono temp for load_audio + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tf: + sf.write(tf.name, wav, sr) + audio_waveform = load_audio(tf.name, model.dtype, model.device) + emissions, stride = generate_emissions(model, audio_waveform, batch_size=1) + tokens_starred, text_starred = preprocess_text(text, romanize=True, language="hin") + segments, scores, blank = get_alignments(emissions, tokens_starred, tokenizer) + spans = get_spans(tokens_starred, segments, blank) + word_ts = postprocess_results(text_starred, spans, stride, scores) + nwords_gt = len(text.split()) + monotonic = all( + word_ts[i]["end"] >= word_ts[i]["start"] for i in range(len(word_ts)) + ) and all( + word_ts[i + 1]["start"] >= word_ts[i]["start"] - 0.05 + for i in range(len(word_ts) - 1) + ) + cov = word_ts[-1]["end"] if word_ts else 0 + print( + f"\n=== {name} dur={dur:.1f}s gt_words={nwords_gt} aligned={len(word_ts)} " + f"last_end={cov:.1f}s coverage={cov/dur*100:.0f}% monotonic={monotonic}", + flush=True, + ) + for w in word_ts[:6] + (["..."] if len(word_ts) > 12 else []) + word_ts[-6:]: + if w == "...": + print(" ...") + continue + print(f" [{w['start']:6.2f}-{w['end']:6.2f}] {w['text']}") diff --git a/extras/asr-services/providers/gemma4/finetune/bench_coshe_12b.py b/extras/asr-services/providers/gemma4/finetune/bench_coshe_12b.py new file mode 100644 index 00000000..1b66b09f --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/bench_coshe_12b.py @@ -0,0 +1,342 @@ +"""Benchmark base google/gemma-4-12B-it on CoSHE-Eval (Hinglish). + +Faithful to the gemma4-asr service path that produced the E4B 26.23% number: + * same diarization prompt + * audio AFTER text, apply_chat_template + generate (model default sampling) + * parse_response -> strip "Speaker N:" labels -> join + * long clips (>30s) split into 30s windows, window texts concatenated + +Selection matches mlexp.runners.dataset.run_coshe(limit=500, seed="coshe-eval"): +all audio_file_names in shard order, shuffled with random.Random(seed), first N. + +Output JSONL (score_coshe schema): {audio_file_name, transcription, hyp, asr_seconds, duration_s} + + python bench_coshe_12b.py --data_dir /home/coshe-data/data --out /home/gemma4ft/out/coshe_12b.jsonl \ + --model google/gemma-4-12B-it --limit 500 --seed coshe-eval +""" + +import argparse +import glob +import io +import json +import os +import random +import re +import time + +import numpy as np +import pyarrow.parquet as pq +import soundfile as sf +import torch +from transformers import AutoModelForMultimodalLM, AutoProcessor + +SR = 16000 +# The gemma4-12B "unified" (encoder-free) model has NO hard 30s audio cap — its +# feature extractor encodes linearly at 25 tok/s with no truncation. CoSHE clips +# are all <=60s, so we feed each clip whole (window_seconds defaults large). +# Keep windowing logic for safety on hypothetical >window_seconds clips. + +# Setup A — HF gemma-4-12B-it model-card ASR prompt (NOT the E4B diarization +# prompt). "in its original language" generalizes the card's {LANGUAGE} +# placeholder for CoSHE's code-switched Hindi/English. Plain transcription. +HF_ASR_PROMPT = ( + "Transcribe the following speech segment in its original language.\n\n" + "Follow these specific instructions for formatting the answer:\n" + "* Only output the transcription, with no newlines.\n" + "* When transcribing numbers, write the digits, i.e. write 1.7 and not " + "one point seven, and write 3 instead of three." +) +# Setup B — strict verbatim transcription prompt (guards against the model +# translating Hinglish to English / adding commentary). +VERBATIM_PROMPT = ( + "You are an expert transcriptionist. Transcribe the provided Hinglish audio " + "verbatim exactly as it is spoken. Output ONLY the transcribed text. Do not " + "translate it to English, do not summarize, and do not add any conversational " + "pleasantries or commentary." +) +PROMPTS = {"hf_asr": HF_ASR_PROMPT, "verbatim": VERBATIM_PROMPT} +SPEAKER_LINE_RE = re.compile(r"^(Speaker \d+):\s*(.+)$", re.MULTILINE) + + +def parse_args(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="google/gemma-4-12B-it") + p.add_argument("--data_dir", default="/home/coshe-data/data") + p.add_argument("--out", default="/home/gemma4ft/out/coshe_12b.jsonl") + p.add_argument("--limit", type=int, default=500) + p.add_argument("--seed", default="coshe-eval") + # CoSHE transcripts are long/dense (mean ~882 chars); 512 truncates and + # inflates WER (deletion trap). 1024 fits the ~60s clips without truncation. + p.add_argument("--max_new_tokens", type=int, default=1024) + p.add_argument( + "--window_seconds", + type=float, + default=600.0, + help="split clips longer than this; default large = whole-clip", + ) + # HF card recommends temp=1.0/top_p=0.95/top_k=64 (general). Lower temp for ASR. + p.add_argument("--temperature", type=float, default=1.0) + p.add_argument("--top_p", type=float, default=0.95) + p.add_argument("--top_k", type=int, default=64) + p.add_argument( + "--greedy", action="store_true", help="do_sample=False (overrides temp)" + ) + p.add_argument("--repetition_penalty", type=float, default=1.0) + p.add_argument( + "--no_repeat_ngram_size", + type=int, + default=0, + help="block repeated n-grams (e.g. 3) to stop greedy loops", + ) + p.add_argument( + "--min_new_tokens", + type=int, + default=0, + help="force >=N generated tokens (stops immediate-EOS empties)", + ) + p.add_argument("--prompt", choices=list(PROMPTS), default="hf_asr") + return p.parse_args() + + +def select_names(data_dir, limit, seed): + """Replicate run_coshe deterministic selection.""" + shards = sorted(glob.glob(os.path.join(data_dir, "eval-*.parquet"))) + names = [] + for s in shards: + names.extend( + pq.read_table(s, columns=["audio_file_name"]).column(0).to_pylist() + ) + if limit and limit < len(names): + random.Random(seed).shuffle(names) + return set(names[:limit]), shards + return set(names), shards + + +def load_done(out_path): + done = set() + if os.path.exists(out_path): + with open(out_path) as f: + for line in f: + try: + rec = json.loads(line) + except Exception: + continue + if "error" in rec: + continue + if rec.get("audio_file_name"): + done.add(rec["audio_file_name"]) + return done + + +def decode_audio(raw): + audio, sr = sf.read(io.BytesIO(raw), dtype="float32") + if audio.ndim > 1: + audio = audio.mean(axis=1) + if sr != SR: + import librosa + + audio = librosa.resample(audio, orig_sr=sr, target_sr=SR) + return np.ascontiguousarray(audio, dtype=np.float32) + + +def windows(audio, window_seconds): + """Whole clip if it fits; else non-overlapping windows of window_seconds.""" + w = int(window_seconds * SR) + if len(audio) <= w: + return [audio] + return [audio[i : i + w] for i in range(0, len(audio), w)] + + +class Model: + def __init__( + self, + model_id, + max_new_tokens, + window_seconds=600.0, + prompt=HF_ASR_PROMPT, + temperature=1.0, + top_p=0.95, + top_k=64, + greedy=False, + repetition_penalty=1.0, + no_repeat_ngram_size=0, + min_new_tokens=0, + ): + self.max_new_tokens = max_new_tokens + self.window_seconds = window_seconds + self.prompt = prompt + self.do_sample = not greedy + self.temperature = temperature + self.top_p = top_p + self.top_k = top_k + self.repetition_penalty = repetition_penalty + self.no_repeat_ngram_size = no_repeat_ngram_size + self.min_new_tokens = min_new_tokens + self.processor = AutoProcessor.from_pretrained(model_id) + self.model = AutoModelForMultimodalLM.from_pretrained( + model_id, dtype=torch.bfloat16, device_map="auto" + ) + self.model.eval() + + def _decode(self, outputs, input_len): + # Decode WITH special tokens so the channel markers survive, then strip + # the thinking channel exactly like the chat template's strip_thinking + # macro: drop everything inside <|channel>thought ... <channel|> blocks. + raw = self.processor.decode(outputs[0][input_len:], skip_special_tokens=False) + parts = [] + for part in raw.split("<channel|>"): + parts.append(part.split("<|channel>")[0] if "<|channel>" in part else part) + text = "".join(parts) + # remove any residual angle-bracket control tokens (<|...|>, <...|>, <eos>, etc.) + text = re.sub(r"<\|?[a-z_]+\|?>", " ", text) + text = re.sub(r"\s+", " ", text).strip() + return text + + def transcribe_window(self, audio): + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": self.prompt}, + {"type": "audio", "audio": audio}, + ], + } + ] + inputs = self.processor.apply_chat_template( + messages, + tokenize=True, + return_dict=True, + return_tensors="pt", + add_generation_prompt=True, + enable_thinking=False, + ).to(self.model.device) + input_len = inputs["input_ids"].shape[-1] + gen_kwargs = { + "max_new_tokens": self.max_new_tokens, + "repetition_penalty": self.repetition_penalty, + } + if self.no_repeat_ngram_size: + gen_kwargs["no_repeat_ngram_size"] = self.no_repeat_ngram_size + if self.min_new_tokens: + # the 12B-it often bails to immediate EOS (empty) on long clips; + # forcing a floor unlocks the transcription it was withholding. + gen_kwargs["min_new_tokens"] = self.min_new_tokens + if self.do_sample: + gen_kwargs.update( + do_sample=True, + temperature=self.temperature, + top_p=self.top_p, + top_k=self.top_k, + ) + else: + gen_kwargs.update(do_sample=False) # greedy / deterministic + with torch.inference_mode(): + outputs = self.model.generate(**inputs, **gen_kwargs) + raw_text = self._decode(outputs, input_len) + matches = list(SPEAKER_LINE_RE.finditer(raw_text)) + clean = ( + " ".join(m.group(2).strip() for m in matches if m.group(2).strip()) + if matches + else raw_text.strip() + ) + if clean == "[NO SPEECH]": + clean = "" + return clean + + def transcribe(self, audio): + parts = [self.transcribe_window(w) for w in windows(audio, self.window_seconds)] + return " ".join(p for p in parts if p).strip() + + +def main(): + args = parse_args() + selected, shards = select_names(args.data_dir, args.limit, args.seed) + done = load_done(args.out) + target = selected - done + print( + f"selected={len(selected)} done={len(selected & done)} todo={len(target)}", + flush=True, + ) + if not target: + print("nothing to do", flush=True) + return + + os.makedirs(os.path.dirname(args.out), exist_ok=True) + m = Model( + args.model, + args.max_new_tokens, + args.window_seconds, + prompt=PROMPTS[args.prompt], + temperature=args.temperature, + top_p=args.top_p, + top_k=args.top_k, + greedy=args.greedy, + repetition_penalty=args.repetition_penalty, + no_repeat_ngram_size=args.no_repeat_ngram_size, + min_new_tokens=args.min_new_tokens, + ) + print( + f"model loaded: {args.model} (prompt={args.prompt} thinking=False " + f"do_sample={not args.greedy} temp={args.temperature} " + f"top_p={args.top_p} top_k={args.top_k} rep_pen={args.repetition_penalty} " + f"no_repeat_ngram={args.no_repeat_ngram_size} min_new={args.min_new_tokens} " + f"max_new={args.max_new_tokens} window={args.window_seconds})", + flush=True, + ) + + n_ok = n_err = 0 + t0 = time.time() + with open(args.out, "a") as out_f: + for shard in shards: + if not target: + break + pf = pq.ParquetFile(shard) + for batch in pf.iter_batches(batch_size=8): + for row in batch.to_pylist(): + name = row["audio_file_name"] + if name not in target: + continue + target.discard(name) + try: + audio = decode_audio(row["audio"]["bytes"]) + dur = len(audio) / SR + ts = time.time() + hyp = m.transcribe(audio) + dt = time.time() - ts + rec = { + "audio_file_name": name, + "transcription": row["transcription"], + "hyp": hyp, + "asr_seconds": round(dt, 2), + "duration_s": round(dur, 2), + } + out_f.write(json.dumps(rec, ensure_ascii=False) + "\n") + out_f.flush() + n_ok += 1 + done_n = n_ok + n_err + print( + f" [{done_n}/{len(selected)}] {name} dur={dur:.0f}s " + f"gen={dt:.1f}s hyp[:80]={hyp[:80]!r}", + flush=True, + ) + except Exception as e: + n_err += 1 + out_f.write( + json.dumps( + {"audio_file_name": name, "error": str(e)}, + ensure_ascii=False, + ) + + "\n" + ) + out_f.flush() + print(f" ERR {name}: {e}", flush=True) + if not target: + break + print( + f"DONE ok={n_ok} err={n_err} elapsed={time.time()-t0:.0f}s out={args.out}", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/gemma4/finetune/build_windowed_dataset.py b/extras/asr-services/providers/gemma4/finetune/build_windowed_dataset.py new file mode 100644 index 00000000..78d25a82 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/build_windowed_dataset.py @@ -0,0 +1,130 @@ +"""Materialize the windowed CoSHE training set: slice each train/val window's audio out +of the CoSHE parquet and write 16k-mono WAVs + a manifest, honoring the clip-level +20/10/70 split (windows inherit their clip's split, so no clip leaks across splits). + +Consumes windows.jsonl from window_coshe.py (honest <=30s GT-accurate windows). Test +clips are NOT sliced here — they are kept whole for windowed-stitch evaluation; we only +emit their clip list + per-window text. + +Out (under --out_dir): + audio/<clip>__w<idx>.wav 16k mono int16, one per train/val window + manifest.jsonl {name, clip, win_idx, split, audio, dur, text} + windowed_split.json {train:[names], val:[names], test_clips:[clips]} +""" + +import argparse +import glob +import io +import json +import wave +from collections import defaultdict +from pathlib import Path + +import numpy as np +import pyarrow.parquet as pq +import soundfile as sf + + +def safe(clip, idx): + return f"{clip.replace('.wav','')}__w{idx}" + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--windows", default="/home/coshe_windowed/windows_v2.jsonl") + ap.add_argument("--split", default="/home/coshe_windowed/split_20_10_70.json") + ap.add_argument("--exclude", default="") + ap.add_argument("--parquet_glob", default="/home/coshe/data/eval-*.parquet") + ap.add_argument("--out_dir", default="/home/coshe_windowed") + ap.add_argument( + "--materialize", + default="train,val", + help="which splits get sliced WAVs (test kept whole for stitch eval)", + ) + args = ap.parse_args() + + drop = {x for x in args.exclude.split(",") if x} + mat = set(args.materialize.split(",")) + split = json.load(open(args.split)) + clip_split = {} + for s in ("train", "val", "test"): + for c in split[s]: + if c not in drop: + clip_split[c] = s + + wins_by_clip = defaultdict(list) + for line in open(args.windows): + r = json.loads(line) + if r["audio_file_name"] in drop: + continue + wins_by_clip[r["audio_file_name"]].append(r) + + out_dir = Path(args.out_dir) + audio_dir = out_dir / "audio" + audio_dir.mkdir(parents=True, exist_ok=True) + man = open(out_dir / "manifest.jsonl", "w") + split_out = { + "train": [], + "val": [], + "test_clips": sorted(c for c, s in clip_split.items() if s == "test"), + } + + want_clips = {c for c, s in clip_split.items() if s in mat} + written = 0 + for f in sorted(glob.glob(args.parquet_glob)): + t = pq.ParquetFile(f).read(columns=["audio_file_name", "audio"]) + names = t["audio_file_name"].to_pylist() + for ci, name in enumerate(names): + if name not in want_clips: + continue + s = clip_split[name] + wav, sr = sf.read(io.BytesIO(t["audio"][ci].as_py()["bytes"])) + if wav.ndim > 1: + wav = wav.mean(1) + if sr != 16000: + import librosa + + wav = librosa.resample( + wav.astype(np.float32), orig_sr=sr, target_sr=16000 + ) + sr = 16000 + pcm = (np.clip(wav, -1, 1) * 32767).astype(np.int16) + for w in wins_by_clip[name]: + nm = safe(name, w["win_idx"]) + a0, a1 = int(w["start"] * sr), int(w["end"] * sr) + seg = pcm[a0:a1] + path = audio_dir / f"{nm}.wav" + with wave.open(str(path), "wb") as wf: + wf.setnchannels(1) + wf.setsampwidth(2) + wf.setframerate(sr) + wf.writeframes(seg.tobytes()) + man.write( + json.dumps( + { + "name": nm, + "clip": name, + "win_idx": w["win_idx"], + "split": s, + "audio": str(path), + "dur": round(len(seg) / sr, 3), + "text": w["text"], + }, + ensure_ascii=False, + ) + + "\n" + ) + split_out[s].append(nm) + written += 1 + print(f" {Path(f).name}: written so far={written}", flush=True) + man.close() + json.dump(split_out, open(out_dir / "windowed_split.json", "w"), ensure_ascii=False) + print( + f"DONE windows materialized={written} " + f"train={len(split_out['train'])} val={len(split_out['val'])} " + f"test_clips={len(split_out['test_clips'])}" + ) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/gemma4/finetune/data.py b/extras/asr-services/providers/gemma4/finetune/data.py new file mode 100644 index 00000000..4af22542 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/data.py @@ -0,0 +1,251 @@ +"""Dataset + collator for Gemma 4 audio QLoRA fine-tuning on CoSHE-Eval sample7. + +The dataset yields raw {audio: float32 16k mono, target: transcript} items. +The collator does the multimodal preprocessing per batch (the HF gemma recipe +pattern): apply_chat_template(tokenize=False) -> processor(text, audio) -> +build labels. + +Label masking (loss only on the transcription): + * everything in the prompt prefix (user turn incl. expanded audio soft + tokens + the `<|turn>model\n` generation marker) -> -100 + * pad tokens and any multimodal special tokens -> -100 +Only the assistant transcription (+ its closing turn token) is supervised. +""" + +import json +import os +import wave +from pathlib import Path + +import numpy as np +import torch + +DEFAULT_PROMPT = ( + "Transcribe the following speech segment in its original language and identify different speakers. " + "Follow these specific instructions for formatting the answer:\n" + "* Label each speaker as Speaker 1, Speaker 2, etc.\n" + "* Format each turn as 'Speaker N: <text>' on its own line.\n" + "* Start a new line when the speaker changes.\n" + "* When transcribing numbers, write the digits, i.e. write 1.7 and not " + "one point seven, and write 3 instead of three.\n" + "* If the audio is silence or contains no speech, respond with exactly: [NO SPEECH]" +) + + +def _load_wav_16k_mono(path: str, max_seconds: float) -> np.ndarray: + with wave.open(path, "rb") as wf: + sr = wf.getframerate() + nch = wf.getnchannels() + raw = wf.readframes(wf.getnframes()) + audio = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0 + if nch > 1: + audio = audio.reshape(-1, nch).mean(axis=1) + if sr != 16000: + import librosa + + audio = librosa.resample(audio, orig_sr=sr, target_sr=16000) + if max_seconds: + audio = audio[: int(16000 * max_seconds)] + return np.ascontiguousarray(audio, dtype=np.float32) + + +class CosheSample7Dataset(torch.utils.data.Dataset): + """Reads coshe-eval/sample7/manifest.json -> {audio, target} items.""" + + def __init__( + self, data_dir: str, max_seconds: float = 30.0, target_max_chars: int = 0 + ): + d = Path(data_dir) + manifest = json.loads((d / "manifest.json").read_text()) + self.items = [] + for row in manifest: + wav = d / "audio" / row["audio_file_name"] + if not wav.exists(): + continue + target = row["transcription"].strip() + # When audio is truncated to the model's 30s window, the back half of + # the full-clip transcript is ungrounded. Truncating the target to ~the + # portion the model can actually hear makes the overfit reproducible at + # generation (CoSHE clips run ~18 chars/sec of speech). + if target_max_chars: + target = target[:target_max_chars] + self.items.append( + { + "audio": _load_wav_16k_mono(str(wav), max_seconds), + "target": target, + "name": row["audio_file_name"], + } + ) + if not self.items: + raise RuntimeError(f"No samples found under {data_dir}") + + def __len__(self): + return len(self.items) + + def __getitem__(self, idx): + return self.items[idx] + + +def _decode_audio_bytes_16k_mono(raw: bytes, max_seconds: float) -> np.ndarray: + import io + + import soundfile as sf + + audio, sr = sf.read(io.BytesIO(raw), dtype="float32") + if audio.ndim > 1: + audio = audio.mean(axis=1) + if sr != 16000: + import librosa + + audio = librosa.resample(audio, orig_sr=sr, target_sr=16000) + if max_seconds: + audio = audio[: int(16000 * max_seconds)] + return np.ascontiguousarray(audio, dtype=np.float32) + + +class CosheParquetDataset(torch.utils.data.Dataset): + """Full CoSHE-Eval: reads parquet shards (audio as embedded wav bytes). + + Pre-decodes everything into RAM once (≈3.8GB for the full 1985 clips at 30s) + so multi-epoch overfitting doesn't re-decode each epoch. `limit` caps samples + for quick recipe checks; `max_seconds` truncates to the model's 30s window. + """ + + def __init__( + self, + parquet_glob: str, + max_seconds: float = 30.0, + target_max_chars: int = 0, + limit: int = 0, + cache_path: str = "", + ): + import glob + import pickle + + import pyarrow.parquet as pq + + # Decoding+resampling all 1985 clips off the HDD is ~40 min; cache the + # decoded (audio, target) items so repeated training runs load instantly. + if cache_path and os.path.exists(cache_path): + with open(cache_path, "rb") as f: + self.items = pickle.load(f) + if target_max_chars: + for it in self.items: + it["target"] = it["target"][:target_max_chars] + if limit: + self.items = self.items[:limit] + return + + paths = sorted(glob.glob(parquet_glob)) + if not paths: + raise RuntimeError(f"No parquet shards match {parquet_glob}") + self.items = [] + for path in paths: + t = pq.read_table( + path, columns=["audio_file_name", "transcription", "audio"] + ) + names = t.column("audio_file_name").to_pylist() + trans = t.column("transcription").to_pylist() + audio_col = t.column("audio").to_pylist() # list of {bytes, path} + for name, tr, au in zip(names, trans, audio_col): + if au is None or au.get("bytes") is None or not tr: + continue + self.items.append( + { + "audio": _decode_audio_bytes_16k_mono(au["bytes"], max_seconds), + "target": tr.strip(), # full target; truncation applied below + "name": name, + } + ) + if limit and len(self.items) >= limit: + break + if limit and len(self.items) >= limit: + break + if not self.items: + raise RuntimeError(f"No samples loaded from {parquet_glob}") + if cache_path: + with open(cache_path, "wb") as f: + pickle.dump(self.items, f) + if target_max_chars: + for it in self.items: + it["target"] = it["target"][:target_max_chars] + + def __len__(self): + return len(self.items) + + def __getitem__(self, idx): + return self.items[idx] + + +class Gemma4AudioCollator: + def __init__( + self, processor, prompt: str = DEFAULT_PROMPT, mask_prompt: bool = True + ): + self.processor = processor + self.prompt = prompt + self.mask_prompt = mask_prompt + tok = processor.tokenizer + tok.padding_side = "right" + self.pad_id = tok.pad_token_id + # multimodal special tokens to exclude from the loss + self.special_ids = [ + t + for t in [ + getattr(tok, "audio_token_id", None), + getattr(tok, "image_token_id", None), + getattr(tok, "boi_token_id", None), + getattr(tok, "eoi_token_id", None), + getattr(tok, "boa_token_id", None), + getattr(tok, "eoa_token_id", None), + ] + if t is not None + ] + + def _user_turn(self, audio): + return { + "role": "user", + "content": [ + {"type": "text", "text": self.prompt}, + {"type": "audio", "audio": audio}, + ], + } + + def __call__(self, features): + audios = [f["audio"] for f in features] + texts_full = [] + for f in features: + msgs = [ + self._user_turn(f["audio"]), + { + "role": "assistant", + "content": [{"type": "text", "text": f["target"]}], + }, + ] + texts_full.append( + self.processor.apply_chat_template( + msgs, tokenize=False, add_generation_prompt=False + ) + ) + batch = self.processor( + text=texts_full, audio=audios, return_tensors="pt", padding=True + ) + + labels = batch["input_ids"].clone() + labels[labels == self.pad_id] = -100 + for tid in self.special_ids: + labels[labels == tid] = -100 + + if self.mask_prompt: + for i, f in enumerate(features): + ptext = self.processor.apply_chat_template( + [self._user_turn(f["audio"])], + tokenize=False, + add_generation_prompt=True, + ) + plen = self.processor( + text=[ptext], audio=[f["audio"]], return_tensors="pt" + )["input_ids"].shape[1] + labels[i, :plen] = -100 + + batch["labels"] = labels + return batch diff --git a/extras/asr-services/providers/gemma4/finetune/data_interleave.py b/extras/asr-services/providers/gemma4/finetune/data_interleave.py new file mode 100644 index 00000000..88a50cf3 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/data_interleave.py @@ -0,0 +1,132 @@ +"""Interleaved multi-chunk dataset + collator for the audio->formatted(pasted_text) task. + +Each example is ONE clip: its audio is fed as N <=28s chunks (multiple audio blocks in a +single user turn), the target is the WHOLE pasted_text, and the prompt is conditioned on the +app the dictation was going into (Wispr formats per-app). No text windowing — formatted text +isn't word-alignable, so the model attends to all chunks and emits the full formatted clip. + +Loss is on the target only (prompt prefix incl. all audio soft-tokens + pad + mm-special +tokens masked to -100), matching the single-audio Gemma4AudioCollator. +""" + +import json +import wave + +import numpy as np +import torch + + +def _load_wav_f32(path: str) -> np.ndarray: + with wave.open(path, "rb") as wf: + raw = wf.readframes(wf.getnframes()) + return np.ascontiguousarray( + np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0 + ) + + +def build_prompt(app: str) -> str: + """App-aware formatting instruction. Wispr formats dictation differently per target app.""" + where = f"the {app} app" if app else "the target app" + return ( + f"You are formatting voice dictation that will be typed into {where}. " + "The audio is provided as one or more consecutive segments of a single dictation. " + "Transcribe all of it and format the result the way it should appear when typed into " + f"{where} (punctuation, capitalization, digits as digits, app-appropriate style). " + "Output only the final formatted text, with no commentary." + ) + + +class ChunkInterleaveDataset(torch.utils.data.Dataset): + """One item per clip: {chunks:[float32 16k], target:str, app:str, clip:str}.""" + + def __init__(self, manifest_path: str, split: str): + self.items = [] + for line in open(manifest_path): + r = json.loads(line) + if r["split"] != split: + continue + self.items.append( + { + "chunks": [_load_wav_f32(p) for p in r["chunks"]], + "target": r["target"].strip(), + "app": r.get("app", ""), + "clip": r["clip"], + } + ) + if not self.items: + raise RuntimeError(f"no '{split}' rows in {manifest_path}") + + def __len__(self): + return len(self.items) + + def __getitem__(self, idx): + return self.items[idx] + + +class InterleaveCollator: + """Batches variable-#chunk examples. Uses the batch processor with audio as a per-example + list of arrays (nested); the chat template emits one audio placeholder per chunk.""" + + def __init__(self, processor, mask_prompt: bool = True): + self.processor = processor + self.mask_prompt = mask_prompt + tok = processor.tokenizer + tok.padding_side = "right" + self.pad_id = tok.pad_token_id + self.special_ids = [ + t + for t in [ + getattr(tok, "audio_token_id", None), + getattr(tok, "image_token_id", None), + getattr(tok, "boi_token_id", None), + getattr(tok, "eoi_token_id", None), + getattr(tok, "boa_token_id", None), + getattr(tok, "eoa_token_id", None), + ] + if t is not None + ] + + def _user_turn(self, chunks, app): + content = [{"type": "text", "text": build_prompt(app)}] + content += [{"type": "audio", "audio": a} for a in chunks] + return {"role": "user", "content": content} + + def __call__(self, features): + # processor wants a FLAT list of all audios across the batch, matched to the audio + # placeholders (one per chunk) in text order. + audios = [a for f in features for a in f["chunks"]] + texts_full = [] + for f in features: + msgs = [ + self._user_turn(f["chunks"], f["app"]), + { + "role": "assistant", + "content": [{"type": "text", "text": f["target"]}], + }, + ] + texts_full.append( + self.processor.apply_chat_template( + msgs, tokenize=False, add_generation_prompt=False + ) + ) + batch = self.processor( + text=texts_full, audio=audios, return_tensors="pt", padding=True + ) + + labels = batch["input_ids"].clone() + labels[labels == self.pad_id] = -100 + for tid in self.special_ids: + labels[labels == tid] = -100 + if self.mask_prompt: + for i, f in enumerate(features): + ptext = self.processor.apply_chat_template( + [self._user_turn(f["chunks"], f["app"])], + tokenize=False, + add_generation_prompt=True, + ) + plen = self.processor( + text=[ptext], audio=f["chunks"], return_tensors="pt" + )["input_ids"].shape[1] + labels[i, :plen] = -100 + batch["labels"] = labels + return batch diff --git a/extras/asr-services/providers/gemma4/finetune/data_windowed.py b/extras/asr-services/providers/gemma4/finetune/data_windowed.py new file mode 100644 index 00000000..fc1ad496 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/data_windowed.py @@ -0,0 +1,57 @@ +"""Dataset over the windowed CoSHE manifest (build_windowed_dataset.py output). + +Each row is one honest <=30s window: a sliced 16k-mono WAV + its EXACT ground-truth text +(forced-alignment-derived, no first-30s truncation hack). Yields the same +{audio, target, name} shape the Gemma4AudioCollator expects, so training reuses the +existing collator unchanged. + +Plain-transcript targets: CoSHE has no speaker labels, so PLAIN_PROMPT asks only for a +verbatim transcript (digits as digits). Train and eval MUST use the same prompt. +""" + +import json +import wave + +import numpy as np +import torch + +PLAIN_PROMPT = ( + "Transcribe the following speech segment verbatim in its original language " + "(Hindi-English code-mixed). Write digits as digits (e.g. 3, not three). " + "Output only the transcript, no speaker labels." +) + + +def _load_wav_f32(path: str) -> np.ndarray: + with wave.open(path, "rb") as wf: + raw = wf.readframes(wf.getnframes()) + return np.ascontiguousarray( + np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0 + ) + + +class WindowedManifestDataset(torch.utils.data.Dataset): + """Loads the windows for one split (train/val). Audio is decoded eagerly into RAM + (train+val is ~1GB) so multi-epoch training doesn't re-read WAVs.""" + + def __init__(self, manifest_path: str, split: str): + self.items = [] + for line in open(manifest_path): + r = json.loads(line) + if r["split"] != split: + continue + self.items.append( + { + "audio": _load_wav_f32(r["audio"]), + "target": r["text"].strip(), + "name": r["name"], + } + ) + if not self.items: + raise RuntimeError(f"no '{split}' rows in {manifest_path}") + + def __len__(self): + return len(self.items) + + def __getitem__(self, idx): + return self.items[idx] diff --git a/extras/asr-services/providers/gemma4/finetune/data_wispr.py b/extras/asr-services/providers/gemma4/finetune/data_wispr.py new file mode 100644 index 00000000..cfe1c395 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/data_wispr.py @@ -0,0 +1,15 @@ +"""Wispr dictation FT: shared English verbatim prompt + windowed dataset. + +The target is asr_text (Wispr's raw verbatim ASR, treated as ground truth) cut into honest +<=28s windows. English single-speaker dictation, so the prompt asks for a plain verbatim +English transcript (no speaker labels, no Hinglish wording). Train, eval, and the base +baseline MUST all use WISPR_PROMPT. +""" + +from data_windowed import WindowedManifestDataset # noqa: F401 (re-exported) + +WISPR_PROMPT = ( + "Transcribe the following speech segment verbatim in English. " + "Write digits as digits (e.g. 3, not three). " + "Only output the transcription text itself, with no commentary or explanation." +) diff --git a/extras/asr-services/providers/gemma4/finetune/debug_gen_nocache.py b/extras/asr-services/providers/gemma4/finetune/debug_gen_nocache.py new file mode 100644 index 00000000..7d7ce4de --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/debug_gen_nocache.py @@ -0,0 +1,66 @@ +"""Test: does generating with use_cache=False make the v3 adapter reproduce targets? + +The cached (use_cache=True) gemma4 multimodal forward gives wrong logits (loss 4.49 +vs 0.10 uncached), so model.generate() — which uses the cache by default — fails. +Force use_cache=False during generation and compare. +""" + +from difflib import SequenceMatcher + +import torch +from data import DEFAULT_PROMPT, CosheSample7Dataset +from peft import PeftModel +from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig + +proc = AutoProcessor.from_pretrained("google/gemma-4-E2B-it") +proc.tokenizer.padding_side = "left" +bnb = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + llm_int8_skip_modules=[ + "model.audio_tower", + "model.vision_tower", + "model.embed_audio", + "model.embed_vision", + "lm_head", + ], +) +m = AutoModelForMultimodalLM.from_pretrained( + "google/gemma-4-E2B-it", + quantization_config=bnb, + dtype=torch.bfloat16, + device_map="auto", + attn_implementation="eager", +) +pm = PeftModel.from_pretrained(m, "/train/out/e2b-overfit-v3") +pm.eval() + +ds = CosheSample7Dataset( + "/data/coshe-eval/sample7", max_seconds=30.0, target_max_chars=540 +) +print("===== generate(use_cache=False) =====", flush=True) +ratios = [] +for item in ds: + msgs = [ + { + "role": "user", + "content": [ + {"type": "text", "text": DEFAULT_PROMPT}, + {"type": "audio", "audio": item["audio"]}, + ], + } + ] + ptext = proc.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True) + inp = proc(text=[ptext], audio=[item["audio"]], return_tensors="pt").to(pm.device) + inlen = inp["input_ids"].shape[-1] + with torch.inference_mode(): + out = pm.generate(**inp, max_new_tokens=320, do_sample=False, use_cache=False) + gen = proc.decode(out[0][inlen:], skip_special_tokens=True).strip() + r = SequenceMatcher(None, gen, item["target"]).ratio() + ratios.append(r) + print(f"\n {item['name']} sim={r:.3f}", flush=True) + print(f" TGT: {item['target'][:220]}", flush=True) + print(f" GEN: {gen[:220]}", flush=True) +print(f"\n MEAN char-sim = {sum(ratios)/len(ratios):.3f}", flush=True) diff --git a/extras/asr-services/providers/gemma4/finetune/debug_inprocess.py b/extras/asr-services/providers/gemma4/finetune/debug_inprocess.py new file mode 100644 index 00000000..ec19759b --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/debug_inprocess.py @@ -0,0 +1,137 @@ +"""Train in-process and eval WITHOUT save/reload, to localize the loss-0-but-wrong bug. + +If in-process teacher-forced loss is ~0 AND generation reproduces -> the bug is in +save/reload (adapter not persisted/loaded correctly). +If in-process loss is ~0 but generation still fails -> teacher-forcing/label issue. +If in-process loss is also high -> the Trainer's reported step loss was the illusion. +""" + +import torch +from data import DEFAULT_PROMPT, CosheSample7Dataset, Gemma4AudioCollator +from peft import LoraConfig, get_peft_model +from transformers import ( + AutoModelForMultimodalLM, + AutoProcessor, + BitsAndBytesConfig, + Trainer, + TrainingArguments, +) + +MODEL = "google/gemma-4-E2B-it" +LORA_TARGETS = ( + r".*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" +) + +processor = AutoProcessor.from_pretrained(MODEL) +bnb = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + llm_int8_skip_modules=[ + "model.audio_tower", + "model.vision_tower", + "model.embed_audio", + "model.embed_vision", + "lm_head", + ], +) +model = AutoModelForMultimodalLM.from_pretrained( + MODEL, + quantization_config=bnb, + dtype=torch.bfloat16, + device_map="auto", + attn_implementation="eager", +) +model.config.use_cache = False +for p in model.parameters(): + p.requires_grad = False +model.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={"use_reentrant": False} +) +model.enable_input_require_grads() +model = get_peft_model( + model, + LoraConfig( + r=16, + lora_alpha=32, + lora_dropout=0.0, + bias="none", + task_type="CAUSAL_LM", + target_modules=LORA_TARGETS, + ), +) +model.print_trainable_parameters() + +ds = CosheSample7Dataset( + "/data/coshe-eval/sample7", max_seconds=30.0, target_max_chars=540 +) +collator = Gemma4AudioCollator(processor) + +trainer = Trainer( + model=model, + args=TrainingArguments( + output_dir="/train/out/_inproc", + per_device_train_batch_size=1, + num_train_epochs=30, + learning_rate=2e-4, + lr_scheduler_type="constant", + bf16=True, + logging_steps=20, + save_strategy="no", + report_to=[], + remove_unused_columns=False, + gradient_checkpointing=False, + ), + train_dataset=ds, + data_collator=collator, +) +trainer.train() + +# ---- eval in the SAME process, no save/reload ---- +model.eval() +print("\n===== IN-PROCESS EVAL (no save/reload) =====", flush=True) +for item in ds: + b = collator([item]) + b = {k: v.to(model.device) for k, v in b.items()} + with torch.inference_mode(): + o = model(**b) + ng = model(**{k: v for k, v in b.items() if k != "labels"}) + lg = ng.logits[0] + lb = b["labels"][0] + p = lg[:-1].argmax(-1) + g = lb[1:] + s = g != -100 + n = int(s.sum()) + m = int((p[s] == g[s]).sum()) + print( + f" {item['name']:16s} model.loss={float(o.loss):.4f} argmax-match={m}/{n} " + f"({100*m/max(n,1):.1f}%)", + flush=True, + ) + +# generation on sample 0 using the SAME collator prompt path +model.config.use_cache = True +item = ds[0] +ptext = processor.apply_chat_template( + [ + { + "role": "user", + "content": [ + {"type": "text", "text": DEFAULT_PROMPT}, + {"type": "audio", "audio": item["audio"]}, + ], + } + ], + tokenize=False, + add_generation_prompt=True, +) +inp = processor(text=[ptext], audio=[item["audio"]], return_tensors="pt").to( + model.device +) +inlen = inp["input_ids"].shape[-1] +with torch.inference_mode(): + out = model.generate(**inp, max_new_tokens=256, do_sample=False) +gen = processor.decode(out[0][inlen:], skip_special_tokens=True) +print(f"\n TARGET[:300]: {item['target'][:300]}", flush=True) +print(f" GEN[:300] : {gen[:300]}", flush=True) diff --git a/extras/asr-services/providers/gemma4/finetune/debug_overfit.py b/extras/asr-services/providers/gemma4/finetune/debug_overfit.py new file mode 100644 index 00000000..9326adff --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/debug_overfit.py @@ -0,0 +1,186 @@ +"""Diagnose why teacher-forced loss -> 0 but generation doesn't reproduce targets. + +Checks two things on the trained v3 adapter (target_max_chars=540): + + 1. TOKEN-FORMAT PARITY: does the training collator's input construction + (apply_chat_template(tokenize=False) -> processor(text=str, audio=)) produce + the SAME prompt token ids as the inference path + (apply_chat_template(tokenize=True, return_dict=True))? A mismatch (e.g. a + duplicated <bos>) would mean we train on one format and generate on another. + + 2. TEACHER-FORCED REPRODUCTION: feed the full (prompt+target) exactly as in + training, take argmax of the logits at each SUPERVISED position, and measure + how many match the gold target token. If this is ~100%, the model truly + memorized and any failure is in the generation path; if it's low, loss never + really reached 0 on the hard (early) tokens. +""" + +import argparse + +import torch +from data import DEFAULT_PROMPT, CosheSample7Dataset, Gemma4AudioCollator +from peft import PeftModel +from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig + + +def parse_args(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="google/gemma-4-E2B-it") + p.add_argument("--adapter", default="/train/out/e2b-overfit-v3") + p.add_argument("--data_dir", default="/data/coshe-eval/sample7") + p.add_argument("--target_max_chars", type=int, default=540) + return p.parse_args() + + +def main(): + args = parse_args() + processor = AutoProcessor.from_pretrained(args.model) + tok = processor.tokenizer + + ds = CosheSample7Dataset( + args.data_dir, max_seconds=30.0, target_max_chars=args.target_max_chars + ) + item = ds[0] + + # ---- 1. TOKEN-FORMAT PARITY ---- + print("===== 1. TOKEN-FORMAT PARITY (prompt only) =====", flush=True) + user_turn = { + "role": "user", + "content": [ + {"type": "text", "text": DEFAULT_PROMPT}, + {"type": "audio", "audio": item["audio"]}, + ], + } + # training-style: tokenize=False -> processor(text=str, audio=) + ptext = processor.apply_chat_template( + [user_turn], tokenize=False, add_generation_prompt=True + ) + train_path = processor(text=[ptext], audio=[item["audio"]], return_tensors="pt") + train_ids = train_path["input_ids"][0] + # inference-style: tokenize=True, return_dict=True + infer_path = processor.apply_chat_template( + [user_turn], + tokenize=True, + return_dict=True, + return_tensors="pt", + add_generation_prompt=True, + ) + infer_ids = infer_path["input_ids"][0] + + print( + f" train-path len={len(train_ids)} first8={train_ids[:8].tolist()}", + flush=True, + ) + print( + f" infer-path len={len(infer_ids)} first8={infer_ids[:8].tolist()}", + flush=True, + ) + bos = tok.bos_token_id + print( + f" bos_token_id={bos} train leading bos count={int((train_ids[:3]==bos).sum())} " + f"infer leading bos count={int((infer_ids[:3]==bos).sum())}", + flush=True, + ) + same = len(train_ids) == len(infer_ids) and bool((train_ids == infer_ids).all()) + print(f" >>> IDENTICAL: {same}", flush=True) + if not same: + # show first divergence + n = min(len(train_ids), len(infer_ids)) + diff = (train_ids[:n] != infer_ids[:n]).nonzero() + first = int(diff[0]) if len(diff) else n + print(f" first divergence at pos {first}", flush=True) + print( + f" train[{first}:{first+6}]={train_ids[first:first+6].tolist()} " + f"-> {tok.convert_ids_to_tokens(train_ids[first:first+6].tolist())}", + flush=True, + ) + print( + f" infer[{first}:{first+6}]={infer_ids[first:first+6].tolist()} " + f"-> {tok.convert_ids_to_tokens(infer_ids[first:first+6].tolist())}", + flush=True, + ) + + # ---- 2. TEACHER-FORCED REPRODUCTION ---- + print( + "\n===== 2. TEACHER-FORCED ARGMAX REPRODUCTION (v3 adapter) =====", flush=True + ) + bnb = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + llm_int8_skip_modules=[ + "model.audio_tower", + "model.vision_tower", + "model.embed_audio", + "model.embed_vision", + "lm_head", + ], + ) + model = AutoModelForMultimodalLM.from_pretrained( + args.model, + quantization_config=bnb, + dtype=torch.bfloat16, + device_map="auto", + attn_implementation="eager", + ) + if args.adapter: + model = PeftModel.from_pretrained(model, args.adapter) + model.eval() + + collator = Gemma4AudioCollator(processor) + item0 = ds[0] + batch = collator([item0]) + batch = {k: v.to(model.device) for k, v in batch.items()} + with torch.inference_mode(): + out_with_labels = model(**batch) # model's OWN loss (its internal shift) + out = model(**{k: v for k, v in batch.items() if k != "labels"}) + logits = out.logits[0] + labels = batch["labels"][0] + # my manual shift: logits[t] predicts token t+1 + pred = logits[:-1].argmax(-1) + gold = labels[1:] + sup = gold != -100 + manual_loss = torch.nn.functional.cross_entropy( + logits[:-1][sup].float(), gold[sup], reduction="mean" + ) + # NO-shift variant: logits[t] vs labels[t] (in case model pre-shifts) + sup0 = labels != -100 + noshift_loss = torch.nn.functional.cross_entropy( + logits[sup0].float(), labels[sup0], reduction="mean" + ) + print( + f" model.forward(labels=...).loss = {float(out_with_labels.loss):.4f}", + flush=True, + ) + print(f" my manual SHIFTED loss = {float(manual_loss):.4f}", flush=True) + print(f" my manual NO-shift loss = {float(noshift_loss):.4f}", flush=True) + + # with vs without adapter (model's own loss) + if args.adapter: + with torch.inference_mode(), model.disable_adapter(): + base_loss = float(model(**batch).loss) + print( + f" model loss WITH adapter={float(out_with_labels.loss):.4f} " + f"WITHOUT adapter (base)={base_loss:.4f}", + flush=True, + ) + + print("\n per-sample argmax-match (shifted):", flush=True) + for item in ds: + b = collator([item]) + b = {k: v.to(model.device) for k, v in b.items()} + with torch.inference_mode(): + o = model(**{k: v for k, v in b.items() if k != "labels"}) + lg = o.logits[0] + lb = b["labels"][0] + p = lg[:-1].argmax(-1) + g = lb[1:] + s = g != -100 + n = int(s.sum()) + m = int((p[s] == g[s]).sum()) + print(f" {item['name']:16s} match={m}/{n} ({100*m/max(n,1):.1f}%)", flush=True) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/gemma4/finetune/debug_reload_loss.py b/extras/asr-services/providers/gemma4/finetune/debug_reload_loss.py new file mode 100644 index 00000000..db431e46 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/debug_reload_loss.py @@ -0,0 +1,51 @@ +import torch +from data import CosheSample7Dataset, Gemma4AudioCollator +from peft import PeftModel +from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig + +proc = AutoProcessor.from_pretrained("google/gemma-4-E2B-it") +bnb = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + llm_int8_skip_modules=[ + "model.audio_tower", + "model.vision_tower", + "model.embed_audio", + "model.embed_vision", + "lm_head", + ], +) +m = AutoModelForMultimodalLM.from_pretrained( + "google/gemma-4-E2B-it", + quantization_config=bnb, + dtype=torch.bfloat16, + device_map="auto", + attn_implementation="eager", +) +pm = PeftModel.from_pretrained(m, "/train/out/e2b-overfit-v3") +pm.eval() + +bnorm = sum(float(p.float().norm()) for n, p in pm.named_parameters() if "lora_B" in n) +print("reloaded lora_B sum norm:", round(bnorm, 3), flush=True) + +ds = CosheSample7Dataset( + "/data/coshe-eval/sample7", max_seconds=30.0, target_max_chars=540 +) +col = Gemma4AudioCollator(proc) + +b0 = col([ds[0]]) +b0 = {k: v.to(pm.device) for k, v in b0.items()} +for uc in [True, False]: + pm.config.use_cache = uc + if hasattr(pm, "base_model"): + pm.base_model.config.use_cache = uc + with torch.inference_mode(): + loss = float(pm(**b0).loss) + print(f" use_cache={uc}: reloaded-adapter loss = {loss:.4f}", flush=True) + +with torch.inference_mode(), pm.disable_adapter(): + b = col([ds[0]]) + b = {k: v.to(pm.device) for k, v in b.items()} + print(" base (adapter disabled) loss =", round(float(pm(**b).loss), 4), flush=True) diff --git a/extras/asr-services/providers/gemma4/finetune/eval_interleave.py b/extras/asr-services/providers/gemma4/finetune/eval_interleave.py new file mode 100644 index 00000000..17e23030 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/eval_interleave.py @@ -0,0 +1,107 @@ +"""Evaluate base or FT Gemma4-E2B on held-out clips for the interleave audio->pasted_text task. +Per test clip: feed all its <=28s chunks in one app-aware prompt (onestep, bf16, sdpa), +generate the formatted text in one shot, score against pasted_text (the target) with jiwer. +Same prompt/decode for base and adapter -> controlled delta. +""" + +import argparse +import json +import re +import statistics + +import jiwer +import torch +from data_interleave import _load_wav_f32, build_prompt +from peft import PeftModel +from transformers import AutoModelForMultimodalLM, AutoProcessor + +_B = re.compile(r"\[[^\]]*\]") +_A = re.compile(r"['’`]") +_P = re.compile(r"[^\w\s]") +_W = re.compile(r"\s+") + + +def norm(t): + t = _B.sub(" ", t).lower() + t = _A.sub("", t) + t = _P.sub(" ", t) + return _W.sub(" ", t).strip() + + +@torch.inference_mode() +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", default="google/gemma-4-E2B-it") + ap.add_argument("--adapter", default="") + ap.add_argument("--manifest", default="/home/wispr_interleave/manifest.jsonl") + ap.add_argument("--out", required=True) + ap.add_argument("--max_new_tokens", type=int, default=512) + args = ap.parse_args() + + test = [ + json.loads(l) for l in open(args.manifest) if json.loads(l)["split"] == "test" + ] + proc = AutoProcessor.from_pretrained(args.model) + proc.tokenizer.padding_side = "left" + model = AutoModelForMultimodalLM.from_pretrained( + args.model, dtype=torch.bfloat16, device_map="auto" + ) # default attn (sdpa), bf16 + if args.adapter: + model = PeftModel.from_pretrained(model, args.adapter) + print(f"adapter: {args.adapter}", flush=True) + model.eval() + model.config.use_cache = True + + rows = [] + fout = open(args.out, "w") + for ci, r in enumerate(test): + chunks = [_load_wav_f32(p) for p in r["chunks"]] + content = [{"type": "text", "text": build_prompt(r.get("app", ""))}] + content += [{"type": "audio", "audio": a} for a in chunks] + inp = proc.apply_chat_template( + [{"role": "user", "content": content}], + tokenize=True, + return_dict=True, + return_tensors="pt", + add_generation_prompt=True, + enable_thinking=False, + ).to(model.device) + g = model.generate( + **inp, max_new_tokens=args.max_new_tokens, do_sample=False, use_cache=True + ) + hyp = proc.decode( + g[0][inp["input_ids"].shape[-1] :], skip_special_tokens=True + ).strip() + ref = r["target"] + rn, hn = norm(ref), norm(hyp) + wer = jiwer.wer(rn, hn) if hn else 1.0 + rows.append( + { + "clip": r["clip"], + "app": r.get("app", ""), + "wer": round(wer, 4), + "ref_words": len(rn.split()), + "n_chunks": len(chunks), + "transcription": ref, + "hyp": hyp, + } + ) + fout.write(json.dumps(rows[-1], ensure_ascii=False) + "\n") + fout.flush() + print(f"[{ci+1}/{len(test)}] {r['clip']} wer={wer:.3f}", flush=True) + fout.close() + + refs_n = [norm(r["transcription"]) for r in rows] + hyps_n = [norm(r["hyp"]) for r in rows] + summary = { + "adapter": args.adapter or "base", + "clips": len(rows), + "corpus_wer": round(jiwer.wer(refs_n, hyps_n), 4), + "median_wer": round(statistics.median(r["wer"] for r in rows), 4), + } + print("SUMMARY " + json.dumps(summary), flush=True) + json.dump(summary, open(args.out.replace(".jsonl", "_summary.json"), "w"), indent=2) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/gemma4/finetune/eval_test_split.py b/extras/asr-services/providers/gemma4/finetune/eval_test_split.py new file mode 100644 index 00000000..492b16ff --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/eval_test_split.py @@ -0,0 +1,152 @@ +"""Transcribe the held-out CoSHE test split with the base model and/or a LoRA adapter, +emitting mlexp-schema JSONL for scoring. + +Using the SAME harness for base and adapters means the base-vs-FT WER delta isolates the +adapter's effect (no service / audio-handling confound). Audio is the cached 30s window +(same as training); the same window + decode settings are applied to every model, so the +delta is controlled even though the absolute WER carries the 30s-truncation tail equally. +""" + +import argparse +import json +import os +import random +import time + +import torch +from data import DEFAULT_PROMPT, CosheParquetDataset +from peft import PeftModel +from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig +from window_target import apply_window_truncation + + +def load_done(path): + done = set() + if os.path.exists(path): + for line in open(path): + try: + done.add(json.loads(line)["audio_file_name"]) + except Exception: + pass + return done + + +@torch.inference_mode() +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", default="google/gemma-4-E4B-it") + ap.add_argument("--adapter", default="") + ap.add_argument("--split_file", default="/home/gemma4ft/split_20_10_70.json") + ap.add_argument("--cache_path", default="/home/gemma4ft/out/coshe_full_cache.pkl") + ap.add_argument("--parquet_glob", default="/home/coshe-data/data/eval-*.parquet") + ap.add_argument("--out", required=True) + ap.add_argument("--batch_size", type=int, default=8) + ap.add_argument("--max_new_tokens", type=int, default=1024) + ap.add_argument("--limit", type=int, default=0, help="cap #clips (smoke test)") + ap.add_argument( + "--window_seconds", + type=float, + default=30.0, + help="proportionally truncate ref to this audio window (0=full)", + ) + ap.add_argument("--durations", default="/home/gemma4ft/durations.json") + args = ap.parse_args() + + proc = AutoProcessor.from_pretrained(args.model) + bnb = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + llm_int8_skip_modules=[ + "model.audio_tower", + "model.vision_tower", + "model.embed_audio", + "model.embed_vision", + "lm_head", + ], + ) + model = AutoModelForMultimodalLM.from_pretrained( + args.model, + quantization_config=bnb, + dtype=torch.bfloat16, + device_map="auto", + attn_implementation="eager", + ) + if args.adapter: + model = PeftModel.from_pretrained(model, args.adapter) + print(f"loaded adapter: {args.adapter}", flush=True) + model.eval() + model.config.use_cache = True + proc.tokenizer.padding_side = "left" + + split = json.load(open(args.split_file)) + test_names = split["test"] + ds = CosheParquetDataset( + args.parquet_glob, max_seconds=30.0, cache_path=args.cache_path + ) + apply_window_truncation(list(ds), args.durations, args.window_seconds) + by = {it["name"]: it for it in ds} + items = [by[n] for n in test_names if n in by] + if args.limit and args.limit < len(items): + # seeded random subset so every model is evaluated on the SAME clips + items = sorted( + random.Random(0).sample(items, args.limit), key=lambda it: it["name"] + ) + done = load_done(args.out) + items = [it for it in items if it["name"] not in done] + print( + f"test split={len(test_names)} to_do={len(items)} already_done={len(done)}", + flush=True, + ) + + fout = open(args.out, "a") + t0 = time.time() + for s in range(0, len(items), args.batch_size): + batch = items[s : s + args.batch_size] + texts, audios = [], [] + for it in batch: + msgs = [ + { + "role": "user", + "content": [ + {"type": "text", "text": DEFAULT_PROMPT}, + {"type": "audio", "audio": it["audio"]}, + ], + } + ] + texts.append( + proc.apply_chat_template( + msgs, tokenize=False, add_generation_prompt=True + ) + ) + audios.append(it["audio"]) + inp = proc(text=texts, audio=audios, return_tensors="pt", padding=True).to( + model.device + ) + ts = time.time() + out = model.generate( + **inp, max_new_tokens=args.max_new_tokens, do_sample=False, use_cache=True + ) + dt = time.time() - ts + for i, it in enumerate(batch): + hyp = proc.decode( + out[i][inp["input_ids"].shape[-1] :], skip_special_tokens=True + ).strip() + rec = { + "audio_file_name": it["name"], + "transcription": it["target"], + "hyp": hyp, + "asr_seconds": round(dt / len(batch), 3), + "duration_s": round(len(it["audio"]) / 16000.0, 2), + "raw": hyp, + } + fout.write(json.dumps(rec, ensure_ascii=False) + "\n") + fout.flush() + print(f"[{s + len(batch)}/{len(items)}] {time.time() - t0:.0f}s", flush=True) + fout.close() + print(f"DONE out={args.out}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/gemma4/finetune/eval_wer.py b/extras/asr-services/providers/gemma4/finetune/eval_wer.py new file mode 100644 index 00000000..5558c288 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/eval_wer.py @@ -0,0 +1,176 @@ +"""Batched WER eval for a Gemma 4 audio LoRA adapter over the full CoSHE dataset. + +No-cache generation is slow (gemma4 cache bug), so we BATCH clips through +generate() with left padding to make evaluating ~2000 clips feasible. Reports +corpus WER (the goal metric) + mean per-clip WER. + + python eval_wer.py --adapter /train/out/full --parquet_glob '/coshe-full/data/eval-*.parquet' \ + --cache_path /train/out/coshe_full_cache.pkl --batch_size 16 [--limit N] +""" + +import argparse + +import jiwer +import torch +from data import DEFAULT_PROMPT, CosheParquetDataset, CosheSample7Dataset +from peft import PeftModel +from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig + +_WER_NORM = jiwer.Compose( + [ + jiwer.ToLowerCase(), + jiwer.RemovePunctuation(), + jiwer.RemoveMultipleSpaces(), + jiwer.Strip(), + jiwer.ReduceToListOfListOfWords(), + ] +) + + +def parse_args(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="google/gemma-4-E2B-it") + p.add_argument("--adapter", default="") + p.add_argument("--parquet_glob", default="") + p.add_argument("--data_dir", default="/data/coshe-eval/sample7") + p.add_argument("--cache_path", default="") + p.add_argument("--limit", type=int, default=0) + p.add_argument("--max_seconds", type=float, default=30.0) + p.add_argument("--target_max_chars", type=int, default=0) + p.add_argument("--max_new_tokens", type=int, default=512) + p.add_argument("--batch_size", type=int, default=16) + p.add_argument( + "--use_cache", + action="store_true", + help="use KV cache (needs --patch for gemma4)", + ) + p.add_argument( + "--patch", action="store_true", help="apply gemma4 use_cache prefill bugfix" + ) + return p.parse_args() + + +def main(): + args = parse_args() + if args.patch: + import gemma4_cache_patch + + gemma4_cache_patch.apply() + print("applied gemma4 use_cache patch", flush=True) + proc = AutoProcessor.from_pretrained(args.model) + proc.tokenizer.padding_side = "left" # left padding for batched generation + + bnb = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + llm_int8_skip_modules=[ + "model.audio_tower", + "model.vision_tower", + "model.embed_audio", + "model.embed_vision", + "lm_head", + ], + ) + model = AutoModelForMultimodalLM.from_pretrained( + args.model, + quantization_config=bnb, + dtype=torch.bfloat16, + device_map="auto", + attn_implementation="eager", + ) + if args.adapter: + model = PeftModel.from_pretrained(model, args.adapter) + model.eval() + # cached gemma4 multimodal forward is buggy unless --patch is applied + model.config.use_cache = args.use_cache + + if args.parquet_glob: + ds = CosheParquetDataset( + args.parquet_glob, + max_seconds=args.max_seconds, + target_max_chars=args.target_max_chars, + limit=args.limit, + cache_path=args.cache_path, + ) + else: + ds = CosheSample7Dataset( + args.data_dir, + max_seconds=args.max_seconds, + target_max_chars=args.target_max_chars, + ) + items = list(ds) + print(f"Evaluating {len(items)} clips, batch_size={args.batch_size}", flush=True) + + refs, hyps, wers = [], [], [] + for start in range(0, len(items), args.batch_size): + batch = items[start : start + args.batch_size] + texts, audios = [], [] + for it in batch: + msgs = [ + { + "role": "user", + "content": [ + {"type": "text", "text": DEFAULT_PROMPT}, + {"type": "audio", "audio": it["audio"]}, + ], + } + ] + texts.append( + proc.apply_chat_template( + msgs, tokenize=False, add_generation_prompt=True + ) + ) + audios.append(it["audio"]) + inp = proc(text=texts, audio=audios, return_tensors="pt", padding=True).to( + model.device + ) + in_len = inp["input_ids"].shape[-1] + with torch.inference_mode(): + out = model.generate( + **inp, + max_new_tokens=args.max_new_tokens, + do_sample=False, + use_cache=args.use_cache, + ) + for i, it in enumerate(batch): + gen = proc.decode(out[i][in_len:], skip_special_tokens=True).strip() + refs.append(it["target"]) + hyps.append(gen) + w = ( + jiwer.wer( + it["target"], + gen, + reference_transform=_WER_NORM, + hypothesis_transform=_WER_NORM, + ) + if it["target"].strip() + else 0.0 + ) + wers.append(w) + done = start + len(batch) + run_corpus = jiwer.wer( + refs, hyps, reference_transform=_WER_NORM, hypothesis_transform=_WER_NORM + ) + print( + f" {done}/{len(items)} running corpus WER={run_corpus*100:.2f}% " + f"batch mean WER={sum(wers[-len(batch):])/len(batch)*100:.2f}%", + flush=True, + ) + + corpus = jiwer.wer( + refs, hyps, reference_transform=_WER_NORM, hypothesis_transform=_WER_NORM + ) + mean = sum(wers) / len(wers) + n_exact = sum(1 for w in wers if w == 0.0) + print(f"\n==== FULL EVAL ====", flush=True) + print( + f"clips={len(wers)} exact(WER=0)={n_exact} " + f"mean per-clip WER={mean*100:.2f}% CORPUS WER={corpus*100:.2f}%", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/gemma4/finetune/eval_windowed_stitch.py b/extras/asr-services/providers/gemma4/finetune/eval_windowed_stitch.py new file mode 100644 index 00000000..2eb5843d --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/eval_windowed_stitch.py @@ -0,0 +1,319 @@ +"""Evaluate a model (base or LoRA adapter) on the held-out CoSHE test clips via the +windowed-stitch path: cut each whole clip into its honest <=28s windows (the SAME +forced-alignment windows used for training, from windows_v2.jsonl), transcribe each window +with PLAIN_PROMPT, then concatenate window hyps in order -> a full-clip hypothesis scored +against the full GT transcript. + +Non-overlapping windows -> stitch is a plain concat (no overlap dedup). The same windowing ++ prompt + decode settings are used for base and every adapter, so the base->FT WER delta +is controlled. use_cache=True is safe (gemma4 cache bug fixed in transformers >=5.10). + +Out: mlexp-schema JSONL keyed by audio_file_name: {transcription (full GT), hyp (stitched), +n_wins, duration_s}. Resumable (skips clips already written). +""" + +import argparse +import glob +import io +import json +import os +import random +import time +from collections import defaultdict + +import numpy as np +import pyarrow.parquet as pq +import soundfile as sf +import torch +from data_windowed import PLAIN_PROMPT +from peft import PeftModel +from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig + +# Exact prod/benchmark prompt (interleave_probe.py PLAIN_PROMPT) that yields E2B ~10.6% +# median on CoSHE-100. gemma4 is prompt-conditioned; an elaborate prompt degrades base. +MINIMAL_PROMPT = ( + "Transcribe the following speech in its original language into text. " + "Only output the transcription text itself, with no commentary or explanation." +) +PROMPTS = {"plain": PLAIN_PROMPT, "minimal": MINIMAL_PROMPT} + + +def load_done(path): + done = set() + if os.path.exists(path): + for line in open(path): + try: + done.add(json.loads(line)["audio_file_name"]) + except Exception: + pass + return done + + +@torch.inference_mode() +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", default="google/gemma-4-E2B-it") + ap.add_argument("--adapter", default="") + ap.add_argument("--windows", default="/home/coshe_windowed/windows_v2.jsonl") + ap.add_argument("--split", default="/home/coshe_windowed/windowed_split.json") + ap.add_argument("--parquet_glob", default="/home/coshe/data/eval-*.parquet") + ap.add_argument("--out", required=True) + ap.add_argument("--batch_size", type=int, default=8) + ap.add_argument("--max_new_tokens", type=int, default=256) + ap.add_argument( + "--no_repeat_ngram", + type=int, + default=0, + help="generate no_repeat_ngram_size; 3 kills greedy repetition-collapse", + ) + ap.add_argument("--repetition_penalty", type=float, default=1.0) + ap.add_argument( + "--limit", type=int, default=0, help="cap #test clips (seeded subset)" + ) + ap.add_argument( + "--quant", + default="bf16", + choices=["bf16", "4bit"], + help="bf16 = prod config (no quant); 4bit only for tiny GPUs", + ) + ap.add_argument( + "--prompt", + default="plain", + choices=["plain", "minimal"], + help="minimal = exact benchmark/prod prompt (E2B ~10.6%); plain = training prompt", + ) + ap.add_argument( + "--names_file", + default="", + help="JSON list of clip names to restrict eval to (∩ held-out)", + ) + ap.add_argument( + "--inject", + default="twostep", + choices=["twostep", "onestep"], + help="onestep = prod path apply_chat_template(tokenize=True,return_dict=True)", + ) + ap.add_argument( + "--fixed_win", + type=float, + default=0.0, + help="ignore fa windows; split each clip into fixed N-sec windows (benchmark style)", + ) + ap.add_argument( + "--attn", + default="", + help="attn_implementation; empty = model default (prod). 'eager' degrades gemma4 gen", + ) + args = ap.parse_args() + prompt_text = PROMPTS[args.prompt] + + test_clips = json.load(open(args.split))["test_clips"] + if args.names_file: + want = set(json.load(open(args.names_file))) + test_clips = [c for c in test_clips if c in want] + if args.limit and args.limit < len(test_clips): + test_clips = sorted(random.Random(0).sample(test_clips, args.limit)) + test_set = set(test_clips) + done = load_done(args.out) + todo = [c for c in test_clips if c not in done] + + wins_by_clip = defaultdict(list) + for line in open(args.windows): + r = json.loads(line) + if r["audio_file_name"] in test_set: + wins_by_clip[r["audio_file_name"]].append(r) + for c in wins_by_clip: + wins_by_clip[c].sort(key=lambda w: w["win_idx"]) + + # decode needed test clips from parquet (audio + full GT), slice window audio + gt, clip_audio = {}, {} + need = set(todo) + for f in sorted(glob.glob(args.parquet_glob)): + if not need: + break + t = pq.ParquetFile(f).read( + columns=["audio_file_name", "transcription", "audio"] + ) + names = t["audio_file_name"].to_pylist() + for ci, name in enumerate(names): + if name not in need: + continue + gt[name] = t["transcription"][ci].as_py().strip() + wav, sr = sf.read(io.BytesIO(t["audio"][ci].as_py()["bytes"])) + if wav.ndim > 1: + wav = wav.mean(1) + if sr != 16000: + import librosa + + wav = librosa.resample( + wav.astype(np.float32), orig_sr=sr, target_sr=16000 + ) + clip_audio[name] = np.ascontiguousarray(wav, dtype=np.float32) + need.discard(name) + print( + f"test clips total={len(test_clips)} todo={len(todo)} decoded={len(clip_audio)}", + flush=True, + ) + + proc = AutoProcessor.from_pretrained(args.model) + load_kw = dict(dtype=torch.bfloat16, device_map="auto") + if ( + args.attn + ): # benchmark/prod uses the MODEL DEFAULT (sdpa); eager degrades gemma4 gen + load_kw["attn_implementation"] = args.attn + if args.quant == "4bit": # QLoRA-training config; degrades small E2B at inference + load_kw["quantization_config"] = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + llm_int8_skip_modules=[ + "model.audio_tower", + "model.vision_tower", + "model.embed_audio", + "model.embed_vision", + "lm_head", + ], + ) + model = AutoModelForMultimodalLM.from_pretrained(args.model, **load_kw) + if args.adapter: + model = PeftModel.from_pretrained(model, args.adapter) + print(f"loaded adapter: {args.adapter}", flush=True) + model.eval() + model.config.use_cache = True + proc.tokenizer.padding_side = "left" + + # optional: ignore fa windows, split each clip into fixed N-sec windows (benchmark style) + if args.fixed_win: + n = int(args.fixed_win * 16000) + wins_by_clip = { + c: [ + { + "win_idx": i, + "start": (i * n) / 16000, + "end": min(len(clip_audio[c]), (i + 1) * n) / 16000, + } + for i in range(max(1, (len(clip_audio[c]) + n - 1) // n)) + ] + for c in todo + } + + # flatten all windows of all todo clips into one work list, batch across clips + work = [] # (clip, win_idx, audio_slice) + for c in todo: + for w in wins_by_clip[c]: + a0, a1 = int(w["start"] * 16000), int(w["end"] * 16000) + work.append((c, w["win_idx"], clip_audio[c][a0:a1])) + print(f"windows to transcribe: {len(work)}", flush=True) + + gen_kw = dict( + max_new_tokens=args.max_new_tokens, + do_sample=False, + use_cache=True, + repetition_penalty=args.repetition_penalty, + ) + if args.no_repeat_ngram: + gen_kw["no_repeat_ngram_size"] = args.no_repeat_ngram + + hyp_parts = defaultdict(dict) # clip -> {win_idx: hyp} + fout = open(args.out, "a") + t0 = time.time() + bs = 1 if args.inject == "onestep" else args.batch_size + for s in range(0, len(work), bs): + batch = work[s : s + bs] + if ( + args.inject == "onestep" + ): # prod path: tokenize=True one-shot, per-window (bs=1) + import os as _os + import tempfile + import wave as _wave + + c, wi, au = batch[0] + # match the benchmark exactly: write a 16k PCM_16 WAV and pass the PATH (not the + # raw array) — the Gemma4 processor conditions differently on array vs file. + tf = tempfile.NamedTemporaryFile(suffix=".wav", delete=False) + with _wave.open(tf.name, "wb") as wf: + wf.setnchannels(1) + wf.setsampwidth(2) + wf.setframerate(16000) + wf.writeframes((np.clip(au, -1, 1) * 32767).astype(np.int16).tobytes()) + msgs = [ + { + "role": "user", + "content": [ + {"type": "text", "text": prompt_text}, + {"type": "audio", "audio": tf.name}, + ], + } + ] + inp = proc.apply_chat_template( + msgs, + tokenize=True, + return_dict=True, + return_tensors="pt", + add_generation_prompt=True, + enable_thinking=False, + ).to(model.device) + out = model.generate(**inp, **gen_kw) + _os.unlink(tf.name) + hyp_parts[c][wi] = proc.decode( + out[0][inp["input_ids"].shape[-1] :], skip_special_tokens=True + ).strip() + else: # legacy two-step path + texts, audios = [], [] + for _, _, au in batch: + msgs = [ + { + "role": "user", + "content": [ + {"type": "text", "text": prompt_text}, + {"type": "audio", "audio": au}, + ], + } + ] + texts.append( + proc.apply_chat_template( + msgs, + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + ) + audios.append(au) + inp = proc(text=texts, audio=audios, return_tensors="pt", padding=True).to( + model.device + ) + out = model.generate(**inp, **gen_kw) + for i, (c, wi, _) in enumerate(batch): + hyp_parts[c][wi] = proc.decode( + out[i][inp["input_ids"].shape[-1] :], skip_special_tokens=True + ).strip() + if (s // bs) % 40 == 0: + print(f"[{s + len(batch)}/{len(work)}] {time.time() - t0:.0f}s", flush=True) + + # stitch + write per fully-completed clip + for c in todo: + parts = hyp_parts.get(c, {}) + if len(parts) != len(wins_by_clip[c]): + continue # incomplete (shouldn't happen); skip to retain resumability + stitched = " ".join(parts[i] for i in sorted(parts)).strip() + dur = round(len(clip_audio[c]) / 16000.0, 2) + fout.write( + json.dumps( + { + "audio_file_name": c, + "transcription": gt[c], + "hyp": stitched, + "n_wins": len(parts), + "duration_s": dur, + }, + ensure_ascii=False, + ) + + "\n" + ) + fout.close() + print(f"DONE out={args.out}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/gemma4/finetune/eval_wispr_stitch.py b/extras/asr-services/providers/gemma4/finetune/eval_wispr_stitch.py new file mode 100644 index 00000000..a18d6781 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/eval_wispr_stitch.py @@ -0,0 +1,148 @@ +"""Evaluate base or LoRA-adapter Gemma4-E2B on the held-out Wispr test clips. + +Test clips are pre-sliced into the SAME fa windows used for training (manifest split=="test"). +Each window is transcribed with WISPR_PROMPT via the prod onestep path +(apply_chat_template(tokenize=True, return_dict=True), bf16, model-default attn=sdpa, +use_cache=True), then window hyps are concatenated per clip (non-overlapping -> plain concat) +into a full-clip hypothesis scored against asr_text (test_refs.jsonl) with jiwer. + +Same windowing + prompt + decode for base and every adapter -> the base->FT WER delta is +controlled. Run once with --adapter "" (base) and once with --adapter <dir> (FT). + +Usage: + python eval_wispr_stitch.py --manifest .../manifest.jsonl --refs .../test_refs.jsonl \ + --adapter "" --out base.jsonl +""" + +import argparse +import json +import re +import statistics +import time +from collections import defaultdict + +import jiwer +import torch +from data_wispr import WISPR_PROMPT +from peft import PeftModel +from transformers import AutoModelForMultimodalLM, AutoProcessor + +_BRACKET = re.compile(r"\[[^\]]*\]") +_APOS = re.compile(r"['’`]") +_PUNCT = re.compile(r"[^\w\s]") +_WS = re.compile(r"\s+") + + +def normalize(t: str) -> str: + t = _BRACKET.sub(" ", t).lower() + t = _APOS.sub("", t) + t = _PUNCT.sub(" ", t) + return _WS.sub(" ", t).strip() + + +@torch.inference_mode() +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", default="google/gemma-4-E2B-it") + ap.add_argument("--adapter", default="") + ap.add_argument("--manifest", default="/home/wispr_windowed/manifest.jsonl") + ap.add_argument("--refs", default="/home/wispr_windowed/test_refs.jsonl") + ap.add_argument("--out", required=True) + ap.add_argument("--max_new_tokens", type=int, default=256) + args = ap.parse_args() + + # test windows grouped by clip + wins_by_clip = defaultdict(list) + for line in open(args.manifest): + r = json.loads(line) + if r["split"] == "test": + wins_by_clip[r["clip"]].append(r) + for c in wins_by_clip: + wins_by_clip[c].sort(key=lambda w: w["win_idx"]) + refs = { + json.loads(l)["audio_file_name"]: json.loads(l)["transcription"] + for l in open(args.refs) + } + + proc = AutoProcessor.from_pretrained(args.model) + proc.tokenizer.padding_side = "left" + model = AutoModelForMultimodalLM.from_pretrained( + args.model, dtype=torch.bfloat16, device_map="auto" + ) # default attn (sdpa), bf16 + if args.adapter: + model = PeftModel.from_pretrained(model, args.adapter) + print(f"adapter: {args.adapter}", flush=True) + model.eval() + model.config.use_cache = True + + gen_kw = dict(max_new_tokens=args.max_new_tokens, do_sample=False, use_cache=True) + t0 = time.time() + rows = [] + fout = open(args.out, "w") + for ci, (clip, wins) in enumerate(sorted(wins_by_clip.items())): + parts = [] + for w in wins: + msgs = [ + { + "role": "user", + "content": [ + {"type": "text", "text": WISPR_PROMPT}, + {"type": "audio", "audio": w["audio"]}, + ], + } + ] + inp = proc.apply_chat_template( + msgs, + tokenize=True, + return_dict=True, + return_tensors="pt", + add_generation_prompt=True, + enable_thinking=False, + ).to(model.device) + out = model.generate(**inp, **gen_kw) + parts.append( + proc.decode( + out[0][inp["input_ids"].shape[-1] :], skip_special_tokens=True + ).strip() + ) + hyp = " ".join(parts).strip() + ref = refs[clip] + rn, hn = normalize(ref), normalize(hyp) + wer = jiwer.wer(rn, hn) if hn else 1.0 + rows.append( + { + "audio_file_name": clip, + "wer": round(wer, 4), + "ref_words": len(rn.split()), + "n_wins": len(wins), + "transcription": ref, + "hyp": hyp, + } + ) + fout.write(json.dumps(rows[-1], ensure_ascii=False) + "\n") + fout.flush() + print( + f"[{ci+1}/{len(wins_by_clip)}] {clip} wer={wer:.3f} ({time.time()-t0:.0f}s)", + flush=True, + ) + fout.close() + + # corpus + median + tot_ref = sum(r["ref_words"] for r in rows) + refs_n = [normalize(refs[r["audio_file_name"]]) for r in rows] + hyps_n = [normalize(r["hyp"]) for r in rows] + corpus = jiwer.wer(refs_n, hyps_n) + median = statistics.median(r["wer"] for r in rows) + summary = { + "adapter": args.adapter or "base", + "clips": len(rows), + "corpus_wer": round(corpus, 4), + "median_wer": round(median, 4), + "total_ref_words": tot_ref, + } + print("SUMMARY " + json.dumps(summary), flush=True) + json.dump(summary, open(args.out.replace(".jsonl", "_summary.json"), "w"), indent=2) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/gemma4/finetune/gemini_transcribe.py b/extras/asr-services/providers/gemma4/finetune/gemini_transcribe.py new file mode 100644 index 00000000..4e633305 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/gemini_transcribe.py @@ -0,0 +1,69 @@ +"""Send an audio window to Gemini for a reference transcription.""" + +import argparse +import base64 +import io +import json +import os +import urllib.error +import urllib.request + +import soundfile as sf + +SR = 16000 + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--audio", default="/w/ankush-yap-fixed.wav") + p.add_argument("--offset", type=float, default=0.0) + p.add_argument("--seconds", type=float, default=60.0) + p.add_argument("--model", default="gemini-3.5-flash") + args = p.parse_args() + + a, sr = sf.read(args.audio, dtype="float32") + if a.ndim > 1: + a = a.mean(1) + if sr != SR: + import librosa + + a = librosa.resample(a, orig_sr=sr, target_sr=SR) + seg = a[int(args.offset * SR) : int((args.offset + args.seconds) * SR)] + buf = io.BytesIO() + sf.write(buf, seg, SR, format="WAV", subtype="PCM_16") + b64 = base64.b64encode(buf.getvalue()).decode() + print( + f"window {args.offset:.0f}-{args.offset+args.seconds:.0f}s, {len(b64)/1e6:.1f}MB b64", + flush=True, + ) + + key = os.environ["GOOGLE_API_KEY"] + url = f"https://generativelanguage.googleapis.com/v1beta/models/{args.model}:generateContent?key={key}" + body = { + "contents": [ + { + "parts": [ + { + "text": "Transcribe this audio verbatim. Output only the transcript text." + }, + {"inline_data": {"mime_type": "audio/wav", "data": b64}}, + ] + } + ] + } + req = urllib.request.Request( + url, + data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}, + ) + try: + resp = json.load(urllib.request.urlopen(req, timeout=120)) + cand = resp["candidates"][0]["content"]["parts"][0]["text"] + print("=== GEMINI TRANSCRIPT ===", flush=True) + print(cand, flush=True) + except urllib.error.HTTPError as e: + print("HTTP", e.code, e.read().decode()[:800], flush=True) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/gemma4/finetune/gemma4_cache_patch.py b/extras/asr-services/providers/gemma4/finetune/gemma4_cache_patch.py new file mode 100644 index 00000000..479a7253 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/gemma4_cache_patch.py @@ -0,0 +1,64 @@ +"""Monkeypatch for the gemma4 use_cache prefill bug (transformers 5.5.0). + +Bug: `create_causal_mask_mapping` only applies the audio/vision BIDIRECTIONAL +attention mask (`or_mask_function`) on the "first iteration" (prefill). It infers +prefill via: + + past_key_values is None or not past_key_values.is_initialized or pixel_values is not None + +But a freshly-created DynamicCache reports `is_initialized=True` *before any token +is cached*, so on the prefill forward (with use_cache=True / generate) this is +False → the audio soft tokens are masked purely causally instead of bidirectionally +→ the audio conditioning is wrong → logits collapse toward the base model +(measured: loss 4.49 with cache vs 0.10 without on a fully-overfit adapter). + +Fix: detect prefill by `get_seq_length() == 0` (no tokens cached yet) instead of +`is_initialized`. This makes generation correct AND fast (KV cache usable), turning +a ~100s/clip no-cache eval into normal cached generation. + +Usage: `import gemma4_cache_patch; gemma4_cache_patch.apply()` BEFORE loading the model. +""" + +import transformers.models.gemma4.modeling_gemma4 as _g4 + +_orig_create_causal_mask_mapping = _g4.create_causal_mask_mapping + + +def _patched_create_causal_mask_mapping( + config, + inputs_embeds, + attention_mask, + past_key_values, + position_ids, + mm_token_type_ids=None, + pixel_values=None, + is_training=False, + is_first_iteration=None, + **kwargs, +): + if is_first_iteration is None: + if past_key_values is None: + is_first_iteration = True + else: + try: + seqlen = past_key_values.get_seq_length() + except Exception: + seqlen = 0 + is_first_iteration = (seqlen == 0) or (pixel_values is not None) + return _orig_create_causal_mask_mapping( + config, + inputs_embeds, + attention_mask, + past_key_values, + position_ids, + mm_token_type_ids=mm_token_type_ids, + pixel_values=pixel_values, + is_training=is_training, + is_first_iteration=is_first_iteration, + **kwargs, + ) + + +def apply(): + _g4.create_causal_mask_mapping = _patched_create_causal_mask_mapping + return True diff --git a/extras/asr-services/providers/gemma4/finetune/infer.py b/extras/asr-services/providers/gemma4/finetune/infer.py new file mode 100644 index 00000000..effd21c3 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/infer.py @@ -0,0 +1,169 @@ +"""Verify a Gemma 4 audio LoRA adapter by transcribing the training clips and +comparing to ground truth (overfit check). + +Loads the 4-bit base + adapter, generates on each sample7 clip, prints the +output next to the manifest target, and reports a crude char-level similarity +(SequenceMatcher ratio) so we can see the loss-near-zero overfit reflected in +the actual decoded text. + +Usage: + python infer.py --model google/gemma-4-E2B-it \ + --adapter /train/out/e2b-overfit --data_dir /data/coshe-eval/sample7 +""" + +import argparse +from difflib import SequenceMatcher + +import jiwer +import torch +from data import DEFAULT_PROMPT, CosheParquetDataset, CosheSample7Dataset +from peft import PeftModel +from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig + +# WER normalization: lowercase, strip punctuation, collapse whitespace. For an +# overfit, the model is trained on the target's exact script, so exact +# reproduction -> WER 0 (no romanization needed; raw mixed-script compares fine +# when hyp and ref use the same script, which they do after training). +_WER_NORM = jiwer.Compose( + [ + jiwer.ToLowerCase(), + jiwer.RemovePunctuation(), + jiwer.RemoveMultipleSpaces(), + jiwer.Strip(), + jiwer.ReduceToListOfListOfWords(), + ] +) + + +def compute_wer(ref: str, hyp: str) -> float: + if not ref.strip(): + return 0.0 if not hyp.strip() else 1.0 + return jiwer.wer( + ref, hyp, reference_transform=_WER_NORM, hypothesis_transform=_WER_NORM + ) + + +def parse_args(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="google/gemma-4-E2B-it") + p.add_argument("--adapter", default="") + p.add_argument("--data_dir", default="/data/coshe-eval/sample7") + p.add_argument("--parquet_glob", default="") + p.add_argument("--limit", type=int, default=0) + p.add_argument("--cache_path", default="") + p.add_argument("--max_seconds", type=float, default=30.0) + p.add_argument("--target_max_chars", type=int, default=0) + p.add_argument("--max_new_tokens", type=int, default=512) + return p.parse_args() + + +def main(): + args = parse_args() + processor = AutoProcessor.from_pretrained(args.model) + processor.tokenizer.padding_side = "left" + + bnb = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + llm_int8_skip_modules=[ + "model.audio_tower", + "model.vision_tower", + "model.embed_audio", + "model.embed_vision", + "lm_head", + ], + ) + model = AutoModelForMultimodalLM.from_pretrained( + args.model, + quantization_config=bnb, + dtype=torch.bfloat16, + device_map="auto", + attn_implementation="eager", + ) + tag = "BASE (no adapter)" + if args.adapter: + model = PeftModel.from_pretrained(model, args.adapter) + tag = f"ADAPTER {args.adapter}" + model.eval() + # CRITICAL: gemma4's multimodal forward produces WRONG audio-conditioned logits + # when the KV cache is enabled (transformers 5.5.0) — with use_cache=True a + # fully-overfit adapter scores loss ~4.5 (≈ base) instead of ~0.1, and + # generation reverts to near-base output. Generate with the cache OFF. + model.config.use_cache = False + + if args.parquet_glob: + dataset = CosheParquetDataset( + args.parquet_glob, + max_seconds=args.max_seconds, + target_max_chars=args.target_max_chars, + limit=args.limit, + cache_path=args.cache_path, + ) + else: + dataset = CosheSample7Dataset( + args.data_dir, + max_seconds=args.max_seconds, + target_max_chars=args.target_max_chars, + ) + print(f"\n===== INFERENCE: {tag} =====", flush=True) + ratios = [] + wers = [] + all_gens = [] + for item in dataset: + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": DEFAULT_PROMPT}, + {"type": "audio", "audio": item["audio"]}, + ], + } + ] + inputs = processor.apply_chat_template( + messages, + tokenize=True, + return_dict=True, + return_tensors="pt", + add_generation_prompt=True, + ).to(model.device) + in_len = inputs["input_ids"].shape[-1] + with torch.inference_mode(): + out = model.generate( + **inputs, + max_new_tokens=args.max_new_tokens, + do_sample=False, + use_cache=False, + ) + gen = processor.decode(out[0][in_len:], skip_special_tokens=True).strip() + target = item["target"] + ratio = SequenceMatcher(None, gen, target).ratio() + wer = compute_wer(target, gen) + ratios.append(ratio) + wers.append(wer) + all_gens.append(gen) + print( + f"\n--- {item['name']} (char-sim {ratio:.3f} WER {wer*100:.2f}%) ---", + flush=True, + ) + print(f" TARGET: {target[:240]}", flush=True) + print(f" OUTPUT: {gen[:240]}", flush=True) + + mean_wer = sum(wers) / len(wers) + # corpus WER (aggregate over all words), the headline metric for the goal + corpus_wer = jiwer.wer( + [d["target"] for d in dataset], + all_gens, + reference_transform=_WER_NORM, + hypothesis_transform=_WER_NORM, + ) + print(f"\nMean char-similarity: {sum(ratios)/len(ratios):.3f}", flush=True) + print( + f"Mean per-clip WER: {mean_wer*100:.2f}% Corpus WER: {corpus_wer*100:.2f}%", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/gemma4/finetune/introspect.py b/extras/asr-services/providers/gemma4/finetune/introspect.py new file mode 100644 index 00000000..87255f76 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/introspect.py @@ -0,0 +1,162 @@ +"""Introspect a Gemma 4 E*B model + processor for QLoRA fine-tuning. + +Loads the model in 4-bit and dumps the facts we need to write the training +script correctly for THIS model (the gemma4 family), since published recipes +target gemma3n and key/module names can differ: + + 1. Processor output keys for a (text + audio) example, with shapes/dtypes. + 2. Audio/image special token ids exposed on the tokenizer/config. + 3. Module-name families for LoRA targets: text-decoder projections, the + audio projector, and embedding/lm_head modules. + +Run inside the gemma4 image (transformers >= 5.5). +""" + +import os +import wave +from collections import Counter + +import numpy as np +import torch +from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig + +MODEL_ID = os.getenv("ASR_MODEL", "google/gemma-4-E2B-it") +AUDIO = os.getenv("PROBE_AUDIO", "/data/coshe-eval/sample7/audio/audio_13.wav") +PROMPT = "Transcribe the following speech segment in its original language." + + +def load_wav_16k_mono(path: str, max_seconds: float = 30.0) -> np.ndarray: + with wave.open(path, "rb") as wf: + sr = wf.getframerate() + n = wf.getnframes() + raw = wf.readframes(n) + audio = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0 + # mono assumed (sample7 is mono). Resample to 16k via librosa. + if sr != 16000: + import librosa + + audio = librosa.resample(audio, orig_sr=sr, target_sr=16000) + audio = audio[: int(16000 * max_seconds)] + return audio + + +def main(): + print(f"=== MODEL: {MODEL_ID} ===", flush=True) + bnb = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + ) + processor = AutoProcessor.from_pretrained(MODEL_ID) + model = AutoModelForMultimodalLM.from_pretrained( + MODEL_ID, + quantization_config=bnb, + dtype=torch.bfloat16, + device_map="auto", + ) + + tok = processor.tokenizer + print("\n=== SPECIAL TOKEN IDS ===", flush=True) + for attr in [ + "pad_token_id", + "eos_token_id", + "bos_token_id", + "audio_token_id", + "image_token_id", + "boi_token_id", + "eoi_token_id", + "boa_token_id", + "eoa_token_id", + ]: + print(f" tokenizer.{attr} = {getattr(tok, attr, '<none>')}", flush=True) + cfg = model.config + for attr in [ + "audio_token_id", + "image_token_id", + "boi_token_id", + "eoi_token_id", + "boa_token_id", + "eoa_token_id", + "eoa_token_index", + ]: + print(f" config.{attr} = {getattr(cfg, attr, '<none>')}", flush=True) + + print("\n=== PROCESSOR OUTPUT (text+audio, TRAINING form) ===", flush=True) + audio = load_wav_16k_mono(AUDIO) + print( + f" probe audio: {AUDIO} samples={len(audio)} ({len(audio)/16000:.1f}s)", + flush=True, + ) + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": PROMPT}, + {"type": "audio", "audio": audio}, + ], + }, + { + "role": "assistant", + "content": [{"type": "text", "text": "Speaker 1: hello world"}], + }, + ] + # tokenize=False then call processor with text+audio (the HF recipe pattern) + text = processor.apply_chat_template( + messages, tokenize=False, add_generation_prompt=False + ) + print(f" templated text (first 300 chars):\n {text[:300]!r}", flush=True) + batch = processor(text=[text], audio=[audio], return_tensors="pt", padding=True) + for k, v in batch.items(): + if hasattr(v, "shape"): + print(f" key {k:24s} shape={tuple(v.shape)} dtype={v.dtype}", flush=True) + else: + print(f" key {k:24s} type={type(v)}", flush=True) + + print("\n=== MODULE NAME FAMILIES (Linear leaf modules) ===", flush=True) + fam = Counter() + audio_like = set() + embed_like = set() + for name, mod in model.named_modules(): + leaf = name.split(".")[-1] + cls = type(mod).__name__ + if "Linear" in cls or "lora" in cls.lower(): + fam[leaf] += 1 + low = name.lower() + if any(t in low for t in ["audio", "speech", "conformer"]): + # record short module suffixes from the audio tower / projector + if "Linear" in cls or "Embedding" in cls or "proj" in low: + audio_like.add(f"{leaf} ({cls})") + if "embed" in low or leaf in ("lm_head",): + embed_like.add(f"{name} ({cls})") + print(" Linear leaf-name -> count:", flush=True) + for leaf, c in fam.most_common(): + print(f" {leaf:28s} x{c}", flush=True) + print("\n audio-tower / projector linear+embedding leaves (sample):", flush=True) + for s in sorted(audio_like)[:40]: + print(f" {s}", flush=True) + print("\n embedding / lm_head modules (full names, sample):", flush=True) + for s in sorted(embed_like)[:25]: + print(f" {s}", flush=True) + + print("\n=== top-level named children ===", flush=True) + for name, _ in model.named_children(): + print(f" {name}", flush=True) + # one level deeper into the LM + audio tower roots + for root in [ + "model", + "language_model", + "audio_tower", + "audio_model", + "multi_modal_projector", + ]: + sub = getattr(model, root, None) + if sub is not None: + kids = [n for n, _ in sub.named_children()] + print(f" {root}.children = {kids}", flush=True) + + print("\nDONE", flush=True) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/gemma4/finetune/poll_until_target.sh b/extras/asr-services/providers/gemma4/finetune/poll_until_target.sh new file mode 100755 index 00000000..09398adc --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/poll_until_target.sh @@ -0,0 +1,32 @@ +#!/bin/bash +# Poll the A100 training log via the Jupyter contents API (SSH-free) until the +# early-stopping loop reports it hit <2% WER ("TARGET REACHED" / "DONE"), or the +# training process dies, or a long timeout. Exits 0 with a summary so the agent is +# re-invoked to retrieve the adapter + pause the instance. +set -u +ENVBKP=/home/ankush/workspaces/friend-lite/.env.bkp +FT=/home/ankush/workspaces/friend-lite/extras/asr-services/providers/gemma4/finetune +JTOKEN=$(grep -hiE "^jarvislabs=" "$ENVBKP" | head -1 | cut -d= -f2- | tr -d '"'\'' \r') + +# Resolve host + jupyter token once. +URL=$(JLTOKEN="$JTOKEN" uv run --with "git+https://github.com/jarvislabsai/jlclient.git" python3 -c " +import os; from jlclient import jarvisclient; from jlclient.jarvisclient import * +jarvisclient.token=os.environ['JLTOKEN']; print(User.get_instances()[0].url)" 2>/dev/null | grep notebooks) +HOST=$(echo "$URL" | sed -E 's#https://([^/]+)/.*#\1#') +TOK=$(echo "$URL" | sed -E 's#.*token=##') +LOG="https://$HOST/api/contents/gemma4ft/out/train_full2.log?token=$TOK&type=file&format=text&content=1" +echo "polling host=$HOST" + +for i in $(seq 1 240); do # up to 240 * 4min = 16h + C=$(curl -s --max-time 90 "$LOG" 2>/dev/null) + # latest epoch-mean signal: last WER eval line and last loss/epoch + EVAL=$(echo "$C" | grep -aoE '\[epoch [0-9]+\][^\\]*(corpus WER = [0-9.]+%|skipping WER eval)' | tail -1) + TGT=$(echo "$C" | grep -aoE 'TARGET REACHED[^\\]*|DONE \(final adapter saved\)') + LOSSEP=$(echo "$C" | grep -aoE "'loss': '[0-9.]+', 'grad_norm'[^}]*'epoch': '[0-9.]+'" | tail -1 | sed -E "s/, 'grad_norm'[^,]*,/,/") + echo "[$(date +%H:%M)] iter=$i | $LOSSEP | eval: ${EVAL:-none}" + if [ -n "$TGT" ]; then echo "STOP_SIGNAL: $TGT"; exit 0; fi + # detect dead training: if process gone (no recent log growth is hard; rely on TARGET/DONE or manual) + sleep 240 +done +echo "poller timed out after 16h without TARGET" +exit 0 diff --git a/extras/asr-services/providers/gemma4/finetune/run_split_pipeline.sh b/extras/asr-services/providers/gemma4/finetune/run_split_pipeline.sh new file mode 100644 index 00000000..594a48c9 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/run_split_pipeline.sh @@ -0,0 +1,49 @@ +#!/bin/bash +# Drives the full split-LoRA generalization experiment on the A100, sequentially +# (single GPU). Each stage logs to its own file and appends a marker to a master +# progress log; a failed stage does not abort the rest. Window=30s "first-30s +# transcription" task throughout (proportional target truncation). +cd /home/gemma4ft || exit 1 +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True +PROG=/home/gemma4ft/out/pipeline_progress.log +echo "PIPELINE_START $(date -u +%H:%M:%S)" > "$PROG" + +stage () { # stage <name> <logfile> <cmd...> + local name="$1" log="$2"; shift 2 + echo "STAGE_START $name $(date -u +%H:%M:%S)" >> "$PROG" + if "$@" > "$log" 2>&1; then echo "STAGE_OK $name $(date -u +%H:%M:%S)" >> "$PROG" + else echo "STAGE_FAILED $name rc=$? $(date -u +%H:%M:%S)" >> "$PROG"; fi +} + +COMMON="--lr 1e-5 --lora_r 16 --lora_alpha 32 --batch_size 2 --grad_accum 2 --patience 4 --epochs 40 --window_seconds 30" +EVAL_COMMON="--limit 700 --batch_size 8 --max_new_tokens 448 --window_seconds 30" + +# 1) decoder-only r16 +stage train_decoder out/split_decoder_r16.log \ + python3 train_lora_split.py --output_dir out/split_decoder_r16 $COMMON + +# 2) include-head r16 +stage train_head out/split_head_r16.log \ + python3 train_lora_split.py --output_dir out/split_head_r16 --include_head $COMMON + +# 3) base eval (no adapter) — the controlled baseline +stage eval_base out/eval_base.log \ + python3 eval_test_split.py --out out/base_test.jsonl $EVAL_COMMON + +# 4) ft decoder eval (guard: only if adapter saved) +if [ -f out/split_decoder_r16/adapter_model.safetensors ]; then + stage eval_ft_decoder out/eval_ft_decoder.log \ + python3 eval_test_split.py --adapter out/split_decoder_r16 --out out/ft_decoder_test.jsonl $EVAL_COMMON +else + echo "SKIP eval_ft_decoder (no adapter)" >> "$PROG" +fi + +# 5) ft head eval (guard) +if [ -f out/split_head_r16/adapter_model.safetensors ]; then + stage eval_ft_head out/eval_ft_head.log \ + python3 eval_test_split.py --adapter out/split_head_r16 --out out/ft_head_test.jsonl $EVAL_COMMON +else + echo "SKIP eval_ft_head (no adapter)" >> "$PROG" +fi + +echo "PIPELINE_DONE $(date -u +%H:%M:%S)" >> "$PROG" diff --git a/extras/asr-services/providers/gemma4/finetune/smoke_interleave.py b/extras/asr-services/providers/gemma4/finetune/smoke_interleave.py new file mode 100644 index 00000000..c3671af7 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/smoke_interleave.py @@ -0,0 +1,93 @@ +"""Smoke test the multi-audio interleave collator + a forward pass BEFORE a full train. +Checks: processor accepts nested multi-audio, labels mask only the target, forward gives a +finite loss. Also runs a base-model generate on one multi-chunk clip to confirm interleave +inference works. +""" + +import sys + +import torch +from data_interleave import ChunkInterleaveDataset, InterleaveCollator, build_prompt +from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig + +MANIFEST = sys.argv[1] if len(sys.argv) > 1 else "/home/wispr_interleave/manifest.jsonl" +MODEL = "google/gemma-4-E2B-it" + +proc = AutoProcessor.from_pretrained(MODEL) +ds = ChunkInterleaveDataset(MANIFEST, "train") +print( + f"train clips={len(ds)}; chunk counts (first 10): {[len(ds[i]['chunks']) for i in range(min(10,len(ds)))]}" +) +# pick a 1-chunk and a multi-chunk example to stress the collator +idx_multi = next((i for i in range(len(ds)) if len(ds[i]["chunks"]) >= 2), 0) +idx_one = next((i for i in range(len(ds)) if len(ds[i]["chunks"]) == 1), 0) +feats = [ds[idx_one], ds[idx_multi]] +print( + f"batch: 1-chunk='{feats[0]['clip']}' ({len(feats[0]['chunks'])}), " + f"multi='{feats[1]['clip']}' ({len(feats[1]['chunks'])} chunks, app={feats[1]['app']})" +) + +coll = InterleaveCollator(proc) +batch = coll(feats) +print("batch keys:", list(batch.keys())) +print( + "input_ids", + tuple(batch["input_ids"].shape), + "| input_features", + tuple(batch["input_features"].shape) if "input_features" in batch else None, +) +sup = (batch["labels"] != -100).sum(dim=1).tolist() +print("supervised tokens/example (should ≈ target lengths):", sup) +# decode the supervised span of example 0 to confirm it's the target, not the prompt +lab = batch["labels"][0] +sup_ids = batch["input_ids"][0][lab != -100] +print("supervised decode[0]:", proc.decode(sup_ids, skip_special_tokens=True)[:120]) + +bnb = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + llm_int8_skip_modules=[ + "model.audio_tower", + "model.vision_tower", + "model.embed_audio", + "model.embed_vision", + "lm_head", + ], +) +model = AutoModelForMultimodalLM.from_pretrained( + MODEL, + quantization_config=bnb, + dtype=torch.bfloat16, + device_map="auto", + attn_implementation="sdpa", +) +model.config.use_cache = False +batch = {k: v.to(model.device) for k, v in batch.items()} +with torch.no_grad(): + out = model(**batch) +print( + "FORWARD loss:", float(out.loss), "(finite:", torch.isfinite(out.loss).item(), ")" +) + +# interleave generate on the multi-chunk clip (base model) +model.config.use_cache = True +f = ds[idx_multi] +content = [{"type": "text", "text": build_prompt(f["app"])}] +content += [{"type": "audio", "audio": a} for a in f["chunks"]] +inp = proc.apply_chat_template( + [{"role": "user", "content": content}], + tokenize=True, + return_dict=True, + return_tensors="pt", + add_generation_prompt=True, + enable_thinking=False, +).to(model.device) +g = model.generate(**inp, max_new_tokens=200, do_sample=False, use_cache=True) +print( + "GEN (base, interleave):", + proc.decode(g[0][inp["input_ids"].shape[-1] :], skip_special_tokens=True)[:200], +) +print("TARGET:", f["target"][:200]) +print("SMOKE OK") diff --git a/extras/asr-services/providers/gemma4/finetune/split_status.py b/extras/asr-services/providers/gemma4/finetune/split_status.py new file mode 100644 index 00000000..ea3e849f --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/split_status.py @@ -0,0 +1,36 @@ +"""Print a compact training status for a split-LoRA run: the val-loss series (from the +latest checkpoint's trainer_state.json), best eval_loss, current step, and any terminal +marker in the run's .log. Used by the progress monitor (the live stdout log is a single +\\r-joined line, so the structured state file is the reliable source).""" + +import glob +import json +import sys + +out = sys.argv[1].rstrip("/") +cks = sorted(glob.glob(out + "/checkpoint-*"), key=lambda p: int(p.split("-")[-1])) +evals, best, step = [], None, None +if cks: + s = json.load(open(cks[-1] + "/trainer_state.json")) + best = s.get("best_metric") + step = s.get("global_step") + for h in s["log_history"]: + if "eval_loss" in h: + evals.append((round(h["epoch"], 2), round(h["eval_loss"], 4))) +term = "" +try: + raw = open(out + ".log").read() + for m in [ + "DONE ", + "Traceback", + "CUDA out of memory", + "RuntimeError", + "OutOfMemory", + "Killed", + ]: + if m in raw: + term = m.strip() + break +except Exception: + pass +print(f"evals={len(evals)} step={step} best={best} series={evals[-6:]} TERM={term}") diff --git a/extras/asr-services/providers/gemma4/finetune/stream_client_test.py b/extras/asr-services/providers/gemma4/finetune/stream_client_test.py new file mode 100644 index 00000000..5b1bd0a5 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/stream_client_test.py @@ -0,0 +1,71 @@ +"""Test the gemma4 /stream WebSocket endpoint end-to-end. + +Streams a clip's PCM (16-bit LE, 16 kHz mono) in chunks, prints interim +messages as they arrive, sends CloseStream, and prints the final result — +exactly the Chronicle stt_stream contract. +""" + +import argparse +import asyncio +import json + +import numpy as np +import soundfile as sf +import websockets + +SR = 16000 + + +async def main(): + p = argparse.ArgumentParser() + p.add_argument("--url", default="ws://localhost:8767/stream") + p.add_argument( + "--audio", + default="extras/test-audios/audio/Cheap-vs-Expensive-Rug-Tufting-for-the-First-Time-original.wav", + ) + p.add_argument("--seconds", type=float, default=30.0) + p.add_argument("--chunk-ms", type=int, default=500) + p.add_argument("--pace", type=float, default=0.05, help="sleep between chunks (s)") + args = p.parse_args() + + a, sr = sf.read(args.audio, dtype="float32", frames=int(args.seconds * 44100)) + if a.ndim > 1: + a = a.mean(1) + if sr != SR: + import librosa + + a = librosa.resample(a, orig_sr=sr, target_sr=SR) + pcm = (np.clip(a, -1, 1) * 32767).astype("<i2").tobytes() + chunk = int(SR * args.chunk_ms / 1000) * 2 + print( + f"streaming {len(pcm)/SR/2:.1f}s of audio in {args.chunk_ms}ms chunks", + flush=True, + ) + + async with websockets.connect(args.url, max_size=None) as ws: + interims = [] + + async def receiver(): + async for raw in ws: + msg = json.loads(raw) + if msg["type"] == "interim": + interims.append(msg["text"]) + print(f" [interim] {msg['text'][:120]}", flush=True) + elif msg["type"] == "final": + print(f"\n[FINAL] {msg['text']}", flush=True) + print( + f"[FINAL] segments={len(msg.get('segments') or [])}", flush=True + ) + return + + recv_task = asyncio.create_task(receiver()) + for i in range(0, len(pcm), chunk): + await ws.send(pcm[i : i + chunk]) + await asyncio.sleep(args.pace) + await ws.send(json.dumps({"type": "CloseStream"})) + await asyncio.wait_for(recv_task, timeout=180) + print(f"\n{len(interims)} interim message(s) received", flush=True) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/extras/asr-services/providers/gemma4/finetune/sweep_12b_v2.py b/extras/asr-services/providers/gemma4/finetune/sweep_12b_v2.py new file mode 100644 index 00000000..e81666bd --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/sweep_12b_v2.py @@ -0,0 +1,254 @@ +"""Fast 12B ASR config sweep on a fixed CoSHE subset, via the REAL service path. + +Loads google/gemma-4-12B-it ONCE (Gemma4Transcriber, the production transcriber) +and runs a small fixed set of CoSHE-Eval clips through several configs, varying +prompt / windowing / decode strategy. Output is one JSONL per config in the +score_coshe schema so evaluate/score_coshe.py can score them directly. + +Unlike bench_coshe_12b.py (manual channel-strip decode, whole-clip 600s window) +this uses the transcriber's parse_response decode + 30s windowing + stitching — +i.e. exactly what the deployed gemma4-asr service does — so a winning config maps +straight to env knobs (GEMMA4_QUANT, TRANSCRIPTION_PROMPT, MAX_NEW_TOKENS, +BATCH_DURATION_SECONDS, BATCH_THRESHOLD_SECONDS). + +Run inside chronicle-asr-gemma4-smoke (transformers 5.10.1 + bnb), GPU, with +/models = HF cache, /coshe-data = CoSHE parquet, /data/coshe-eval mounted. +""" + +import argparse +import glob +import io +import json +import os +import tempfile +import time +import wave + +import numpy as np +import pyarrow.parquet as pq +import soundfile as sf + +SR = 16000 + +# 11 fixed samples chosen from the existing 500-sample run to span the +# distribution: catastrophic-tail (loops/empties), median, and easy clips. +SAMPLES = [ + "audio_1400.wav", + "audio_299.wav", + "audio_353.wav", + "audio_131.wav", + "audio_564.wav", + "audio_1404.wav", + "audio_243.wav", + "audio_666.wav", + "audio_1543.wav", + "audio_650.wav", + "audio_1938.wav", +] + +# Prompts under test ---------------------------------------------------------- +PROD_PROMPT = ( # current service default (DEFAULT_TRANSCRIPTION_PROMPT) + "Transcribe the following speech segment in its original language into text. " + "Only output the transcription text itself, with no commentary or explanation. " + "When transcribing numbers, write the digits, i.e. write 1.7 and not " + "one point seven, and write 3 instead of three." +) +HF_PROMPT = ( + "Transcribe the following speech segment in its original language.\n\n" + "Follow these specific instructions for formatting the answer:\n" + "* Only output the transcription, with no newlines.\n" + "* When transcribing numbers, write the digits, i.e. write 1.7 and not " + "one point seven, and write 3 instead of three." +) +VERBATIM_PROMPT = ( + "You are an expert transcriptionist. Transcribe the provided Hinglish audio " + "verbatim exactly as it is spoken. Output ONLY the transcribed text. Do not " + "translate it to English, do not summarize, and do not add any conversational " + "pleasantries or commentary." +) +HINGLISH_PROMPT = ( # explicit code-switch steer + "Transcribe the following Hindi-English code-mixed (Hinglish) speech verbatim. " + "Write each word in the language it is spoken. Only output the transcription " + "text itself, with no commentary. When transcribing numbers, write the digits." +) +# prod prompt WITHOUT the "write digits" clause (CoSHE refs spell numbers out, so +# the digit instruction inflates WER) and without code-switch wording. +NODIGIT_PROMPT = ( + "Transcribe the following speech segment in its original language into text. " + "Only output the transcription text itself, with no commentary or explanation." +) + +# Each config: (prompt, threshold_s, window_s, overlap_s, max_new, extra_gen) --- +WHOLE = 100000.0 +CONFIGS = { + # production defaults: 30s windows + overlap, greedy, 512 tokens + "prod": (PROD_PROMPT, 30.0, 30.0, 5.0, 512, {}), + # isolate windowing: production prompt but whole-clip (matches old bench geometry) + "prod_whole": (PROD_PROMPT, WHOLE, 30.0, 5.0, 1024, {}), + # anti-loop safety on the windowed path + "prod_norep3": (PROD_PROMPT, 30.0, 30.0, 5.0, 512, {"no_repeat_ngram_size": 3}), + # prompt variants on the windowed path + "hfprompt": (HF_PROMPT, 30.0, 30.0, 5.0, 512, {}), + "verbatim": (VERBATIM_PROMPT, 30.0, 30.0, 5.0, 512, {}), + "hinglish": (HINGLISH_PROMPT, 30.0, 30.0, 5.0, 512, {}), + # low-temp sampling on the windowed path + "sample_t03": ( + PROD_PROMPT, + 30.0, + 30.0, + 5.0, + 512, + {"do_sample": True, "temperature": 0.3, "top_p": 0.95, "top_k": 64}, + ), + # round 2: best prompts + anti-loop n-gram (tail-safe), windowed greedy + "hf_norep3": (HF_PROMPT, 30.0, 30.0, 5.0, 512, {"no_repeat_ngram_size": 3}), + "hinglish_norep3": ( + HINGLISH_PROMPT, + 30.0, + 30.0, + 5.0, + 512, + {"no_repeat_ngram_size": 3}, + ), + "nodigit_norep3": ( + NODIGIT_PROMPT, + 30.0, + 30.0, + 5.0, + 512, + {"no_repeat_ngram_size": 3}, + ), +} + + +def select_rows(data_dir, names): + """Pull the target rows from the parquet shards.""" + want = set(names) + found = {} + for s in sorted(glob.glob(os.path.join(data_dir, "eval-*.parquet"))): + pf = pq.ParquetFile(s) + for batch in pf.iter_batches(batch_size=64): + for row in batch.to_pylist(): + if row["audio_file_name"] in want: + found[row["audio_file_name"]] = row + if len(found) == len(want): + break + return found + + +def write_wav(raw_bytes, path): + audio, sr = sf.read(io.BytesIO(raw_bytes), dtype="float32") + if audio.ndim > 1: + audio = audio.mean(axis=1) + if sr != SR: + import librosa + + audio = librosa.resample(audio, orig_sr=sr, target_sr=SR) + pcm = np.clip(audio, -1.0, 1.0) + pcm16 = (pcm * 32767.0).astype(np.int16) + with wave.open(path, "wb") as wf: + wf.setnchannels(1) + wf.setsampwidth(2) + wf.setframerate(SR) + wf.writeframes(pcm16.tobytes()) + return len(audio) / SR + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--data_dir", default="/coshe-data") + ap.add_argument("--out_dir", default="/data/coshe-eval/results/12b_sweep") + ap.add_argument("--model", default="google/gemma-4-12B-it") + ap.add_argument("--quant", default="4bit") + ap.add_argument("--configs", default="", help="comma list; default = all") + args = ap.parse_args() + + # Force the transcriber into the desired load mode BEFORE import/instantiation. + os.environ["GEMMA4_QUANT"] = args.quant + os.environ["GEMMA4_MTP"] = "0" # off: clean target distribution, faster load + os.environ["ASR_MODEL"] = args.model + os.environ["TORCH_DTYPE"] = "bfloat16" + + from providers.gemma4.transcriber import Gemma4Transcriber + + os.makedirs(args.out_dir, exist_ok=True) + + print(f"selecting {len(SAMPLES)} samples from {args.data_dir} ...", flush=True) + rows = select_rows(args.data_dir, SAMPLES) + missing = set(SAMPLES) - set(rows) + if missing: + print(f"WARNING missing samples: {missing}", flush=True) + + # Decode all target clips to temp WAVs once. + tmp = tempfile.mkdtemp(prefix="coshe_") + clips = {} # name -> (wav_path, ref_text, duration_s) + for name in SAMPLES: + if name not in rows: + continue + wav_path = os.path.join(tmp, name) + dur = write_wav(rows[name]["audio"]["bytes"], wav_path) + clips[name] = (wav_path, rows[name]["transcription"], dur) + print(f"decoded {len(clips)} clips", flush=True) + + # Load model once. + t = Gemma4Transcriber(model_id=args.model) + t0 = time.time() + t.load_model() + print(f"model loaded quant={args.quant} in {time.time()-t0:.0f}s", flush=True) + + # Monkeypatch _generate to merge per-config extra gen kwargs (anti-loop/sampling). + orig_generate = t._generate + t._extra_gen = {} + + def patched_generate(inputs, **kw): + kw.update(t._extra_gen) + return orig_generate(inputs, **kw) + + t._generate = patched_generate + + selected = args.configs.split(",") if args.configs else list(CONFIGS) + for cfg_name in selected: + prompt, thr, win, ov, max_new, extra = CONFIGS[cfg_name] + t.batch_threshold = thr + t.batch_duration = win + t.batch_overlap = ov + t.max_new_tokens = max_new + t._extra_gen = dict(extra) + out_path = os.path.join(args.out_dir, f"{cfg_name}.jsonl") + print( + f"\n=== CONFIG {cfg_name} (thr={thr} win={win} max_new={max_new} " + f"extra={extra}) -> {out_path}", + flush=True, + ) + with open(out_path, "w") as f: + for name in SAMPLES: + if name not in clips: + continue + wav_path, ref, dur = clips[name] + ts = time.time() + try: + res = t.transcribe(wav_path, prompt_override=prompt) + hyp = res.text + except Exception as e: + hyp = "" + print(f" ERR {name}: {e}", flush=True) + dt = time.time() - ts + rec = { + "audio_file_name": name, + "transcription": ref, + "hyp": hyp, + "asr_seconds": round(dt, 2), + "duration_s": round(dur, 2), + } + f.write(json.dumps(rec, ensure_ascii=False) + "\n") + f.flush() + print( + f" [{name}] dur={dur:.0f}s gen={dt:.1f}s " + f"hyp[:70]={hyp[:70]!r}", + flush=True, + ) + print("\nDONE", flush=True) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/gemma4/finetune/test_cache57.py b/extras/asr-services/providers/gemma4/finetune/test_cache57.py new file mode 100644 index 00000000..3b316bfc --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/test_cache57.py @@ -0,0 +1,90 @@ +"""On the A100/transformers-5.7: is the use_cache bug still present? + +A correct implementation gives IDENTICAL logits/loss for a single full forward +regardless of use_cache. We compare use_cache True vs False on one batch with the +current adapter. If equal -> cache is fine -> fast cached generation/eval. If they +differ -> bug persists -> use_cache=False for eval. +""" + +import argparse +import time + +import torch +from data import DEFAULT_PROMPT, CosheParquetDataset, Gemma4AudioCollator +from peft import PeftModel +from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig + +p = argparse.ArgumentParser() +p.add_argument("--adapter", default="/home/gemma4ft/out/full/checkpoint-994") +p.add_argument("--parquet_glob", default="/home/coshe-data/data/eval-*.parquet") +p.add_argument("--cache_path", default="/home/gemma4ft/out/coshe_full_cache.pkl") +args = p.parse_args() + +proc = AutoProcessor.from_pretrained("google/gemma-4-E2B-it") +bnb = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + llm_int8_skip_modules=[ + "model.audio_tower", + "model.vision_tower", + "model.embed_audio", + "model.embed_vision", + "lm_head", + ], +) +model = AutoModelForMultimodalLM.from_pretrained( + "google/gemma-4-E2B-it", + quantization_config=bnb, + dtype=torch.bfloat16, + device_map="auto", + attn_implementation="eager", +) +model = PeftModel.from_pretrained(model, args.adapter) +model.eval() + +ds = CosheParquetDataset( + args.parquet_glob, + max_seconds=30.0, + target_max_chars=0, + cache_path=args.cache_path, + limit=8, +) +col = Gemma4AudioCollator(proc) +b = col([ds[0]]) +b = {k: v.to(model.device) for k, v in b.items()} + +print("=== single-forward loss: use_cache True vs False ===", flush=True) +for uc in [False, True]: + with torch.inference_mode(): + loss = float(model(**b, use_cache=uc).loss) + print(f" use_cache={uc}: loss={loss:.4f}", flush=True) + +print("\n=== generation speed: cache vs no-cache (4 clips) ===", flush=True) +proc.tokenizer.padding_side = "left" +texts, audios = [], [] +for it in [ds[i] for i in range(4)]: + msgs = [ + { + "role": "user", + "content": [ + {"type": "text", "text": DEFAULT_PROMPT}, + {"type": "audio", "audio": it["audio"]}, + ], + } + ] + texts.append( + proc.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True) + ) + audios.append(it["audio"]) +inp = proc(text=texts, audio=audios, return_tensors="pt", padding=True).to(model.device) +for uc in [True, False]: + t = time.time() + with torch.inference_mode(): + out = model.generate(**inp, max_new_tokens=200, do_sample=False, use_cache=uc) + dt = time.time() - t + txt = proc.decode(out[0][inp["input_ids"].shape[-1] :], skip_special_tokens=True)[ + :120 + ] + print(f" use_cache={uc}: {dt:.1f}s for 4 clips | sample: {txt!r}", flush=True) diff --git a/extras/asr-services/providers/gemma4/finetune/test_longaudio.py b/extras/asr-services/providers/gemma4/finetune/test_longaudio.py new file mode 100644 index 00000000..f4ef0ece --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/test_longaudio.py @@ -0,0 +1,64 @@ +"""Does gemma4 accept a full ~55s clip in one forward (vs the 30s window)? + +Prints token/feature counts and runs a forward at 30s and ~60s to see if it +errors or how VRAM scales. Decides whether full 1-min clips are usable for +training without chunking. +""" + +import torch +from data import DEFAULT_PROMPT, _load_wav_16k_mono +from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig + +WAV = "/data/coshe-eval/sample7/audio/audio_13.wav" + +proc = AutoProcessor.from_pretrained("google/gemma-4-E2B-it") +bnb = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + llm_int8_skip_modules=[ + "model.audio_tower", + "model.vision_tower", + "model.embed_audio", + "model.embed_vision", + "lm_head", + ], +) +model = AutoModelForMultimodalLM.from_pretrained( + "google/gemma-4-E2B-it", + quantization_config=bnb, + dtype=torch.bfloat16, + device_map="auto", + attn_implementation="eager", +) +model.config.use_cache = False +model.eval() + +for secs in [30.0, 60.0]: + audio = _load_wav_16k_mono(WAV, secs) + msgs = [ + { + "role": "user", + "content": [ + {"type": "text", "text": DEFAULT_PROMPT}, + {"type": "audio", "audio": audio}, + ], + } + ] + text = proc.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True) + batch = proc(text=[text], audio=[audio], return_tensors="pt").to(model.device) + torch.cuda.reset_peak_memory_stats() + try: + with torch.inference_mode(): + out = model(**batch) + peak = torch.cuda.max_memory_allocated() / 1e9 + print( + f"audio={secs:.0f}s samples={len(audio)} " + f"input_ids={batch['input_ids'].shape[-1]} " + f"audio_frames={batch['input_features'].shape[1]} " + f"forward OK logits={tuple(out.logits.shape)} peakVRAM={peak:.1f}GB", + flush=True, + ) + except Exception as e: + print(f"audio={secs:.0f}s FAILED: {type(e).__name__}: {e}", flush=True) diff --git a/extras/asr-services/providers/gemma4/finetune/test_patch_loss.py b/extras/asr-services/providers/gemma4/finetune/test_patch_loss.py new file mode 100644 index 00000000..3867da7f --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/test_patch_loss.py @@ -0,0 +1,131 @@ +"""Instrument the gemma4 cache-mask logic to fix the use_cache prefill bug. + +Runs a single full forward (use_cache=True) on a memorized sample7 clip with the +s7-full adapter and prints what create_causal_mask_mapping sees (cache type, +get_seq_length, is_initialized, chosen is_first_iteration) + the resulting loss. +Goal: find the right prefill test so loss drops 4.49 -> ~0.1 with the cache on. +""" + +import torch +import transformers.models.gemma4.modeling_gemma4 as g4 +from data import DEFAULT_PROMPT, CosheSample7Dataset, Gemma4AudioCollator +from peft import PeftModel +from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig + +_orig = g4.create_causal_mask_mapping +_seen = {} + + +def _instrumented( + config, + inputs_embeds, + attention_mask, + past_key_values, + position_ids, + mm_token_type_ids=None, + pixel_values=None, + is_training=False, + is_first_iteration=None, + **kwargs, +): + pkv = past_key_values + info = { + "pkv_type": type(pkv).__name__ if pkv is not None else None, + "is_initialized": ( + getattr(pkv, "is_initialized", "n/a") if pkv is not None else None + ), + "get_seq_length": (pkv.get_seq_length() if pkv is not None else None), + "mm_token_type_ids_none": mm_token_type_ids is None, + "q_len": inputs_embeds.shape[1], + "pixel_values_none": pixel_values is None, + } + if not _seen: + _seen.update(info) + print("MASK CALL INFO:", info, flush=True) + return _orig( + config, + inputs_embeds, + attention_mask, + past_key_values, + position_ids, + mm_token_type_ids=mm_token_type_ids, + pixel_values=pixel_values, + is_training=is_training, + is_first_iteration=is_first_iteration, + **kwargs, + ) + + +g4.create_causal_mask_mapping = _instrumented + +proc = AutoProcessor.from_pretrained("google/gemma-4-E2B-it") +bnb = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + llm_int8_skip_modules=[ + "model.audio_tower", + "model.vision_tower", + "model.embed_audio", + "model.embed_vision", + "lm_head", + ], +) +model = AutoModelForMultimodalLM.from_pretrained( + "google/gemma-4-E2B-it", + quantization_config=bnb, + dtype=torch.bfloat16, + device_map="auto", + attn_implementation="eager", +) +model = PeftModel.from_pretrained(model, "/train/out/s7-full") +model.eval() + +ds = CosheSample7Dataset( + "/data/coshe-eval/sample7", max_seconds=30.0, target_max_chars=0 +) +col = Gemma4AudioCollator(proc) +b = col([ds[0]]) +b = {k: v.to(model.device) for k, v in b.items()} + +for uc in [False, True]: + model.config.use_cache = uc + _seen.clear() + with torch.inference_mode(): + loss = float(model(**b, use_cache=uc).loss) + print(f"use_cache={uc}: loss={loss:.4f}", flush=True) + + +# now try forcing is_first_iteration=True always (prefill-style) to confirm theory +def _force_true( + config, + inputs_embeds, + attention_mask, + past_key_values, + position_ids, + mm_token_type_ids=None, + pixel_values=None, + is_training=False, + is_first_iteration=None, + **kwargs, +): + return _orig( + config, + inputs_embeds, + attention_mask, + past_key_values, + position_ids, + mm_token_type_ids=mm_token_type_ids, + pixel_values=pixel_values, + is_training=is_training, + is_first_iteration=True, + **kwargs, + ) + + +g4.create_causal_mask_mapping = _force_true +model.config.use_cache = True +with torch.inference_mode(): + loss = float(model(**b, use_cache=True).loss) +print(f"use_cache=True + force is_first_iteration=True: loss={loss:.4f}", flush=True) diff --git a/extras/asr-services/providers/gemma4/finetune/train.py b/extras/asr-services/providers/gemma4/finetune/train.py new file mode 100644 index 00000000..a21b8dc8 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/train.py @@ -0,0 +1,203 @@ +"""QLoRA fine-tuning for Gemma 4 E*B on CoSHE-Eval sample7 (overfit smoke test). + +4-bit NF4 base (double quant, bf16 compute) + LoRA on the text decoder only +(audio tower + multimodal embedders frozen). Plain HF Trainer with the custom +multimodal collator in data.py. + +Usage (inside the gemma4 image venv): + python train.py \ + --model google/gemma-4-E2B-it \ + --data_dir /data/coshe-eval/sample7 \ + --output_dir /train/out/e2b-overfit \ + --epochs 40 --lr 2e-4 --batch_size 1 --grad_accum 1 +""" + +import argparse + +import torch +from data import CosheParquetDataset, CosheSample7Dataset, Gemma4AudioCollator +from peft import LoraConfig, get_peft_model +from transformers import ( + AutoModelForMultimodalLM, + AutoProcessor, + BitsAndBytesConfig, + Trainer, + TrainingArguments, +) + +# LoRA scoped to the text decoder (language_model) attention + MLP projections. +# Regex (PEFT re.fullmatch) keeps it off the identically-named audio-tower linears. +# NOTE: an experiment adding lm_head + embed_tokens here (rank 64) made free-running +# generation DEGENERATE into "Speaker 1: Speaker 1:..." loops — the blocker is target +# grounding, not output capacity, so we keep the lean attention+MLP set. +LORA_TARGETS = ( + r".*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" +) +# With --include_head, also adapt lm_head + input embeddings for more memorization +# capacity (safe now that generation uses use_cache=False; the earlier "degenerate" +# result with these was a cache-bug artifact, not a real failure). +LORA_TARGETS_HEAD = ( + r"(.*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" + r"|.*language_model\.embed_tokens|lm_head)" +) + + +def parse_args(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="google/gemma-4-E2B-it") + p.add_argument("--data_dir", default="/data/coshe-eval/sample7") + p.add_argument( + "--parquet_glob", + default="", + help="if set, use full CoSHE parquet instead of sample7", + ) + p.add_argument("--limit", type=int, default=0, help="cap #samples (parquet only)") + p.add_argument( + "--cache_path", default="", help="pickle cache for decoded parquet audio" + ) + p.add_argument("--output_dir", default="/train/out/e2b-overfit") + p.add_argument("--epochs", type=float, default=40.0) + p.add_argument("--lr", type=float, default=2e-4) + p.add_argument("--batch_size", type=int, default=1) + p.add_argument("--grad_accum", type=int, default=1) + p.add_argument("--max_seconds", type=float, default=30.0) + p.add_argument("--target_max_chars", type=int, default=0) + p.add_argument("--lora_r", type=int, default=16) + p.add_argument("--lora_alpha", type=int, default=32) + p.add_argument("--lora_dropout", type=float, default=0.0) + p.add_argument( + "--optim", + default="adamw_torch", + help="e.g. adamw_8bit / paged_adamw_8bit to save VRAM", + ) + p.add_argument( + "--include_head", action="store_true", help="also LoRA lm_head + embed_tokens" + ) + p.add_argument( + "--resume", + action="store_true", + help="resume from latest checkpoint in output_dir", + ) + return p.parse_args() + + +def main(): + args = parse_args() + torch.manual_seed(0) + + print(f"Loading processor + 4-bit model: {args.model}", flush=True) + processor = AutoProcessor.from_pretrained(args.model) + + bnb = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + # Keep the multimodal towers / embedders / lm_head in bf16. The audio + # tower's ClippableLinear calls torch.finfo(weight.dtype) for gradient + # clipping, which breaks on 4-bit (uint8-stored) weights. We only train + # LoRA on the text decoder anyway. + llm_int8_skip_modules=[ + "model.audio_tower", + "model.vision_tower", + "model.embed_audio", + "model.embed_vision", + "lm_head", + ], + ) + model = AutoModelForMultimodalLM.from_pretrained( + args.model, + quantization_config=bnb, + dtype=torch.bfloat16, + device_map="auto", + attn_implementation="eager", + ) + model.config.use_cache = False + + # Manual k-bit prep (NOT peft's prepare_model_for_kbit_training): that helper + # upcasts every non-4bit param to fp32, which makes the text embeddings fp32 + # while the (quant-skipped, bf16) audio tower stays bf16 -> Gemma4's + # multimodal masked_scatter merge then errors on the dtype mismatch. Keeping + # everything bf16 avoids that and matches the recipe's bf16-throughout advice. + for p in model.parameters(): + p.requires_grad = False + model.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={"use_reentrant": False} + ) + model.enable_input_require_grads() + + lora = LoraConfig( + r=args.lora_r, + lora_alpha=args.lora_alpha, + lora_dropout=args.lora_dropout, + bias="none", + task_type="CAUSAL_LM", + target_modules=LORA_TARGETS_HEAD if args.include_head else LORA_TARGETS, + ) + model = get_peft_model(model, lora) + model.print_trainable_parameters() + + if args.parquet_glob: + dataset = CosheParquetDataset( + args.parquet_glob, + max_seconds=args.max_seconds, + target_max_chars=args.target_max_chars, + limit=args.limit, + cache_path=args.cache_path, + ) + else: + dataset = CosheSample7Dataset( + args.data_dir, + max_seconds=args.max_seconds, + target_max_chars=args.target_max_chars, + ) + print(f"Dataset: {len(dataset)} samples", flush=True) + collator = Gemma4AudioCollator(processor) + + targs = TrainingArguments( + output_dir=args.output_dir, + per_device_train_batch_size=args.batch_size, + gradient_accumulation_steps=args.grad_accum, + num_train_epochs=args.epochs, + learning_rate=args.lr, + lr_scheduler_type="constant", + warmup_steps=0, + bf16=True, + fp16=False, + logging_steps=1, + save_strategy="epoch", + save_total_limit=2, + optim=args.optim, + report_to=[], + remove_unused_columns=False, + dataloader_num_workers=0, + gradient_checkpointing=False, # already enabled via prepare_model_for_kbit_training + ) + + trainer = Trainer( + model=model, + args=targs, + train_dataset=dataset, + data_collator=collator, + ) + + # Resume from the latest checkpoint in output_dir if present and --resume set. + # /home is persisted on Jarvis Labs but the process dies on VM pause/restart, so + # resuming from the last per-epoch checkpoint avoids wasting progress. + import os as _os + + resume = False + if args.resume and _os.path.isdir(args.output_dir): + ckpts = [d for d in _os.listdir(args.output_dir) if d.startswith("checkpoint-")] + resume = len(ckpts) > 0 + print(f"Starting training... (resume={resume})", flush=True) + trainer.train(resume_from_checkpoint=resume) + + print(f"Saving adapter to {args.output_dir}", flush=True) + trainer.save_model(args.output_dir) + processor.save_pretrained(args.output_dir) + print("DONE", flush=True) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/gemma4/finetune/train_interleave.py b/extras/asr-services/providers/gemma4/finetune/train_interleave.py new file mode 100644 index 00000000..291bd39d --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/train_interleave.py @@ -0,0 +1,150 @@ +"""LoRA finetune Gemma4-E2B on the interleaved audio->pasted_text (app-aware formatting) task. +One example = one clip (N<=28s audio chunks in a single prompt) -> full formatted pasted_text. +Decoder-only LoRA, val-loss early stop. Mirrors train_wispr.py but uses the interleave +dataset/collator (multi-audio per example, formatted target). +""" + +import argparse +import json +import os + +import torch +from data_interleave import ChunkInterleaveDataset, InterleaveCollator +from peft import LoraConfig, get_peft_model +from transformers import ( + AutoModelForMultimodalLM, + AutoProcessor, + BitsAndBytesConfig, + EarlyStoppingCallback, + Trainer, + TrainingArguments, +) + +LORA_TARGETS_HEAD = ( + r"(.*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" + r"|.*language_model\.embed_tokens|lm_head)" +) +LORA_TARGETS = ( + r".*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" +) + + +def parse_args(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="google/gemma-4-E2B-it") + p.add_argument("--manifest", default="/home/wispr_interleave/manifest.jsonl") + p.add_argument("--output_dir", required=True) + p.add_argument("--epochs", type=float, default=40.0) + p.add_argument("--lr", type=float, default=1e-5) + p.add_argument("--batch_size", type=int, default=1) + p.add_argument("--grad_accum", type=int, default=4) + p.add_argument("--eval_batch_size", type=int, default=1) + p.add_argument("--lora_r", type=int, default=16) + p.add_argument("--lora_alpha", type=int, default=32) + p.add_argument("--include_head", action="store_true") + p.add_argument("--optim", default="adamw_8bit") + p.add_argument("--patience", type=int, default=5) + p.add_argument("--attn", default="sdpa") + return p.parse_args() + + +def main(): + args = parse_args() + torch.manual_seed(0) + proc = AutoProcessor.from_pretrained(args.model) + bnb = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + llm_int8_skip_modules=[ + "model.audio_tower", + "model.vision_tower", + "model.embed_audio", + "model.embed_vision", + "lm_head", + ], + ) + model = AutoModelForMultimodalLM.from_pretrained( + args.model, + quantization_config=bnb, + dtype=torch.bfloat16, + device_map="auto", + attn_implementation=args.attn, + ) + model.config.use_cache = False + for pm in model.parameters(): + pm.requires_grad = False + model.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={"use_reentrant": False} + ) + model.enable_input_require_grads() + model = get_peft_model( + model, + LoraConfig( + r=args.lora_r, + lora_alpha=args.lora_alpha, + lora_dropout=0.0, + bias="none", + task_type="CAUSAL_LM", + target_modules=LORA_TARGETS_HEAD if args.include_head else LORA_TARGETS, + ), + ) + model.print_trainable_parameters() + + train_ds = ChunkInterleaveDataset(args.manifest, "train") + val_ds = ChunkInterleaveDataset(args.manifest, "val") + print(f"train clips={len(train_ds)} val clips={len(val_ds)}", flush=True) + collator = InterleaveCollator(proc) + + targs = TrainingArguments( + output_dir=args.output_dir, + per_device_train_batch_size=args.batch_size, + gradient_accumulation_steps=args.grad_accum, + per_device_eval_batch_size=args.eval_batch_size, + num_train_epochs=args.epochs, + learning_rate=args.lr, + lr_scheduler_type="constant", + warmup_steps=0, + bf16=True, + fp16=False, + logging_steps=10, + eval_strategy="epoch", + save_strategy="epoch", + save_total_limit=2, + load_best_model_at_end=True, + metric_for_best_model="eval_loss", + greater_is_better=False, + optim=args.optim, + report_to=[], + remove_unused_columns=False, + dataloader_num_workers=0, + gradient_checkpointing=False, + seed=0, + data_seed=0, + ) + trainer = Trainer( + model=model, + args=targs, + train_dataset=train_ds, + eval_dataset=val_ds, + data_collator=collator, + callbacks=[EarlyStoppingCallback(early_stopping_patience=args.patience)], + ) + print("Starting training...", flush=True) + trainer.train() + trainer.save_model(args.output_dir) + proc.save_pretrained(args.output_dir) + json.dump( + trainer.state.log_history, + open(os.path.join(args.output_dir, "log_history.json"), "w"), + ) + print( + f"DONE best_checkpoint={trainer.state.best_model_checkpoint} " + f"best_eval_loss={trainer.state.best_metric}", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/gemma4/finetune/train_lora_split.py b/extras/asr-services/providers/gemma4/finetune/train_lora_split.py new file mode 100644 index 00000000..2a19bdf4 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/train_lora_split.py @@ -0,0 +1,199 @@ +"""Split-aware LoRA finetune of Gemma4-E4B on CoSHE — a *generalization* experiment. + +Unlike `train_until_wer.py` (which memorizes the full set to <2% WER), this trains on a +random 20% slice and selects the checkpoint with the lowest **validation loss** on a +held-out 10% slice. The remaining 70% (the test split) is never touched here — it is +scored separately for WER against the base-model baseline. + +Split comes from a JSON file ({"train":[names], "val":[names], "test":[names]}) so the +exact same partition is used here (training) and locally (scoring). Items are matched to +the decoded-audio cache by `name` (== audio_file_name). + +Strategy (per the user's spec): very low constant LR, LoRA adapter, early-stop when the +validation loss stops improving (HF EarlyStoppingCallback + load_best_model_at_end on +eval_loss). Two configs are supported via --include_head: + * decoder-only : LoRA on language_model q/k/v/o/gate/up/down (audio tower frozen) + * include_head : the above + LoRA on embed_tokens and lm_head +""" + +import argparse +import json +import os + +import torch +from data import CosheParquetDataset, Gemma4AudioCollator +from peft import LoraConfig, get_peft_model +from transformers import ( + AutoModelForMultimodalLM, + AutoProcessor, + BitsAndBytesConfig, + EarlyStoppingCallback, + Trainer, + TrainingArguments, +) +from window_target import apply_window_truncation + +LORA_TARGETS_HEAD = ( + r"(.*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" + r"|.*language_model\.embed_tokens|lm_head)" +) +LORA_TARGETS = ( + r".*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" +) + + +class ListDataset(torch.utils.data.Dataset): + def __init__(self, items): + self.items = items + + def __len__(self): + return len(self.items) + + def __getitem__(self, idx): + return self.items[idx] + + +def parse_args(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="google/gemma-4-E4B-it") + p.add_argument("--parquet_glob", default="/home/coshe-data/data/eval-*.parquet") + p.add_argument("--cache_path", default="/home/gemma4ft/out/coshe_full_cache.pkl") + p.add_argument("--split_file", default="/home/gemma4ft/split_20_10_70.json") + p.add_argument("--output_dir", required=True) + p.add_argument("--epochs", type=float, default=40.0) + p.add_argument("--lr", type=float, default=1e-5) + p.add_argument("--batch_size", type=int, default=2) + p.add_argument("--grad_accum", type=int, default=2) + p.add_argument("--eval_batch_size", type=int, default=4) + p.add_argument("--lora_r", type=int, default=16) + p.add_argument("--lora_alpha", type=int, default=32) + p.add_argument("--include_head", action="store_true") + p.add_argument("--optim", default="adamw_8bit") + p.add_argument("--patience", type=int, default=4) + p.add_argument( + "--window_seconds", + type=float, + default=30.0, + help="proportionally truncate target to this audio window (0=full)", + ) + p.add_argument("--durations", default="/home/gemma4ft/durations.json") + return p.parse_args() + + +def main(): + args = parse_args() + torch.manual_seed(0) + proc = AutoProcessor.from_pretrained(args.model) + bnb = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + llm_int8_skip_modules=[ + "model.audio_tower", + "model.vision_tower", + "model.embed_audio", + "model.embed_vision", + "lm_head", + ], + ) + model = AutoModelForMultimodalLM.from_pretrained( + args.model, + quantization_config=bnb, + dtype=torch.bfloat16, + device_map="auto", + attn_implementation="eager", + ) + model.config.use_cache = False + for pm in model.parameters(): + pm.requires_grad = False + model.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={"use_reentrant": False} + ) + model.enable_input_require_grads() + model = get_peft_model( + model, + LoraConfig( + r=args.lora_r, + lora_alpha=args.lora_alpha, + lora_dropout=0.0, + bias="none", + task_type="CAUSAL_LM", + target_modules=LORA_TARGETS_HEAD if args.include_head else LORA_TARGETS, + ), + ) + model.print_trainable_parameters() + + split = json.load(open(args.split_file)) + ds = CosheParquetDataset( + args.parquet_glob, + max_seconds=30.0, + target_max_chars=0, + cache_path=args.cache_path, + ) + apply_window_truncation(list(ds), args.durations, args.window_seconds) + by = {it["name"]: it for it in ds} + train_items = [by[n] for n in split["train"] if n in by] + val_items = [by[n] for n in split["val"] if n in by] + missing_tr = len(split["train"]) - len(train_items) + missing_va = len(split["val"]) - len(val_items) + print( + f"train={len(train_items)} (missing {missing_tr}) " + f"val={len(val_items)} (missing {missing_va}) " + f"test(held out)={len(split['test'])}", + flush=True, + ) + collator = Gemma4AudioCollator(proc) + + targs = TrainingArguments( + output_dir=args.output_dir, + per_device_train_batch_size=args.batch_size, + gradient_accumulation_steps=args.grad_accum, + per_device_eval_batch_size=args.eval_batch_size, + num_train_epochs=args.epochs, + learning_rate=args.lr, + lr_scheduler_type="constant", + warmup_steps=0, + bf16=True, + fp16=False, + logging_steps=10, + eval_strategy="epoch", + save_strategy="epoch", + save_total_limit=2, + load_best_model_at_end=True, + metric_for_best_model="eval_loss", + greater_is_better=False, + optim=args.optim, + report_to=[], + remove_unused_columns=False, + dataloader_num_workers=0, + gradient_checkpointing=False, + seed=0, + data_seed=0, + ) + + trainer = Trainer( + model=model, + args=targs, + train_dataset=ListDataset(train_items), + eval_dataset=ListDataset(val_items), + data_collator=collator, + callbacks=[EarlyStoppingCallback(early_stopping_patience=args.patience)], + ) + + print("Starting training...", flush=True) + trainer.train() + trainer.save_model(args.output_dir) + proc.save_pretrained(args.output_dir) + json.dump( + trainer.state.log_history, + open(os.path.join(args.output_dir, "log_history.json"), "w"), + ) + # surface the chosen (best) checkpoint + its val loss for the writeup + best = getattr(trainer.state, "best_model_checkpoint", None) + best_loss = getattr(trainer.state, "best_metric", None) + print(f"DONE best_checkpoint={best} best_eval_loss={best_loss}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/gemma4/finetune/train_lora_windowed.py b/extras/asr-services/providers/gemma4/finetune/train_lora_windowed.py new file mode 100644 index 00000000..471dd8c3 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/train_lora_windowed.py @@ -0,0 +1,160 @@ +"""LoRA finetune of Gemma4-E2B on the HONEST windowed CoSHE set (20% train / 10% val), +val-loss early stopping. Successor to train_lora_split.py: instead of the "first-30s + +proportional char-truncation" hack, it trains on real <=30s (audio, GT-text) windows cut +at forced-alignment word timings (build_windowed_dataset.py). The held-out 70% test clips +are never touched here; they are scored separately via the windowed-stitch path. + +decoder-only vs --include_head as before. Same low constant LR + LoRA recipe. +""" + +import argparse +import json +import os + +import torch +from data import Gemma4AudioCollator +from data_windowed import PLAIN_PROMPT, WindowedManifestDataset +from peft import LoraConfig, get_peft_model +from transformers import ( + AutoModelForMultimodalLM, + AutoProcessor, + BitsAndBytesConfig, + EarlyStoppingCallback, + Trainer, + TrainingArguments, +) + +LORA_TARGETS_HEAD = ( + r"(.*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" + r"|.*language_model\.embed_tokens|lm_head)" +) +LORA_TARGETS = ( + r".*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" +) + + +def parse_args(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="google/gemma-4-E2B-it") + p.add_argument("--manifest", default="/home/coshe_windowed/manifest.jsonl") + p.add_argument("--output_dir", required=True) + p.add_argument("--epochs", type=float, default=40.0) + p.add_argument("--lr", type=float, default=1e-5) + p.add_argument("--batch_size", type=int, default=2) + p.add_argument("--grad_accum", type=int, default=2) + p.add_argument("--eval_batch_size", type=int, default=4) + p.add_argument("--lora_r", type=int, default=16) + p.add_argument("--lora_alpha", type=int, default=32) + p.add_argument("--include_head", action="store_true") + p.add_argument("--optim", default="adamw_8bit") + p.add_argument("--patience", type=int, default=4) + p.add_argument( + "--attn", + default="sdpa", + help="sdpa matches prod/eval (gemma4 default); eager was the old buggy-at-eval choice", + ) + return p.parse_args() + + +def main(): + args = parse_args() + torch.manual_seed(0) + proc = AutoProcessor.from_pretrained(args.model) + bnb = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + llm_int8_skip_modules=[ + "model.audio_tower", + "model.vision_tower", + "model.embed_audio", + "model.embed_vision", + "lm_head", + ], + ) + model = AutoModelForMultimodalLM.from_pretrained( + args.model, + quantization_config=bnb, + dtype=torch.bfloat16, + device_map="auto", + attn_implementation=args.attn, + ) + model.config.use_cache = False + for pm in model.parameters(): + pm.requires_grad = False + model.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={"use_reentrant": False} + ) + model.enable_input_require_grads() + model = get_peft_model( + model, + LoraConfig( + r=args.lora_r, + lora_alpha=args.lora_alpha, + lora_dropout=0.0, + bias="none", + task_type="CAUSAL_LM", + target_modules=LORA_TARGETS_HEAD if args.include_head else LORA_TARGETS, + ), + ) + model.print_trainable_parameters() + + train_ds = WindowedManifestDataset(args.manifest, "train") + val_ds = WindowedManifestDataset(args.manifest, "val") + print(f"train windows={len(train_ds)} val windows={len(val_ds)}", flush=True) + collator = Gemma4AudioCollator(proc, prompt=PLAIN_PROMPT) + + targs = TrainingArguments( + output_dir=args.output_dir, + per_device_train_batch_size=args.batch_size, + gradient_accumulation_steps=args.grad_accum, + per_device_eval_batch_size=args.eval_batch_size, + num_train_epochs=args.epochs, + learning_rate=args.lr, + lr_scheduler_type="constant", + warmup_steps=0, + bf16=True, + fp16=False, + logging_steps=10, + eval_strategy="epoch", + save_strategy="epoch", + save_total_limit=2, + load_best_model_at_end=True, + metric_for_best_model="eval_loss", + greater_is_better=False, + optim=args.optim, + report_to=[], + remove_unused_columns=False, + dataloader_num_workers=0, + gradient_checkpointing=False, + seed=0, + data_seed=0, + ) + + trainer = Trainer( + model=model, + args=targs, + train_dataset=train_ds, + eval_dataset=val_ds, + data_collator=collator, + callbacks=[EarlyStoppingCallback(early_stopping_patience=args.patience)], + ) + + print("Starting training...", flush=True) + trainer.train() + trainer.save_model(args.output_dir) + proc.save_pretrained(args.output_dir) + json.dump( + trainer.state.log_history, + open(os.path.join(args.output_dir, "log_history.json"), "w"), + ) + print( + f"DONE best_checkpoint={trainer.state.best_model_checkpoint} " + f"best_eval_loss={trainer.state.best_metric}", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/gemma4/finetune/train_until_wer.py b/extras/asr-services/providers/gemma4/finetune/train_until_wer.py new file mode 100644 index 00000000..6f4b93c1 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/train_until_wer.py @@ -0,0 +1,302 @@ +"""Overfit CoSHE, evaluating WER every N epochs and stopping when corpus WER < target. + +transformers 5.7 has the use_cache bug fixed, so eval uses fast cached generation. +Every `--eval_every` epochs: eval a fast subset; if subset WER < target, eval the +FULL set and stop only if that's also < target. Saves the adapter on stop. +""" + +import argparse +import time + +import jiwer +import torch +from data import DEFAULT_PROMPT, CosheParquetDataset, Gemma4AudioCollator +from peft import LoraConfig, PeftModel, get_peft_model +from transformers import ( + AutoModelForMultimodalLM, + AutoProcessor, + BitsAndBytesConfig, + Trainer, + TrainerCallback, + TrainingArguments, +) + +LORA_TARGETS_HEAD = ( + r"(.*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" + r"|.*language_model\.embed_tokens|lm_head)" +) +LORA_TARGETS = ( + r".*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" +) +_WER_NORM = jiwer.Compose( + [ + jiwer.ToLowerCase(), + jiwer.RemovePunctuation(), + jiwer.RemoveMultipleSpaces(), + jiwer.Strip(), + jiwer.ReduceToListOfListOfWords(), + ] +) + + +def corpus_wer(refs, hyps): + return jiwer.wer( + refs, hyps, reference_transform=_WER_NORM, hypothesis_transform=_WER_NORM + ) + + +@torch.inference_mode() +def eval_wer(model, proc, items, batch_size, max_new_tokens): + model.eval() + prev_uc = model.config.use_cache + model.config.use_cache = True + proc.tokenizer.padding_side = "left" + refs, hyps = [], [] + for s in range(0, len(items), batch_size): + batch = items[s : s + batch_size] + texts, audios = [], [] + for it in batch: + msgs = [ + { + "role": "user", + "content": [ + {"type": "text", "text": DEFAULT_PROMPT}, + {"type": "audio", "audio": it["audio"]}, + ], + } + ] + texts.append( + proc.apply_chat_template( + msgs, tokenize=False, add_generation_prompt=True + ) + ) + audios.append(it["audio"]) + inp = proc(text=texts, audio=audios, return_tensors="pt", padding=True).to( + model.device + ) + out = model.generate( + **inp, max_new_tokens=max_new_tokens, do_sample=False, use_cache=True + ) + for i, it in enumerate(batch): + hyps.append( + proc.decode( + out[i][inp["input_ids"].shape[-1] :], skip_special_tokens=True + ).strip() + ) + refs.append(it["target"]) + model.config.use_cache = prev_uc + model.train() + return corpus_wer(refs, hyps), refs, hyps + + +class WERStopCallback(TrainerCallback): + def __init__( + self, + proc, + subset, + full, + every, + target, + batch_size, + max_new_tokens, + out_dir, + loss_gate=0.15, + ): + self.proc, self.subset, self.full = proc, subset, full + self.every, self.target, self.bs, self.mnt = ( + every, + target, + batch_size, + max_new_tokens, + ) + self.out_dir = out_dir + self.loss_gate = loss_gate + + def _recent_loss(self, state): + # Mean of the last ~30 logged losses — the single last value is too noisy + # (per-batch variance 0.1–0.27) and was erratically skipping evals. + vals = [r["loss"] for r in state.log_history if "loss" in r] + if not vals: + return None + tail = vals[-30:] + return sum(tail) / len(tail) + + def on_epoch_end(self, args, state, control, model=None, **kwargs): + ep = int(round(state.epoch)) + if ep == 0 or ep % self.every != 0: + return control + # Free-running generation can't be good while teacher-forced loss is high + # (sample7: WER only collapsed once loss -> ~0). Skip the expensive WER eval + # until loss drops below the gate, to avoid wasting ~10min/eval early on. + rl = self._recent_loss(state) + if rl is not None and rl > self.loss_gate: + print( + f"[epoch {ep}] loss={rl:.3f} > {self.loss_gate}; skipping WER eval", + flush=True, + ) + return control + t = time.time() + wer, _, _ = eval_wer(model, self.proc, self.subset, self.bs, self.mnt) + print( + f"[epoch {ep}] subset({len(self.subset)}) corpus WER = {wer*100:.2f}% ({time.time()-t:.0f}s)", + flush=True, + ) + if wer < self.target: + fw, _, _ = eval_wer(model, self.proc, self.full, self.bs, self.mnt) + print( + f"[epoch {ep}] FULL({len(self.full)}) corpus WER = {fw*100:.2f}%", + flush=True, + ) + if fw < self.target: + print( + f"[epoch {ep}] TARGET REACHED (<{self.target*100:.0f}% on full). Stopping.", + flush=True, + ) + control.should_training_stop = True + return control + + +def parse_args(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="google/gemma-4-E2B-it") + p.add_argument("--parquet_glob", default="/home/coshe-data/data/eval-*.parquet") + p.add_argument("--cache_path", default="/home/gemma4ft/out/coshe_full_cache.pkl") + p.add_argument("--output_dir", default="/home/gemma4ft/out/full") + p.add_argument("--epochs", type=float, default=150.0) + p.add_argument("--lr", type=float, default=2e-4) + p.add_argument("--batch_size", type=int, default=4) + p.add_argument("--lora_r", type=int, default=256) + p.add_argument("--lora_alpha", type=int, default=512) + p.add_argument("--optim", default="adamw_8bit") + p.add_argument("--include_head", action="store_true") + p.add_argument("--resume", action="store_true") + p.add_argument( + "--init_adapter", + default="", + help="load this adapter as trainable start (fresh optimizer)", + ) + p.add_argument("--eval_every", type=int, default=2) + p.add_argument("--wer_target", type=float, default=0.02) + p.add_argument("--eval_subset", type=int, default=300) + p.add_argument("--loss_gate", type=float, default=0.15) + p.add_argument("--eval_batch_size", type=int, default=24) + p.add_argument("--eval_max_new_tokens", type=int, default=512) + return p.parse_args() + + +def main(): + import os + + args = parse_args() + torch.manual_seed(0) + proc = AutoProcessor.from_pretrained(args.model) + bnb = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + llm_int8_skip_modules=[ + "model.audio_tower", + "model.vision_tower", + "model.embed_audio", + "model.embed_vision", + "lm_head", + ], + ) + model = AutoModelForMultimodalLM.from_pretrained( + args.model, + quantization_config=bnb, + dtype=torch.bfloat16, + device_map="auto", + attn_implementation="eager", + ) + model.config.use_cache = False + for pm in model.parameters(): + pm.requires_grad = False + model.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={"use_reentrant": False} + ) + model.enable_input_require_grads() + if args.init_adapter: + # Load a previously-trained adapter as the starting point but with a FRESH + # optimizer (so a new/higher LR takes effect) — used to accelerate a run + # whose loss descent has slowed, without losing learned weights. + print(f"Loading init adapter (trainable): {args.init_adapter}", flush=True) + model = PeftModel.from_pretrained(model, args.init_adapter, is_trainable=True) + else: + model = get_peft_model( + model, + LoraConfig( + r=args.lora_r, + lora_alpha=args.lora_alpha, + lora_dropout=0.0, + bias="none", + task_type="CAUSAL_LM", + target_modules=LORA_TARGETS_HEAD if args.include_head else LORA_TARGETS, + ), + ) + model.print_trainable_parameters() + + ds = CosheParquetDataset( + args.parquet_glob, + max_seconds=30.0, + target_max_chars=0, + cache_path=args.cache_path, + ) + items = list(ds) + subset = items[: args.eval_subset] + print(f"Dataset: {len(items)} samples; eval subset={len(subset)}", flush=True) + collator = Gemma4AudioCollator(proc) + + targs = TrainingArguments( + output_dir=args.output_dir, + per_device_train_batch_size=args.batch_size, + num_train_epochs=args.epochs, + learning_rate=args.lr, + lr_scheduler_type="constant", + warmup_steps=0, + bf16=True, + fp16=False, + logging_steps=10, + save_strategy="epoch", + save_total_limit=2, + optim=args.optim, + report_to=[], + remove_unused_columns=False, + dataloader_num_workers=0, + gradient_checkpointing=False, + ) + + cb = WERStopCallback( + proc, + subset, + items, + args.eval_every, + args.wer_target, + args.eval_batch_size, + args.eval_max_new_tokens, + args.output_dir, + loss_gate=args.loss_gate, + ) + trainer = Trainer( + model=model, + args=targs, + train_dataset=ds, + data_collator=collator, + callbacks=[cb], + ) + + resume = ( + args.resume + and os.path.isdir(args.output_dir) + and any(d.startswith("checkpoint-") for d in os.listdir(args.output_dir)) + ) + print(f"Starting training... (resume={resume})", flush=True) + trainer.train(resume_from_checkpoint=resume) + trainer.save_model(args.output_dir) + proc.save_pretrained(args.output_dir) + print("DONE (final adapter saved)", flush=True) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/gemma4/finetune/train_wispr.py b/extras/asr-services/providers/gemma4/finetune/train_wispr.py new file mode 100644 index 00000000..ffb0257b --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/train_wispr.py @@ -0,0 +1,168 @@ +"""LoRA finetune of Gemma4-E2B on the Wispr dictation windows (60% train / 20% val), +val-loss early stopping. Mirrors train_lora_windowed.py but uses WISPR_PROMPT (English +verbatim) and the Wispr windowed manifest. Held-out 20% test clips are scored separately +by eval_wispr_stitch.py. decoder-only LoRA by default (--include_head optional). +""" + +import argparse +import json +import os + +import torch +from data import Gemma4AudioCollator +from data_wispr import WISPR_PROMPT, WindowedManifestDataset +from peft import LoraConfig, get_peft_model +from transformers import ( + AutoModelForMultimodalLM, + AutoProcessor, + BitsAndBytesConfig, + EarlyStoppingCallback, + Trainer, + TrainingArguments, +) + +LORA_TARGETS_HEAD = ( + r"(.*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" + r"|.*language_model\.embed_tokens|lm_head)" +) +LORA_TARGETS = ( + r".*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" +) +# Audio conformer linears live INSIDE Gemma4ClippableLinear wrappers, so target the inner +# `.linear` of each attention/FFN projection. Lets the model adapt its *hearing*, not just +# the text decoder (which a frozen encoder can't fix). +AUDIO_TARGETS = ( + r".*audio_tower.*\.(q_proj|k_proj|v_proj|post|relative_k_proj" + r"|ffw_layer_1|ffw_layer_2)\.linear" +) + + +def parse_args(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="google/gemma-4-E2B-it") + p.add_argument("--manifest", default="/home/wispr_windowed/manifest.jsonl") + p.add_argument("--output_dir", required=True) + p.add_argument("--epochs", type=float, default=40.0) + p.add_argument("--lr", type=float, default=1e-5) + p.add_argument("--batch_size", type=int, default=2) + p.add_argument("--grad_accum", type=int, default=2) + p.add_argument("--eval_batch_size", type=int, default=4) + p.add_argument("--lora_r", type=int, default=16) + p.add_argument("--lora_alpha", type=int, default=32) + p.add_argument("--include_head", action="store_true") + p.add_argument( + "--audio_lora", + action="store_true", + help="also LoRA the audio conformer (adapt hearing, not just decoder)", + ) + p.add_argument("--optim", default="adamw_8bit") + p.add_argument("--patience", type=int, default=4) + p.add_argument("--attn", default="sdpa") + return p.parse_args() + + +def main(): + args = parse_args() + torch.manual_seed(0) + proc = AutoProcessor.from_pretrained(args.model) + bnb = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + llm_int8_skip_modules=[ + "model.audio_tower", + "model.vision_tower", + "model.embed_audio", + "model.embed_vision", + "lm_head", + ], + ) + model = AutoModelForMultimodalLM.from_pretrained( + args.model, + quantization_config=bnb, + dtype=torch.bfloat16, + device_map="auto", + attn_implementation=args.attn, + ) + model.config.use_cache = False + for pm in model.parameters(): + pm.requires_grad = False + model.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={"use_reentrant": False} + ) + model.enable_input_require_grads() + targets = LORA_TARGETS_HEAD if args.include_head else LORA_TARGETS + if args.audio_lora: + targets = f"({targets}|{AUDIO_TARGETS})" + model = get_peft_model( + model, + LoraConfig( + r=args.lora_r, + lora_alpha=args.lora_alpha, + lora_dropout=0.0, + bias="none", + task_type="CAUSAL_LM", + target_modules=targets, + ), + ) + model.print_trainable_parameters() + + train_ds = WindowedManifestDataset(args.manifest, "train") + val_ds = WindowedManifestDataset(args.manifest, "val") + print(f"train windows={len(train_ds)} val windows={len(val_ds)}", flush=True) + collator = Gemma4AudioCollator(proc, prompt=WISPR_PROMPT) + + targs = TrainingArguments( + output_dir=args.output_dir, + per_device_train_batch_size=args.batch_size, + gradient_accumulation_steps=args.grad_accum, + per_device_eval_batch_size=args.eval_batch_size, + num_train_epochs=args.epochs, + learning_rate=args.lr, + lr_scheduler_type="constant", + warmup_steps=0, + bf16=True, + fp16=False, + logging_steps=10, + eval_strategy="epoch", + save_strategy="epoch", + save_total_limit=2, + load_best_model_at_end=True, + metric_for_best_model="eval_loss", + greater_is_better=False, + optim=args.optim, + report_to=[], + remove_unused_columns=False, + dataloader_num_workers=0, + gradient_checkpointing=False, + seed=0, + data_seed=0, + ) + + trainer = Trainer( + model=model, + args=targs, + train_dataset=train_ds, + eval_dataset=val_ds, + data_collator=collator, + callbacks=[EarlyStoppingCallback(early_stopping_patience=args.patience)], + ) + + print("Starting training...", flush=True) + trainer.train() + trainer.save_model(args.output_dir) + proc.save_pretrained(args.output_dir) + json.dump( + trainer.state.log_history, + open(os.path.join(args.output_dir, "log_history.json"), "w"), + ) + print( + f"DONE best_checkpoint={trainer.state.best_model_checkpoint} " + f"best_eval_loss={trainer.state.best_metric}", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/gemma4/finetune/verify_fix.py b/extras/asr-services/providers/gemma4/finetune/verify_fix.py new file mode 100644 index 00000000..2c6df283 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/verify_fix.py @@ -0,0 +1,53 @@ +"""Verify enable_thinking=False + no_repeat_ngram fixes the bad clips.""" + +import glob +import os +import sys + +sys.path.insert(0, "/home/gemma4ft") +import pyarrow.parquet as pq +from bench_coshe_12b import SR, VERBATIM_PROMPT, Model, decode_audio + +DATA = "/home/coshe-data/data" +TARGETS = [ + "audio_1400.wav", + "audio_299.wav", + "audio_160.wav", + "audio_1404.wav", + "audio_176.wav", + "audio_12.wav", +] + +want = {n: None for n in TARGETS} +for s in sorted(glob.glob(os.path.join(DATA, "eval-*.parquet"))): + t = pq.read_table(s, columns=["audio_file_name", "transcription", "audio"]) + for nm, tr, au in zip( + t.column(0).to_pylist(), t.column(1).to_pylist(), t.column(2).to_pylist() + ): + if nm in want and want[nm] is None and au and au.get("bytes"): + want[nm] = (decode_audio(au["bytes"]), tr) + if all(v is not None for v in want.values()): + break + +m = Model( + "google/gemma-4-12B-it", + max_new_tokens=1024, + prompt=VERBATIM_PROMPT, + greedy=True, + no_repeat_ngram_size=3, +) +print("model loaded (greedy, thinking=False, no_repeat_ngram=3)\n", flush=True) + +for nm in TARGETS: + if want[nm] is None: + print(f"{nm}: NOT FOUND") + continue + audio, gt = want[nm] + hyp = m.transcribe(audio) + print( + f"=== {nm} dur={len(audio)/SR:.0f}s hyp_words={len(hyp.split())} ===", + flush=True, + ) + print("GT :", gt[:200].replace(chr(10), " "), flush=True) + print("HYP:", hyp[:200].replace(chr(10), " ") if hyp else "(EMPTY)", flush=True) + print(flush=True) diff --git a/extras/asr-services/providers/gemma4/finetune/vm_exec.py b/extras/asr-services/providers/gemma4/finetune/vm_exec.py new file mode 100644 index 00000000..64576526 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/vm_exec.py @@ -0,0 +1,90 @@ +"""Run a shell command on the Jarvis VM via the Jupyter kernel API (SSH-free). + +SSH to the instance is currently throttled/blocked, but the Jupyter server is +reachable. This opens a kernel, executes a python snippet that shells out, streams +stdout back, and tears the kernel down. + +Usage: + JLTOKEN=... python vm_exec.py "<shell command>" +It resolves the instance URL+token via jlclient automatically. +""" + +import json +import os +import ssl +import sys +import time +import uuid + +import requests +from jlclient import jarvisclient +from jlclient.jarvisclient import * +from websocket import create_connection + +jarvisclient.token = os.environ["JLTOKEN"] +url = User.get_instances()[0].url # https://<host>/lab?token=<tok> +host = url.split("//")[1].split("/")[0] +tok = url.split("token=")[1] +base = f"https://{host}" + +cmd = sys.argv[1] +code = ( + "import subprocess;" + f"r=subprocess.run({cmd!r}, shell=True, capture_output=True, text=True);" + "print(r.stdout);" + "print(r.stderr) if r.stderr else None" +) + +r = requests.post(f"{base}/api/kernels", params={"token": tok}, timeout=30) +kid = r.json()["id"] +try: + ws = create_connection( + f"wss://{host}/api/kernels/{kid}/channels?token={tok}", + sslopt={"cert_reqs": ssl.CERT_NONE}, + timeout=30, + ) + msg_id = uuid.uuid4().hex + hdr = { + "msg_id": msg_id, + "username": "u", + "session": uuid.uuid4().hex, + "msg_type": "execute_request", + "version": "5.3", + } + ws.send( + json.dumps( + { + "header": hdr, + "parent_header": {}, + "metadata": {}, + "content": { + "code": code, + "silent": False, + "store_history": False, + "user_expressions": {}, + "allow_stdin": False, + "stop_on_error": True, + }, + "channel": "shell", + } + ) + ) + out = [] + t0 = time.time() + while time.time() - t0 < 600: + m = json.loads(ws.recv()) + if m.get("parent_header", {}).get("msg_id") != msg_id: + continue + mt = m.get("msg_type") + if mt == "stream": + out.append(m["content"]["text"]) + elif mt in ("execute_result", "display_data"): + out.append(str(m["content"]["data"].get("text/plain", ""))) + elif mt == "error": + out.append("\n".join(m["content"]["traceback"])) + elif mt == "status" and m["content"]["execution_state"] == "idle": + break + ws.close() + print("".join(out)) +finally: + requests.delete(f"{base}/api/kernels/{kid}", params={"token": tok}, timeout=30) diff --git a/extras/asr-services/providers/gemma4/finetune/window_coshe.py b/extras/asr-services/providers/gemma4/finetune/window_coshe.py new file mode 100644 index 00000000..3387233b --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/window_coshe.py @@ -0,0 +1,129 @@ +"""Window CoSHE clips into <=max-seconds segments with EXACT ground-truth targets. + +Input is `fa_words.jsonl` — a faithful forced alignment of the CoSHE ground-truth +transcript (per clip: the GT word sequence, each with start/end seconds). This is the +VibeVoice-ASR data recipe: a timestamped source gives word timings, and we map the exact +GT text onto that timeline. Because the alignment IS the GT (word sequence == GT.split()), +windowing is a lossless contiguous partition — concatenating window texts reproduces the +full transcript. + +Why window: Gemma4's audio encoder hears <=~30s (E2B hard-caps audio at ~31s / 786 tok), +and CoSHE clips run ~57s median. The prior finetune used a "first-30s + proportional +char-truncation" hack (no timestamps were available then); with real word timings we can +cut honest <=30s (audio, text) pairs and train/eval on the true windowed task. + +Cuts fall at WORD boundaries, preferring an inter-word silence (gap >= --min-gap) near the +window tail so boundaries land in pauses rather than mid-utterance. + +Out: windows.jsonl, one row per window: + {audio_file_name, win_idx, n_wins, start, end, dur, n_words, text, cut_gap} +""" + +import argparse +import json +import statistics as st +from pathlib import Path + + +def window_clip(words, max_s, min_s, min_gap): + """Partition `words` (list of {word,start,end}) into contiguous [i,j] index spans, + each spanning <= max_s seconds, preferring to end at a silence >= min_gap.""" + spans = [] + i, n = 0, len(words) + while i < n: + wstart = words[i]["start"] + j = i + while j + 1 < n and words[j + 1]["end"] - wstart <= max_s: + j += 1 + cut_gap = 0.0 + if j + 1 < n: # more words remain -> consider cutting at a pause + best_k, best_gap = j, -1.0 + for k in range(j, i, -1): + if words[k]["end"] - wstart < min_s: + break + gap = words[k + 1]["start"] - words[k]["end"] + if gap > best_gap: + best_gap, best_k = gap, k + if best_gap >= min_gap: + j, cut_gap = best_k, best_gap + spans.append((i, j, cut_gap)) + i = j + 1 + return spans + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--fa", default="/home/ft/fa_words.jsonl") + ap.add_argument("--exclude", default="", help="comma-sep audio_file_names to drop") + ap.add_argument("--out", default="/home/coshe_windowed/windows.jsonl") + ap.add_argument("--max_seconds", type=float, default=28.0) + ap.add_argument("--min_seconds", type=float, default=5.0) + ap.add_argument("--min_gap", type=float, default=0.25) + ap.add_argument("--pad", type=float, default=0.1) + args = ap.parse_args() + + drop = {x for x in args.exclude.split(",") if x} + rows = [json.loads(l) for l in open(args.fa)] + Path(args.out).parent.mkdir(parents=True, exist_ok=True) + + out = open(args.out, "w") + n_clips = n_wins = n_dropped = 0 + durs, wins_per_clip, mismatches = [], [], 0 + for r in rows: + name = r["audio_file_name"] + if name in drop: + n_dropped += 1 + continue + words = r["words"] + if not words: + continue + spans = window_clip(words, args.max_seconds, args.min_seconds, args.min_gap) + n_clips += 1 + wins_per_clip.append(len(spans)) + # lossless check: contiguous partition reproduces the GT word sequence + flat = [words[k]["word"] for a, b, _ in spans for k in range(a, b + 1)] + if flat != [w["word"] for w in words]: + mismatches += 1 + for wi, (a, b, cut_gap) in enumerate(spans): + start = max(0.0, words[a]["start"] - args.pad) + end = words[b]["end"] + args.pad + text = " ".join(words[k]["word"] for k in range(a, b + 1)) + dur = round(end - start, 3) + durs.append(dur) + out.write( + json.dumps( + { + "audio_file_name": name, + "win_idx": wi, + "n_wins": len(spans), + "start": round(start, 3), + "end": round(end, 3), + "dur": dur, + "n_words": b - a + 1, + "text": text, + "cut_gap": round(cut_gap, 3), + }, + ensure_ascii=False, + ) + + "\n" + ) + n_wins += 1 + out.close() + + from collections import Counter + + c = Counter(wins_per_clip) + print(f"clips kept={n_clips} dropped(foreign)={n_dropped} windows={n_wins}") + print(f"partition-mismatch clips (should be 0): {mismatches}") + print(f"windows/clip: {dict(sorted(c.items()))}") + print( + f"window dur s: min {min(durs):.1f} med {st.median(durs):.1f} " + f"mean {st.mean(durs):.1f} max {max(durs):.1f} p95 {sorted(durs)[int(len(durs)*0.95)]:.1f}" + ) + over = sum(1 for d in durs if d > 30.0) + print(f"windows > 30s (should be ~0): {over}") + print(f"out -> {args.out}") + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/gemma4/finetune/window_target.py b/extras/asr-services/providers/gemma4/finetune/window_target.py new file mode 100644 index 00000000..ec7a0744 --- /dev/null +++ b/extras/asr-services/providers/gemma4/finetune/window_target.py @@ -0,0 +1,25 @@ +"""Proportionally truncate transcripts to the audio window the model actually hears. + +CoSHE clips are mostly ~57s but Gemma4's audio encoder caps at 30s, and the dataset has +no word timestamps. So for a coherent "transcribe the first W seconds" task we truncate +each clip's transcript to the first `W/duration` fraction of its characters (≈ the words +spoken in the heard window, assuming roughly uniform speech rate). Clips <= W keep the +full transcript. Applied identically in training and eval so base/FT and ref/hyp stay +consistent. +""" + +import json + + +def apply_window_truncation(items, durations_path, window_seconds): + """Mutate `items` (list of {name,target,...}) in place; return it. No-op if + window_seconds is falsy.""" + if not window_seconds: + return items + dur = json.load(open(durations_path)) + for it in items: + d = dur.get(it["name"], 0.0) + if d and d > window_seconds: + n = max(1, int(len(it["target"]) * (window_seconds / d))) + it["target"] = it["target"][:n].rstrip() + return items diff --git a/extras/asr-services/providers/nemo/finetune/cut_segments.py b/extras/asr-services/providers/nemo/finetune/cut_segments.py new file mode 100644 index 00000000..26a3c1c8 --- /dev/null +++ b/extras/asr-services/providers/nemo/finetune/cut_segments.py @@ -0,0 +1,94 @@ +"""Slice source WAVs at the segment boundaries from segment_by_alignment.py and write a +NeMo training manifest of the short, correctly-targeted clips. + +Tool-agnostic: consumes segments.jsonl ({audio_file_name, seg_start, seg_end, text, +review}) + the full-clip WAV dir (from make_manifest.py --all). Emits one WAV per segment +and a NeMo manifest line per segment ({audio_filepath, text, duration, target_lang}). + +By default skips segments flagged review=true (low GT<->hyp match) — pass --keep-flagged +to include them. + +Usage: + python cut_segments.py --segments segments.jsonl --wav-dir /home/ft/data_full/wav \ + --out-dir /home/ft/data_seg --target-lang hi-IN +""" + +import argparse +import json +from pathlib import Path + +import soundfile as sf + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--segments", required=True) + ap.add_argument( + "--wav-dir", + required=True, + help="dir of full-clip WAVs (stem = audio_file_name stem)", + ) + ap.add_argument("--out-dir", required=True) + ap.add_argument("--target-lang", default="hi-IN") + ap.add_argument("--min-seconds", type=float, default=0.5) + ap.add_argument("--keep-flagged", action="store_true") + args = ap.parse_args() + + out = Path(args.out_dir) + seg_wav_dir = out / "wav" + seg_wav_dir.mkdir(parents=True, exist_ok=True) + wav_dir = Path(args.wav_dir) + + # cache decoded source audio so we read each full clip once + cache: dict[str, tuple] = {} + n_seg = n_skip = 0 + with open(out / "train.json", "w") as mf: + for line in open(args.segments): + s = json.loads(line) + if s.get("review") and not args.keep_flagged: + n_skip += 1 + continue + dur = s["seg_end"] - s["seg_start"] + if dur < args.min_seconds or not s["text"].strip(): + n_skip += 1 + continue + name = s["audio_file_name"] + stem = Path(name).stem + if stem not in cache: + data, sr = sf.read(str(wav_dir / f"{stem}.wav"), dtype="float32") + if data.ndim > 1: + data = data.mean(axis=1) + cache[stem] = (data, sr) + data, sr = cache[stem] + a = max(0, int(s["seg_start"] * sr)) + b = min(len(data), int(s["seg_end"] * sr)) + if b - a < int(args.min_seconds * sr): + n_skip += 1 + continue + seg_name = f"{stem}_{int(s['seg_start']*1000):07d}_{int(s['seg_end']*1000):07d}.wav" + seg_path = seg_wav_dir / seg_name + sf.write(str(seg_path), data[a:b], sr, subtype="PCM_16") + mf.write( + json.dumps( + { + "audio_filepath": str(seg_path), + "text": s["text"], + "duration": round((b - a) / sr, 3), + "target_lang": args.target_lang, + "audio_file_name": seg_name, + "source_clip": name, + }, + ensure_ascii=False, + ) + + "\n" + ) + n_seg += 1 + print( + f"wrote {n_seg} segment WAVs + manifest, skipped {n_skip} " + f"-> {out / 'train.json'}", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/nemo/finetune/eval_full.py b/extras/asr-services/providers/nemo/finetune/eval_full.py new file mode 100644 index 00000000..ae1dc9cb --- /dev/null +++ b/extras/asr-services/providers/nemo/finetune/eval_full.py @@ -0,0 +1,120 @@ +"""Evaluate a fine-tuned nemotron .nemo checkpoint on a CoSHE manifest. + +Loads a saved .nemo (restore_from), transcribes every clip in the manifest, and +writes one JSONL line per clip in the schema mlexp score_coshe.py consumes +({audio_file_name, transcription, hyp, asr_seconds, duration_s}). Score locally: + + uv run --with "jiwer,requests,tqdm" python3 \ + extras/ml-experiments/src/mlexp/evaluate/score_coshe.py \ + --result ft=ft.jsonl --result base=nemotron_auto_full.jsonl --out-dir report/ + +Resumable: clips already present in --out are skipped. + +Usage: + python eval_full.py --init /home/ft/out/warm/final.nemo \ + --manifest /home/ft/data_full/all.json --out /home/ft/results/ft_warm.jsonl +""" + +import argparse +import json +import time +from pathlib import Path + +import nemo.collections.asr as nemo_asr +import torch +from nemo.collections.asr.data.audio_to_text_lhotse_prompt_index import ( + LhotseSpeechToTextBpeDatasetWithPromptIndex, +) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--init", required=True, help="path to fine-tuned .nemo") + ap.add_argument("--manifest", required=True) + ap.add_argument("--out", required=True) + ap.add_argument("--target-lang", default="hi-IN") + ap.add_argument("--batch-size", type=int, default=1) + ap.add_argument( + "--max-symbols", + type=int, + default=None, + help="override RNN-T greedy max_symbols per frame (default cfg=10)", + ) + args = ap.parse_args() + + rows = [json.loads(l) for l in open(args.manifest)] + out_path = Path(args.out) + out_path.parent.mkdir(parents=True, exist_ok=True) + done = set() + if out_path.exists(): + for l in open(out_path): + try: + done.add(json.loads(l)["audio_file_name"]) + except Exception: + pass + todo = [r for r in rows if r["audio_file_name"] not in done] + print(f"{len(done)} done, {len(todo)} to eval", flush=True) + + # --init can be a local .nemo (fine-tuned) OR a HF model id (base, for comparison) + if Path(args.init).is_file(): + model = nemo_asr.models.ASRModel.restore_from(args.init) + else: + model = nemo_asr.models.ASRModel.from_pretrained(model_name=args.init) + model.eval() + + # RNN-T greedy decode caps at `max_symbols` tokens per encoder frame (default 10); + # too low truncates long / token-dense (Devanagari) segments. Bump to test. + if args.max_symbols: + from omegaconf import open_dict + + dec = model.cfg.decoding + with open_dict(dec): + dec.strategy = "greedy_batch" + dec.greedy.max_symbols = args.max_symbols + model.change_decoding_strategy(dec) + + # force the training target_lang prompt; num_workers=0 so this in-process + # override applies (worker subprocesses wouldn't see it). See train_overfit.py. + LhotseSpeechToTextBpeDatasetWithPromptIndex._get_prompt_index_for_cut = ( + lambda self, cut, _tl=args.target_lang: self._get_prompt_index(_tl) + ) + + t0 = time.time() + with open(out_path, "a") as out_f, torch.no_grad(): + for i in range(0, len(todo), args.batch_size): + batch = todo[i : i + args.batch_size] + wavs = [r["audio_filepath"] for r in batch] + t1 = time.time() + hyps = model.transcribe( + wavs, + batch_size=args.batch_size, + target_lang=args.target_lang, + num_workers=0, + verbose=False, + ) + dt = (time.time() - t1) / len(batch) + for r, h in zip(batch, hyps): + text = h.text if hasattr(h, "text") else str(h) + out_f.write( + json.dumps( + { + "audio_file_name": r["audio_file_name"], + "transcription": r["text"], + "hyp": text, + "asr_seconds": round(dt, 3), + "duration_s": r.get("duration", 0), + }, + ensure_ascii=False, + ) + + "\n" + ) + out_f.flush() + n = i + len(batch) + if n % 50 == 0 or n == len(todo): + el = time.time() - t0 + print(f" {n}/{len(todo)} ({el:.0f}s, {n/el:.2f} clip/s)", flush=True) + print(f"DONE -> {out_path}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/nemo/finetune/fa_align.py b/extras/asr-services/providers/nemo/finetune/fa_align.py new file mode 100644 index 00000000..fd7c441d --- /dev/null +++ b/extras/asr-services/providers/nemo/finetune/fa_align.py @@ -0,0 +1,114 @@ +"""FREE local forced alignment of CoSHE audio to its ground-truth transcript, on the +4090, via ctc-forced-aligner (MMS uroman CTC). Times the GROUND-TRUTH words directly — +no hyp<->GT reconciliation. + +IMPORTANT: the high-level get_word_stamps() does NOT romanize, so it silently drops +Devanagari (keeps only Latin words). We use the lower-level ONNX pipeline with +preprocess_text(romanize=True) so Hinglish (Devanagari + Latin) aligns fully. Needs +onnxruntime-gpu for CUDA. + +Emits the normalized JSONL segment_by_alignment.py consumes (words = GT words + times): + {"audio_file_name": "...", "words": [{"word": "<gt>", "start": 1.2, "end": 1.4}, ...]} + +Usage: + python fa_align.py --manifest overfit.json --out fa_words.jsonl --language hin +""" + +import argparse +import json +import os +import time +from pathlib import Path + +import onnxruntime +from ctc_forced_aligner import ( + MODEL_URL, + Tokenizer, + ensure_onnx_model, + generate_emissions, + get_alignments, + get_spans, + load_audio, + postprocess_results, + preprocess_text, +) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--manifest", required=True) + ap.add_argument("--out", required=True) + ap.add_argument( + "--language", + default="hin", + help="uroman iso (hin romanizes Devanagari; Latin passes through)", + ) + ap.add_argument("--batch-size", type=int, default=8) + ap.add_argument( + "--model-path", default=os.path.expanduser("~/.cache/ctc_fa/mms_fa.onnx") + ) + args = ap.parse_args() + + Path(args.model_path).parent.mkdir(parents=True, exist_ok=True) + ensure_onnx_model(args.model_path, MODEL_URL) + providers = onnxruntime.get_available_providers() + use = ( + ["CUDAExecutionProvider", "CPUExecutionProvider"] + if "CUDAExecutionProvider" in providers + else ["CPUExecutionProvider"] + ) + print(f"onnx providers: {providers} -> using {use}", flush=True) + session = onnxruntime.InferenceSession(args.model_path, providers=use) + tokenizer = Tokenizer() + + rows = [json.loads(l) for l in open(args.manifest)] + out_path = Path(args.out) + out_path.parent.mkdir(parents=True, exist_ok=True) + + t0 = time.time() + with open(out_path, "w") as out_f: + for r in rows: + audio = load_audio(r["audio_filepath"], ret_type="np") + emissions, stride = generate_emissions( + session, audio, batch_size=args.batch_size + ) + tokens_starred, text_starred = preprocess_text( + r["text"], + romanize=True, + language=args.language, + ) + segments, scores, blank = get_alignments( + emissions, tokens_starred, tokenizer + ) + spans = get_spans(tokens_starred, segments, blank) + word_ts = postprocess_results(text_starred, spans, stride, scores) + words = [ + { + "word": w["text"], + "start": round(w["start"], 3), + "end": round(w["end"], 3), + } + for w in word_ts + ] + out_f.write( + json.dumps( + { + "audio_file_name": r["audio_file_name"], + "words": words, + }, + ensure_ascii=False, + ) + + "\n" + ) + out_f.flush() + gt_n = len(r["text"].split()) + print( + f" {r['audio_file_name']}: {len(words)} aligned / {gt_n} GT words " + f"[{words[0]['start'] if words else 0:.1f}..{words[-1]['end'] if words else 0:.1f}s]", + flush=True, + ) + print(f"DONE {len(rows)} clips in {time.time()-t0:.0f}s -> {out_path}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/nemo/finetune/make_manifest.py b/extras/asr-services/providers/nemo/finetune/make_manifest.py new file mode 100644 index 00000000..9bd63c73 --- /dev/null +++ b/extras/asr-services/providers/nemo/finetune/make_manifest.py @@ -0,0 +1,167 @@ +"""Build NeMo manifests + 16 kHz mono WAVs from the CoSHE-Eval parquet shards. + +Nemotron-3.5-asr-streaming is ``EncDecRNNTBPEModelWithPrompt`` (NeMo 2.8.0rc0): a +cache-aware FastConformer-RNNT whose Lhotse dataloader reads a per-utterance +language from the manifest ``prompt_field`` (``target_lang``) and maps it through +the model's ``prompt_dictionary`` (hi-IN -> 6, en-US -> 0, ...). So every manifest +line needs ``audio_filepath``, ``text``, ``duration`` and ``target_lang``. + +This mirrors bench_coshe.decode_to_wav16k so train/eval audio is byte-identical to +the benchmark. It powers all three FT stages: + + Stage 1 (overfit smoke): --indices 0 1 (1-2 clips) + Stage 2 (full overfit): --all (1985 clips, one split) + Stage 3 (generalization): --split 0.2 --seed 0 (train/val/test jsonls) + +Usage: + python make_manifest.py --dataset /home/coshe/data --out-dir /home/ft/data \ + --target-lang hi-IN --indices 0 1 + python make_manifest.py --dataset /home/coshe/data --out-dir /home/ft/data \ + --target-lang hi-IN --split 0.2 --seed 0 +""" + +import argparse +import io +import json +import random +from pathlib import Path + +import librosa +import pyarrow.parquet as pq +import soundfile as sf + + +def decode_to_wav16k(audio_bytes: bytes, dst: str) -> float: + """Decode embedded audio bytes -> 16 kHz mono PCM16 WAV. Returns duration (s). + + Identical to bench_coshe.decode_to_wav16k so FT audio matches the benchmark. + """ + data, sr = sf.read(io.BytesIO(audio_bytes), dtype="float32", always_2d=False) + if data.ndim > 1: + data = data.mean(axis=1) + if sr != 16000: + data = librosa.resample(data, orig_sr=sr, target_sr=16000) + sr = 16000 + sf.write(dst, data, sr, subtype="PCM_16") + return len(data) / sr + + +def iter_rows(dataset: str): + """Yield (global_index, row_dict) over all eval-*.parquet shards in order.""" + shards = sorted(Path(dataset).glob("eval-*.parquet")) + if not shards: + raise SystemExit(f"No eval-*.parquet shards under {dataset}") + gi = 0 + for shard in shards: + pf = pq.ParquetFile(shard) + for batch in pf.iter_batches(batch_size=64): + for row in batch.to_pylist(): + yield gi, row + gi += 1 + + +def write_manifest(rows, wav_dir: Path, manifest_path: Path, target_lang: str) -> int: + wav_dir.mkdir(parents=True, exist_ok=True) + manifest_path.parent.mkdir(parents=True, exist_ok=True) + n = 0 + with open(manifest_path, "w") as mf: + for _gi, row in rows: + name = row["audio_file_name"] + wav_path = wav_dir / f"{Path(name).stem}.wav" + dur = decode_to_wav16k(row["audio"]["bytes"], str(wav_path)) + mf.write( + json.dumps( + { + "audio_filepath": str(wav_path), + "text": row["transcription"], + "duration": round(dur, 3), + "target_lang": target_lang, + "audio_file_name": name, + }, + ensure_ascii=False, + ) + + "\n" + ) + n += 1 + print(f" wrote {n} -> {manifest_path}", flush=True) + return n + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--dataset", default="/home/coshe/data") + ap.add_argument("--out-dir", required=True) + ap.add_argument( + "--target-lang", + default="hi-IN", + help="prompt_dictionary key written to every manifest line", + ) + g = ap.add_mutually_exclusive_group(required=True) + g.add_argument("--indices", type=int, nargs="+", help="explicit global row indices") + g.add_argument("--names", nargs="+", help="explicit audio_file_name values") + g.add_argument("--all", action="store_true", help="every clip into one manifest") + g.add_argument( + "--split", type=float, help="held-out fraction, e.g. 0.2 (train=this frac)" + ) + ap.add_argument("--seed", type=int, default=0) + args = ap.parse_args() + + out = Path(args.out_dir) + wav_dir = out / "wav" + + if args.indices is not None: + want = set(args.indices) + rows = [(gi, r) for gi, r in iter_rows(args.dataset) if gi in want] + write_manifest(rows, wav_dir, out / "overfit.json", args.target_lang) + + elif args.names is not None: + want = set(args.names) + rows = [ + (gi, r) for gi, r in iter_rows(args.dataset) if r["audio_file_name"] in want + ] + got = {r["audio_file_name"] for _gi, r in rows} + missing = want - got + if missing: + raise SystemExit(f"names not found: {sorted(missing)}") + write_manifest(rows, wav_dir, out / "overfit.json", args.target_lang) + + elif args.all: + rows = list(iter_rows(args.dataset)) + write_manifest(rows, wav_dir, out / "all.json", args.target_lang) + + else: # --split: train = args.split fraction, then half of remainder = val, half = test + all_rows = list(iter_rows(args.dataset)) + idx = list(range(len(all_rows))) + random.Random(args.seed).shuffle(idx) + n = len(idx) + n_train = int(round(args.split * n)) + n_val = (n - n_train) // 2 + train_i = set(idx[:n_train]) + val_i = set(idx[n_train : n_train + n_val]) + test_i = set(idx[n_train + n_val :]) + print( + f"split seed={args.seed}: train={len(train_i)} val={len(val_i)} test={len(test_i)}", + flush=True, + ) + write_manifest( + [(gi, r) for gi, r in enumerate(all_rows) if gi in train_i], + wav_dir, + out / "train.json", + args.target_lang, + ) + write_manifest( + [(gi, r) for gi, r in enumerate(all_rows) if gi in val_i], + wav_dir, + out / "val.json", + args.target_lang, + ) + write_manifest( + [(gi, r) for gi, r in enumerate(all_rows) if gi in test_i], + wav_dir, + out / "test.json", + args.target_lang, + ) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/nemo/finetune/run_deepgram.py b/extras/asr-services/providers/nemo/finetune/run_deepgram.py new file mode 100644 index 00000000..c1d2e146 --- /dev/null +++ b/extras/asr-services/providers/nemo/finetune/run_deepgram.py @@ -0,0 +1,125 @@ +"""Get word-level timestamps for CoSHE clips from Deepgram Nova-3 (multilingual), +emitting the normalized JSONL that segment_by_alignment.py consumes: + + {"audio_file_name": "...", "dg_text": "<deepgram hypothesis>", + "words": [{"word": "...", "start": 1.2, "end": 1.4}, ...]} + +Deepgram Nova-3 multilingual is the only option that advertises Hindi<->English +intra-utterance code-switching (CoSHE is exactly that) and returns precise word +timestamps. Already the project's transcription provider, so the key is in hand. +~$0.0092/min -> ~$16 for all 1985 clips (likely free under the $200 credit). + +Start small to de-risk alignment quality on real Hinglish before spending on all 1985: + DEEPGRAM_API_KEY=... python run_deepgram.py --manifest all.json \ + --out hyp_words.jsonl --limit 5 + +Then the full run (drop --limit). Resumable: clips already in --out are skipped. + +`dg_text` is kept so we also get a free Deepgram-vs-base WER on CoSHE later. +""" + +import argparse +import json +import os +import time +from pathlib import Path + +import requests + +DG_URL = "https://api.deepgram.com/v1/listen" + + +def transcribe(wav_path: str, key: str, model: str, language: str) -> dict: + """POST one WAV to Deepgram prerecorded; return its first alternative dict.""" + params = { + "model": model, + "language": language, + "punctuate": "true", + "smart_format": "false", + } + with open(wav_path, "rb") as f: + resp = requests.post( + DG_URL, + params=params, + headers={"Authorization": f"Token {key}", "Content-Type": "audio/wav"}, + data=f.read(), + timeout=120, + ) + resp.raise_for_status() + alts = resp.json()["results"]["channels"][0]["alternatives"] + return alts[0] if alts else {"transcript": "", "words": []} + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument( + "--manifest", required=True, help="all.json (audio_filepath + audio_file_name)" + ) + ap.add_argument("--out", required=True) + ap.add_argument("--model", default="nova-3") + ap.add_argument( + "--language", default="multi", help="'multi' = Nova-3 code-switching" + ) + ap.add_argument( + "--limit", type=int, default=None, help="only the first N clips (validation)" + ) + ap.add_argument("--names", nargs="+", help="only these audio_file_name values") + args = ap.parse_args() + + key = os.environ.get("DEEPGRAM_API_KEY") + if not key: + raise SystemExit("set DEEPGRAM_API_KEY") + + rows = [json.loads(l) for l in open(args.manifest)] + if args.names: + want = set(args.names) + rows = [r for r in rows if r["audio_file_name"] in want] + if args.limit: + rows = rows[: args.limit] + + out_path = Path(args.out) + out_path.parent.mkdir(parents=True, exist_ok=True) + done = set() + if out_path.exists(): + for l in open(out_path): + try: + done.add(json.loads(l)["audio_file_name"]) + except Exception: + pass + todo = [r for r in rows if r["audio_file_name"] not in done] + print( + f"{len(done)} done, {len(todo)} to transcribe (model={args.model} lang={args.language})", + flush=True, + ) + + t0 = time.time() + with open(out_path, "a") as out_f: + for i, r in enumerate(todo, 1): + try: + alt = transcribe(r["audio_filepath"], key, args.model, args.language) + except Exception as e: + print(f" ERR {r['audio_file_name']}: {str(e)[:160]}", flush=True) + continue + words = [ + {"word": w["word"], "start": w["start"], "end": w["end"]} + for w in alt.get("words", []) + ] + out_f.write( + json.dumps( + { + "audio_file_name": r["audio_file_name"], + "dg_text": alt.get("transcript", ""), + "words": words, + }, + ensure_ascii=False, + ) + + "\n" + ) + out_f.flush() + if i % 25 == 0 or i == len(todo): + print(f" {i}/{len(todo)} ({time.time()-t0:.0f}s)", flush=True) + print(f"DONE -> {out_path}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/nemo/finetune/segment_by_alignment.py b/extras/asr-services/providers/nemo/finetune/segment_by_alignment.py new file mode 100644 index 00000000..be823047 --- /dev/null +++ b/extras/asr-services/providers/nemo/finetune/segment_by_alignment.py @@ -0,0 +1,206 @@ +"""Segment long CoSHE clips into <=N-second windows with CORRECT ground-truth targets, +using a timestamped ASR hypothesis as the timing source (forced-alignment substitute). + +Why: RNN-T loss is O(T*U*V) -> a 54s clip OOMs even at batch=1 (see +[[nemotron-coshe-finetune]]). The fix is short clips, but CoSHE has no word timestamps, +so we can't cut the transcript to match a shorter audio window. Solution: run a +timestamped ASR (Deepgram Nova-3 multilingual / Google Chirp), use its WORD TIMINGS to +find cut points, and map our exact GT text onto that timeline. + +This stage is TOOL-AGNOSTIC: it consumes a normalized per-clip word-timestamp JSONL: + {"audio_file_name": "...", "words": [{"word": "...", "start": 1.2, "end": 1.4}, ...]} +A thin adapter (see deepgram_words.py / chirp_words.py, TODO) converts each provider's +response into that shape. GT text comes from the CoSHE benchmark jsonl ("transcription"). + +Pipeline per clip: + 1. romanize GT words and hyp words (IndicXlit) so Devanagari-GT aligns to a hyp in + either script; + 2. sequence-align GT<->hyp (difflib) and assign each GT word an approx time (matched + words take the hyp time; unmatched GT words interpolate between matched anchors); + 3. greedily pack GT words into <=--max-seconds segments, preferring to cut at a hyp + PAUSE (word gap >= --min-gap) so boundaries fall in silence; + 4. emit one segment row per window: {audio_file_name, seg_start, seg_end, text, ...} + plus a confidence flag (match_ratio) so low-confidence segments can be reviewed. + +A later step (cut_segments.py, TODO) slices the WAVs at [seg_start, seg_end] and writes a +NeMo manifest. Nothing here is provider-specific or destructive. + +Usage: + python segment_by_alignment.py --timestamps hyp_words.jsonl \ + --gt nemotron_auto_full.jsonl --out segments.jsonl --max-seconds 18 --min-gap 0.3 +""" + +import argparse +import json +from difflib import SequenceMatcher +from pathlib import Path + +try: + from mlexp.utils.indicxlit import romanize as _romanize +except Exception: # script may run where mlexp/IndicXlit isn't importable + _romanize = None + + +def romanize_words(words: list[str]) -> list[str]: + """Lowercase romanized form per word for cross-script alignment. Falls back to a + plain lowercase when IndicXlit isn't available (Latin-only hyps still align).""" + if _romanize is None: + return [w.lower() for w in words] + # romanize the joined text once (IndicXlit is per-word internally + cached), re-split + rom = _romanize(" ".join(words)).lower().split() + # keep length aligned with input; fall back per-word if the join changed token count + if len(rom) == len(words): + return rom + return [_romanize(w).lower() for w in words] + + +def assign_times(gt_words, hyp_words): + """Return a list of (start, end) per GT word by aligning to the timed hyp words. + + Matched GT words inherit the hyp word's [start,end]; runs of unmatched GT words are + linearly interpolated between the surrounding matched anchors. Returns (times, + match_ratio) where match_ratio is the fraction of GT words that matched a hyp word. + """ + rom_gt = romanize_words([w["w"] for w in gt_words]) + rom_hyp = romanize_words([w["word"] for w in hyp_words]) + sm = SequenceMatcher(a=rom_gt, b=rom_hyp, autojunk=False) + + times = [None] * len(gt_words) + matched = 0 + for tag, i1, i2, j1, _j2 in sm.get_opcodes(): + if tag == "equal": + for k in range(i2 - i1): + h = hyp_words[j1 + k] + times[i1 + k] = (float(h["start"]), float(h["end"])) + matched += 1 + + # interpolate the None runs between anchors + n = len(times) + anchors = [(i, t) for i, t in enumerate(times) if t is not None] + if not anchors: + return None, 0.0 + # head/tail extrapolation: clamp to first/last anchor time + first_i, first_t = anchors[0] + last_i, last_t = anchors[-1] + for i in range(first_i): + times[i] = (first_t[0], first_t[0]) + for i in range(last_i + 1, n): + times[i] = (last_t[1], last_t[1]) + # interior gaps + for (ia, ta), (ib, tb) in zip(anchors, anchors[1:]): + gap = ib - ia + if gap <= 1: + continue + t0, t1 = ta[1], tb[0] + for k in range(1, gap): + frac = k / gap + t = t0 + (t1 - t0) * frac + times[ia + k] = (t, t) + return times, matched / max(1, len(gt_words)) + + +def pack_segments(gt_words, times, hyp_words, max_seconds, min_gap): + """Greedily group GT words into <=max_seconds windows, preferring cut points that + coincide with a hyp pause (a gap >= min_gap between consecutive hyp words).""" + # precompute pause times (end of a hyp word that is followed by a >=min_gap silence) + pauses = [] + for a, b in zip(hyp_words, hyp_words[1:]): + if float(b["start"]) - float(a["end"]) >= min_gap: + pauses.append((float(a["end"]) + float(b["start"])) / 2.0) + + segs = [] + cur_start_idx = 0 + seg_t0 = times[0][0] + for i in range(len(gt_words)): + cur_t1 = times[i][1] + if cur_t1 - seg_t0 >= max_seconds and i > cur_start_idx: + # find nearest pause at/before cur_t1 within this window, else cut here + window_pauses = [p for p in pauses if seg_t0 < p <= cur_t1] + cut_t = window_pauses[-1] if window_pauses else cur_t1 + segs.append((cur_start_idx, i, seg_t0, cut_t)) + cur_start_idx = i + seg_t0 = cut_t # contiguous: next segment starts where this one cut + # final segment + segs.append((cur_start_idx, len(gt_words), seg_t0, times[-1][1])) + + # merge a too-short final fragment (< 1/3 max) back into the previous segment + if len(segs) >= 2 and (segs[-1][3] - segs[-1][2]) < max_seconds / 3: + s_i, _e_i, t0, _t1 = segs[-2] + last = segs[-1] + segs[-2] = (s_i, last[1], t0, last[3]) + segs.pop() + return segs + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument( + "--timestamps", required=True, help="normalized word-timestamp JSONL" + ) + ap.add_argument( + "--gt", required=True, help="CoSHE jsonl with audio_file_name + transcription" + ) + ap.add_argument("--out", required=True) + ap.add_argument("--max-seconds", type=float, default=18.0) + ap.add_argument("--min-gap", type=float, default=0.3) + ap.add_argument( + "--min-match-ratio", + type=float, + default=0.4, + help="flag clips whose GT<->hyp match ratio is below this for review", + ) + args = ap.parse_args() + + gt = {} + for line in open(args.gt): + r = json.loads(line) + if "transcription" in r: + gt[r["audio_file_name"]] = r["transcription"] + + out_path = Path(args.out) + out_path.parent.mkdir(parents=True, exist_ok=True) + n_clip = n_seg = n_flag = 0 + with open(out_path, "w") as out_f: + for line in open(args.timestamps): + ts = json.loads(line) + name = ts["audio_file_name"] + if name not in gt or not ts.get("words"): + continue + gt_words = [{"w": w} for w in gt[name].split()] + hyp_words = ts["words"] + times, ratio = assign_times(gt_words, hyp_words) + if times is None: + continue + flagged = ratio < args.min_match_ratio + n_flag += flagged + for s_i, e_i, t0, t1 in pack_segments( + gt_words, times, hyp_words, args.max_seconds, args.min_gap + ): + text = " ".join(gt_words[k]["w"] for k in range(s_i, e_i)) + out_f.write( + json.dumps( + { + "audio_file_name": name, + "seg_start": round(t0, 3), + "seg_end": round(t1, 3), + "duration": round(t1 - t0, 3), + "text": text, + "match_ratio": round(ratio, 3), + "review": flagged, + }, + ensure_ascii=False, + ) + + "\n" + ) + n_seg += 1 + n_clip += 1 + print( + f"{n_clip} clips -> {n_seg} segments " + f"({n_seg/max(1,n_clip):.1f}/clip), {n_flag} clips flagged for review " + f"-> {out_path}", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/nemo/finetune/train_full.py b/extras/asr-services/providers/nemo/finetune/train_full.py new file mode 100644 index 00000000..a88931cc --- /dev/null +++ b/extras/asr-services/providers/nemo/finetune/train_full.py @@ -0,0 +1,154 @@ +"""Stage 2 full-CoSHE overfit for nemotron-3.5-asr-streaming (NeMo RNNT-prompt). + +Memorize all 1985 CoSHE clips to <2% WER (capacity proof), via the warm->anneal +recipe proven on gemma4/qwen3 ([[gemma4_qlora_finetune]], [[qwen3_asr_coshe_overfit]]): + + warm: --model <hf id> --lr 2e-4 train until train-loss plateaus + anneal: --init <warm.nemo> --lr 3e-5 fresh optimizer -> loss collapses + +Full fine-tune (mode=full) already trains the RNNT joint's 13087-way output +projection, so it's the max-capacity setting (no separate "include_head" needed, +unlike the HF decoder-LM models). + +Runs full-length on correct full targets (CoSHE has no word timestamps, so audio +must NOT be truncated) — requires A100-80GB: one ~60s clip's RNNT joint is ~34 GB. + +Checkpoints every --save-every steps to --save-dir so the run is stop/resume/anneal +-able. WER eval is a SEPARATE step (eval_full.py) — in-training transcribe competes +with the ~34 GB training footprint and risks OOM. + +Usage (A100-80GB, NeMo-main venv): + python train_full.py --manifest /home/ft/data_full/all.json \ + --lr 2e-4 --steps 200000 --batch-size 1 --save-dir /home/ft/out/warm + python train_full.py --init /home/ft/out/warm/step120000.nemo \ + --manifest /home/ft/data_full/all.json --lr 3e-5 --steps 20000 \ + --save-dir /home/ft/out/anneal +""" + +import argparse +import time +from pathlib import Path + +import lightning.pytorch as pl +import nemo.collections.asr as nemo_asr +import torch +from lightning.pytorch.callbacks import Callback +from omegaconf import open_dict + + +class PeriodicSave(Callback): + """Save a .nemo every `every` global steps (NeMo .nemo, not a Lightning ckpt).""" + + def __init__(self, save_dir: str, every: int): + self.save_dir = Path(save_dir) + self.every = every + self.save_dir.mkdir(parents=True, exist_ok=True) + + def on_train_batch_end(self, trainer, pl_module, *args, **kwargs): + step = trainer.global_step + if step > 0 and step % self.every == 0: + path = self.save_dir / f"step{step}.nemo" + pl_module.save_to(str(path)) + print(f"[ckpt] step {step} -> {path}", flush=True) + + +def build_train_cfg(model, manifest: str, batch_size: int, max_dur: float): + cfg = model.cfg.train_ds + with open_dict(cfg): + cfg.manifest_filepath = manifest + cfg.is_tarred = False + cfg.tarred_audio_filepaths = None + cfg.shard_manifests = False + cfg.use_lhotse = True + cfg.shuffle = True + cfg.batch_size = batch_size + cfg.batch_duration = None + cfg.bucketing_batch_size = None + cfg.num_buckets = 0 + cfg.max_duration = max_dur + cfg.min_duration = 0.1 + cfg.num_workers = 4 + cfg.prompt_field = "target_lang" + return cfg + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", default="nvidia/nemotron-3.5-asr-streaming-0.6b") + ap.add_argument( + "--init", default=None, help="resume from a saved .nemo (fresh optimizer)" + ) + ap.add_argument("--manifest", required=True) + ap.add_argument("--steps", type=int, default=200000) + ap.add_argument("--lr", type=float, default=2e-4) + ap.add_argument("--batch-size", type=int, default=1) + ap.add_argument("--fused-batch-size", type=int, default=1) + ap.add_argument("--max-dur", type=float, default=65.0) + ap.add_argument("--mode", choices=["full", "decoder"], default="full") + ap.add_argument("--save-dir", required=True) + ap.add_argument("--save-every", type=int, default=20000) + args = ap.parse_args() + + torch.set_float32_matmul_precision("high") + + if args.init: + print(f"Resuming from {args.init} (fresh optimizer, lr={args.lr})", flush=True) + model = nemo_asr.models.ASRModel.restore_from(args.init) + else: + print(f"Loading base {args.model} (lr={args.lr})", flush=True) + model = nemo_asr.models.ASRModel.from_pretrained(model_name=args.model) + model.set_trainer(None) + + with open_dict(model.cfg.optim): + model.cfg.optim.name = "adamw" + model.cfg.optim.lr = args.lr + model.cfg.optim.weight_decay = 0.0 + model.cfg.optim.pop("sched", None) + + model.setup_training_data( + build_train_cfg(model, args.manifest, args.batch_size, args.max_dur) + ) + + if args.fused_batch_size > 0 and hasattr(model.joint, "set_fuse_loss_wer"): + model.joint.set_fuse_loss_wer(True, loss=model.loss, metric=model.wer) + model.joint.set_fused_batch_size(args.fused_batch_size) + print( + f"fused RNNT loss/wer enabled, fused_batch_size={args.fused_batch_size}", + flush=True, + ) + + if args.mode == "decoder": + for p in model.encoder.parameters(): + p.requires_grad = False + + tr = sum(p.numel() for p in model.parameters() if p.requires_grad) + tot = sum(p.numel() for p in model.parameters()) + print( + f"mode={args.mode} trainable={tr/1e6:.1f}M / {tot/1e6:.1f}M ({100*tr/tot:.1f}%)", + flush=True, + ) + + trainer = pl.Trainer( + max_steps=args.steps, + accelerator="gpu", + devices=1, + precision="bf16-mixed", + enable_checkpointing=False, + logger=False, + enable_progress_bar=True, + log_every_n_steps=50, + limit_val_batches=0, + num_sanity_val_steps=0, + callbacks=[PeriodicSave(args.save_dir, args.save_every)], + ) + model.set_trainer(trainer) + + t0 = time.time() + trainer.fit(model) + final = Path(args.save_dir) / "final.nemo" + model.save_to(str(final)) + print(f"train done in {time.time()-t0:.0f}s -> {final}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/nemo/finetune/train_overfit.py b/extras/asr-services/providers/nemo/finetune/train_overfit.py new file mode 100644 index 00000000..2c64b704 --- /dev/null +++ b/extras/asr-services/providers/nemo/finetune/train_overfit.py @@ -0,0 +1,204 @@ +"""Stage 1 overfit smoke test for nemotron-3.5-asr-streaming (NeMo RNNT-prompt). + +Goal: prove the hardware + data pipeline + NeMo training loop end-to-end by +memorizing 1-2 CoSHE clips to WER -> 0. Mirrors the gemma4/qwen3 `sample7` stage, +but the mechanics are NeMo-native (Lhotse manifest dataloader + Lightning), not +HF PEFT. + +Model: EncDecRNNTBPEModelWithPrompt (NeMo main). 24-layer Conformer encoder +(chunked_limited streaming) -> RNNT decoder+joint (13087 BPE). Per-utterance +language comes from the manifest `target_lang` field via the model's +prompt_dictionary (hi-IN -> 6). + +Tuning surface (--mode): + full : unfreeze everything (most reliable for a 1-2 clip overfit) + decoder : freeze encoder, train decoder + joint (+ prompt) only (lighter) + +Trains, then evaluates IN-PROCESS (no save/reload) so a WER->0 result can't be a +load-path artifact -- the exact trap that cost 3 runs on gemma4. + +Usage (on the GPU box, NeMo-main venv): + python train_overfit.py --model nvidia/nemotron-3.5-asr-streaming-0.6b \ + --manifest /home/ft/data/overfit.json --steps 400 --lr 1e-4 --mode full +""" + +import argparse +import json +import time + +import jiwer +import lightning.pytorch as pl # NeMo 2.x uses the `lightning` namespace +import nemo.collections.asr as nemo_asr +import torch +from nemo.collections.asr.data.audio_to_text_lhotse_prompt_index import ( + LhotseSpeechToTextBpeDatasetWithPromptIndex, +) +from omegaconf import open_dict + + +def build_train_cfg(model, manifest: str, batch_size: int, max_dur: float): + """Clone the model's train_ds cfg and point it at our tiny non-tarred manifest.""" + cfg = model.cfg.train_ds + with open_dict(cfg): + cfg.manifest_filepath = manifest + cfg.is_tarred = False + cfg.tarred_audio_filepaths = None + cfg.shard_manifests = False + cfg.use_lhotse = True + cfg.shuffle = True + cfg.batch_size = batch_size + # drop dynamic bucketing/batch_duration so a handful of clips form fixed batches + cfg.batch_duration = None + cfg.bucketing_batch_size = None + cfg.num_buckets = 0 + cfg.max_duration = max_dur # CoSHE clips run ~54s; default cap is 20 + cfg.min_duration = 0.1 + cfg.num_workers = 2 + cfg.prompt_field = "target_lang" # already the default, kept explicit + return cfg + + +def apply_freeze(model, mode: str): + if mode == "full": + for p in model.parameters(): + p.requires_grad = True + return + if mode == "decoder": + model.encoder.freeze() + for p in model.encoder.parameters(): + p.requires_grad = False + for mod in (model.decoder, model.joint): + for p in mod.parameters(): + p.requires_grad = True + return + raise SystemExit(f"unknown --mode {mode}") + + +def n_trainable(model) -> tuple[int, int]: + tot = sum(p.numel() for p in model.parameters()) + tr = sum(p.numel() for p in model.parameters() if p.requires_grad) + return tr, tot + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", default="nvidia/nemotron-3.5-asr-streaming-0.6b") + ap.add_argument("--manifest", required=True) + ap.add_argument("--steps", type=int, default=400) + ap.add_argument("--lr", type=float, default=1e-4) + ap.add_argument("--batch-size", type=int, default=2) + ap.add_argument( + "--fused-batch-size", + type=int, + default=1, + help="RNNT joint fused loss/wer sub-batch; keeps the [B,T,U,V] " + "joint off-GPU on long CoSHE clips (0 disables fusion)", + ) + ap.add_argument("--max-dur", type=float, default=65.0) + ap.add_argument("--mode", choices=["full", "decoder"], default="full") + ap.add_argument("--target-lang", default="hi-IN") + ap.add_argument("--save", default=None, help="optional .nemo save path") + args = ap.parse_args() + + torch.set_float32_matmul_precision("high") + rows = [json.loads(l) for l in open(args.manifest)] + print( + f"Overfit on {len(rows)} clip(s): " f"{[r['audio_file_name'] for r in rows]}", + flush=True, + ) + + model = nemo_asr.models.ASRModel.from_pretrained(model_name=args.model) + model.set_trainer(None) + + # constant LR, no scheduler -> pure overfit signal (pretrain default lr=0.5/Noam). + # NeMo's setup_optimization does `optim['sched']['max_steps']=...` whenever a + # `sched` key is present, so the key must be REMOVED, not set to None. + with open_dict(model.cfg.optim): + model.cfg.optim.name = "adamw" + model.cfg.optim.lr = args.lr + model.cfg.optim.weight_decay = 0.0 + model.cfg.optim.pop("sched", None) + + model.setup_training_data( + build_train_cfg(model, args.manifest, args.batch_size, args.max_dur) + ) + + # RNN-T loss on long CoSHE clips (~59s -> T~737 frames, ~480 BPE tokens) would + # materialize a [B,T,U,13087] joint (~34 GB) and OOM a 40 GB GPU. Fuse loss+WER + # into the joint so it's computed in `fused_batch_size` sub-batches without ever + # holding the full tensor. + if args.fused_batch_size > 0 and hasattr(model.joint, "set_fuse_loss_wer"): + model.joint.set_fuse_loss_wer(True, loss=model.loss, metric=model.wer) + model.joint.set_fused_batch_size(args.fused_batch_size) + print( + f"fused RNNT loss/wer enabled, fused_batch_size={args.fused_batch_size}", + flush=True, + ) + + apply_freeze(model, args.mode) + tr, tot = n_trainable(model) + print( + f"mode={args.mode} trainable={tr/1e6:.1f}M / {tot/1e6:.1f}M ({100*tr/tot:.1f}%)", + flush=True, + ) + + trainer = pl.Trainer( + max_steps=args.steps, + accelerator="gpu", + devices=1, + precision="bf16-mixed", + enable_checkpointing=False, + logger=False, + enable_progress_bar=True, + log_every_n_steps=10, + limit_val_batches=0, + num_sanity_val_steps=0, + ) + model.set_trainer(trainer) + + t0 = time.time() + trainer.fit(model) + print(f"train done in {time.time()-t0:.0f}s", flush=True) + + # ---- in-process eval (no reload) ---- + model.eval() + wavs = [r["audio_filepath"] for r in rows] + refs = [r["text"] for r in rows] + + # transcribe() builds cuts with supervision.language=None. The prompt-index + # dataset's default "unified" mode reads that None language -> "Unknown prompt + # key: 'None'". Force every eval cut to our training target_lang index so the + # prompt matches training and no None lookup happens. + # num_workers=0 keeps the dataset in-process so this class override actually + # applies (worker subprocesses wouldn't see an in-process monkeypatch). + LhotseSpeechToTextBpeDatasetWithPromptIndex._get_prompt_index_for_cut = ( + lambda self, cut, _tl=args.target_lang: self._get_prompt_index(_tl) + ) + + with torch.no_grad(): + hyps = model.transcribe( + wavs, + batch_size=1, + target_lang=args.target_lang, + num_workers=0, + verbose=False, + ) + hyps = [h.text if hasattr(h, "text") else str(h) for h in hyps] + + for r, h in zip(rows, hyps): + w = jiwer.wer(r["text"], h) + print(f"\n[{r['audio_file_name']}] WER={w*100:.2f}%", flush=True) + print(f" REF: {r['text'][:160]}", flush=True) + print(f" HYP: {h[:160]}", flush=True) + corpus = jiwer.wer(refs, hyps) + print( + f"\n=== OVERFIT corpus WER = {corpus*100:.2f}% (target: ~0%) ===", flush=True + ) + + if args.save: + model.save_to(args.save) + print(f"saved -> {args.save}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/nemo/finetune/train_split.py b/extras/asr-services/providers/nemo/finetune/train_split.py new file mode 100644 index 00000000..d0c53c0a --- /dev/null +++ b/extras/asr-services/providers/nemo/finetune/train_split.py @@ -0,0 +1,148 @@ +"""Stage 3 generalization for nemotron-3.5-asr-streaming (NeMo RNNT-prompt). + +Fine-tune on a 20% split of the CoSHE SEGMENTS (leakage-free by source clip), with +val-loss early stopping, then eval the held-out test split vs base — does FT generalize? +Mirrors the gemma4 20% experiment ([[gemma4_coshe_20pct_generalization]]): very low LR, +EarlyStopping(patience), keep the best-val checkpoint. + +Unlike train_full.py (overfit, no val), this sets up validation_data + compute_eval_wer, +saves the best-val .nemo, and stops when val_wer stops improving. + +Usage: + python train_split.py --train s3_train.json --val s3_val.json \ + --lr 1e-5 --batch-size 8 --max-epochs 40 --patience 5 \ + --save-dir /home/ft/out/s3 --mode decoder +""" + +import argparse +import time +from pathlib import Path + +import lightning.pytorch as pl +import nemo.collections.asr as nemo_asr +import torch +from lightning.pytorch.callbacks import Callback, EarlyStopping +from omegaconf import open_dict + + +class BestNemoSaver(Callback): + """Save a .nemo whenever val_wer improves (NeMo .nemo, not a Lightning ckpt).""" + + def __init__(self, path: str, monitor: str = "val_wer"): + self.path = path + self.monitor = monitor + self.best = float("inf") + + def on_validation_end(self, trainer, pl_module): + v = trainer.callback_metrics.get(self.monitor) + if v is None: + return + v = float(v) + if v < self.best: + self.best = v + pl_module.save_to(self.path) + print(f"[best] {self.monitor}={v:.4f} -> saved {self.path}", flush=True) + + +def ds_cfg(model, manifest, batch_size, max_dur, shuffle): + cfg = model.cfg.train_ds + with open_dict(cfg): + cfg.manifest_filepath = manifest + cfg.is_tarred = False + cfg.tarred_audio_filepaths = None + cfg.shard_manifests = False + cfg.use_lhotse = True + cfg.shuffle = shuffle + cfg.batch_size = batch_size + cfg.batch_duration = None + cfg.bucketing_batch_size = None + cfg.num_buckets = 0 + cfg.max_duration = max_dur + cfg.min_duration = 0.1 + cfg.num_workers = 4 + cfg.prompt_field = "target_lang" + return cfg + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", default="nvidia/nemotron-3.5-asr-streaming-0.6b") + ap.add_argument("--train", required=True) + ap.add_argument("--val", required=True) + ap.add_argument("--lr", type=float, default=1e-5) + ap.add_argument("--batch-size", type=int, default=8) + ap.add_argument("--fused-batch-size", type=int, default=1) + ap.add_argument("--max-dur", type=float, default=30.0) + ap.add_argument("--mode", choices=["full", "decoder"], default="decoder") + ap.add_argument("--max-epochs", type=int, default=40) + ap.add_argument("--patience", type=int, default=5) + ap.add_argument("--save-dir", required=True) + args = ap.parse_args() + + torch.set_float32_matmul_precision("high") + model = nemo_asr.models.ASRModel.from_pretrained(model_name=args.model) + model.set_trainer(None) + + # enable validation loss (off by default in the pretrain config) + with open_dict(model.cfg): + model.cfg.compute_eval_wer = True + with open_dict(model.cfg.optim): + model.cfg.optim.name = "adamw" + model.cfg.optim.lr = args.lr + model.cfg.optim.weight_decay = 0.0 + model.cfg.optim.pop("sched", None) + + model.setup_training_data( + ds_cfg(model, args.train, args.batch_size, args.max_dur, True) + ) + model.setup_validation_data( + ds_cfg(model, args.val, args.batch_size, args.max_dur, False) + ) + + if args.fused_batch_size > 0 and hasattr(model.joint, "set_fuse_loss_wer"): + model.joint.set_fuse_loss_wer(True, loss=model.loss, metric=model.wer) + model.joint.set_fused_batch_size(args.fused_batch_size) + + if args.mode == "decoder": + for p in model.encoder.parameters(): + p.requires_grad = False + + tr = sum(p.numel() for p in model.parameters() if p.requires_grad) + tot = sum(p.numel() for p in model.parameters()) + print( + f"mode={args.mode} lr={args.lr} trainable={tr/1e6:.1f}M/{tot/1e6:.1f}M", + flush=True, + ) + + save_dir = Path(args.save_dir) + save_dir.mkdir(parents=True, exist_ok=True) + best_path = str(save_dir / "best.nemo") + + trainer = pl.Trainer( + max_epochs=args.max_epochs, + accelerator="gpu", + devices=1, + precision="bf16-mixed", + enable_checkpointing=False, + logger=False, + enable_progress_bar=True, + log_every_n_steps=50, + num_sanity_val_steps=0, + check_val_every_n_epoch=1, + limit_val_batches=120, # ~960 val segs — fast, stable early-stop signal (full set decode/epoch is slow) + callbacks=[ + EarlyStopping( + monitor="val_wer", patience=args.patience, mode="min", verbose=True + ), + BestNemoSaver(best_path), + ], + ) + model.set_trainer(trainer) + + t0 = time.time() + trainer.fit(model) + print(f"train done in {time.time()-t0:.0f}s -> best {best_path}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/qwen3_asr/finetune/README.md b/extras/asr-services/providers/qwen3_asr/finetune/README.md new file mode 100644 index 00000000..46b5a58d --- /dev/null +++ b/extras/asr-services/providers/qwen3_asr/finetune/README.md @@ -0,0 +1,94 @@ +# Qwen3-ASR LoRA fine-tuning on CoSHE (Hinglish) + +LoRA fine-tuning of `Qwen/Qwen3-ASR-0.6B` (AuT audio encoder + Qwen3 text decoder, built on +Qwen3-Omni) on the **CoSHE-Eval** Hinglish dataset (1985 clips, mixed Devanagari+Latin +code-switched transcripts). Mirrors the gemma4 finetune dir layout but uses Qwen's official SFT +pipeline + a gemma4-style anneal/early-stop loop. + +## Result — overfit goal achieved + +**Full-CoSHE corpus WER = 1.18%** (1355/1985 clips reproduced exactly) — overfit to <2% on all +1985 clips. The sample7 smoke test reproduces 7/7 clips exactly (0.00% WER). + +Qwen3-ASR is multimodal (audio→text) and natively produces **mixed Devanagari+Latin Hinglish** +(e.g. `complain कर लो technology evolve…`), with intra-sentence code-switching — confirmed +empirically on CoSHE. + +## Files + +- `coshe_infer.py` — transcribe CoSHE parquet shards with any Qwen3-ASR model → mlexp-schema JSONL + (used for the base-model + srota baselines). **Use `--max_new_tokens 2048`**: the 512/1024 + default truncates CoSHE's long transcripts and inflates WER (the deletion-heavy artifact seen on srota). +- `make_manifest.py` / `make_manifest_full.py` — build `train.jsonl` (sample7) / `train_full.jsonl` + (all 1985) in the official SFT schema: `{"audio": "/abs.wav", "text": "language None<asr_text>" + transcript}`. +- `introspect.py` — dump module names; LoRA targets the Qwen3 decoder `thinker.model.layers.*` + (q/k/v/o/gate/up/down), keeping off the frozen AuT `thinker.audio_tower`. +- `qwen3_asr_sft.py` — vendored official trainer (`QwenLM/Qwen3-ASR/finetuning`); reused for its + data/collator/`patch_outer_forward` (model.forward → `thinker.forward`) + label masking. +- `train_qwen3_until_wer.py` — LoRA training with gemma4-style anneal + loss-gated WER<target + early-stop (subset gate → full-1985 confirm). Reuses the sft collator. +- `verify.py` — load base + adapter, transcribe clips, report char-sim / exact / WER. + +## Recipe (what worked, and the gotchas) + +Runs on the Jarvis Labs A100 (torch 2.11/cu130; `pip install qwen-asr datasets jiwer peft`; note +qwen-asr pins transformers 4.57). Single GPU. + +1. **Smoke test (sample7, 7 clips):** `train_qwen3_until_wer.py --lora_r 16 --lr 2e-4 --epochs 50` + → loss → 0.0001, **7/7 exact** once eval uses `--eval_max_new_tokens 2048` (512 truncates → false + 12.6% WER). +2. **Full CoSHE (1985):** + - **Capacity matters.** Decoder-only LoRA (even r256) **plateaus at loss ~3.0** — it can't + memorize 1985 long transcripts with a frozen output layer. `--include_head` (LoRA also on + `lm_head` + `embed_tokens`, 239.8M params / 23.5%) breaks the wall → loss descends to ~0.08. + (`Qwen3ASRForConditionalGeneration` doesn't expose `get_input_embeddings`; the script patches + it to delegate to `thinker.model.embed_tokens` so PEFT can adapt the embeddings.) + - **Warm then anneal (the gemma4 lesson).** Constant `lr 2e-4` drives loss to ~0.08 but WER + **plateaus/bounces 11–22%** (memorizes some clips, too hot to settle the rest). Restart from + the plateaued checkpoint with a **fresh optimizer at `lr 3e-5`** (`--init_adapter <ckpt>`) → + loss collapses and **WER → 1.18% in a single anneal epoch.** + - **Eval OOM.** The in-training WER eval (`wrapper.transcribe`) competes with training memory; + use a small `--eval_chunk` (4–16) + `gc.collect()/empty_cache()` before eval, and + `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`. + +## Run + +```bash +# build full manifest (decodes 1985 clips → WAVs) +python3 make_manifest_full.py + +# warm phase +python3 train_qwen3_until_wer.py --model Qwen/Qwen3-ASR-0.6B \ + --train_file train_full.jsonl --output_dir out/coshe-full-head \ + --lr 2e-4 --batch_size 2 --lora_r 256 --lora_alpha 512 --include_head \ + --eval_every 2 --eval_subset 200 --eval_chunk 16 --wer_target 0.02 --eval_max_new_tokens 2048 + +# anneal from the plateaued checkpoint (fresh optimizer, lower LR) until WER < 2% +python3 train_qwen3_until_wer.py --model Qwen/Qwen3-ASR-0.6B \ + --train_file train_full.jsonl --output_dir out/coshe-anneal \ + --init_adapter out/coshe-full-head/checkpoint-<latest> \ + --lr 3e-5 --batch_size 2 --eval_every 1 --eval_subset 200 --eval_chunk 16 --wer_target 0.02 + +# verify +python3 verify.py --adapter out/coshe-anneal --train_file train_full.jsonl --max_new_tokens 2048 +``` + +## CoSHE WER comparison (IndicXlit-romanized, via mlexp `score_coshe`, 500 common clips) + +| model | corpus WER | note | +|-------|-----------|------| +| Qwen3-ASR-0.6B + CoSHE LoRA (this) | **1.18%** | overfit goal (training-callback jiwer, full 1985) | +| Qwen3-ASR-0.6B (base, @2048 tok) | **25.16%** | base 0.6B — competitive with the 4B gemma4 | +| gemma4-E4B (base) | 26.23% | 4B base model | +| srota = qwen3-asr-0.6b-hinglish (out-of-domain FT) | 48.85% | trained on HiACC+OpenSLR, not CoSHE | + +Findings: +- **The 1.18% is a memorization/overfit result** on the training set — it shows the model + LoRA + recipe can fully fit CoSHE, not a generalization number. +- **Base Qwen3-ASR-0.6B (25.2%) ≈ gemma4-E4B (26.2%)** on CoSHE despite being ~7× smaller — a strong + base model. (It has occasional greedy repetition-collapse: per-sample WER max 621%, so corpus WER + has a heavy tail; a repetition guard would help.) +- **srota (the community Hinglish finetune) is *worse* than base (48.8% vs 25.2%)** on CoSHE — the + out-of-domain SFT on HiACC/OpenSLR hurt generalization to this distribution (plus repetition-collapse + and, in the original 1024-token run, truncation). A cautionary data point: a same-language finetune + isn't automatically better on a different corpus. diff --git a/extras/asr-services/providers/qwen3_asr/finetune/coshe_infer.py b/extras/asr-services/providers/qwen3_asr/finetune/coshe_infer.py new file mode 100644 index 00000000..97682b8c --- /dev/null +++ b/extras/asr-services/providers/qwen3_asr/finetune/coshe_infer.py @@ -0,0 +1,168 @@ +"""Transcribe the CoSHE-Eval dataset with a Qwen3-ASR model and emit mlexp-schema JSONL. + +Runs on the A100 (qwen-asr installed). Reads the ``eval-*.parquet`` shards, decodes each +row's audio bytes, batch-transcribes with ``Qwen3ASRModel`` (``language=None`` → language- +agnostic), and appends one JSONL line per clip in the exact schema mlexp's ``score_coshe`` +expects: + + {"audio_file_name", "transcription" (ground truth), "hyp", "asr_seconds", + "duration_s", "raw"} + +Resumable: rows whose ``audio_file_name`` already has a non-error line in the output are +skipped (mirrors ``mlexp.runners.dataset.load_done``). WER scoring (IndicXlit romanization + +jiwer) happens locally back in extras/ml-experiments, not here — this script only produces +hypotheses. +""" + +import argparse +import io +import json +import time +from glob import glob +from pathlib import Path + +import numpy as np +import pyarrow.parquet as pq +import soundfile as sf +from qwen_asr import Qwen3ASRModel + + +def load_done(out_path: Path) -> set[str]: + done: set[str] = set() + if not out_path.exists(): + return done + with open(out_path) as f: + for line in f: + try: + rec = json.loads(line) + except Exception: + continue + if "error" in rec: + continue + name = rec.get("audio_file_name") + if name: + done.add(name) + return done + + +def decode_audio(raw_bytes: bytes) -> tuple[np.ndarray, int]: + """WAV/audio bytes -> (float32 mono, sr). qwen_asr handles resampling internally.""" + data, sr = sf.read(io.BytesIO(raw_bytes), dtype="float32") + if data.ndim > 1: + data = data.mean(axis=1) + return np.ascontiguousarray(data), int(sr) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument( + "--model", required=True, help="HF id or local path of the Qwen3-ASR model" + ) + p.add_argument("--out", required=True, help="output JSONL path") + p.add_argument("--parquet_glob", default="/home/coshe-data/data/eval-*.parquet") + p.add_argument( + "--limit", type=int, default=0, help="cap #clips (shard order); 0 = all" + ) + p.add_argument("--batch_size", type=int, default=16) + p.add_argument("--max_new_tokens", type=int, default=1024) + args = p.parse_args() + + out_path = Path(args.out) + out_path.parent.mkdir(parents=True, exist_ok=True) + done = load_done(out_path) + print(f"Model: {args.model}", flush=True) + print(f"Out: {out_path} (resuming, {len(done)} already done)", flush=True) + + import torch + + model = Qwen3ASRModel.from_pretrained( + args.model, + dtype=torch.bfloat16, + device_map="cuda:0", + max_new_tokens=args.max_new_tokens, + max_inference_batch_size=args.batch_size, + ) + + shards = sorted(glob(args.parquet_glob)) + n_ok = n_err = n_seen = 0 + t_start = time.time() + + def flush_batch(batch, out_f): + nonlocal n_ok, n_err + if not batch: + return + audios = [b["audio"] for b in batch] + try: + t0 = time.time() + results = model.transcribe(audios, language=[None] * len(audios)) + dt = (time.time() - t0) / len(audios) + except Exception as e: # batch-level failure -> record per-clip errors + for b in batch: + out_f.write( + json.dumps( + {"audio_file_name": b["name"], "error": str(e)}, + ensure_ascii=False, + ) + + "\n" + ) + n_err += 1 + out_f.flush() + return + for b, res in zip(batch, results): + rec = { + "audio_file_name": b["name"], + "transcription": b["gt"], + "hyp": res.text, + "asr_seconds": round(dt, 2), + "duration_s": b["duration_s"], + "raw": {"text": res.text, "language": getattr(res, "language", None)}, + } + out_f.write(json.dumps(rec, ensure_ascii=False) + "\n") + n_ok += 1 + out_f.flush() + + with open(out_path, "a") as out_f: + batch = [] + for shard in shards: + pf = pq.ParquetFile(shard) + for rb in pf.iter_batches(batch_size=64): + for row in rb.to_pylist(): + if args.limit and n_seen >= args.limit: + break + name = row["audio_file_name"] + n_seen += 1 + if name in done: + continue + audio, sr = decode_audio(row["audio"]["bytes"]) + batch.append( + { + "name": name, + "gt": row["transcription"], + "audio": (audio, sr), + "duration_s": round(len(audio) / sr, 2), + } + ) + if len(batch) >= args.batch_size: + flush_batch(batch, out_f) + batch = [] + if n_ok % 64 == 0: + el = time.time() - t_start + print( + f" done={n_ok} err={n_err} seen={n_seen} " + f"({el:.0f}s, {n_ok / max(el, 1):.2f} clip/s)", + flush=True, + ) + if args.limit and n_seen >= args.limit: + break + if args.limit and n_seen >= args.limit: + break + flush_batch(batch, out_f) + + print( + f"DONE: ok={n_ok} err={n_err} seen={n_seen} elapsed={time.time() - t_start:.0f}s", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/qwen3_asr/finetune/introspect.py b/extras/asr-services/providers/qwen3_asr/finetune/introspect.py new file mode 100644 index 00000000..0c1eade4 --- /dev/null +++ b/extras/asr-services/providers/qwen3_asr/finetune/introspect.py @@ -0,0 +1,62 @@ +"""Dump Qwen3-ASR module structure to choose LoRA target modules (runs on the A100). + +Prints the top-level layout of ``model.thinker`` and the unique nn.Linear *leaf names*, +split into "audio/encoder" vs "decoder/text" so we can scope the LoRA target regex to the +Qwen3 text decoder and keep it off the AuT audio encoder (mirrors the gemma4 approach of +adapting the text decoder only). + +Loads on CPU (``device_map=None``, no .cuda()) so it can run while the GPU is busy. +""" + +import argparse +import re + +import torch +from qwen_asr import Qwen3ASRModel + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="Qwen/Qwen3-ASR-0.6B") + args = p.parse_args() + + wrapper = Qwen3ASRModel.from_pretrained( + args.model, dtype=torch.bfloat16, device_map=None + ) + model = wrapper.model + print("=== model class ===", model.__class__.__name__) + print("=== thinker children ===") + for n, _ in model.thinker.named_children(): + print(" ", n) + + # Collect every nn.Linear's full path; classify by which subtree it lives in. + audio_kw = ("audio", "encoder", "conv", "merger", "proj_audio") + lin_paths = [] + for name, mod in model.named_modules(): + if mod.__class__.__name__.endswith("Linear"): + lin_paths.append(name) + + def is_audio(path: str) -> bool: + return any(k in path.lower() for k in audio_kw) + + decoder = sorted({p.split(".")[-1] for p in lin_paths if not is_audio(p)}) + audio = sorted({p.split(".")[-1] for p in lin_paths if is_audio(p)}) + print("\n=== Linear leaf names: DECODER/text subtree ===", decoder) + print("=== Linear leaf names: AUDIO/encoder subtree ===", audio) + + # Show a few representative full paths so the regex can be anchored correctly. + print("\n=== sample decoder Linear paths ===") + for p_ in [x for x in lin_paths if not is_audio(x)][:8]: + print(" ", p_) + print("=== sample audio Linear paths ===") + for p_ in [x for x in lin_paths if is_audio(x)][:8]: + print(" ", p_) + + # Candidate regex (q/k/v/o/gate/up/down on the decoder layers, excluding audio). + cand = r"^(?!.*(audio|encoder)).*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)$" + n_match = sum(1 for p_ in lin_paths if re.search(cand, p_)) + print(f"\n=== candidate regex matches {n_match} Linear modules ===\n {cand}") + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/qwen3_asr/finetune/make_manifest.py b/extras/asr-services/providers/qwen3_asr/finetune/make_manifest.py new file mode 100644 index 00000000..8a826f73 --- /dev/null +++ b/extras/asr-services/providers/qwen3_asr/finetune/make_manifest.py @@ -0,0 +1,80 @@ +"""Build the sample7 overfit train.jsonl for Qwen3-ASR SFT (runs on the A100). + +Extracts the 7 CoSHE sample7 clips from the parquet shards by ``audio_file_name``, writes +each as a 16 kHz mono WAV, and emits a JSONL in the schema the official ``qwen3_asr_sft.py`` +expects: + + {"audio": "/abs/path.wav", "text": "language None<asr_text>" + <ground-truth transcript>} + +The ``language None`` prefix matches how we evaluate (``transcribe(language=None)``); WER is +computed on the parsed ``<asr_text>`` payload, so the language token itself is irrelevant. +""" + +import argparse +import io +import json +from glob import glob +from pathlib import Path + +import numpy as np +import pyarrow.parquet as pq +import soundfile as sf + +SAMPLE7 = [ + "audio_13.wav", + "audio_278.wav", + "audio_1320.wav", + "audio_58.wav", + "audio_1197.wav", + "audio_1149.wav", + "audio_1059.wav", +] + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--parquet_glob", default="/home/coshe-data/data/eval-*.parquet") + p.add_argument("--audio_dir", default="/home/qwen3ft/sample7_audio") + p.add_argument("--out", default="/home/qwen3ft/train.jsonl") + p.add_argument("--prefix", default="language None<asr_text>") + args = p.parse_args() + + audio_dir = Path(args.audio_dir) + audio_dir.mkdir(parents=True, exist_ok=True) + want = set(SAMPLE7) + found: dict[str, dict] = {} + + for shard in sorted(glob(args.parquet_glob)): + if len(found) == len(want): + break + pf = pq.ParquetFile(shard) + for rb in pf.iter_batches(batch_size=64): + for row in rb.to_pylist(): + name = row["audio_file_name"] + if name not in want or name in found: + continue + data, sr = sf.read(io.BytesIO(row["audio"]["bytes"]), dtype="float32") + if data.ndim > 1: + data = data.mean(axis=1) + wav_path = audio_dir / name + sf.write(str(wav_path), np.ascontiguousarray(data), sr) + found[name] = { + "audio": str(wav_path), + "text": args.prefix + row["transcription"], + } + if len(found) == len(want): + break + + missing = want - set(found) + if missing: + raise SystemExit(f"Missing clips not found in parquet: {sorted(missing)}") + + with open(args.out, "w") as f: + for name in SAMPLE7: + f.write(json.dumps(found[name], ensure_ascii=False) + "\n") + print(f"Wrote {len(found)} examples -> {args.out}") + print(f"Audio in {audio_dir}") + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/qwen3_asr/finetune/make_manifest_full.py b/extras/asr-services/providers/qwen3_asr/finetune/make_manifest_full.py new file mode 100644 index 00000000..0a2b38d8 --- /dev/null +++ b/extras/asr-services/providers/qwen3_asr/finetune/make_manifest_full.py @@ -0,0 +1,70 @@ +"""Build the FULL CoSHE train.jsonl for the <2% WER overfit (runs on the A100). + +Decodes every clip in the parquet shards to a 16 kHz mono WAV under ``--audio_dir`` and emits +a JSONL in the official ``qwen3_asr_sft.py`` schema: + + {"audio": "/abs/path.wav", "text": "language None<asr_text>" + <ground-truth transcript>} + +~1985 clips → a few GB of WAVs. Resumable: skips clips whose WAV already exists. Mirrors the +gemma4 CosheParquetDataset coverage (all shards, full transcripts). +""" + +import argparse +import io +import json +from glob import glob +from pathlib import Path + +import numpy as np +import pyarrow.parquet as pq +import soundfile as sf + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--parquet_glob", default="/home/coshe-data/data/eval-*.parquet") + p.add_argument("--audio_dir", default="/home/qwen3ft/coshe_audio") + p.add_argument("--out", default="/home/qwen3ft/train_full.jsonl") + p.add_argument("--prefix", default="language None<asr_text>") + p.add_argument("--limit", type=int, default=0, help="cap #clips (0 = all)") + args = p.parse_args() + + audio_dir = Path(args.audio_dir) + audio_dir.mkdir(parents=True, exist_ok=True) + n = 0 + with open(args.out, "w") as out_f: + for shard in sorted(glob(args.parquet_glob)): + pf = pq.ParquetFile(shard) + for rb in pf.iter_batches(batch_size=64): + for row in rb.to_pylist(): + if args.limit and n >= args.limit: + break + name = row["audio_file_name"] + wav_path = audio_dir / name + if not wav_path.exists(): + data, sr = sf.read( + io.BytesIO(row["audio"]["bytes"]), dtype="float32" + ) + if data.ndim > 1: + data = data.mean(axis=1) + sf.write(str(wav_path), np.ascontiguousarray(data), sr) + out_f.write( + json.dumps( + { + "audio": str(wav_path), + "text": args.prefix + row["transcription"], + }, + ensure_ascii=False, + ) + + "\n" + ) + n += 1 + if args.limit and n >= args.limit: + break + if args.limit and n >= args.limit: + break + print(f"Wrote {n} examples -> {args.out}; audio in {audio_dir}") + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/qwen3_asr/finetune/qwen3_asr_sft.py b/extras/asr-services/providers/qwen3_asr/finetune/qwen3_asr_sft.py new file mode 100644 index 00000000..fc953131 --- /dev/null +++ b/extras/asr-services/providers/qwen3_asr/finetune/qwen3_asr_sft.py @@ -0,0 +1,334 @@ +# coding=utf-8 +# Copyright 2026 The Alibaba Qwen team. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import argparse +import os +import re +import shutil +from dataclasses import dataclass +from typing import Any, Dict, List, Optional + +import librosa +import torch +from datasets import load_dataset +from qwen_asr import Qwen3ASRModel +from transformers import GenerationConfig, Trainer, TrainerCallback, TrainingArguments + + +def patch_outer_forward(model): + cls = model.__class__ + if getattr(cls, "_forward_patched", False): + return + + if not hasattr(model, "thinker") or not hasattr(model.thinker, "forward"): + raise RuntimeError( + "Cannot patch forward: model has no `.thinker.forward`. " + "Your qwen3_asr model may be incompatible." + ) + + def forward( + self, + input_ids=None, + attention_mask=None, + input_features=None, + feature_attention_mask=None, + labels=None, + **kwargs, + ): + return self.thinker.forward( + input_ids=input_ids, + attention_mask=attention_mask, + input_features=input_features, + feature_attention_mask=feature_attention_mask, + labels=labels, + **kwargs, + ) + + cls.forward = forward + cls._forward_patched = True + + +_CKPT_RE = re.compile(r"^checkpoint-(\d+)$") + + +def find_latest_checkpoint(output_dir: str) -> Optional[str]: + if not output_dir or not os.path.isdir(output_dir): + return None + best_step = None + best_path = None + for name in os.listdir(output_dir): + m = _CKPT_RE.match(name) + if not m: + continue + step = int(m.group(1)) + path = os.path.join(output_dir, name) + if os.path.isdir(path) and (best_step is None or step > best_step): + best_step = step + best_path = path + return best_path + + +def load_audio(path: str, sr: int = 16000): + wav, _ = librosa.load(path, sr=sr, mono=True) + return wav + + +def build_prefix_messages(prompt: str, audio_array): + return [ + {"role": "system", "content": prompt or ""}, + {"role": "user", "content": [{"type": "audio", "audio": audio_array}]}, + ] + + +def make_preprocess_fn_prefix_only(processor): + def _preprocess(ex: Dict[str, Any]) -> Dict[str, Any]: + prompt = ex.get("prompt", "") + dummy_audio = None + prefix_msgs = build_prefix_messages(prompt, dummy_audio) + prefix_text = processor.apply_chat_template( + [prefix_msgs], add_generation_prompt=True, tokenize=False + )[0] + return { + "prompt": prompt, + "audio": ex["audio"], + "target": ex["text"], + "prefix_text": prefix_text, + } + + return _preprocess + + +@dataclass +class DataCollatorForQwen3ASRFinetuning: + processor: Any + sampling_rate: int = 16000 + + def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]: + audio_paths = [f["audio"] for f in features] + prefix_texts = [f["prefix_text"] for f in features] + targets = [f["target"] for f in features] + + eos = self.processor.tokenizer.eos_token or "" + full_texts = [pfx + tgt + eos for pfx, tgt in zip(prefix_texts, targets)] + audios = [load_audio(p, sr=self.sampling_rate) for p in audio_paths] + + full_inputs = self.processor( + text=full_texts, + audio=audios, + return_tensors="pt", + padding=True, + truncation=False, + ) + prefix_inputs = self.processor( + text=prefix_texts, + audio=audios, + return_tensors="pt", + padding=True, + truncation=False, + ) + + prefix_lens = prefix_inputs["attention_mask"].sum(dim=1).tolist() + labels = full_inputs["input_ids"].clone() + for i, pl in enumerate(prefix_lens): + labels[i, :pl] = -100 + + pad_id = self.processor.tokenizer.pad_token_id + if pad_id is not None: + labels[labels == pad_id] = -100 + + full_inputs["labels"] = labels + return full_inputs + + +class CastFloatInputsTrainer(Trainer): + def _prepare_inputs(self, inputs): + inputs = super()._prepare_inputs(inputs) + model_dtype = getattr(self.model, "dtype", None) + if model_dtype is not None: + for k, v in list(inputs.items()): + if torch.is_tensor(v) and v.is_floating_point(): + inputs[k] = v.to(dtype=model_dtype) + return inputs + + +def copy_required_hf_files_for_qwen_asr(src_dir: str, dst_dir: str): + os.makedirs(dst_dir, exist_ok=True) + required = [ + "config.json", + "generation_config.json", + "preprocessor_config.json", + "processor_config.json", + "tokenizer_config.json", + "tokenizer.json", + "special_tokens_map.json", + "chat_template.json", + "merges.txt", + "vocab.json", + ] + for fn in required: + src = os.path.join(src_dir, fn) + if os.path.exists(src): + shutil.copy2(src, os.path.join(dst_dir, fn)) + + +class MakeEveryCheckpointInferableCallback(TrainerCallback): + def __init__(self, base_model_path: str): + self.base_model_path = base_model_path + + def on_save(self, args: TrainingArguments, state, control, **kwargs): + if args.process_index != 0: + return control + + ckpt_dir = os.path.join(args.output_dir, f"checkpoint-{state.global_step}") + if not os.path.isdir(ckpt_dir): + ckpt_dir = kwargs.get("checkpoint", ckpt_dir) + + copy_required_hf_files_for_qwen_asr(self.base_model_path, ckpt_dir) + return control + + +def parse_args(): + p = argparse.ArgumentParser("Qwen3-ASR Finetuning") + + # Paths + p.add_argument("--model_path", type=str, default="Qwen/Qwen3-ASR-1.7B") + p.add_argument("--train_file", type=str, default="train.jsonl") + p.add_argument("--eval_file", type=str, default="") + p.add_argument("--output_dir", type=str, default="./qwen3-asr-finetuning-out") + + # Audio + p.add_argument("--sr", type=int, default=16000) + + # Train hyper-params + p.add_argument("--batch_size", type=int, default=32) + p.add_argument("--grad_acc", type=int, default=4) + p.add_argument("--lr", type=float, default=2e-5) + p.add_argument("--epochs", type=float, default=1) + p.add_argument("--log_steps", type=int, default=10) + p.add_argument("--lr_scheduler_type", type=str, default="linear") + p.add_argument("--warmup_ratio", type=float, default=0.02) + + # DataLoader + p.add_argument("--num_workers", type=int, default=4) + p.add_argument("--pin_memory", type=int, default=1) + p.add_argument("--persistent_workers", type=int, default=1) + p.add_argument("--prefetch_factor", type=int, default=2) + + # Save + p.add_argument("--save_strategy", type=str, default="steps") + p.add_argument("--save_steps", type=int, default=200) + p.add_argument("--save_total_limit", type=int, default=5) + + # Resume + p.add_argument("--resume_from", type=str, default="") + p.add_argument("--resume", type=int, default=0) + + return p.parse_args() + + +def main(): + args_cli = parse_args() + + if not args_cli.train_file: + raise ValueError( + "TRAIN_FILE is required (json/jsonl). Needs fields: audio, text, optional prompt" + ) + + use_bf16 = torch.cuda.is_available() and torch.cuda.get_device_capability(0)[0] >= 8 + asr_wrapper = Qwen3ASRModel.from_pretrained( + args_cli.model_path, + dtype=torch.bfloat16 if use_bf16 else torch.float16, + device_map=None, + ) + model = asr_wrapper.model + processor = asr_wrapper.processor + + patch_outer_forward(model) + model.generation_config = GenerationConfig.from_model_config(model.config) + + raw_ds = load_dataset( + "json", + data_files={ + "train": args_cli.train_file, + **({"validation": args_cli.eval_file} if args_cli.eval_file else {}), + }, + ) + ds = raw_ds.map(make_preprocess_fn_prefix_only(processor), num_proc=1) + + keep = {"prompt", "audio", "target", "prefix_text"} + for split in ds.keys(): + drop = [c for c in ds[split].column_names if c not in keep] + if drop: + ds[split] = ds[split].remove_columns(drop) + + collator = DataCollatorForQwen3ASRFinetuning( + processor=processor, sampling_rate=args_cli.sr + ) + + training_args = TrainingArguments( + output_dir=args_cli.output_dir, + per_device_train_batch_size=args_cli.batch_size, + gradient_accumulation_steps=args_cli.grad_acc, + learning_rate=args_cli.lr, + num_train_epochs=args_cli.epochs, + logging_steps=args_cli.log_steps, + lr_scheduler_type=args_cli.lr_scheduler_type, + warmup_ratio=args_cli.warmup_ratio, + dataloader_num_workers=args_cli.num_workers, + dataloader_pin_memory=(args_cli.pin_memory == 1), + dataloader_persistent_workers=(args_cli.persistent_workers == 1), + dataloader_prefetch_factor=( + args_cli.prefetch_factor if args_cli.num_workers > 0 else None + ), + save_strategy=args_cli.save_strategy, + save_steps=args_cli.save_steps, + save_total_limit=args_cli.save_total_limit, + save_safetensors=True, + eval_strategy="steps", + eval_steps=args_cli.save_steps, + do_eval=bool(args_cli.eval_file), + bf16=use_bf16, + fp16=not use_bf16, + ddp_find_unused_parameters=False, + remove_unused_columns=False, + report_to="none", + ) + + trainer = CastFloatInputsTrainer( + model=model, + args=training_args, + train_dataset=ds["train"], + eval_dataset=ds.get("validation", None), + data_collator=collator, + tokenizer=processor.tokenizer, + callbacks=[ + MakeEveryCheckpointInferableCallback(base_model_path=args_cli.model_path) + ], + ) + + resume_from = (args_cli.resume_from or "").strip() + if not resume_from and args_cli.resume == 1: + resume_from = find_latest_checkpoint(training_args.output_dir) or "" + + if resume_from: + if trainer.args.process_index == 0: + print(f"[resume] resume_from_checkpoint = {resume_from}") + trainer.train(resume_from_checkpoint=resume_from) + else: + trainer.train() + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/qwen3_asr/finetune/train_qwen3_until_wer.py b/extras/asr-services/providers/qwen3_asr/finetune/train_qwen3_until_wer.py new file mode 100644 index 00000000..7c5c89e8 --- /dev/null +++ b/extras/asr-services/providers/qwen3_asr/finetune/train_qwen3_until_wer.py @@ -0,0 +1,314 @@ +"""LoRA overfit of Qwen3-ASR on CoSHE sample7, with gemma4's anneal + WER<target early-stop. + +Reuses the official ``qwen3_asr_sft.py`` data/forward/label pipeline (prefix-masked target, +audio-path collator, ``patch_outer_forward`` → ``thinker.forward``) but: + - trains **LoRA on the Qwen3 text decoder only** (``thinker.model.layers.*`` q/k/v/o/gate/up/down), + AuT audio encoder (``thinker.audio_tower``) frozen — mirrors the gemma4 text-decoder-only recipe; + - runs `lr_scheduler_type="constant"` and supports a fresh-optimizer **anneal restart** via + ``--init_adapter`` (warm at high LR, then re-launch from the checkpoint at a lower LR to settle); + - a loss-gated ``WERStopCallback`` evaluates corpus WER on the 7 clips each ``--eval_every`` epochs + and stops at ``--wer_target`` (default 2%). transformers 4.57 has no use_cache bug, so eval uses + normal cached generation via ``wrapper.transcribe``. +""" + +import argparse +import time + +import jiwer +import torch +from peft import LoraConfig, PeftModel, get_peft_model +from qwen3_asr_sft import ( + DataCollatorForQwen3ASRFinetuning, + find_latest_checkpoint, + make_preprocess_fn_prefix_only, + patch_outer_forward, +) +from transformers import GenerationConfig, Trainer, TrainerCallback, TrainingArguments + +# LoRA on the Qwen3 text decoder only (excludes thinker.audio_tower.* and lm_head). +LORA_TARGETS = r".*thinker\.model\.layers\..*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" +# With --include_head, also adapt the output head + input embeddings for more memorization +# capacity on the full dataset (gemma4 needed this to drive full-CoSHE WER below 2%). +LORA_TARGETS_HEAD = ( + r"(.*thinker\.model\.layers\..*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" + r"|.*thinker\.model\.embed_tokens|.*thinker\.lm_head)" +) + +_ASR_TAG = "<asr_text>" +_WER_NORM = jiwer.Compose( + [ + jiwer.ToLowerCase(), + jiwer.RemovePunctuation(), + jiwer.RemoveMultipleSpaces(), + jiwer.Strip(), + jiwer.ReduceToListOfListOfWords(), + ] +) + + +def strip_target(text: str) -> str: + """'language None<asr_text>foo' -> 'foo' (the payload we compare against).""" + return text.split(_ASR_TAG, 1)[1] if _ASR_TAG in text else text + + +def corpus_wer(refs, hyps): + return jiwer.wer( + refs, hyps, reference_transform=_WER_NORM, hypothesis_transform=_WER_NORM + ) + + +def char_sim(a: str, b: str) -> float: + from difflib import SequenceMatcher + + return SequenceMatcher(None, a, b).ratio() + + +@torch.inference_mode() +def eval_clips(wrapper, items, chunk=4): + """Transcribe clips (language=None), chunked, and return (wer, exact, refs, hyps).""" + refs = [strip_target(it["text"]) for it in items] + hyps = [] + for s in range(0, len(items), chunk): + paths = [it["audio"] for it in items[s : s + chunk]] + results = wrapper.transcribe(paths, language=[None] * len(paths)) + hyps.extend(r.text for r in results) + exact = sum(1 for r, h in zip(refs, hyps) if r.strip() == h.strip()) + return corpus_wer(refs, hyps), exact, refs, hyps + + +class WERStopCallback(TrainerCallback): + """Eval corpus WER each `every` epochs (loss-gated). Subset gate -> full confirm (gemma4 style).""" + + def __init__( + self, wrapper, items, every, target, loss_gate, eval_subset=0, chunk=4 + ): + self.wrapper, self.items = wrapper, items + self.every, self.target, self.loss_gate = every, target, loss_gate + self.chunk = chunk + self.subset = ( + items[:eval_subset] if eval_subset and eval_subset < len(items) else items + ) + + def _recent_loss(self, state): + vals = [r["loss"] for r in state.log_history if "loss" in r] + return sum(vals[-30:]) / len(vals[-30:]) if vals else None + + def on_epoch_end(self, args, state, control, model=None, **kwargs): + ep = int(round(state.epoch)) + if ep == 0 or ep % self.every != 0: + return control + rl = self._recent_loss(state) + if rl is not None and rl > self.loss_gate: + print( + f"[epoch {ep}] loss={rl:.3f} > gate {self.loss_gate}; skip WER eval", + flush=True, + ) + return control + was_training = model.training + model.eval() + import gc + + gc.collect() + torch.cuda.empty_cache() # free training caches before generation + t = time.time() + wer, exact, _, _ = eval_clips(self.wrapper, self.subset, chunk=self.chunk) + print( + f"[epoch {ep}] subset({len(self.subset)}) corpus WER = {wer*100:.2f}% " + f"exact={exact}/{len(self.subset)} ({time.time()-t:.0f}s)", + flush=True, + ) + if wer < self.target and len(self.subset) < len(self.items): + fw, fex, _, _ = eval_clips(self.wrapper, self.items, chunk=self.chunk) + print( + f"[epoch {ep}] FULL({len(self.items)}) corpus WER = {fw*100:.2f}% " + f"exact={fex}/{len(self.items)}", + flush=True, + ) + wer = fw + gc.collect() + torch.cuda.empty_cache() + if was_training: + model.train() + if wer < self.target: + print( + f"[epoch {ep}] TARGET REACHED (corpus WER < {self.target*100:.0f}%). Stopping.", + flush=True, + ) + control.should_training_stop = True + return control + + +def parse_args(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="Qwen/Qwen3-ASR-0.6B") + p.add_argument("--train_file", default="/home/qwen3ft/train.jsonl") + p.add_argument("--output_dir", default="/home/qwen3ft/out/sample7-overfit") + p.add_argument("--epochs", type=float, default=150.0) + p.add_argument("--lr", type=float, default=2e-4) + p.add_argument("--batch_size", type=int, default=1) + p.add_argument("--grad_acc", type=int, default=1) + p.add_argument("--lora_r", type=int, default=16) + p.add_argument("--lora_alpha", type=int, default=32) + p.add_argument( + "--include_head", action="store_true", help="also LoRA lm_head + embed_tokens" + ) + p.add_argument("--optim", default="adamw_torch") + p.add_argument( + "--init_adapter", + default="", + help="load this adapter as trainable start (fresh optimizer) for LR anneal", + ) + p.add_argument("--resume", action="store_true") + p.add_argument("--eval_every", type=int, default=2) + p.add_argument( + "--eval_subset", type=int, default=0, help="WER-gate on first N clips; 0 = all" + ) + p.add_argument( + "--eval_chunk", + type=int, + default=4, + help="transcribe batch size during WER eval (small to avoid OOM atop training memory)", + ) + p.add_argument("--wer_target", type=float, default=0.02) + p.add_argument("--loss_gate", type=float, default=0.15) + p.add_argument( + "--eval_max_new_tokens", + type=int, + default=2048, + help="generation cap for WER eval; 512 truncates long CoSHE transcripts -> false high WER", + ) + p.add_argument("--sr", type=int, default=16000) + return p.parse_args() + + +def main(): + import json + import os + + args = parse_args() + torch.manual_seed(0) + + from qwen_asr import Qwen3ASRModel + + wrapper = Qwen3ASRModel.from_pretrained( + args.model, + dtype=torch.bfloat16, + device_map=None, + max_new_tokens=args.eval_max_new_tokens, + ) + model = wrapper.model + processor = wrapper.processor + patch_outer_forward(model) + # Qwen3ASRForConditionalGeneration doesn't expose (get|set)_input_embeddings at the top + # level; PEFT needs them when LoRA targets embed_tokens (--include_head). Delegate to thinker. + cls = type(model) + if not getattr(cls, "_emb_patched", False): + cls.get_input_embeddings = lambda self: self.thinker.model.embed_tokens + cls.set_input_embeddings = lambda self, v: setattr( + self.thinker.model, "embed_tokens", v + ) + cls._emb_patched = True + model.generation_config = GenerationConfig.from_model_config(model.config) + model.to("cuda") + model.config.use_cache = False + + for pm in model.parameters(): + pm.requires_grad = False + # No gradient checkpointing: a 0.6B + LoRA fits a 40GB A100 with full activations, and + # Qwen3ASRForConditionalGeneration doesn't expose get_input_embeddings at the top level + # (so enable_input_require_grads / grad-checkpointing error). LoRA on the decoder layers + # introduces the grad-requiring path on its own with the base frozen. + + if args.init_adapter: + print(f"Loading init adapter (trainable): {args.init_adapter}", flush=True) + model = PeftModel.from_pretrained(model, args.init_adapter, is_trainable=True) + else: + model = get_peft_model( + model, + LoraConfig( + r=args.lora_r, + lora_alpha=args.lora_alpha, + lora_dropout=0.0, + bias="none", + task_type="CAUSAL_LM", + target_modules=LORA_TARGETS_HEAD if args.include_head else LORA_TARGETS, + ), + ) + model.print_trainable_parameters() + + from datasets import load_dataset + + raw = load_dataset("json", data_files={"train": args.train_file}) + ds = raw.map(make_preprocess_fn_prefix_only(processor), num_proc=1) + keep = {"prompt", "audio", "target", "prefix_text"} + drop = [c for c in ds["train"].column_names if c not in keep] + if drop: + ds["train"] = ds["train"].remove_columns(drop) + collator = DataCollatorForQwen3ASRFinetuning( + processor=processor, sampling_rate=args.sr + ) + + # raw items (audio path + full target) for the WER callback + with open(args.train_file) as f: + items = [json.loads(line) for line in f if line.strip()] + + targs = TrainingArguments( + output_dir=args.output_dir, + per_device_train_batch_size=args.batch_size, + gradient_accumulation_steps=args.grad_acc, + num_train_epochs=args.epochs, + learning_rate=args.lr, + lr_scheduler_type="constant", + warmup_steps=0, + bf16=True, + fp16=False, + logging_steps=5, + save_strategy="epoch", + save_total_limit=2, + optim=args.optim, + report_to="none", + remove_unused_columns=False, + dataloader_num_workers=0, + gradient_checkpointing=False, + ) + + cb = WERStopCallback( + wrapper, + items, + args.eval_every, + args.wer_target, + args.loss_gate, + eval_subset=args.eval_subset, + chunk=args.eval_chunk, + ) + + class BF16CastTrainer(Trainer): + def _prepare_inputs(self, inputs): + inputs = super()._prepare_inputs(inputs) + for k, v in list(inputs.items()): + if torch.is_tensor(v) and v.is_floating_point(): + inputs[k] = v.to(dtype=torch.bfloat16) + return inputs + + trainer = BF16CastTrainer( + model=model, + args=targs, + train_dataset=ds["train"], + data_collator=collator, + tokenizer=processor.tokenizer, + callbacks=[cb], + ) + + resume = ( + args.resume + and os.path.isdir(args.output_dir) + and find_latest_checkpoint(args.output_dir) is not None + ) + print(f"Starting training... (resume={resume})", flush=True) + trainer.train(resume_from_checkpoint=resume) + trainer.save_model(args.output_dir) + print("DONE (adapter saved)", flush=True) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/providers/qwen3_asr/finetune/verify.py b/extras/asr-services/providers/qwen3_asr/finetune/verify.py new file mode 100644 index 00000000..5757e7a2 --- /dev/null +++ b/extras/asr-services/providers/qwen3_asr/finetune/verify.py @@ -0,0 +1,76 @@ +"""Verify a Qwen3-ASR LoRA overfit: transcribe the sample7 clips and compare to ground truth. + +Loads the base model + the trained LoRA adapter (PeftModel), transcribes each train clip with +``language=None``, and prints hyp vs ground-truth with char-similarity + an exact-match count and +corpus WER (mirrors the gemma4 ``infer.py`` smoke check). +""" + +import argparse +import json +from difflib import SequenceMatcher + +import jiwer +import torch +from peft import PeftModel +from qwen_asr import Qwen3ASRModel + +_ASR_TAG = "<asr_text>" +_NORM = jiwer.Compose( + [ + jiwer.ToLowerCase(), + jiwer.RemovePunctuation(), + jiwer.RemoveMultipleSpaces(), + jiwer.Strip(), + jiwer.ReduceToListOfListOfWords(), + ] +) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="Qwen/Qwen3-ASR-0.6B") + p.add_argument("--adapter", default="/home/qwen3ft/out/sample7-overfit") + p.add_argument("--train_file", default="/home/qwen3ft/train.jsonl") + p.add_argument("--max_new_tokens", type=int, default=2048) + args = p.parse_args() + + wrapper = Qwen3ASRModel.from_pretrained( + args.model, + dtype=torch.bfloat16, + device_map="cuda:0", + max_new_tokens=args.max_new_tokens, + ) + if args.adapter: + wrapper.model = PeftModel.from_pretrained(wrapper.model, args.adapter) + print(f"Loaded adapter: {args.adapter}", flush=True) + + with open(args.train_file) as f: + items = [json.loads(line) for line in f if line.strip()] + + refs = [it["text"].split(_ASR_TAG, 1)[-1] for it in items] + paths = [it["audio"] for it in items] + results = wrapper.transcribe(paths, language=[None] * len(paths)) + hyps = [r.text for r in results] + + exact = 0 + for it, ref, hyp in zip(items, refs, hyps): + cs = SequenceMatcher(None, ref.strip(), hyp.strip()).ratio() + ok = ref.strip() == hyp.strip() + exact += ok + name = it["audio"].split("/")[-1] + print( + f"\n=== {name} char-sim={cs:.3f} exact={ok} len(ref)={len(ref)} len(hyp)={len(hyp)} ===" + ) + print(f" REF: {ref[:200]}") + print(f" HYP: {hyp[:200]}") + print(f" REF tail: ...{ref[-120:]}") + print(f" HYP tail: ...{hyp[-120:]}") + + wer = jiwer.wer(refs, hyps, reference_transform=_NORM, hypothesis_transform=_NORM) + print( + f"\nSUMMARY: exact={exact}/{len(items)} corpus WER={wer*100:.2f}%", flush=True + ) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/pyproject.toml b/extras/asr-services/pyproject.toml index 968e8f1e..ecb06c8f 100644 --- a/extras/asr-services/pyproject.toml +++ b/extras/asr-services/pyproject.toml @@ -202,6 +202,11 @@ dev = [ "pytest>=8.0.0", ] +test = [ + "pytest>=8.0.0", + "pytest-cov>=6.0.0", +] + demo = [ "fastrtc>=0.0.23", "gradio>=5.29.0", @@ -214,3 +219,27 @@ build-backend = "setuptools.build_meta" [tool.setuptools] packages = ["common", "providers", "providers.faster_whisper", "providers.transformers", "providers.nemo", "providers.vibevoice", "providers.qwen3_asr", "providers.gemma4", "providers.af_next", "providers.granite"] + +[tool.pytest.ini_options] +minversion = "8.0" +testpaths = ["tests"] +addopts = ["-ra", "--strict-markers", "--strict-config"] +markers = [ + "asr: requires the containerized ASR service", +] + +[tool.coverage.run] +branch = true +relative_files = true +source = ["common", "providers"] + +[tool.coverage.report] +precision = 1 +show_missing = true +skip_empty = true + +[tool.coverage.xml] +output = "coverage-reports/coverage.xml" + +[tool.coverage.html] +directory = "coverage-reports/html" diff --git a/extras/asr-services/scripts/validate_nemotron.py b/extras/asr-services/scripts/validate_nemotron.py new file mode 100644 index 00000000..f4cf09fa --- /dev/null +++ b/extras/asr-services/scripts/validate_nemotron.py @@ -0,0 +1,143 @@ +"""Validate the Nemotron streaming ASR service on a WAV file. + +Exercises both endpoints the Chronicle backend uses: + * POST /transcribe — offline full-file decode + * WS /stream — cache-aware streaming (binary PCM in, interim/final out), + driven exactly like backend RegistryStreamingProvider: + raw PCM frames, then {"type":"CloseStream"} to finalize. + +Usage: + uv run python scripts/validate_nemotron.py <wav_path> [--port 8771] +""" + +import argparse +import asyncio +import json +import sys +import time +import wave + +import httpx +import websockets + +CHUNK_SECONDS = 0.25 # mirrors AudioStreamProducer's 0.25s chunks + + +def read_pcm(path: str): + with wave.open(path, "rb") as wf: + assert wf.getnchannels() == 1, "expected mono" + assert wf.getsampwidth() == 2, "expected 16-bit" + rate = wf.getframerate() + pcm = wf.readframes(wf.getnframes()) + return pcm, rate + + +async def test_offline(base_url: str, path: str): + print(f"\n=== /transcribe (offline) ===") + t0 = time.time() + with open(path, "rb") as f: + files = {"file": ("audio.wav", f, "audio/wav")} + async with httpx.AsyncClient(timeout=300) as client: + resp = await client.post(f"{base_url}/transcribe", files=files) + resp.raise_for_status() + data = resp.json() + dt = time.time() - t0 + print(f" time: {dt:.1f}s") + print(f" words: {len(data.get('words') or [])}") + print(f" text: {data.get('text', '')!r}") + return data.get("text", "") + + +async def test_streaming(ws_url: str, pcm: bytes, rate: int): + print(f"\n=== /stream (cache-aware streaming) ===") + chunk_bytes = int(rate * 2 * CHUNK_SECONDS) + interims = [] + final = None + t0 = time.time() + first_interim_at = None + # ping_interval=None: a client that continuously pushes audio doesn't rely on + # library-level keepalive (GIL-heavy RNNT decode can starve the server loop + # past a 20s ping window while it drains its backlog after CloseStream). + async with websockets.connect(ws_url, max_size=None, ping_interval=None) as ws: + + async def sender(): + for i in range(0, len(pcm), chunk_bytes): + await ws.send(pcm[i : i + chunk_bytes]) + await asyncio.sleep(CHUNK_SECONDS) # real-time pacing + await ws.send(json.dumps({"type": "CloseStream"})) + + send_task = asyncio.create_task(sender()) + try: + while True: + msg = await asyncio.wait_for(ws.recv(), timeout=120) + data = json.loads(msg) + if data.get("type") == "interim": + if first_interim_at is None: + first_interim_at = time.time() - t0 + interims.append(data.get("text", "")) + elif data.get("type") == "final": + final = data + break + finally: + send_task.cancel() + dt = time.time() - t0 + print(f" interims received: {len(interims)}") + if first_interim_at is not None: + print(f" first interim at: {first_interim_at:.2f}s into stream") + if interims: + print(f" sample interims:") + for t in interims[: min(3, len(interims))]: + print(f" - {t!r}") + print(f" - (last) {interims[-1]!r}") + print(f" final words: {len(final.get('words') or []) if final else 0}") + print(f" final text: {final.get('text', '') if final else None!r}") + print(f" total stream time: {dt:.1f}s") + return (final or {}).get("text", ""), interims + + +async def main(): + ap = argparse.ArgumentParser() + ap.add_argument("wav") + ap.add_argument("--port", type=int, default=8772) + ap.add_argument("--host", default="localhost") + args = ap.parse_args() + + base_url = f"http://{args.host}:{args.port}" + ws_url = f"ws://{args.host}:{args.port}/stream" + + # Wait for health + print(f"Waiting for {base_url}/health ...") + async with httpx.AsyncClient(timeout=10) as client: + for _ in range(120): + try: + r = await client.get(f"{base_url}/health") + if r.status_code == 200 and r.json().get("status") == "healthy": + print(f" healthy: {r.json()}") + break + except Exception: + pass + await asyncio.sleep(5) + else: + print("Service did not become healthy in time") + sys.exit(1) + + pcm, rate = read_pcm(args.wav) + print(f"audio: {len(pcm)} bytes @ {rate}Hz ({len(pcm)/(rate*2):.1f}s)") + + offline_text = await test_offline(base_url, args.wav) + stream_text, interims = await test_streaming(ws_url, pcm, rate) + + print("\n=== SUMMARY ===") + print(f"streaming produced {len(interims)} interim updates") + print(f"offline final text len: {len(offline_text)}") + print(f"streaming final text len: {len(stream_text)}") + ok = bool(stream_text.strip()) and len(interims) > 1 + print( + f"RESULT: {'PASS' if ok else 'FAIL'} " + f"(needs non-empty streaming final + multiple interims)" + ) + sys.exit(0 if ok else 2) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/extras/asr-services/tests/capture_vibevoice_ground_truth.py b/extras/asr-services/tests/capture_vibevoice_ground_truth.py new file mode 100644 index 00000000..41ff338b --- /dev/null +++ b/extras/asr-services/tests/capture_vibevoice_ground_truth.py @@ -0,0 +1,168 @@ +""" +Capture VibeVoice transcription output as ground truth for regression tests. + +Sends the 4-minute test audio to a running VibeVoice service and saves the +full response (segments, text, words) as a JSON fixture. + +To exercise batching/stitching on the 4-min file, restart the service with +small batch windows first: + + cd extras/asr-services + BATCH_THRESHOLD_SECONDS=60 BATCH_DURATION_SECONDS=90 BATCH_OVERLAP_SECONDS=15 \ + docker compose up vibevoice-asr -d + +Then run the capture: + + uv run python tests/capture_vibevoice_ground_truth.py + +To restore normal settings afterwards: + + docker compose up vibevoice-asr -d # env vars unset → uses config defaults + +The output is saved to tests/fixtures/vibevoice_4min_ground_truth.json. +Review the output manually before committing — it becomes the reference +for segment overlap and stitching regression tests. +""" + +import argparse +import json +import sys +from pathlib import Path + +import httpx + +DEFAULT_AUDIO = ( + Path(__file__).parent.parent.parent.parent + / "tests" + / "test_assets" + / "DIY_Experts_Glass_Blowing_16khz_mono_4min.wav" +) +DEFAULT_OUTPUT = Path(__file__).parent / "fixtures" / "vibevoice_4min_ground_truth.json" + + +def capture(service_url: str, audio_path: str, output_path: str) -> dict: + """Send audio to VibeVoice /transcribe and return the parsed result.""" + audio_path = Path(audio_path) + if not audio_path.exists(): + print(f"ERROR: Audio file not found: {audio_path}") + sys.exit(1) + + print(f"Sending {audio_path.name} to {service_url}/transcribe ...") + + with open(audio_path, "rb") as f: + # VibeVoice may return NDJSON for long audio — read the full + # response and extract the final "result" event. + with httpx.Client(timeout=httpx.Timeout(600.0)) as client: + resp = client.post( + f"{service_url}/transcribe", + files={"file": (audio_path.name, f, "audio/wav")}, + ) + resp.raise_for_status() + + content_type = resp.headers.get("content-type", "") + + if "application/x-ndjson" in content_type: + # Parse NDJSON — last "result" line is the transcription + data = None + for line in resp.text.strip().split("\n"): + line = line.strip() + if not line: + continue + event = json.loads(line) + if event.get("type") == "progress": + current = event.get("current", "?") + total = event.get("total", "?") + print(f" Progress: batch {current}/{total}") + elif event.get("type") == "result": + data = event + if data is None: + print("ERROR: NDJSON stream ended without a result event") + sys.exit(1) + else: + data = resp.json() + + # Summarize + segments = data.get("segments", []) + words = data.get("words", []) + text = data.get("text", "") + + print(f"\nResult summary:") + print(f" Text length: {len(text)} chars") + print(f" Segments: {len(segments)}") + print(f" Words: {len(words)}") + + if segments: + last_seg = segments[-1] + print(f" Duration: {last_seg.get('end', 0):.1f}s") + + # Check for overlaps + overlaps = [] + for i in range(len(segments) - 1): + overlap = segments[i]["end"] - segments[i + 1]["start"] + if overlap > 0.5: + overlaps.append( + ( + i, + i + 1, + overlap, + segments[i]["text"][:40], + segments[i + 1]["text"][:40], + ) + ) + + if overlaps: + print(f"\n WARNING: {len(overlaps)} segment overlaps detected!") + for idx_a, idx_b, dur, text_a, text_b in overlaps: + print( + f' [{idx_a}→{idx_b}] {dur:.2f}s overlap: "{text_a}" / "{text_b}"' + ) + else: + print(f" No segment overlaps detected") + + # Save + output = Path(output_path) + output.parent.mkdir(parents=True, exist_ok=True) + with open(output, "w") as f: + json.dump(data, f, indent=2) + + print(f"\nSaved to {output}") + return data + + +def main(): + parser = argparse.ArgumentParser( + description="Capture VibeVoice ground truth", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +To force batching on the 4-min test file, restart the service first: + + cd extras/asr-services + BATCH_THRESHOLD_SECONDS=60 BATCH_DURATION_SECONDS=90 BATCH_OVERLAP_SECONDS=15 \\ + docker compose up vibevoice-asr -d + +Then run this script. Restore defaults after with: + + docker compose up vibevoice-asr -d +""", + ) + parser.add_argument( + "--url", + default="http://localhost:8767", + help="VibeVoice service URL (default: http://localhost:8767)", + ) + parser.add_argument( + "--audio", + default=str(DEFAULT_AUDIO), + help="Path to test audio file", + ) + parser.add_argument( + "--output", + default=str(DEFAULT_OUTPUT), + help="Output JSON path", + ) + args = parser.parse_args() + capture(args.url, args.audio, args.output) + + +if __name__ == "__main__": + main() diff --git a/extras/asr-services/uv.lock b/extras/asr-services/uv.lock index 81e86640..a9e4f485 100644 --- a/extras/asr-services/uv.lock +++ b/extras/asr-services/uv.lock @@ -1105,6 +1105,10 @@ qwen3-asr-full = [ { name = "soundfile" }, { name = "websockets" }, ] +test = [ + { name = "pytest" }, + { name = "pytest-cov" }, +] transformers = [ { name = "accelerate" }, { name = "diffusers" }, @@ -1229,6 +1233,10 @@ qwen3-asr-full = [ { name = "soundfile", specifier = ">=0.12" }, { name = "websockets", specifier = ">=13.0" }, ] +test = [ + { name = "pytest", specifier = ">=8.0.0" }, + { name = "pytest-cov", specifier = ">=6.0.0" }, +] transformers = [ { name = "accelerate", specifier = ">=0.30.0" }, { name = "diffusers", specifier = ">=0.30.0" }, @@ -2076,6 +2084,109 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/87/68/7f46fb537958e87427d98a4074bcde4b67a70b04900cfc5ce29bc2f556c1/contourpy-1.3.2-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:8c5acb8dddb0752bf252e01a3035b21443158910ac16a3b0d20e7fed7d534ce5", size = 221791, upload_time = "2025-04-15T17:45:24.794Z" }, ] +[[package]] +name = "coverage" +version = "7.15.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4a/83/6051c2a2feab48ae5bd27c84ef047191d2d4a3172f689e38eaa48ed17db1/coverage-7.15.1.tar.gz", hash = "sha256:165e9949eaf222ef1f018635d0d7f368a23bfe0212af558534c40d8c04686d67", size = 927640, upload_time = "2026-07-12T20:58:19.908Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ef/76/eef242d6bb95fb737fefd9bbf3a318897e4e82cf543b53a61c2a6e8588eb/coverage-7.15.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:05d87c2a43373ad6b976d0a99ad58c48633633bcdeb896dc645a006472cc4a71", size = 221195, upload_time = "2026-07-12T20:55:55.82Z" }, + { url = "https://files.pythonhosted.org/packages/89/73/41cd50da59d513a30fd3a34b5516b962ac17514defdb6705948446a5c0b1/coverage-7.15.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:2afce82f2cf8f4c9002746a42755e1dc61baff33d9f7ab5569b4c9101f8f4d1d", size = 221714, upload_time = "2026-07-12T20:55:58.104Z" }, + { url = "https://files.pythonhosted.org/packages/22/9d/6df6ac6677719a087c839208186525b7b1f3653c50196c43246482ada65e/coverage-7.15.1-cp310-cp310-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:0a545ef5384d787d0fcac6c349afc2c5f99dcc39e13ed3c191b2c06305f64c04", size = 248454, upload_time = "2026-07-12T20:55:59.336Z" }, + { url = "https://files.pythonhosted.org/packages/9f/f5/86f92c429fbf942986d01c3bb7b0b6c3ad06b09f7fa745023877413dde83/coverage-7.15.1-cp310-cp310-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:afaa144b8f5b3bc69fe0ce50d401c46b01ab264782553bfd05a3f98804524ecb", size = 250282, upload_time = "2026-07-12T20:56:00.589Z" }, + { url = "https://files.pythonhosted.org/packages/47/c6/99d0f1c1dd1583aaf8c1deca447e44426be54a76027a8c55e29970f22b15/coverage-7.15.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e1b6099490e5f88569c46b18605f556c3b30acc9a0a219cf7ef8fab8f7161ec4", size = 252149, upload_time = "2026-07-12T20:56:01.853Z" }, + { url = "https://files.pythonhosted.org/packages/b2/c3/5d30740ec96f5fd8a659e46b6ee9224197f87aa66ce4a20e93dcf46221ac/coverage-7.15.1-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:bb5f6c2cf1ffd0bf2bd925c7cdcae9b4f208e9696d453ed51eb1f5fa0cc5b45b", size = 254061, upload_time = "2026-07-12T20:56:03.277Z" }, + { url = "https://files.pythonhosted.org/packages/71/e5/3afb2950673ca6cd22bc9ad6a9d6058c853bef8fe27a6d1df3adc274fd88/coverage-7.15.1-cp310-cp310-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:81572011fc1fc271317da35da593944daef7bfd507085e35751abbe702b74f69", size = 249152, upload_time = "2026-07-12T20:56:04.618Z" }, + { url = "https://files.pythonhosted.org/packages/ca/53/25ab3c0a7cf517ed4d30cd4dd82331e005d4b57fec8c19c0ffd94fb50baa/coverage-7.15.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:8f765d13c08497687d0780cca66115c6aa4ba6703ad43b61e94fab9db689e3a3", size = 250187, upload_time = "2026-07-12T20:56:06Z" }, + { url = "https://files.pythonhosted.org/packages/ad/2a/8afb5f994e26e919e5dca5164f1c7b6b7ce6090aab1aac32b1ae1782f047/coverage-7.15.1-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:771caba880fee96493d18dfc465c318e08ab74e3bc2a3e4089e52514be6a6e54", size = 248193, upload_time = "2026-07-12T20:56:07.376Z" }, + { url = "https://files.pythonhosted.org/packages/d4/62/b6616e15288c9a7b628f7745e832fd8af2c03ec1f7dc11da25947f5ab16f/coverage-7.15.1-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:433a73200848e80f27712fc113b6ff5311f29b479a7d3bd4b1106138a77f9674", size = 252004, upload_time = "2026-07-12T20:56:08.787Z" }, + { url = "https://files.pythonhosted.org/packages/9e/21/5d80d65e4df3018dedb23a44f276f20ce26e429ffd76df03de03bcf9a767/coverage-7.15.1-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:5d6fa45079db9fbeba0a69e3d91189f05301d6ac918162a53179d32fc9ed4910", size = 248463, upload_time = "2026-07-12T20:56:10.245Z" }, + { url = "https://files.pythonhosted.org/packages/53/4b/04a22708691561c44d77e1d5a60857b351310fdfbf253533780bb31c8b6f/coverage-7.15.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:00b6703c6640075cdce5124e9335dfdf9167272475301828acfdd09c0e5ee731", size = 249064, upload_time = "2026-07-12T20:56:11.614Z" }, + { url = "https://files.pythonhosted.org/packages/82/c9/a5c1e8954e657559de478acc1ef73e9393d3fc9b4a2ab4427501d765aca2/coverage-7.15.1-cp310-cp310-win32.whl", hash = "sha256:3ad9a0eac4728327fd870d52f74d2e631d176c5f178eaea2d9983ab5b9755a55", size = 223248, upload_time = "2026-07-12T20:56:13.068Z" }, + { url = "https://files.pythonhosted.org/packages/54/49/5f6566e14611a58e93a926f3a8a7b22e4e0702ebb4c467ac38f452c7f4cf/coverage-7.15.1-cp310-cp310-win_amd64.whl", hash = "sha256:eb5fa75dc3d30e3a1b75da97973479b20ffa9b0641ff56d6e94b5f3e210daa54", size = 223874, upload_time = "2026-07-12T20:56:14.396Z" }, + { url = "https://files.pythonhosted.org/packages/67/ee/5095e07a0986930174fc153fd0bb115f7f2724f5344cc672e016d4f23a4e/coverage-7.15.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:6506330b4a8dcf53b95bd84d8d0e817107cdb3fc1438e835029cdf0bc6612eb0", size = 221320, upload_time = "2026-07-12T20:56:16.036Z" }, + { url = "https://files.pythonhosted.org/packages/43/7a/7e2598ce98433a214020a61c0755d5961f9f26df0c630dc73e06b86d3416/coverage-7.15.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:f8a51a8ec382c39d939cba0ab07ae949077ae4e842343bd4eed22d432358cea9", size = 221823, upload_time = "2026-07-12T20:56:17.54Z" }, + { url = "https://files.pythonhosted.org/packages/8b/4f/67dac6a69139b490a239d91b6d5ef127982ffb89159769241656ab879ee5/coverage-7.15.1-cp311-cp311-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:18ea20e3922d7f8ca9e0ef1084408d08c4ad62d5e531cb9c1f6896a99297ebea", size = 252242, upload_time = "2026-07-12T20:56:18.868Z" }, + { url = "https://files.pythonhosted.org/packages/05/37/5e439725a96de592309725550c8df639c359c209dcb4b152ed46032d4c0d/coverage-7.15.1-cp311-cp311-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:aab9902a64b8390e3b56e539fddae1d79a267807fe5cb0c18d7d2f544ce867e2", size = 254153, upload_time = "2026-07-12T20:56:20.118Z" }, + { url = "https://files.pythonhosted.org/packages/0f/1c/671ba9e15f9651a99962fb588a1fe4fb267cae4bb85f9bb4ac4413c6f7ef/coverage-7.15.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8cc316264317b07a9e90d7f2b4188a15e36e9b54e651081b791b0515fa612a29", size = 256261, upload_time = "2026-07-12T20:56:21.682Z" }, + { url = "https://files.pythonhosted.org/packages/d6/3c/88305322b4d015b4536b7174b8f9b87f850f01af9d89008cdbf984b65ff2/coverage-7.15.1-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:d4e47e7eea81a8ccf060a07627654151d929da62c7b715738387c200905cec89", size = 258222, upload_time = "2026-07-12T20:56:22.962Z" }, + { url = "https://files.pythonhosted.org/packages/8d/20/d02e42333a57f266ded61ed4fb3821da1f2dc143734aaa3dcc9c117204c5/coverage-7.15.1-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:c6829d9a3b55ad2b73ef5fda8302e5be03683789e88b1a079dcf4a773229c21d", size = 252368, upload_time = "2026-07-12T20:56:24.475Z" }, + { url = "https://files.pythonhosted.org/packages/f3/d0/94ebc010926a191d83e83b1f1236f8bd0d14d0eb98b5c324aa4be2589a8e/coverage-7.15.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:764e045811f9c8cda436641f3f088283351d331a519b5807f19041cd0a68da1c", size = 253953, upload_time = "2026-07-12T20:56:26.036Z" }, + { url = "https://files.pythonhosted.org/packages/2e/a6/2c8553191da0c2da2c43ff294fb9bab57a5b0fae03560304a82a8b5addc3/coverage-7.15.1-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:09b3b088aa24489c4082bcc35fcc8224281ab94a653dfb6d3f0c8165b0d628ab", size = 252016, upload_time = "2026-07-12T20:56:27.374Z" }, + { url = "https://files.pythonhosted.org/packages/17/99/20f47986d8360fa1a31d9858dd2ba0be009d44ecfa8d02120b1eb93c028b/coverage-7.15.1-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:7911b02f57053adf8164ae63edb1c26574d24dfccabadc5268cf69310a69a358", size = 255783, upload_time = "2026-07-12T20:56:28.921Z" }, + { url = "https://files.pythonhosted.org/packages/17/36/fb61983b5259f78fa5e62ee74e91fe4de4c556fcbc9a3b9c9021c2f8b57b/coverage-7.15.1-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:664e279ed40599b8ed16f4db18d92a7e212c73129672bec8f5d96d4da48d2404", size = 251735, upload_time = "2026-07-12T20:56:30.465Z" }, + { url = "https://files.pythonhosted.org/packages/6c/2e/d109d8d2d6c726ba2e78125297073918a2a91464f0ea777e5398fa3cf75c/coverage-7.15.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:34fe7cf79d5f1f87f2e8ce7dc1c32950841f50e10d0120f263856acfad66de34", size = 252645, upload_time = "2026-07-12T20:56:32.076Z" }, + { url = "https://files.pythonhosted.org/packages/bc/a0/5b3219cbb46135e14673427f2bf5890c4da4e6f655243692f3448aa3f42c/coverage-7.15.1-cp311-cp311-win32.whl", hash = "sha256:5e2d2536d2f57a354aa382ed303ac0e2e5c9522a508c05b998d26181b94163a7", size = 223414, upload_time = "2026-07-12T20:56:33.432Z" }, + { url = "https://files.pythonhosted.org/packages/5b/80/24e13ade0b6df8090e2492f9c55044bf1423f7ef28c76fc1af8c827a61cb/coverage-7.15.1-cp311-cp311-win_amd64.whl", hash = "sha256:c337da8fca7ea93ab43f3868cfcde6cf6dad32c3906b273cfbad5d7390bc423b", size = 223888, upload_time = "2026-07-12T20:56:34.784Z" }, + { url = "https://files.pythonhosted.org/packages/2d/9e/34659930ae380491af48e2f5f5f9a2e99cd9fb7de3d30f27be5ac5425137/coverage-7.15.1-cp311-cp311-win_arm64.whl", hash = "sha256:db3403fdb7a94d5eb73e099befad8104d2a7d110a0f0d99df0de61c5d1fa756c", size = 223435, upload_time = "2026-07-12T20:56:36.302Z" }, + { url = "https://files.pythonhosted.org/packages/d9/76/32c1826309beaf4604c54accef108fdd611e5e5e93f2f5192f050cd5f6bd/coverage-7.15.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d9476292594309db922cc841dd13b303b3c388f4c25d279884f7e2341c681f80", size = 221497, upload_time = "2026-07-12T20:56:37.628Z" }, + { url = "https://files.pythonhosted.org/packages/db/5c/b88ce0d68fa550c7f3b58617fbf363bce64df5bf8295a01b627e4696e022/coverage-7.15.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:0c579056b0de461b3a62318b63d0b6ce90aed7f8158d3f00da094df82f29d189", size = 221854, upload_time = "2026-07-12T20:56:39.033Z" }, + { url = "https://files.pythonhosted.org/packages/0e/fe/8509fd2a66fc4e0a829f76a0f0b1dc3cc163368352435b5f243168658077/coverage-7.15.1-cp312-cp312-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:23214bdbe226f2b0e9c66a7d6a1d59d4a88045dcf86e702cf0fe0d0935e3d615", size = 253359, upload_time = "2026-07-12T20:56:40.423Z" }, + { url = "https://files.pythonhosted.org/packages/e5/81/c7009ed7ea9765adb2b9d095054d748266fae5f07ac6c5f925f33715fcde/coverage-7.15.1-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:df164be93b46b4825cc39339440a05edc54c4d1d865ba4a60fd43d151a2a1cd3", size = 256096, upload_time = "2026-07-12T20:56:42.115Z" }, + { url = "https://files.pythonhosted.org/packages/21/52/dc8ee03968a5ba86e2da5aa48ddc9e3747bd65d63825fdce2d96acb9c5ff/coverage-7.15.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a524fca1a6f08927d9dc2d4c873cfb7bd7202c247f08b14bdc02424071b8b304", size = 257211, upload_time = "2026-07-12T20:56:43.513Z" }, + { url = "https://files.pythonhosted.org/packages/b8/27/95d7623908da8937deb53d48efcdbf423907a47540e63c62fa21372c652b/coverage-7.15.1-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:d70f3542cd38de85a9e257dcb1ac4c1ab4b6d7d2c2a645809207556628755d1c", size = 259473, upload_time = "2026-07-12T20:56:44.974Z" }, + { url = "https://files.pythonhosted.org/packages/f1/3b/730d761928de97d585465680b568ae69622fb40716babadeabffe75cb51b/coverage-7.15.1-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:d78aa537237212c4313aabe5e964b66acc86350ed19ebc56a3e202df33b6077b", size = 253759, upload_time = "2026-07-12T20:56:46.615Z" }, + { url = "https://files.pythonhosted.org/packages/f1/fc/6b9277acff1f9484b6c12857af5774689d1a6a95e13265f7405329d2f5da/coverage-7.15.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:a318112bb4f79d9d04766196d5a3388caa825908a6a9b052aa87de3d9aea7c61", size = 255131, upload_time = "2026-07-12T20:56:48.073Z" }, + { url = "https://files.pythonhosted.org/packages/3d/f2/c704f86129594ba34e25a64695d2068c71d51c2b98907184d716c94f4aec/coverage-7.15.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:e55d24cada901963eed5bc89fa562aa033f0d84b9d3de4ecf363737c13aed11e", size = 253275, upload_time = "2026-07-12T20:56:49.538Z" }, + { url = "https://files.pythonhosted.org/packages/f6/29/80fee8af47de4a6dce71ccf2938491f444687a756af258a56d8469b8f1b0/coverage-7.15.1-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:3c78f0cea7275342cf2adc2ad5fdd0aafa106ad91e66d573568f2fcf62c41df5", size = 257345, upload_time = "2026-07-12T20:56:51.038Z" }, + { url = "https://files.pythonhosted.org/packages/20/21/a1e7d7ed1b48a8adf8fd5154d9e83fcc5ad8e6ff20ae00e44865057dce8d/coverage-7.15.1-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:86bd37eabe39977216f630a7fc1b698e7f5e81a191c7186013245c6c3d313f9d", size = 252844, upload_time = "2026-07-12T20:56:52.535Z" }, + { url = "https://files.pythonhosted.org/packages/a7/8c/a4bc26e6ee207d412f3678f04d74be1550e83140563ca0e4997510579712/coverage-7.15.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:c6db15c217693bdc3ca0b84de1ba9afafe1c14c26a8a29d77f4ed0de2b6132e2", size = 254716, upload_time = "2026-07-12T20:56:53.968Z" }, + { url = "https://files.pythonhosted.org/packages/11/9d/8ad0266ecfada6353cf6627a1a02294cf55a907521b6ee0bd7b770cfd659/coverage-7.15.1-cp312-cp312-win32.whl", hash = "sha256:359f3fbe09a51500c51966596ee4ee4070b356552c70b3b2420eb200d68e0f76", size = 223554, upload_time = "2026-07-12T20:56:55.583Z" }, + { url = "https://files.pythonhosted.org/packages/81/6d/24224929e06c6e05a93f738bc5f9e8e6ab658f8f1d9b823e7b85430e28b8/coverage-7.15.1-cp312-cp312-win_amd64.whl", hash = "sha256:fa75dc099c126e941a9c0baa8ebd2cbc78bd778687534fe410baf754f6d9e374", size = 224087, upload_time = "2026-07-12T20:56:57.041Z" }, + { url = "https://files.pythonhosted.org/packages/35/23/f81441dd01de88e53c97842e706907b307d9078918c3f4998b11e9ac7250/coverage-7.15.1-cp312-cp312-win_arm64.whl", hash = "sha256:26f89cf6d0634375f454fa71057945ad18edb0f1607a90fecf22c57dc3dc289a", size = 223472, upload_time = "2026-07-12T20:56:58.594Z" }, + { url = "https://files.pythonhosted.org/packages/ca/1e/6fa289d7993a2a39f1b283ddb58c4bfec80f7800be654b8ba8a9f6a07c63/coverage-7.15.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:71ac4ca1658ca99160fd58cc6967110e989c34b04627f24ed6ec9f70fb24571a", size = 221519, upload_time = "2026-07-12T20:57:00.081Z" }, + { url = "https://files.pythonhosted.org/packages/6f/e1/0db4902a0588234a70ab0218073c0b20fbc5c740aa35f91d360160a2ebc9/coverage-7.15.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:26a40cbf2b13bd94af53ee02a424cb3bb96a9edfac0d00834bd068512a62714b", size = 221895, upload_time = "2026-07-12T20:57:01.867Z" }, + { url = "https://files.pythonhosted.org/packages/b4/cb/3719783865092dac5e08df842730305ee9ab1973ae7ddb6fbdf27d401f30/coverage-7.15.1-cp313-cp313-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:f4c5a5eff4ad4f9f7088fd3fc7a66d98d06566ee294b3b053309fb0a3b45be1e", size = 252882, upload_time = "2026-07-12T20:57:03.459Z" }, + { url = "https://files.pythonhosted.org/packages/1c/5e/caf3abbdbb22629626160ffc9c017eb995b7cb11c0be46b974834cef1792/coverage-7.15.1-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:962aa56c1c9b016d681265880eb6acc9966029d2c4c559319cc43a1abbb9b59a", size = 255479, upload_time = "2026-07-12T20:57:04.984Z" }, + { url = "https://files.pythonhosted.org/packages/d6/f1/d60f375bfe095fef944f0f19427aefdbf9bdd5a9571c41a4bf6e2f5fdb81/coverage-7.15.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1678eb2dc57a8ce67601b029582ef6d41e9e6ca22692aaeccd4107e40f27386c", size = 256715, upload_time = "2026-07-12T20:57:06.446Z" }, + { url = "https://files.pythonhosted.org/packages/d7/17/8b0cbc90d02dc5adad4d9034c1824ec3fa567771b4c39d9c1e3f9b1431b8/coverage-7.15.1-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:1174900a43f6f8c425fee10d7dbddc308adefcdc78aaced32357f5ab750a0e90", size = 258845, upload_time = "2026-07-12T20:57:08.092Z" }, + { url = "https://files.pythonhosted.org/packages/92/29/c5e69f5fb75c322e9a3e4ef64d02eebfc3d66efceccc8514ff80a3c13a56/coverage-7.15.1-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:98847557a6859cadf693792ce89f440cb89692993f60dc6d3a7e35f3d340216f", size = 253098, upload_time = "2026-07-12T20:57:09.636Z" }, + { url = "https://files.pythonhosted.org/packages/64/57/21144252fdd0c01d707d48fbcea13a80b0b7c42ced3f299f885ab8978c3a/coverage-7.15.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:8697b2edb57143546a24389efc11e1b000cd5800fc20d84f04edb601e4a7cfb8", size = 254844, upload_time = "2026-07-12T20:57:11.141Z" }, + { url = "https://files.pythonhosted.org/packages/59/2a/499a28a322b0ce6768328e6c5bb2e2ad00ac068a7c7adb2ecd8533c8c5d9/coverage-7.15.1-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:6827ac0519be3fe91bf96b4060eb00d1d24f82649b29862cd75a3cfca248b02a", size = 252807, upload_time = "2026-07-12T20:57:12.678Z" }, + { url = "https://files.pythonhosted.org/packages/2c/6f/928a95da5da8b60f2b00e1482c7787b3316188e6d2d227fb8e124ada43a1/coverage-7.15.1-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:2de8ecbbc77c7e4d22572779920ed8979c69168675e96be3a548c996568c6c31", size = 256965, upload_time = "2026-07-12T20:57:14.326Z" }, + { url = "https://files.pythonhosted.org/packages/16/10/889adbc1b8c9f866ed51e18a98bcafc0259fb9d29b81f50a719407c64ea8/coverage-7.15.1-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:2b25f0f0fa5260df9d7bb55d47c8bdc23fa3382c1a18f7c9cae122e6c320b1ad", size = 252628, upload_time = "2026-07-12T20:57:15.892Z" }, + { url = "https://files.pythonhosted.org/packages/1a/30/a5e1871e5d93416511f8e359d1ccebfe0cbb050a1bbf7dd20228533ec0cf/coverage-7.15.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:4a2effcbd93ae340a58db718fe4181d967f84d352c4cefeaab4ff82ce813901a", size = 254399, upload_time = "2026-07-12T20:57:17.703Z" }, + { url = "https://files.pythonhosted.org/packages/2d/26/c36fbffd549dadbdd1a75827528fb00a4c46aa3187b007b750b1e2cebbf2/coverage-7.15.1-cp313-cp313-win32.whl", hash = "sha256:895e65c96aef0cecea250f6e35e9a32f11375514e1a0cb5210e0fda128c04e8e", size = 223564, upload_time = "2026-07-12T20:57:19.253Z" }, + { url = "https://files.pythonhosted.org/packages/16/fc/becbb9d2c4206d242b9b1e1e8e24a42f7926c0200dd3c788b9fab4bb96d5/coverage-7.15.1-cp313-cp313-win_amd64.whl", hash = "sha256:f6d0a28b63a0d75f9ed5118105d1154fc3aa40a8605a30d5d87e3d043ad90fe7", size = 224106, upload_time = "2026-07-12T20:57:21.108Z" }, + { url = "https://files.pythonhosted.org/packages/d3/30/1cfc641461369b6858799fca61c0a8b5edc490c519bf7c636ffa6bbf556f/coverage-7.15.1-cp313-cp313-win_arm64.whl", hash = "sha256:b4ee9818e8bae3544379ad2c09b851c4fb886aaa8860d57a1c1316ddcb16db49", size = 223497, upload_time = "2026-07-12T20:57:22.734Z" }, + { url = "https://files.pythonhosted.org/packages/f0/46/81961952e7aebfb38ad0ae4264e8954cc607a7af9e7ac111f9fa986595cc/coverage-7.15.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:a886af95f59edf67d5770fd3564d53f4a8af93f25f8c1d60d27e00d7f5674ee8", size = 221560, upload_time = "2026-07-12T20:57:24.282Z" }, + { url = "https://files.pythonhosted.org/packages/13/d2/ee14d715889f216baf47301d9f469e08fff6995552aaf67e897b282865f6/coverage-7.15.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:985657ebd707941de90d488d1cbb5efac20bdf81f7b91eba771624ccda4d36f4", size = 221894, upload_time = "2026-07-12T20:57:25.87Z" }, + { url = "https://files.pythonhosted.org/packages/f3/38/f830bc6e6c2c5f23f43847125e6c650d378872f7eeba8d49f1d42193e8a9/coverage-7.15.1-cp314-cp314-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:5bbe2a06e0a5e1404d9ffbdb49b819bbd6a3bb198ebea4c8dfe7ad9f1e1c2e81", size = 252938, upload_time = "2026-07-12T20:57:27.506Z" }, + { url = "https://files.pythonhosted.org/packages/e1/53/0d3dd963631259d794c898735d5436e68d6a8d40749c419a07ff7c171469/coverage-7.15.1-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:bde0fe24083d0b7b3dbafa7a09f0796410af1afa2523f28f5f208d8340a4aaca", size = 255445, upload_time = "2026-07-12T20:57:29.234Z" }, + { url = "https://files.pythonhosted.org/packages/b1/fd/aabed228557565c958259251b89bab8c5669b31291fa63b3e2154ebb017a/coverage-7.15.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f89f7453d6d46db14cf233e2cd8edcd78de2b9c49d4f1dc109590b4e5dbfbb74", size = 256790, upload_time = "2026-07-12T20:57:30.826Z" }, + { url = "https://files.pythonhosted.org/packages/bc/aa/1cc888e5d3623e603c4e5399653cb25728bb2b40d7519188a3e293d24620/coverage-7.15.1-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:cc3656c9ecc27b36bd0907455b77f83c0069ca9ad4a66dec892b76c696eb6047", size = 259104, upload_time = "2026-07-12T20:57:32.63Z" }, + { url = "https://files.pythonhosted.org/packages/5f/61/fc16d5f5e53098dae41efa21e8ccc611a9b4fe922750dd03dc56db552182/coverage-7.15.1-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:24d8e85a2a45e44883b488c2659f51fa761dad5353fdb319b672a93facbd2ca9", size = 252956, upload_time = "2026-07-12T20:57:34.316Z" }, + { url = "https://files.pythonhosted.org/packages/a1/f3/52384668c3de4519ca770bf1975a89e4d6eb5aa2faf0da0577a14008cba4/coverage-7.15.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:68931b5fe746ed4fdaa8892989cab9e6c35781eeb3b0ab2ded893d561e1b3652", size = 254797, upload_time = "2026-07-12T20:57:35.947Z" }, + { url = "https://files.pythonhosted.org/packages/ce/68/54b807e7c1868178e902fd8360b5d4e559394462f97285c50edf1c4608db/coverage-7.15.1-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:1ce6947e2a95534ecaa5a15e73c21e550514c980d80eda204d064d789a95f6a4", size = 252762, upload_time = "2026-07-12T20:57:37.856Z" }, + { url = "https://files.pythonhosted.org/packages/7b/48/dde8adf0338e3ace738757dccf1ce817e5fdcadfae77e1b48a77e5a3b265/coverage-7.15.1-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:841befdbc89b9c82435fc25b0f4f41858b6238693e45af758bec4cfc1968171c", size = 257037, upload_time = "2026-07-12T20:57:39.488Z" }, + { url = "https://files.pythonhosted.org/packages/07/f2/179dd88cf60a0aeeee16a970ffe250dccea8b80ed4beab4c5d3f6c41ad4b/coverage-7.15.1-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:5d3de58b837375e7f4c0e1a088ccab5f655efb2fd7427b729df02c862a559633", size = 252577, upload_time = "2026-07-12T20:57:41.363Z" }, + { url = "https://files.pythonhosted.org/packages/2c/37/8a593d69ab521beb6a105a2017cac4ba94425ee0a8349e29c3c0b522d24f/coverage-7.15.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:b1801963f9f44ae0c0f6d737bc7aeb2bbcde7d1fe7e3b43cddc1961af42d3b41", size = 254235, upload_time = "2026-07-12T20:57:43.025Z" }, + { url = "https://files.pythonhosted.org/packages/6d/34/bc9b3bced66f2cdad4bf5e57ae51c54ea226e8aaaebfc9370a9a11877bf3/coverage-7.15.1-cp314-cp314-win32.whl", hash = "sha256:8c7953c4128ef53b6ffb5f90d87c87d4ce26731df294760bb2314eb0e069e44b", size = 223771, upload_time = "2026-07-12T20:57:44.662Z" }, + { url = "https://files.pythonhosted.org/packages/4c/f3/4d20337bed61915d14349e62b88d5e4144d5a9872b64adbe90e9906db6db/coverage-7.15.1-cp314-cp314-win_amd64.whl", hash = "sha256:6f0bab60a582d415f0fb535ccff13ba334a47a1538f98913330a525d23bd535a", size = 224257, upload_time = "2026-07-12T20:57:46.412Z" }, + { url = "https://files.pythonhosted.org/packages/7b/df/bbfeae4948f3ded516f92b32f2d57952427fc5ecfc0924487bb6ee6a5f38/coverage-7.15.1-cp314-cp314-win_arm64.whl", hash = "sha256:0f410ee8f0ac4ec7db71bc0b7632a8b9994e1cad2755bd1566c17e6a162caa74", size = 223683, upload_time = "2026-07-12T20:57:48.106Z" }, + { url = "https://files.pythonhosted.org/packages/35/65/0b431856064e387d1f5cf474625e4a0465e907024d42f35de6af19ced0be/coverage-7.15.1-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:fc868bab88e049d41fcd41766810d790a8b960053be2a45e060f5ce0d31d258b", size = 222298, upload_time = "2026-07-12T20:57:49.882Z" }, + { url = "https://files.pythonhosted.org/packages/a6/96/50eac9bd49df8a3df5f3d38746d1bf332299dffb554486c94ebd55c9dc49/coverage-7.15.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:206d4ec6028f2773b40932d09f074539d6bcdd8f6b318d40cb04bdbd68ed0b49", size = 222561, upload_time = "2026-07-12T20:57:51.688Z" }, + { url = "https://files.pythonhosted.org/packages/e0/5b/6ba1c4a27e10b8816fd2622b98162c83d3bdf1185097360373611bf96364/coverage-7.15.1-cp314-cp314t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:620482ef1c9f4e61f962e159325fe77dea59d16e39d9c9470d069053b244d864", size = 263923, upload_time = "2026-07-12T20:57:53.392Z" }, + { url = "https://files.pythonhosted.org/packages/e4/59/fe03ade97a3ca2d890e98c572cf48a99fda9adba85757c34b823f41efe1e/coverage-7.15.1-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:d385fc9b054e309ad3cecdc77b586d2af0c98aeec2fdb3773544586f366e817c", size = 266043, upload_time = "2026-07-12T20:57:55.095Z" }, + { url = "https://files.pythonhosted.org/packages/16/e0/55c4b1217a572a43e13b39e1eb78d0da29fb23679003bd0cdf22c50b1978/coverage-7.15.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f1198bca9c0dd7c188aae1f185b0c0b5fc4f0a2b6909000858c29550320bdb07", size = 268465, upload_time = "2026-07-12T20:57:57.017Z" }, + { url = "https://files.pythonhosted.org/packages/b1/e2/ee47944f76afc03909119b036fe9e0da8cbd274a5141287de79791a0fb6d/coverage-7.15.1-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:5d0297e6a070eadb49df7cddd0ab6f420b8b689dd8904c7dd815a323168fa57e", size = 269584, upload_time = "2026-07-12T20:57:58.958Z" }, + { url = "https://files.pythonhosted.org/packages/cb/8a/6b4d9779c7b2e21c3d12c3425e3261aa7411399319e27aa402dfec4db5d0/coverage-7.15.1-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:916fcf2214f56960e409561b37fc32a160a42b6e85483d0652d7b70fa55d707e", size = 263019, upload_time = "2026-07-12T20:58:00.979Z" }, + { url = "https://files.pythonhosted.org/packages/c4/1e/db5c7fa0c8ba5ece390a1e1a3f30db71d440240a80589df28e66a7503c40/coverage-7.15.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:f837bae572c7869ffaa502e604c87e182543012831cf87aae4586ad090ac6dcf", size = 265916, upload_time = "2026-07-12T20:58:03.005Z" }, + { url = "https://files.pythonhosted.org/packages/83/53/fe5176682b00709b13fab36addd26883139d0dea430816fea412e69255e2/coverage-7.15.1-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:3ea65e3ee6c7c32349fd00559927a9e577bdd72386087eeed1c42b62dfce9b82", size = 263520, upload_time = "2026-07-12T20:58:04.994Z" }, + { url = "https://files.pythonhosted.org/packages/f9/e7/16f15127be93fbc70c667df5ec5dce934fc76c9b0888d84969a5d5341e2c/coverage-7.15.1-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:345034976f46a1c54bd17f4e43eb30bb92cb7082fcddff03250cff136cc4eb82", size = 267254, upload_time = "2026-07-12T20:58:06.824Z" }, + { url = "https://files.pythonhosted.org/packages/cd/73/e5119111f6f065376395a525f7ce6e9174d83f3db6d217ea0211a61cca4d/coverage-7.15.1-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:4f051a64eb8f8addb4661c2b41d6eea5b7ebc68ad4b2baea8d9bc54e1956e5f7", size = 262366, upload_time = "2026-07-12T20:58:08.555Z" }, + { url = "https://files.pythonhosted.org/packages/d7/9c/6d0a81182df18a73b081e7a8630f0e2a52b12dfd7898c6ab839551a454d2/coverage-7.15.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:a7625770f7720b49bb30d194ad2f8d50fab3c5177874af3d2399676f95f9c594", size = 264680, upload_time = "2026-07-12T20:58:10.359Z" }, + { url = "https://files.pythonhosted.org/packages/d9/a2/bac0cbd4450638f1be2041a464b1766c8cc94abf705a2df6f1c8d4be870d/coverage-7.15.1-cp314-cp314t-win32.whl", hash = "sha256:81e503d130a472ad1bd38199ecd35116b40d92bcd31e27a2cacde035381f2070", size = 224077, upload_time = "2026-07-12T20:58:12.065Z" }, + { url = "https://files.pythonhosted.org/packages/f0/b2/d83c5403155172a43ba47c08641bad3f89822d8405102423a41339d2c857/coverage-7.15.1-cp314-cp314t-win_amd64.whl", hash = "sha256:724e878b213b302ad46e9f2fc872d386613f20ebfc492a211482d917ea76c14f", size = 224908, upload_time = "2026-07-12T20:58:13.956Z" }, + { url = "https://files.pythonhosted.org/packages/cc/41/442b74cad832cc77712080585455482e7cc4f4a9a13192f65731dcd18231/coverage-7.15.1-cp314-cp314t-win_arm64.whl", hash = "sha256:ce2f05c14d077f406fefc4fa5e4f093ad0e0787549f6582535d6e28766f0361b", size = 224219, upload_time = "2026-07-12T20:58:16.315Z" }, + { url = "https://files.pythonhosted.org/packages/34/98/07a67cf1a26e795d617ed5c540c042b0ac87b72f810c30c07f076cf334f3/coverage-7.15.1-py3-none-any.whl", hash = "sha256:717d01e6e00bed56ad13306f19e0dd2f4f645ee8159d2c72c72301d6cfc7090c", size = 213284, upload_time = "2026-07-12T20:58:18.079Z" }, +] + +[package.optional-dependencies] +toml = [ + { name = "tomli", marker = "python_full_version <= '3.11' or (extra == 'extra-12-asr-services-cu126' and extra == 'extra-12-asr-services-cu128') or (extra == 'extra-12-asr-services-cu126' and extra == 'extra-12-asr-services-strixhalo') or (extra == 'extra-12-asr-services-cu128' and extra == 'extra-12-asr-services-strixhalo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-faster-whisper') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-transformers' and extra == 'group-12-asr-services-vibevoice')" }, +] + [[package]] name = "cryptography" version = "45.0.3" @@ -2170,7 +2281,7 @@ name = "cuda-bindings" version = "12.9.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cuda-pathfinder" }, + { name = "cuda-pathfinder", marker = "sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/37/31/bfcc870f69c6a017c4ad5c42316207fc7551940db6f3639aa4466ec5faf3/cuda_bindings-12.9.4-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a022c96b8bd847e8dc0675523431149a4c3e872f440e3002213dbb9e08f0331a", size = 11800959, upload_time = "2025-10-21T14:51:26.458Z" }, @@ -2554,7 +2665,7 @@ name = "exceptiongroup" version = "1.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.12' or (python_full_version == '3.12.*' and extra == 'extra-12-asr-services-cu126' and extra == 'extra-12-asr-services-cu128') or (python_full_version == '3.12.*' and extra == 'extra-12-asr-services-cu126' and extra == 'extra-12-asr-services-strixhalo') or (python_full_version == '3.12.*' and extra == 'extra-12-asr-services-cu128' and extra == 'extra-12-asr-services-strixhalo') or (python_full_version >= '3.13' and extra == 'extra-12-asr-services-cu126' and extra == 'extra-12-asr-services-cu128') or (python_full_version >= '3.13' and extra == 'extra-12-asr-services-cu126' and extra == 'extra-12-asr-services-strixhalo') or (python_full_version >= '3.13' and extra == 'extra-12-asr-services-cu128' and extra == 'extra-12-asr-services-strixhalo') or (python_full_version == '3.12.*' and extra != 'extra-12-asr-services-cu126' and extra != 'extra-12-asr-services-cu128' and extra != 'extra-12-asr-services-strixhalo')" }, + { name = "typing-extensions", marker = "python_full_version < '3.11' or (extra == 'extra-12-asr-services-cu126' and extra == 'extra-12-asr-services-cu128') or (extra == 'extra-12-asr-services-cu126' and extra == 'extra-12-asr-services-strixhalo') or (extra == 'extra-12-asr-services-cu128' and extra == 'extra-12-asr-services-strixhalo')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0b/9f/a65090624ecf468cdca03533906e7c69ed7588582240cfe7cc9e770b50eb/exceptiongroup-1.3.0.tar.gz", hash = "sha256:b241f5885f560bc56a59ee63ca4c6a8bfa46ae4ad651af316d4e81817bb9fd88", size = 29749, upload_time = "2025-05-10T17:42:51.123Z" } wheels = [ @@ -6558,8 +6669,8 @@ name = "nvidia-cudnn-cu12" version = "9.10.2.21" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas-cu12", version = "12.6.4.1", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-12-asr-services-cu126' or (extra == 'extra-12-asr-services-cu128' and extra == 'extra-12-asr-services-strixhalo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-faster-whisper') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-transformers' and extra == 'group-12-asr-services-vibevoice')" }, - { name = "nvidia-cublas-cu12", version = "12.8.4.1", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-12-asr-services-cu128' or extra == 'extra-12-asr-services-strixhalo' or extra != 'extra-12-asr-services-cu126' or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-faster-whisper') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-transformers' and extra == 'group-12-asr-services-vibevoice')" }, + { name = "nvidia-cublas-cu12", version = "12.6.4.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-12-asr-services-cu126') or (extra == 'extra-12-asr-services-cu126' and extra == 'extra-12-asr-services-cu128') or (extra == 'extra-12-asr-services-cu126' and extra == 'extra-12-asr-services-strixhalo') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-af-next') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-faster-whisper') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-gemma4') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-granite') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-transformers') or (extra == 'extra-12-asr-services-cu128' and extra == 'extra-12-asr-services-strixhalo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-faster-whisper') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-transformers' and extra == 'group-12-asr-services-vibevoice') or (extra == 'extra-12-asr-services-cu126' and extra != 'group-12-asr-services-nemo' and extra != 'group-12-asr-services-nemo-rocm' and extra != 'group-12-asr-services-parakeet' and extra != 'group-12-asr-services-vibevoice')" }, + { name = "nvidia-cublas-cu12", version = "12.8.4.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-12-asr-services-cu126') or (extra == 'extra-12-asr-services-cu128' and extra == 'extra-12-asr-services-strixhalo') or (extra == 'extra-12-asr-services-cu128' and extra == 'group-12-asr-services-af-next') or (extra == 'extra-12-asr-services-cu128' and extra == 'group-12-asr-services-faster-whisper') or (extra == 'extra-12-asr-services-cu128' and extra == 'group-12-asr-services-gemma4') or (extra == 'extra-12-asr-services-cu128' and extra == 'group-12-asr-services-granite') or (extra == 'extra-12-asr-services-cu128' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'extra-12-asr-services-cu128' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'extra-12-asr-services-cu128' and extra == 'group-12-asr-services-transformers') or (extra != 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-af-next') or (extra != 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-faster-whisper') or (extra != 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-gemma4') or (extra != 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-granite') or (extra == 'extra-12-asr-services-cu126' and extra == 'extra-12-asr-services-cu128') or (extra == 'extra-12-asr-services-cu126' and extra == 'extra-12-asr-services-strixhalo') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-faster-whisper') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-gemma4') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-granite') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-parakeet') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-transformers') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-vibevoice') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-gemma4') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-granite') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-parakeet') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-transformers') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-vibevoice') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-granite') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-parakeet') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-transformers') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-vibevoice') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-parakeet') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-transformers') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-vibevoice') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-parakeet') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-transformers') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-vibevoice') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-parakeet') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-transformers') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-vibevoice') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-transformers') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-vibevoice') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-transformers') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-vibevoice') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-transformers') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-vibevoice') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-transformers' and extra == 'group-12-asr-services-vibevoice') or (extra == 'extra-12-asr-services-cu128' and extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'extra-12-asr-services-cu128' and extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-parakeet') or (extra == 'extra-12-asr-services-cu128' and extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-vibevoice') or (extra == 'extra-12-asr-services-cu128' and extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-parakeet') or (extra == 'extra-12-asr-services-cu128' and extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-vibevoice') or (extra == 'extra-12-asr-services-cu128' and extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-vibevoice') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-faster-whisper') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-gemma4') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-granite') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo-rocm') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-parakeet') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-transformers') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-vibevoice') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-gemma4') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-granite') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo-rocm') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-parakeet') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-transformers') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-vibevoice') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-granite') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo-rocm') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-parakeet') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-transformers') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-vibevoice') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo-rocm') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-parakeet') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-transformers') or (extra != 'extra-12-asr-services-strixhalo' and extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-vibevoice') or (extra == 'extra-12-asr-services-cu128' and extra != 'group-12-asr-services-nemo' and extra != 'group-12-asr-services-nemo-rocm' and extra != 'group-12-asr-services-parakeet' and extra != 'group-12-asr-services-vibevoice') or (extra != 'extra-12-asr-services-cu126' and extra != 'group-12-asr-services-nemo' and extra != 'group-12-asr-services-nemo-rocm' and extra != 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr') or (extra != 'extra-12-asr-services-cu126' and extra != 'group-12-asr-services-nemo' and extra != 'group-12-asr-services-nemo-rocm' and extra != 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra != 'extra-12-asr-services-cu126' and extra != 'group-12-asr-services-nemo' and extra != 'group-12-asr-services-nemo-rocm' and extra != 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-transformers') or (extra != 'extra-12-asr-services-cu126' and extra != 'group-12-asr-services-nemo' and extra != 'group-12-asr-services-nemo-rocm' and extra != 'group-12-asr-services-parakeet' and extra != 'group-12-asr-services-vibevoice') or (extra != 'group-12-asr-services-af-next' and extra != 'group-12-asr-services-faster-whisper' and extra != 'group-12-asr-services-gemma4' and extra != 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-nemo-rocm') or (extra != 'group-12-asr-services-af-next' and extra != 'group-12-asr-services-faster-whisper' and extra != 'group-12-asr-services-gemma4' and extra != 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-parakeet') or (extra != 'group-12-asr-services-af-next' and extra != 'group-12-asr-services-faster-whisper' and extra != 'group-12-asr-services-gemma4' and extra != 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr') or (extra != 'group-12-asr-services-af-next' and extra != 'group-12-asr-services-faster-whisper' and extra != 'group-12-asr-services-gemma4' and extra != 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra != 'group-12-asr-services-af-next' and extra != 'group-12-asr-services-faster-whisper' and extra != 'group-12-asr-services-gemma4' and extra != 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-transformers') or (extra != 'group-12-asr-services-af-next' and extra != 'group-12-asr-services-faster-whisper' and extra != 'group-12-asr-services-gemma4' and extra != 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-vibevoice') or (extra != 'group-12-asr-services-af-next' and extra != 'group-12-asr-services-faster-whisper' and extra != 'group-12-asr-services-gemma4' and extra != 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-parakeet') or (extra != 'group-12-asr-services-af-next' and extra != 'group-12-asr-services-faster-whisper' and extra != 'group-12-asr-services-gemma4' and extra != 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr') or (extra != 'group-12-asr-services-af-next' and extra != 'group-12-asr-services-faster-whisper' and extra != 'group-12-asr-services-gemma4' and extra != 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra != 'group-12-asr-services-af-next' and extra != 'group-12-asr-services-faster-whisper' and extra != 'group-12-asr-services-gemma4' and extra != 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-transformers') or (extra != 'group-12-asr-services-af-next' and extra != 'group-12-asr-services-faster-whisper' and extra != 'group-12-asr-services-gemma4' and extra != 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-vibevoice') or (extra != 'group-12-asr-services-af-next' and extra != 'group-12-asr-services-faster-whisper' and extra != 'group-12-asr-services-gemma4' and extra != 'group-12-asr-services-granite' and extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr') or (extra != 'group-12-asr-services-af-next' and extra != 'group-12-asr-services-faster-whisper' and extra != 'group-12-asr-services-gemma4' and extra != 'group-12-asr-services-granite' and extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra != 'group-12-asr-services-af-next' and extra != 'group-12-asr-services-faster-whisper' and extra != 'group-12-asr-services-gemma4' and extra != 'group-12-asr-services-granite' and extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-transformers') or (extra != 'group-12-asr-services-af-next' and extra != 'group-12-asr-services-faster-whisper' and extra != 'group-12-asr-services-gemma4' and extra != 'group-12-asr-services-granite' and extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-vibevoice') or (extra != 'extra-12-asr-services-strixhalo' and extra != 'group-12-asr-services-nemo' and extra != 'group-12-asr-services-nemo-rocm' and extra != 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra != 'extra-12-asr-services-strixhalo' and extra != 'group-12-asr-services-nemo' and extra != 'group-12-asr-services-nemo-rocm' and extra != 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-transformers') or (extra != 'extra-12-asr-services-strixhalo' and extra != 'group-12-asr-services-nemo' and extra != 'group-12-asr-services-nemo-rocm' and extra != 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-vibevoice') or (extra != 'extra-12-asr-services-strixhalo' and extra != 'group-12-asr-services-nemo' and extra != 'group-12-asr-services-nemo-rocm' and extra != 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-transformers') or (extra != 'extra-12-asr-services-strixhalo' and extra != 'group-12-asr-services-nemo' and extra != 'group-12-asr-services-nemo-rocm' and extra != 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-vibevoice') or (extra != 'extra-12-asr-services-strixhalo' and extra != 'group-12-asr-services-nemo' and extra != 'group-12-asr-services-nemo-rocm' and extra != 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-transformers' and extra == 'group-12-asr-services-vibevoice')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/fa/41/e79269ce215c857c935fd86bcfe91a451a584dfc27f1e068f568b9ad1ab7/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:c9132cc3f8958447b4910a1720036d9eff5928cc3179b0a51fb6d167c6cc87d8", size = 705026878, upload_time = "2025-06-06T21:52:51.348Z" }, @@ -6670,7 +6781,7 @@ resolution-markers = [ "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-12-asr-services-cu126' and extra != 'extra-12-asr-services-cu128' and extra != 'extra-12-asr-services-strixhalo' and extra != 'group-12-asr-services-af-next' and extra != 'group-12-asr-services-faster-whisper' and extra != 'group-12-asr-services-gemma4' and extra != 'group-12-asr-services-granite' and extra != 'group-12-asr-services-nemo' and extra != 'group-12-asr-services-nemo-rocm' and extra != 'group-12-asr-services-parakeet' and extra != 'group-12-asr-services-qwen3-asr' and extra != 'group-12-asr-services-qwen3-asr-full' and extra != 'group-12-asr-services-transformers' and extra != 'group-12-asr-services-vibevoice'", ] dependencies = [ - { name = "nvidia-nvjitlink-cu12", version = "12.6.85", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-12-asr-services-cu126' or (extra == 'extra-12-asr-services-cu128' and extra == 'extra-12-asr-services-strixhalo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-faster-whisper') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-transformers' and extra == 'group-12-asr-services-vibevoice')" }, + { name = "nvidia-nvjitlink-cu12", version = "12.6.85", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-12-asr-services-cu126') or (extra == 'extra-12-asr-services-cu126' and extra == 'extra-12-asr-services-cu128') or (extra == 'extra-12-asr-services-cu126' and extra == 'extra-12-asr-services-strixhalo') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-af-next') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-faster-whisper') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-gemma4') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-granite') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-transformers') or (extra == 'extra-12-asr-services-cu128' and extra == 'extra-12-asr-services-strixhalo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-faster-whisper') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-transformers' and extra == 'group-12-asr-services-vibevoice') or (extra == 'extra-12-asr-services-cu126' and extra != 'group-12-asr-services-nemo' and extra != 'group-12-asr-services-nemo-rocm' and extra != 'group-12-asr-services-parakeet' and extra != 'group-12-asr-services-vibevoice')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/1f/37/c50d2b2f2c07e146776389e3080f4faf70bcc4fa6e19d65bb54ca174ebc3/nvidia_cufft_cu12-11.3.0.4-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d16079550df460376455cba121db6564089176d9bac9e4f360493ca4741b22a6", size = 200164144, upload_time = "2024-11-20T17:40:58.288Z" }, @@ -6879,7 +6990,7 @@ resolution-markers = [ "python_full_version < '3.11' and extra != 'extra-12-asr-services-cu126' and extra != 'extra-12-asr-services-cu128' and extra != 'extra-12-asr-services-strixhalo' and extra != 'group-12-asr-services-af-next' and extra != 'group-12-asr-services-faster-whisper' and extra != 'group-12-asr-services-gemma4' and extra != 'group-12-asr-services-granite' and extra != 'group-12-asr-services-nemo' and extra != 'group-12-asr-services-nemo-rocm' and extra != 'group-12-asr-services-parakeet' and extra != 'group-12-asr-services-qwen3-asr' and extra != 'group-12-asr-services-qwen3-asr-full' and extra != 'group-12-asr-services-transformers' and extra != 'group-12-asr-services-vibevoice'", ] dependencies = [ - { name = "nvidia-nvjitlink-cu12", version = "12.8.93", source = { registry = "https://pypi.org/simple" } }, + { name = "nvidia-nvjitlink-cu12", version = "12.8.93", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/60/bc/7771846d3a0272026c416fbb7e5f4c1f146d6d80704534d0b187dd6f4800/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:848ef7224d6305cdb2a4df928759dca7b1201874787083b6e7550dd6765ce69a", size = 193109211, upload_time = "2025-03-07T01:44:56.873Z" }, @@ -7614,9 +7725,9 @@ resolution-markers = [ "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-12-asr-services-cu126' and extra != 'extra-12-asr-services-cu128' and extra != 'extra-12-asr-services-strixhalo' and extra != 'group-12-asr-services-af-next' and extra != 'group-12-asr-services-faster-whisper' and extra != 'group-12-asr-services-gemma4' and extra != 'group-12-asr-services-granite' and extra != 'group-12-asr-services-nemo' and extra != 'group-12-asr-services-nemo-rocm' and extra != 'group-12-asr-services-parakeet' and extra != 'group-12-asr-services-qwen3-asr' and extra != 'group-12-asr-services-qwen3-asr-full' and extra != 'group-12-asr-services-transformers' and extra != 'group-12-asr-services-vibevoice'", ] dependencies = [ - { name = "nvidia-cublas-cu12", version = "12.6.4.1", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-12-asr-services-cu126' or (extra == 'extra-12-asr-services-cu128' and extra == 'extra-12-asr-services-strixhalo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-faster-whisper') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-transformers' and extra == 'group-12-asr-services-vibevoice')" }, - { name = "nvidia-cusparse-cu12", version = "12.5.4.2", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-12-asr-services-cu126' or (extra == 'extra-12-asr-services-cu128' and extra == 'extra-12-asr-services-strixhalo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-faster-whisper') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-transformers' and extra == 'group-12-asr-services-vibevoice')" }, - { name = "nvidia-nvjitlink-cu12", version = "12.6.85", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-12-asr-services-cu126' or (extra == 'extra-12-asr-services-cu128' and extra == 'extra-12-asr-services-strixhalo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-faster-whisper') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-transformers' and extra == 'group-12-asr-services-vibevoice')" }, + { name = "nvidia-cublas-cu12", version = "12.6.4.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-12-asr-services-cu126') or (extra == 'extra-12-asr-services-cu126' and extra == 'extra-12-asr-services-cu128') or (extra == 'extra-12-asr-services-cu126' and extra == 'extra-12-asr-services-strixhalo') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-af-next') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-faster-whisper') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-gemma4') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-granite') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-transformers') or (extra == 'extra-12-asr-services-cu128' and extra == 'extra-12-asr-services-strixhalo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-faster-whisper') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-transformers' and extra == 'group-12-asr-services-vibevoice') or (extra == 'extra-12-asr-services-cu126' and extra != 'group-12-asr-services-nemo' and extra != 'group-12-asr-services-nemo-rocm' and extra != 'group-12-asr-services-parakeet' and extra != 'group-12-asr-services-vibevoice')" }, + { name = "nvidia-cusparse-cu12", version = "12.5.4.2", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-12-asr-services-cu126') or (extra == 'extra-12-asr-services-cu126' and extra == 'extra-12-asr-services-cu128') or (extra == 'extra-12-asr-services-cu126' and extra == 'extra-12-asr-services-strixhalo') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-af-next') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-faster-whisper') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-gemma4') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-granite') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-transformers') or (extra == 'extra-12-asr-services-cu128' and extra == 'extra-12-asr-services-strixhalo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-faster-whisper') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-transformers' and extra == 'group-12-asr-services-vibevoice') or (extra == 'extra-12-asr-services-cu126' and extra != 'group-12-asr-services-nemo' and extra != 'group-12-asr-services-nemo-rocm' and extra != 'group-12-asr-services-parakeet' and extra != 'group-12-asr-services-vibevoice')" }, + { name = "nvidia-nvjitlink-cu12", version = "12.6.85", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-12-asr-services-cu126') or (extra == 'extra-12-asr-services-cu126' and extra == 'extra-12-asr-services-cu128') or (extra == 'extra-12-asr-services-cu126' and extra == 'extra-12-asr-services-strixhalo') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-af-next') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-faster-whisper') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-gemma4') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-granite') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-transformers') or (extra == 'extra-12-asr-services-cu128' and extra == 'extra-12-asr-services-strixhalo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-faster-whisper') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-transformers' and extra == 'group-12-asr-services-vibevoice') or (extra == 'extra-12-asr-services-cu126' and extra != 'group-12-asr-services-nemo' and extra != 'group-12-asr-services-nemo-rocm' and extra != 'group-12-asr-services-parakeet' and extra != 'group-12-asr-services-vibevoice')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/93/17/dbe1aa865e4fdc7b6d4d0dd308fdd5aaab60f939abfc0ea1954eac4fb113/nvidia_cusolver_cu12-11.7.1.2-py3-none-manylinux2014_aarch64.whl", hash = "sha256:0ce237ef60acde1efc457335a2ddadfd7610b892d94efee7b776c64bb1cac9e0", size = 157833628, upload_time = "2024-10-01T17:05:05.591Z" }, @@ -7825,9 +7936,9 @@ resolution-markers = [ "python_full_version < '3.11' and extra != 'extra-12-asr-services-cu126' and extra != 'extra-12-asr-services-cu128' and extra != 'extra-12-asr-services-strixhalo' and extra != 'group-12-asr-services-af-next' and extra != 'group-12-asr-services-faster-whisper' and extra != 'group-12-asr-services-gemma4' and extra != 'group-12-asr-services-granite' and extra != 'group-12-asr-services-nemo' and extra != 'group-12-asr-services-nemo-rocm' and extra != 'group-12-asr-services-parakeet' and extra != 'group-12-asr-services-qwen3-asr' and extra != 'group-12-asr-services-qwen3-asr-full' and extra != 'group-12-asr-services-transformers' and extra != 'group-12-asr-services-vibevoice'", ] dependencies = [ - { name = "nvidia-cublas-cu12", version = "12.8.4.1", source = { registry = "https://pypi.org/simple" } }, - { name = "nvidia-cusparse-cu12", version = "12.5.8.93", source = { registry = "https://pypi.org/simple" } }, - { name = "nvidia-nvjitlink-cu12", version = "12.8.93", source = { registry = "https://pypi.org/simple" } }, + { name = "nvidia-cublas-cu12", version = "12.8.4.1", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'linux'" }, + { name = "nvidia-cusparse-cu12", version = "12.5.8.93", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", version = "12.8.93", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/c8/32/f7cd6ce8a7690544d084ea21c26e910a97e077c9b7f07bf5de623ee19981/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:db9ed69dbef9715071232caa9b69c52ac7de3a95773c2db65bdba85916e4e5c0", size = 267229841, upload_time = "2025-03-07T01:46:54.356Z" }, @@ -7938,7 +8049,7 @@ resolution-markers = [ "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-12-asr-services-cu126' and extra != 'extra-12-asr-services-cu128' and extra != 'extra-12-asr-services-strixhalo' and extra != 'group-12-asr-services-af-next' and extra != 'group-12-asr-services-faster-whisper' and extra != 'group-12-asr-services-gemma4' and extra != 'group-12-asr-services-granite' and extra != 'group-12-asr-services-nemo' and extra != 'group-12-asr-services-nemo-rocm' and extra != 'group-12-asr-services-parakeet' and extra != 'group-12-asr-services-qwen3-asr' and extra != 'group-12-asr-services-qwen3-asr-full' and extra != 'group-12-asr-services-transformers' and extra != 'group-12-asr-services-vibevoice'", ] dependencies = [ - { name = "nvidia-nvjitlink-cu12", version = "12.6.85", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-12-asr-services-cu126' or (extra == 'extra-12-asr-services-cu128' and extra == 'extra-12-asr-services-strixhalo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-faster-whisper') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-transformers' and extra == 'group-12-asr-services-vibevoice')" }, + { name = "nvidia-nvjitlink-cu12", version = "12.6.85", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-12-asr-services-cu126') or (extra == 'extra-12-asr-services-cu126' and extra == 'extra-12-asr-services-cu128') or (extra == 'extra-12-asr-services-cu126' and extra == 'extra-12-asr-services-strixhalo') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-af-next') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-faster-whisper') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-gemma4') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-granite') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'extra-12-asr-services-cu126' and extra == 'group-12-asr-services-transformers') or (extra == 'extra-12-asr-services-cu128' and extra == 'extra-12-asr-services-strixhalo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-faster-whisper') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-af-next' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-gemma4') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-faster-whisper' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-granite') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-gemma4' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-granite' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-nemo-rocm') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-parakeet') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-nemo-rocm' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-parakeet' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-qwen3-asr-full') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-transformers') or (extra == 'group-12-asr-services-qwen3-asr-full' and extra == 'group-12-asr-services-vibevoice') or (extra == 'group-12-asr-services-transformers' and extra == 'group-12-asr-services-vibevoice') or (extra == 'extra-12-asr-services-cu126' and extra != 'group-12-asr-services-nemo' and extra != 'group-12-asr-services-nemo-rocm' and extra != 'group-12-asr-services-parakeet' and extra != 'group-12-asr-services-vibevoice')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/eb/eb/6681efd0aa7df96b4f8067b3ce7246833dd36830bb4cec8896182773db7d/nvidia_cusparse_cu12-12.5.4.2-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d25b62fb18751758fe3c93a4a08eff08effedfe4edf1c6bb5afd0890fe88f887", size = 216451147, upload_time = "2024-11-20T17:44:18.055Z" }, @@ -8147,7 +8258,7 @@ resolution-markers = [ "python_full_version < '3.11' and extra != 'extra-12-asr-services-cu126' and extra != 'extra-12-asr-services-cu128' and extra != 'extra-12-asr-services-strixhalo' and extra != 'group-12-asr-services-af-next' and extra != 'group-12-asr-services-faster-whisper' and extra != 'group-12-asr-services-gemma4' and extra != 'group-12-asr-services-granite' and extra != 'group-12-asr-services-nemo' and extra != 'group-12-asr-services-nemo-rocm' and extra != 'group-12-asr-services-parakeet' and extra != 'group-12-asr-services-qwen3-asr' and extra != 'group-12-asr-services-qwen3-asr-full' and extra != 'group-12-asr-services-transformers' and extra != 'group-12-asr-services-vibevoice'", ] dependencies = [ - { name = "nvidia-nvjitlink-cu12", version = "12.8.93", source = { registry = "https://pypi.org/simple" } }, + { name = "nvidia-nvjitlink-cu12", version = "12.8.93", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/bc/f7/cd777c4109681367721b00a106f491e0d0d15cfa1fd59672ce580ce42a97/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:9b6c161cb130be1a07a27ea6923df8141f3c295852f4b260c65f18f3e0a091dc", size = 288117129, upload_time = "2025-03-07T01:47:40.407Z" }, @@ -9996,6 +10107,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/29/16/c8a903f4c4dffe7a12843191437d7cd8e32751d5de349d45d3fe69544e87/pytest-8.4.1-py3-none-any.whl", hash = "sha256:539c70ba6fcead8e78eebbf1115e8b589e7565830d7d006a8723f19ac8a0afb7", size = 365474, upload_time = "2025-06-18T05:48:03.955Z" }, ] +[[package]] +name = "pytest-cov" +version = "7.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "coverage", extra = ["toml"] }, + { name = "pluggy" }, + { name = "pytest" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b1/51/a849f96e117386044471c8ec2bd6cfebacda285da9525c9106aeb28da671/pytest_cov-7.1.0.tar.gz", hash = "sha256:30674f2b5f6351aa09702a9c8c364f6a01c27aae0c1366ae8016160d1efc56b2", size = 55592, upload_time = "2026-03-21T20:11:16.284Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9d/7a/d968e294073affff457b041c2be9868a40c1c71f4a35fcc1e45e5493067b/pytest_cov-7.1.0-py3-none-any.whl", hash = "sha256:a0461110b7865f9a271aa1b51e516c9a95de9d696734a2f71e3e78f46e1d4678", size = 22876, upload_time = "2026-03-21T20:11:14.438Z" }, +] + [[package]] name = "python-dateutil" version = "2.9.0.post0" From 43b25ae7f06e2142c8822b2ecaeb6a3c6632c80d Mon Sep 17 00:00:00 2001 From: Ankush Malaker <43288948+AnkushMalaker@users.noreply.github.com> Date: Fri, 17 Jul 2026 14:10:50 +0000 Subject: [PATCH 03/16] feat(speaker): add enrollment audit workflows --- extras/speaker-recognition/.env.template | 2 + extras/speaker-recognition/docker-compose.yml | 2 +- .../scripts/backfill_segment_embeddings.py | 121 ++++++ .../api/routers/enrollment_audit.py | 383 +++++++++++++++++- .../core/audio_backend.py | 4 +- .../core/enrollment_audit.py | 12 +- .../webui/src/components/ConnectionStatus.tsx | 14 +- .../webui/src/contexts/UserContext.tsx | 15 +- .../webui/src/pages/EnrollmentHealth.tsx | 83 +++- .../webui/src/services/api.ts | 35 +- .../speaker-recognition/webui/vite.config.ts | 25 ++ 11 files changed, 635 insertions(+), 61 deletions(-) create mode 100644 extras/speaker-recognition/scripts/backfill_segment_embeddings.py diff --git a/extras/speaker-recognition/.env.template b/extras/speaker-recognition/.env.template index db795be8..dde6901d 100644 --- a/extras/speaker-recognition/.env.template +++ b/extras/speaker-recognition/.env.template @@ -38,6 +38,8 @@ GROQ_API_KEY= # Service binding configuration # SPEAKER_SERVICE_HOST=0.0.0.0 # SPEAKER_SERVICE_PORT=8085 +# Internal hostname used by the web UI proxy (usually the compose service alias) +# SPEAKER_SERVICE_PROXY_HOST=speaker-service # React UI configuration # REACT_UI_HOST=0.0.0.0 diff --git a/extras/speaker-recognition/docker-compose.yml b/extras/speaker-recognition/docker-compose.yml index d76af898..e2f995ad 100644 --- a/extras/speaker-recognition/docker-compose.yml +++ b/extras/speaker-recognition/docker-compose.yml @@ -120,7 +120,7 @@ services: - REACT_UI_PORT=${REACT_UI_PORT} - REACT_UI_HTTPS=${REACT_UI_HTTPS} - VITE_ALLOWED_HOSTS=${VITE_ALLOWED_HOSTS:-localhost 127.0.0.1} - - SPEAKER_SERVICE_HOST=${SPEAKER_SERVICE_HOST} + - SPEAKER_SERVICE_PROXY_HOST=${SPEAKER_SERVICE_PROXY_HOST:-speaker-service} - SPEAKER_SERVICE_PORT=${SPEAKER_SERVICE_PORT} restart: unless-stopped healthcheck: diff --git a/extras/speaker-recognition/scripts/backfill_segment_embeddings.py b/extras/speaker-recognition/scripts/backfill_segment_embeddings.py new file mode 100644 index 00000000..2afc4e97 --- /dev/null +++ b/extras/speaker-recognition/scripts/backfill_segment_embeddings.py @@ -0,0 +1,121 @@ +"""Backfill per-clip embeddings (SpeakerAudioSegment rows) from enrollment audio. + +Historically the service stored only a single averaged centroid per speaker +(`speakers.embedding_data`); the `speaker_audio_segments.embedding` column was never +populated. The enrollment-health audit (and a future multi-vector gallery) need one +embedding per enrolled clip, so this one-time importer walks the on-disk enrollment +audio, embeds each clip with the wespeaker model, and inserts a SpeakerAudioSegment row. + +Idempotent: skips any (speaker_id, filename) that already has a row, so it's safe to +re-run after new enrollments. New enrollments persist their own segments going forward +(see enrollment.py:save_segment_record), so this is only for pre-existing data. + +Run inside the speaker-service container: + podman exec speaker-recognition_speaker-service-gpu_1 \ + python3 /app/scripts/backfill_segment_embeddings.py +""" + +import glob +import json +import os +import sys +from pathlib import Path + +import numpy as np +import torch + +# Make the package importable when run as a plain script inside the container. +sys.path.insert(0, "/app/src") + +from pyannote.audio import Audio # noqa: E402 +from pyannote.audio.pipelines.speaker_verification import ( # noqa: E402 + PretrainedSpeakerEmbedding, +) + +from simple_speaker_recognition.database import get_db_session # noqa: E402 +from simple_speaker_recognition.database.models import ( # noqa: E402 + Speaker, + SpeakerAudioSegment, +) + +DATA_DIR = Path(os.getenv("SPEAKER_DATA_DIR", "/app/data")) +ENROLL_DIR = DATA_DIR / "enrollment_audio" +MIN_SAMPLES = 400 # wespeaker fbank window (25 ms @ 16 kHz) + + +def main() -> None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + embedder = PretrainedSpeakerEmbedding( + "pyannote/wespeaker-voxceleb-resnet34-LM", device=device + ) + loader = Audio(sample_rate=16000, mono="downmix") + + def embed(path: str): + wav, _ = loader(path) + wav = wav.unsqueeze(0) + if wav.shape[-1] < MIN_SAMPLES: + wav = torch.nn.functional.pad(wav, (0, MIN_SAMPLES - wav.shape[-1])) + with torch.inference_mode(): + e = embedder(wav.to(device)) + if isinstance(e, torch.Tensor): + e = e.cpu().numpy() + e = np.asarray(e).reshape(-1) + n = np.linalg.norm(e) + return (e / n) if np.isfinite(n) and n > 0 else None + + session = get_db_session() + try: + speaker_ids = {s.id for s in session.query(Speaker).all()} + # Existing (speaker_id, filename) pairs for idempotency. + existing = { + (row.speaker_id, os.path.basename(row.audio_file_path)) + for row in session.query(SpeakerAudioSegment).all() + } + + added = skipped = orphan = 0 + for user_dir in sorted(glob.glob(str(ENROLL_DIR / "*"))): + for sdir in sorted(glob.glob(f"{user_dir}/*")): + speaker_id = os.path.basename(sdir) + if speaker_id not in speaker_ids: + orphan += 1 + continue + for wav in sorted(glob.glob(f"{sdir}/*.wav")): + fname = os.path.basename(wav) + if (speaker_id, fname) in existing: + skipped += 1 + continue + try: + dur = float(loader.get_duration(wav)) + vec = embed(wav) + except Exception as e: # noqa: BLE001 + print(f" SKIP {wav}: {e}") + continue + if vec is None: + print(f" NULL embedding {wav}") + continue + rel = str(Path(wav).resolve().relative_to(ENROLL_DIR.resolve())) + session.add( + SpeakerAudioSegment( + speaker_id=speaker_id, + audio_file_path=rel, + original_file_path=fname, + start_time=0.0, + end_time=dur, + duration_seconds=dur, + embedding=json.dumps(vec.astype(np.float32).tolist()), + ) + ) + added += 1 + session.commit() + + session.commit() + print( + f"Backfill done: added={added} skipped(existing)={skipped} " + f"orphan_dirs(no speaker row)={orphan}" + ) + finally: + session.close() + + +if __name__ == "__main__": + main() diff --git a/extras/speaker-recognition/src/simple_speaker_recognition/api/routers/enrollment_audit.py b/extras/speaker-recognition/src/simple_speaker_recognition/api/routers/enrollment_audit.py index 9d9a2b43..e69500e7 100644 --- a/extras/speaker-recognition/src/simple_speaker_recognition/api/routers/enrollment_audit.py +++ b/extras/speaker-recognition/src/simple_speaker_recognition/api/routers/enrollment_audit.py @@ -6,23 +6,38 @@ fix takes effect on the next identification. """ +import asyncio +import json import logging +import os import shutil +from datetime import datetime from pathlib import Path from typing import Optional -from fastapi import APIRouter, Depends, Form, HTTPException, Query +import numpy as np +from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, UploadFile from fastapi.responses import FileResponse +from pydantic import BaseModel, Field + +from simple_speaker_recognition.api.core.utils import secure_temp_file from simple_speaker_recognition.core.enrollment_audit import ( compute_audit, recompute_speaker_centroid, ) from simple_speaker_recognition.core.unified_speaker_db import UnifiedSpeakerDB from simple_speaker_recognition.database import get_db_session -from simple_speaker_recognition.database.models import Speaker, SpeakerAudioSegment +from simple_speaker_recognition.database.models import ( + ProcessingJob, + Speaker, + SpeakerAudioSegment, +) +from simple_speaker_recognition.utils.audio_processing import get_audio_info router = APIRouter() log = logging.getLogger("speaker_service") +BACKFILL_JOB_TYPE = "enrollment_segment_backfill" +_backfill_tasks: set[asyncio.Task] = set() async def get_db(): @@ -59,11 +74,168 @@ async def enrollment_health( user_id: Optional[int] = Query( None, description="Scope audit to a user's speakers" ), + before: Optional[datetime] = Query( + None, description="Only include clips created before this timestamp" + ), ): """Contamination audit over per-clip enrollment embeddings.""" session = get_db_session() try: - return compute_audit(session, user_id) + return compute_audit(session, user_id, before) + finally: + session.close() + + +def _backfill_job_response(job: ProcessingJob) -> dict: + output = json.loads(job.output_data) if job.output_data else None + return { + "id": job.id, + "status": job.status, + "progress": job.progress, + "result": output, + "error": job.error_message, + "created_at": job.created_at, + "started_at": job.started_at, + "completed_at": job.completed_at, + } + + +async def _run_segment_backfill(job_id: int, user_id: int) -> None: + """Populate missing per-clip embeddings using the service's loaded GPU model.""" + from .. import service + + session = get_db_session() + try: + job = session.query(ProcessingJob).filter(ProcessingJob.id == job_id).one() + job.status = "running" + job.started_at = datetime.utcnow() + session.commit() + + speakers = session.query(Speaker).filter(Speaker.user_id == user_id).all() + speaker_ids = {speaker.id for speaker in speakers} + existing = { + (row.speaker_id, os.path.basename(row.audio_file_path)) + for row in session.query(SpeakerAudioSegment) + .join(Speaker) + .filter(Speaker.user_id == user_id) + .all() + } + enrollment_dir = get_auth().enrollment_audio_dir.resolve() + user_dir = enrollment_dir / str(user_id) + candidates = ( + [ + (speaker_dir.name, wav) + for speaker_dir in sorted(user_dir.iterdir()) + if speaker_dir.is_dir() and speaker_dir.name in speaker_ids + for wav in sorted(speaker_dir.glob("*.wav")) + if (speaker_dir.name, wav.name) not in existing + ] + if user_dir.exists() + else [] + ) + + added = failed = 0 + total = len(candidates) + for index, (speaker_id, wav) in enumerate(candidates, start=1): + try: + duration = float(service.audio_backend.loader.get_duration(str(wav))) + wave = service.audio_backend.load_wave(wav) + vector = np.asarray( + await service.audio_backend.async_embed(wave), dtype=np.float32 + ).reshape(-1) + norm = np.linalg.norm(vector) + if not np.isfinite(norm) or norm <= 0: + raise ValueError("embedding was not finite") + vector /= norm + relative_path = str(wav.resolve().relative_to(enrollment_dir)) + session.add( + SpeakerAudioSegment( + speaker_id=speaker_id, + audio_file_path=relative_path, + original_file_path=wav.name, + start_time=0.0, + end_time=duration, + duration_seconds=duration, + embedding=json.dumps(vector.tolist()), + ) + ) + added += 1 + except Exception: + failed += 1 + log.exception("Could not backfill enrollment clip %s", wav) + + job.progress = round(index * 100 / total, 1) if total else 100.0 + session.commit() + + job.status = "completed" + job.progress = 100.0 + job.output_data = json.dumps( + {"added": added, "failed": failed, "already_available": len(existing)} + ) + job.completed_at = datetime.utcnow() + session.commit() + except Exception as exc: + session.rollback() + job = session.query(ProcessingJob).filter(ProcessingJob.id == job_id).first() + if job: + job.status = "failed" + job.error_message = str(exc) + job.completed_at = datetime.utcnow() + session.commit() + log.exception("Enrollment segment backfill job %s failed", job_id) + finally: + session.close() + + +@router.post("/enrollment/segments/backfill", status_code=202) +async def start_segment_backfill(user_id: int = Query(...)): + """Start or attach to the idempotent per-clip embedding backfill.""" + session = get_db_session() + try: + active = ( + session.query(ProcessingJob) + .filter( + ProcessingJob.job_type == BACKFILL_JOB_TYPE, + ProcessingJob.status.in_(["pending", "running"]), + ProcessingJob.input_data == json.dumps({"user_id": user_id}), + ) + .order_by(ProcessingJob.id.desc()) + .first() + ) + if active: + return _backfill_job_response(active) + + job = ProcessingJob( + job_type=BACKFILL_JOB_TYPE, + input_data=json.dumps({"user_id": user_id}), + ) + session.add(job) + session.commit() + session.refresh(job) + response = _backfill_job_response(job) + task = asyncio.create_task(_run_segment_backfill(job.id, user_id)) + _backfill_tasks.add(task) + task.add_done_callback(_backfill_tasks.discard) + return response + finally: + session.close() + + +@router.get("/enrollment/segments/backfill") +async def get_segment_backfill(user_id: int = Query(...)): + """Return the latest backfill job so refreshed pages reattach to it.""" + session = get_db_session() + try: + job = ( + session.query(ProcessingJob) + .filter( + ProcessingJob.job_type == BACKFILL_JOB_TYPE, + ProcessingJob.input_data == json.dumps({"user_id": user_id}), + ) + .order_by(ProcessingJob.id.desc()) + .first() + ) + return _backfill_job_response(job) if job else None finally: session.close() @@ -89,6 +261,211 @@ async def segment_audio(segment_id: int): return FileResponse(str(path), media_type="audio/wav", filename=path.name) +def _unit(v: np.ndarray) -> Optional[np.ndarray]: + v = np.asarray(v, dtype=np.float32).reshape(-1) + n = np.linalg.norm(v) + if v.size == 0 or not np.all(np.isfinite(v)) or n == 0: + return None + return v / n + + +class EmbeddingScoreRequest(BaseModel): + speaker_id: str + embeddings: list[list[float]] = Field(..., min_length=1, max_length=5000) + + +def _score_embeddings(embeddings: list[list[float]], speaker_id: str) -> list[dict]: + """Compare precomputed unit embeddings with one live speaker gallery.""" + session = get_db_session() + try: + target = session.query(Speaker).filter(Speaker.id == speaker_id).first() + if not target: + raise HTTPException(404, "Target speaker not found") + centroid = ( + _unit(json.loads(target.embedding_data)) if target.embedding_data else None + ) + if centroid is None: + raise HTTPException(422, "Target speaker has no valid centroid") + + clip_vectors = [] + for seg in ( + session.query(SpeakerAudioSegment) + .filter(SpeakerAudioSegment.speaker_id == speaker_id) + .all() + ): + if seg.embedding: + vector = _unit(json.loads(seg.embedding)) + if vector is not None: + clip_vectors.append(vector) + + others = [] + for other in ( + session.query(Speaker) + .filter(Speaker.id != speaker_id, Speaker.user_id == target.user_id) + .all() + ): + if other.embedding_data: + vector = _unit(json.loads(other.embedding_data)) + if vector is not None: + others.append((other, vector)) + + results = [] + for values in embeddings: + emb = _unit(values) + if emb is None or emb.shape != centroid.shape: + results.append({"error": "invalid_embedding"}) + continue + best_other = None + for other, vector in others: + score = float(np.dot(emb, vector)) + if best_other is None or score > best_other["score"]: + best_other = { + "speaker_id": other.id, + "name": other.name, + "score": round(score, 4), + } + results.append( + { + "sim_centroid": round(float(np.dot(emb, centroid)), 4), + "max_clip_sim": ( + round( + max(float(np.dot(emb, vector)) for vector in clip_vectors), + 4, + ) + if clip_vectors + else None + ), + "n_gallery_clips": len(clip_vectors), + "best_other": best_other, + } + ) + return results + finally: + session.close() + + +@router.post("/enrollment/candidates/score-embeddings") +async def score_enrollment_embeddings(body: EmbeddingScoreRequest): + """Score cached corpus embeddings without decoding or embedding audio again.""" + return {"scores": _score_embeddings(body.embeddings, body.speaker_id)} + + +@router.post("/enrollment/candidates/score") +async def score_enrollment_candidate( + file: UploadFile = File(..., description="Candidate clip (single speaker)"), + speaker_id: str = Form(..., description="Target enrolled speaker"), +): + """Score a candidate clip's enrollment value for one target speaker. + + Returns the clip's cosine to the target centroid, its redundancy against the + target's existing per-clip gallery (max cosine — high means the gallery already + covers this acoustic condition), and the closest *other* enrolled speaker, so a + caller can rank candidates by marginal information instead of blind similarity. + """ + with secure_temp_file() as tmp: + tmp.write(await file.read()) + tmp_path = Path(tmp.name) + + try: + duration = get_audio_info(str(tmp_path)).get("duration_seconds") + from .. import service + + wav = service.audio_backend.load_wave(tmp_path) + emb = _unit(await service.audio_backend.async_embed(wav)) + if emb is None: + raise HTTPException(422, "Could not extract a valid embedding from clip") + finally: + tmp_path.unlink(missing_ok=True) + + session = get_db_session() + try: + target = session.query(Speaker).filter(Speaker.id == speaker_id).first() + if not target: + raise HTTPException(404, "Target speaker not found") + centroid = ( + _unit(json.loads(target.embedding_data)) if target.embedding_data else None + ) + if centroid is None: + raise HTTPException(422, "Target speaker has no valid centroid") + + clip_sims = [] + for seg in ( + session.query(SpeakerAudioSegment) + .filter(SpeakerAudioSegment.speaker_id == speaker_id) + .all() + ): + if not seg.embedding: + continue + v = _unit(json.loads(seg.embedding)) + if v is not None: + clip_sims.append(float(np.dot(emb, v))) + + best_other = None + others = ( + session.query(Speaker) + .filter(Speaker.id != speaker_id, Speaker.user_id == target.user_id) + .all() + ) + for other in others: + if not other.embedding_data: + continue + v = _unit(json.loads(other.embedding_data)) + if v is None: + continue + score = float(np.dot(emb, v)) + if best_other is None or score > best_other["score"]: + best_other = { + "speaker_id": other.id, + "name": other.name, + "score": round(score, 4), + } + + return { + "duration": round(float(duration), 3) if duration else None, + "sim_centroid": round(float(np.dot(emb, centroid)), 4), + "max_clip_sim": round(max(clip_sims), 4) if clip_sims else None, + "n_gallery_clips": len(clip_sims), + "best_other": best_other, + } + finally: + session.close() + + +@router.post("/enrollment/candidates/embed") +async def embed_enrollment_candidate( + file: UploadFile = File(..., description="Human-labeled evaluation clip"), +): + """Extract a unit speaker embedding without changing the live gallery.""" + with secure_temp_file() as tmp: + tmp.write(await file.read()) + tmp_path = Path(tmp.name) + try: + duration = get_audio_info(str(tmp_path)).get("duration_seconds") + from .. import service + + wav = service.audio_backend.load_wave(tmp_path) + emb = _unit(await service.audio_backend.async_embed(wav)) + if emb is None: + raise HTTPException(422, "Could not extract a valid embedding from clip") + return { + "embedding": emb.tolist(), + "duration": round(float(duration), 3) if duration else None, + "embedding_model": service.audio_backend.EMBEDDING_MODEL_ID, + } + finally: + tmp_path.unlink(missing_ok=True) + + +@router.get("/enrollment/candidates/embedding-info") +async def enrollment_embedding_info(): + """Fingerprint used to invalidate cached evaluation embeddings.""" + from .. import service + + return { + "embedding_model": service.audio_backend.EMBEDDING_MODEL_ID, + } + + @router.post("/enrollment/segments/{segment_id}/relabel") async def relabel_segment( segment_id: int, diff --git a/extras/speaker-recognition/src/simple_speaker_recognition/core/audio_backend.py b/extras/speaker-recognition/src/simple_speaker_recognition/core/audio_backend.py index d6dd173b..00354638 100644 --- a/extras/speaker-recognition/src/simple_speaker_recognition/core/audio_backend.py +++ b/extras/speaker-recognition/src/simple_speaker_recognition/core/audio_backend.py @@ -20,6 +20,8 @@ class AudioBackend: """Wrapper around PyAnnote & SpeechBrain components.""" + EMBEDDING_MODEL_ID = "pyannote/wespeaker-voxceleb-resnet34-LM" + def __init__(self, hf_token: str, device: torch.device): self.device = device self.diar = Pipeline.from_pretrained( @@ -38,7 +40,7 @@ def __init__(self, hf_token: str, device: torch.device): # Use the EXACT same embedding model that the diarization pipeline uses internally self.embedder = PretrainedSpeakerEmbedding( - "pyannote/wespeaker-voxceleb-resnet34-LM", device=device + self.EMBEDDING_MODEL_ID, device=device ) self.loader = Audio(sample_rate=16_000, mono="downmix") diff --git a/extras/speaker-recognition/src/simple_speaker_recognition/core/enrollment_audit.py b/extras/speaker-recognition/src/simple_speaker_recognition/core/enrollment_audit.py index 21685828..33b358a3 100644 --- a/extras/speaker-recognition/src/simple_speaker_recognition/core/enrollment_audit.py +++ b/extras/speaker-recognition/src/simple_speaker_recognition/core/enrollment_audit.py @@ -19,9 +19,11 @@ import json import logging import os +from datetime import datetime from typing import Optional import numpy as np + from simple_speaker_recognition.database.models import Speaker, SpeakerAudioSegment log = logging.getLogger("speaker_service") @@ -41,13 +43,15 @@ def _norm(v: np.ndarray) -> np.ndarray: return v / n if n > 0 else v -def _load_segments(session, user_id: Optional[int]): +def _load_segments(session, user_id: Optional[int], before: Optional[datetime] = None): """Return [(SpeakerAudioSegment, Speaker, unit_vec)] with valid embeddings.""" q = session.query(SpeakerAudioSegment, Speaker).join( Speaker, SpeakerAudioSegment.speaker_id == Speaker.id ) if user_id is not None: q = q.filter(Speaker.user_id == user_id) + if before is not None: + q = q.filter(SpeakerAudioSegment.created_at < before) rows = [] for seg, spk in q.all(): @@ -63,9 +67,11 @@ def _load_segments(session, user_id: Optional[int]): return rows -def compute_audit(session, user_id: Optional[int] = None) -> dict: +def compute_audit( + session, user_id: Optional[int] = None, before: Optional[datetime] = None +) -> dict: """Build the enrollment-health report for a user (or all users).""" - rows = _load_segments(session, user_id) + rows = _load_segments(session, user_id, before) by_spk: dict = collections.defaultdict(list) # speaker_id -> [(seg, vec)] name_by_id: dict = {} diff --git a/extras/speaker-recognition/webui/src/components/ConnectionStatus.tsx b/extras/speaker-recognition/webui/src/components/ConnectionStatus.tsx index beec5350..5b7e929f 100644 --- a/extras/speaker-recognition/webui/src/components/ConnectionStatus.tsx +++ b/extras/speaker-recognition/webui/src/components/ConnectionStatus.tsx @@ -1,4 +1,4 @@ -import React, { useState, useEffect } from 'react' +import { useState, useEffect } from 'react' import { Circle } from 'lucide-react' import axios from 'axios' @@ -15,17 +15,7 @@ export default function ConnectionStatus() { const [status, setStatus] = useState<BackendStatus>({ status: 'disconnected' }) const getBackendUrl = () => { - // For display purposes, show the URL that the browser can actually access - // In development, this would be localhost:8085 (via port-forward) - // In production with ingress, this would be the ingress URL - const isDevelopment = process.env.NODE_ENV === 'development' - - if (isDevelopment) { - return 'http://localhost:8085' - } else { - // In production, use the current window location but with /api prefix - return `${window.location.protocol}//${window.location.host}/api` - } + return `${window.location.origin}/api` } const checkBackendConnection = async () => { diff --git a/extras/speaker-recognition/webui/src/contexts/UserContext.tsx b/extras/speaker-recognition/webui/src/contexts/UserContext.tsx index 1534e5a9..3ce82758 100644 --- a/extras/speaker-recognition/webui/src/contexts/UserContext.tsx +++ b/extras/speaker-recognition/webui/src/contexts/UserContext.tsx @@ -59,17 +59,20 @@ export function UserProvider({ children }: { children: ReactNode }) { const initializeUser = async () => { setIsLoading(true) try { - // First, refresh users list - await refreshUsers() + const userList = await apiService.getUsers() + setUsers(userList) - // Check if there's a saved user const savedUser = localStorage.getItem('selectedUser') if (savedUser) { const parsedUser = JSON.parse(savedUser) - setUser(parsedUser) + const canonicalUser = userList.find(u => u.username === parsedUser.username) + if (canonicalUser) { + setUser(canonicalUser) + localStorage.setItem('selectedUser', JSON.stringify(canonicalUser)) + } else { + localStorage.removeItem('selectedUser') + } } else { - // Auto-create admin user if no users exist - const userList = await apiService.getUsers() if (userList.length === 0) { await createUser('admin') } diff --git a/extras/speaker-recognition/webui/src/pages/EnrollmentHealth.tsx b/extras/speaker-recognition/webui/src/pages/EnrollmentHealth.tsx index 4b3d34bb..ee46ca9a 100644 --- a/extras/speaker-recognition/webui/src/pages/EnrollmentHealth.tsx +++ b/extras/speaker-recognition/webui/src/pages/EnrollmentHealth.tsx @@ -1,7 +1,7 @@ import { useState, useEffect, useCallback } from 'react' import { ShieldCheck, ShieldAlert, ChevronDown, ChevronRight, - RefreshCw, Archive, ArrowRightLeft, HelpCircle, + RefreshCw, Archive, ArrowRightLeft, HelpCircle, Database, } from 'lucide-react' import { useUser } from '../contexts/UserContext' import { apiService } from '../services/api' @@ -32,6 +32,13 @@ interface AuditReport { total_clips: number speakers_without_segments: { speaker_id: string; name: string }[] } +interface BackfillJob { + id: number + status: 'pending' | 'running' | 'completed' | 'failed' + progress: number + result: { added: number; failed: number; already_available: number } | null + error: string | null +} const VERDICT_STYLES: Record<string, string> = { contaminated: 'bg-red-100 text-red-800 dark:bg-red-900 dark:text-red-200', @@ -59,6 +66,8 @@ export default function EnrollmentHealth() { const [error, setError] = useState<string | null>(null) const [expanded, setExpanded] = useState<Set<string>>(new Set()) const [busy, setBusy] = useState<number | null>(null) + const [backfill, setBackfill] = useState<BackfillJob | null>(null) + const [backfillError, setBackfillError] = useState<string | null>(null) const load = useCallback(async () => { if (!user) return @@ -79,6 +88,44 @@ export default function EnrollmentHealth() { useEffect(() => { load() }, [load]) + const checkBackfill = useCallback(async () => { + if (!user) return + try { + const res = await apiService.get('/enrollment/segments/backfill', { + params: { user_id: user.id }, + }) + setBackfill(res.data) + return res.data as BackfillJob | null + } catch (e: any) { + setBackfillError(e?.response?.data?.detail || e?.message || 'Could not check import status') + return null + } + }, [user]) + + useEffect(() => { checkBackfill() }, [checkBackfill]) + + useEffect(() => { + if (!backfill || !['pending', 'running'].includes(backfill.status)) return + const interval = setInterval(async () => { + const job = await checkBackfill() + if (job?.status === 'completed') await load() + }, 1500) + return () => clearInterval(interval) + }, [backfill?.id, backfill?.status, checkBackfill, load]) + + const startBackfill = async () => { + if (!user) return + setBackfillError(null) + try { + const res = await apiService.post('/enrollment/segments/backfill', undefined, { + params: { user_id: user.id }, + }) + setBackfill(res.data) + } catch (e: any) { + setBackfillError(e?.response?.data?.detail || e?.message || 'Could not start enrollment import') + } + } + const toggle = (id: string) => setExpanded(prev => { const next = new Set(prev) @@ -160,11 +207,35 @@ export default function EnrollmentHealth() { </div> {report.total_clips === 0 && ( - <div className="bg-yellow-50 dark:bg-yellow-900/30 text-yellow-800 dark:text-yellow-200 p-4 rounded-md text-sm"> - No per-clip embeddings found. Run the one-time backfill to import existing enrollment audio: - <code className="block mt-2 text-xs"> - podman exec speaker-recognition_speaker-service-gpu_1 python3 /app/scripts/backfill_segment_embeddings.py - </code> + <div className="bg-yellow-50 dark:bg-yellow-900/30 text-yellow-900 dark:text-yellow-100 p-4 rounded-md"> + <div className="font-medium">Existing enrollment clips need analysis</div> + <p className="mt-1 text-sm"> + Build the per-clip embeddings used to find mislabeled and low-quality audio. + Existing speaker voiceprints are not changed. + </p> + <div className="mt-3 flex items-center gap-3"> + <button + onClick={startBackfill} + disabled={backfill?.status === 'pending' || backfill?.status === 'running'} + className="flex items-center gap-2 px-3 py-2 text-sm font-medium bg-blue-600 text-white rounded-md hover:bg-blue-700 disabled:opacity-60" + > + <Database className={`h-4 w-4 ${backfill?.status === 'running' ? 'animate-pulse' : ''}`} /> + {backfill?.status === 'pending' || backfill?.status === 'running' + ? `Analyzing clips · ${Math.round(backfill.progress)}%` + : backfill?.status === 'failed' ? 'Try again' : 'Analyze enrollment clips'} + </button> + {backfill?.status === 'completed' && backfill.result && ( + <span className="text-sm"> + Added {backfill.result.added} clips + {backfill.result.failed > 0 ? ` · ${backfill.result.failed} failed` : ''} + </span> + )} + </div> + {(backfillError || backfill?.error) && ( + <div className="mt-3 text-sm text-red-700 dark:text-red-300"> + Analysis failed: {backfillError || backfill?.error} + </div> + )} </div> )} diff --git a/extras/speaker-recognition/webui/src/services/api.ts b/extras/speaker-recognition/webui/src/services/api.ts index 31318393..64229549 100644 --- a/extras/speaker-recognition/webui/src/services/api.ts +++ b/extras/speaker-recognition/webui/src/services/api.ts @@ -71,39 +71,16 @@ class ApiService { // User management async getUsers(): Promise<User[]> { - try { - const response = await api.get('/users') - return response.data - } catch (error) { - // If endpoint doesn't exist, return empty array and log - console.warn('Users endpoint not available, using local storage') - return [] + const response = await api.get('/users') + if (!Array.isArray(response.data)) { + throw new Error('Users endpoint returned an invalid response') } + return response.data } async getOrCreateUser(username: string): Promise<User> { - try { - const response = await api.post('/users', { username }) - return response.data - } catch (error) { - // Fallback to local user creation - const existingUsers = JSON.parse(localStorage.getItem('users') || '[]') - const existingUser = existingUsers.find((u: User) => u.username === username) - - if (existingUser) { - return existingUser - } - - const newUser: User = { - id: Date.now(), - username, - created_at: new Date().toISOString() - } - - existingUsers.push(newUser) - localStorage.setItem('users', JSON.stringify(existingUsers)) - return newUser - } + const response = await api.post('/users', { username }) + return response.data } // Speaker management diff --git a/extras/speaker-recognition/webui/vite.config.ts b/extras/speaker-recognition/webui/vite.config.ts index af46c533..5680580b 100644 --- a/extras/speaker-recognition/webui/vite.config.ts +++ b/extras/speaker-recognition/webui/vite.config.ts @@ -1,6 +1,10 @@ import { defineConfig } from 'vite' import react from '@vitejs/plugin-react' +const speakerServiceHost = process.env.SPEAKER_SERVICE_PROXY_HOST || 'speaker-service' +const speakerServicePort = process.env.SPEAKER_SERVICE_PORT || '8085' +const speakerServiceTarget = `http://${speakerServiceHost}:${speakerServicePort}` + // https://vitejs.dev/config/ // The dev server runs plain HTTP. When HTTPS is enabled, Caddy terminates TLS // (Tailscale/Let's Encrypt cert) and reverse-proxies to this server over HTTP — @@ -17,6 +21,27 @@ export default defineConfig({ '127.0.0.1', '.nip.io' ], + proxy: { + '/api': { + target: speakerServiceTarget, + changeOrigin: true, + rewrite: path => path.replace(/^\/api/, ''), + }, + '/health': { + target: speakerServiceTarget, + changeOrigin: true, + }, + '/ws': { + target: speakerServiceTarget, + changeOrigin: true, + ws: true, + }, + '/v1': { + target: speakerServiceTarget, + changeOrigin: true, + ws: true, + }, + }, }, define: { global: 'globalThis', From df99bb92ffcb29f3ae6ed4ea72365a7e227ebf79 Mon Sep 17 00:00:00 2001 From: Ankush Malaker <43288948+AnkushMalaker@users.noreply.github.com> Date: Fri, 17 Jul 2026 14:11:05 +0000 Subject: [PATCH 04/16] docs: record repository cleanup inventory --- REPOSITORY_CLEANUP.md | 98 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 98 insertions(+) create mode 100644 REPOSITORY_CLEANUP.md diff --git a/REPOSITORY_CLEANUP.md b/REPOSITORY_CLEANUP.md new file mode 100644 index 00000000..fae8aaad --- /dev/null +++ b/REPOSITORY_CLEANUP.md @@ -0,0 +1,98 @@ +# Repository cleanup inventory + +This document tracks the cleanup of the working tree started on 2026-07-17. + +## Initial state + +- Branch at inventory time: `dev` (matching `origin/dev`) +- Cleanup branch: `git-cleanup` +- Modified tracked files: 84 +- Untracked files: approximately 1,304 +- Untracked disk usage: approximately 25 GB +- Staged files at inventory time: none + +## Categories + +### Project work to review and commit + +- Advanced backend features and tests +- Advanced backend web UI changes +- ASR service source, configuration, and experiments that are suitable for the repository +- Speaker-recognition service and web UI changes +- CI and test infrastructure + +### Generated outputs likely needing repository ignore rules + +- Model checkpoints, adapters, and optimizer state +- Fine-tuning output directories +- Evaluation outputs and experiment logs +- Model and dataset caches +- Coverage reports + +These generated artifacts account for most of the untracked disk usage. Individual +source files inside experiment directories must be separated from generated outputs +before ignore rules are added. + +### Potentially personal or sensitive material + +- Root-level audio, screenshots, QR codes, transcripts, and notes +- `golden_data/` +- Conversation-vault evaluation artifacts under `artifacts/` +- Voice-cloning experiments and recordings +- Personal planning and research documents + +Do not move these without explicit confirmation. + +### Local-only files + +- Environment-file backups +- Caddy and TTS configuration backups +- Local logs, debug data, and model caches + +Use `.git/info/exclude` when a file is specific to this checkout and cannot be moved +under `untracked/`. Use repository `.gitignore` only for outputs routinely generated +for multiple contributors. + +## Privacy review + +The names `ankush`, `anushpa`, and `jahnvi` occur in both tracked and untracked +content. Tracked occurrences include examples, tests, documentation, and sample-vault +data. Untracked occurrences include evaluation artifacts, experiment logs, datasets, +research notes, and tests. The evaluation vault under `artifacts/` should be treated +as sensitive until reviewed. + +PII already present in Git history is out of scope. Current tracked and proposed new +content must be reviewed before the cleanup branch is finalized. + +## Rules for this cleanup + +- Confirm before moving any file. +- Do not delete or overwrite personal artifacts during categorization. +- Commit project work in coherent groups. +- Keep generated, local-only, and potentially sensitive material out of project + commits until explicitly classified. +- Do not add backward-compatibility code as part of cleanup. + +## Progress + +- [x] Initial read-only inventory +- [x] Create cleanup branch +- [x] Commit advanced backend project work (`4d52a0c1`) +- [x] Commit ASR service project work (`667f5716`) +- [x] Commit speaker-recognition project work (`43b25ae7`) +- [ ] Classify remaining untracked content +- [ ] Apply approved ignore and relocation decisions +- [ ] Clean PII from the current repository state + +## Verification notes + +- Repository pre-commit hooks passed for all three project commits. Black and isort + reformatted newly added Python files before the successful commits. +- The broad advanced-backend pytest run could not complete because existing MongoDB + persistence tests require unavailable local infrastructure. +- `backends/advanced/tests/test_vad_analysis.py` was deliberately left untracked. It + imports `silence_gaps_from_regions`, but that implementation is absent from the + working tree and Git history, so the test is currently incomplete. +- ASR checkpoints, optimizer state, adapter outputs, caches, and logs were not + committed. +- Speaker-recognition debug data and environment backups were not committed. From e2ab9ccd534a9fbfd326c64521f67196e89589f3 Mon Sep 17 00:00:00 2001 From: Ankush Malaker <43288948+AnkushMalaker@users.noreply.github.com> Date: Fri, 17 Jul 2026 14:12:30 +0000 Subject: [PATCH 05/16] chore(asr): untrack ad-hoc experiment files --- .../gemma4/UNIFIED_ENDPOINT_DESIGN.md | 178 ----------------- .../gemma4/finetune/debug_gen_nocache.py | 66 ------- .../gemma4/finetune/debug_inprocess.py | 137 ------------- .../gemma4/finetune/debug_overfit.py | 186 ------------------ .../gemma4/finetune/debug_reload_loss.py | 51 ----- .../gemma4/finetune/stream_client_test.py | 71 ------- .../providers/gemma4/finetune/test_cache57.py | 90 --------- .../gemma4/finetune/test_longaudio.py | 64 ------ .../gemma4/finetune/test_patch_loss.py | 131 ------------ .../tests/capture_vibevoice_ground_truth.py | 168 ---------------- 10 files changed, 1142 deletions(-) delete mode 100644 extras/asr-services/providers/gemma4/UNIFIED_ENDPOINT_DESIGN.md delete mode 100644 extras/asr-services/providers/gemma4/finetune/debug_gen_nocache.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/debug_inprocess.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/debug_overfit.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/debug_reload_loss.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/stream_client_test.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/test_cache57.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/test_longaudio.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/test_patch_loss.py delete mode 100644 extras/asr-services/tests/capture_vibevoice_ground_truth.py diff --git a/extras/asr-services/providers/gemma4/UNIFIED_ENDPOINT_DESIGN.md b/extras/asr-services/providers/gemma4/UNIFIED_ENDPOINT_DESIGN.md deleted file mode 100644 index 4d773b42..00000000 --- a/extras/asr-services/providers/gemma4/UNIFIED_ENDPOINT_DESIGN.md +++ /dev/null @@ -1,178 +0,0 @@ -# Gemma 4 unified multimodal endpoint — design - -## Problem - -Gemma 4 E2B is **one combined multimodal model**: audio understanding and text generation -are the *same* forward pass over interleaved text+audio tokens. But the service exposes it -as a set of task-specific wrappers, and none of them accept **more than one audio clip per -request**: - -| endpoint | audio in | limitation | -|----------|----------|-----------| -| `/transcribe` | 1 file (multipart) | single audio; ASR-only framing | -| `/v1/chat/completions` | none (text-only) | `ChatMessage.content: str`; no audio at all | -| `/judge` | 1 file | single audio + transcript | -| `/stream` (ws) | PCM frames | streaming ASR only | - -Consequences: -1. **Multi-audio prompting is impossible.** Few-shot wake-word classification (show K example - clips + a candidate in one prompt) can't be expressed → today it requires loading a *second* - copy of the model in a one-off container, which OOMs the 24 GB GPU, so we **stop the live ASR - service, run, and restart** every time. Redundant model load + ASR downtime + GPU contention. -2. **Redundant encode/decode round-trips.** The OpenAI `input_audio` transport is base64. For a - *local* caller that already has a wav on disk, the path is `file → base64 → (wire) → base64-decode - → temp file → processor re-reads+decodes`. Two of those steps are pure overhead. - -## Core insight - -Everything the model does is one operation: - -``` -generate(messages) -> text # messages = interleaved [text | audio] parts -``` - -Transcribe, classify, judge, and chat are all **presets** of this. The clean shape is: - -- a single **canonical multimodal generate** path on the loaded model, -- the **OpenAI-compatible chat/completions** endpoint as the *general* interface (this is the - industry-standard way to serve a combined model — same schema vLLM/Gemini/OpenRouter use for - audio: `content: [{type:"text"}, {type:"input_audio", input_audio:{data,format}}]`), -- the ASR-specific wrappers kept thin on top. - -"Use the combined model as a general endpoint for both chat and ASR" = make `/v1/chat/completions` -multimodal; transcription becomes "chat with a transcribe prompt + one audio". - -## Design — three layers - -``` -HTTP adapters (thin; format + presets) - POST /v1/chat/completions GENERAL multimodal chat (text+audio messages) ← chat, classify, judge, 1-shot ASR - POST /transcribe ASR preset: silence-gate + >30s window/stitch + diarization-parse + NDJSON - POST /classify (optional) convenience: example clips+labels+candidate -> label - POST /judge, WS /stream existing presets - │ normalize_audio() ── decode EVERY transport to np.float32@16k mono, ONCE - ▼ -Transcriber.generate(messages_with_array_audio) -> (text, in_tok, out_tok) ← single canonical path, loaded model -``` - -### Layer 1 — canonical generate (transcriber.py) -`generate(messages, *, max_tokens, temperature=0, top_p, top_k) -> (text, in_tok, out_tok)` -- `messages`: `[{role, content:[parts]}]`; parts are `{"type":"text",...}` or - `{"type":"audio","audio": <np.ndarray>}` (mono float32 @ 16 kHz). -- The ONLY place that calls `apply_chat_template` + `model.generate` + `_decode_response`. -- No silence gate, no batching, no temp files — pure model access on the in-memory weights. -- `generate_chat()` (today text-only) becomes a thin caller → full back-compat for text chat. - -### Layer 2 — audio normalization at the boundary (the round-trip fix) -One helper `normalize_audio_parts(content) -> content` that accepts audio in any transport and -decodes **exactly once to an array** (`soundfile.read(BytesIO)` / `load_audio_file`), never to disk: - -| transport | when | cost | -|-----------|------|------| -| `input_audio` base64 | remote / cloud-parity | b64→BytesIO→array (no temp file) | -| multipart raw bytes | local, avoid +33% base64 tax | bytes→array | -| `audio_path` ref (shared RO volume) | trusted local batch | `load_audio_file(path)` → array, **zero re-encode** | - -All converge to `{"type":"audio","audio": <array>}`. The `file→b64→file→decode` chain collapses to -a single decode regardless of caller. (Validation note: confirm the installed processor accepts an -ndarray under the `audio` key; today `_transcribe_single` passes a path. If a given transformers -version wants a path/URL, normalize to the processor's accepted form once — still one decode, in a -single place.) - -### Layer 3 — HTTP endpoints (thin adapters) -- **`/v1/chat/completions`** — widen `ChatMessage.content` from `str` to `str | list[part]`; parts may - include `input_audio`. Pass through `normalize_audio_parts` → `generate`. This is the general - endpoint; classification/judge/one-shot-ASR are just prompts over it. -- **`/transcribe`** — unchanged contract; still owns the ASR-only concerns the general endpoint - shouldn't carry (silence gate, >30 s windowing+stitch, diarization parse, NDJSON progress). - Internally routed through `normalize_audio_parts` + `generate`. -- **`/classify`** (optional sugar) — multipart `examples[]` + `labels[]` + `candidate` → builds the - few-shot messages server-side → returns `{label, raw}`. Saves callers from hand-building messages. -- `/judge`, `/stream` — refactor to call `generate`; behavior unchanged. - -## Consistency: OpenAI ↔ Gemma - -- The OpenAI `input_audio` part maps **deterministically** to Gemma's `{"type":"audio"}` part — a - faithful 1:1 adapter, *if normalized at the boundary*. Same payload works against OpenRouter/Gemini - **and** the local service → identical client code, just swap `base_url`. (This is what lets - `threeway_probe.py` target cloud or local unchanged.) -- One Gemma-specific nuance: it prefers the audio **after** its text label in a turn. We preserve - client-provided part order; presets emit the correct order. So consistency holds. -- The chat path deliberately **skips the silence gate** (that's an ASR-output concern) — correct for - classification, where we want a verdict even on near-silent clips. - -## Transport guidance (answering "just random round trips") - -- **Remote/cloud:** base64 `input_audio` is unavoidable (no shared FS) — but we decode once to an - array, never to a temp file. -- **Local batch (triage over pending clips):** prefer **`audio_path` refs** with the wakeword - `data/samples` volume mounted read-only into the asr container → the service reads each wav - straight to an array. No base64, no temp files, no copies. Multipart raw is the middle option. - -## Impact / migration -- **Back-compat:** `/transcribe` and text-only `/v1/chat` behave exactly as today. `ChatMessage.content` - widened to `str | list` (additive). One image rebuild. -- **Deletes the stop/restart workflow:** multi-audio few-shot classification is served by the *live* - model — no second load, no ASR downtime, no extra idle VRAM (same weights). -- **Less code, one behavior:** transcribe/classify/judge/chat unified on one generate path. -- **Concurrency:** GPU generate is serial; a batch labeling job shares the card with live ASR/wake - transcription. Fine for background triage; if it matters, route batch on the existing ASR "normal" - lane (priority lane keeps wake clips ahead). See `asr-priority-lane`. - -## Status — IMPLEMENTED + validated (2026-06-14) - -Done (minimal, backward-compatible): -- `common/audio_utils.py`: added `load_audio_bytes(bytes)` — in-memory WAV→array (no temp file), - reuses `convert_audio_to_numpy`/`convert_to_mono`/`resample_audio`. -- `providers/gemma4/transcriber.py`: added `_normalize_chat_content()` and wired it into - `generate_chat()`. `input_audio` base64 → decode once to ndarray; **path refs + ndarrays + text - pass straight through** (processor loads paths natively). No new endpoint, no schema change — - `/v1/chat/completions` already forwards `messages: list[dict]`, so audio parts just flow through. -- Verified the processor accepts an **ndarray** under `audio` (not only a path), so no temp files. - -Validated on the LIVE service (no stop/restart, no second model load): -- Multi-audio few-shot through `/v1/chat/completions` works on the running model. -- 3-way on the 27 held-out clips via the endpoint = **96.3% (26/27)** vs **100% in-container**. - The single flip is a borderline `hermes`: in-container passed **file paths** (transformers' own - loader), the endpoint sent **base64** decoded by our `wave` loader (`/32767` scaling) — slightly - different decode flips one boundary case (both greedy). - -Decode-path takeaway: for **local** callers, prefer **path refs** (mount `data/samples` RO into the -asr container, send `{"type":"audio","audio":"<container path>"}`) → exact `/transcribe` parity, no -base64, no decode. base64 stays for remote/cloud where there's no shared FS (accept the 1-in-27 -boundary wobble, or align the loaders for bit-parity later). - -NOTE: the running image was built before the final "path refs pass through" tweak; rebuild to -activate path-ref parity. base64 flow is already live and correct. - -## Conclusion (where we stopped, 2026-06-14) - -**Validated + serving primitive built; the wakeword-service integration is deferred.** -- Approach proven: 3-way few-shot ("hey hermes" / "hermes" / "nothing"), gemma4 E2B best - (98% labeled, 100%/96% on the user's 27 held-out hand-labels), local + free. -- `/v1/chat/completions` now serves multi-audio on the live model — no stop/restart. -- NOT built: any loop inside `extras/wakeword-service/` (no trigger, no pending triage, no - suggestion write/surface). Probe scripts under `ml-experiments/.../wakeword_gemma_classify/` - are eval harnesses only. - -**Eventual shape (per user):** mostly a **button** in the review UI ("suggest labels") that calls the -classifier over `pending/` and pre-fills suggestions for confirm/correct — NOT an autonomous labeler. - -**Key open problem before that ships: a reliable "unsure"/abstain.** The classifier currently makes a -forced 3-way choice; for a triage button it must *defer* hard cases to the human instead of guessing. -Candidate signals (none built yet): -- explicit 4th label "unsure" in the prompt (cheap, but models rarely self-abstain reliably); -- self-consistency — run N times with shuffled few-shot order / different ref sets, abstain on - disagreement (needs sampling or perturbation since greedy is deterministic); -- two-stage classify→verify, abstain when they disagree; -- token-logprob confidence if we expose it from `generate` (top-token prob as a calibrated score); -- cheap energy/VAD pre-gate to auto-route near-silence to "nothing" (note: a real wake existed at - rms 0.0023, so a global floor is unsafe — use only as a soft signal). -The abstain mechanism, not the classifier, is the remaining design work. - -## Build order -1. `transcriber.generate()` canonical path + refactor `generate_chat` to delegate. -2. `normalize_audio_parts()` (base64 + multipart + path), decode-once-to-array. -3. Widen `/v1/chat/completions` schema + wire normalization. -4. (optional) `/classify` convenience preset. -5. Rebuild image. Point `threeway_probe.py` / a new `classify_pending.py` at `localhost:8767` — zero restarts. diff --git a/extras/asr-services/providers/gemma4/finetune/debug_gen_nocache.py b/extras/asr-services/providers/gemma4/finetune/debug_gen_nocache.py deleted file mode 100644 index 7d7ce4de..00000000 --- a/extras/asr-services/providers/gemma4/finetune/debug_gen_nocache.py +++ /dev/null @@ -1,66 +0,0 @@ -"""Test: does generating with use_cache=False make the v3 adapter reproduce targets? - -The cached (use_cache=True) gemma4 multimodal forward gives wrong logits (loss 4.49 -vs 0.10 uncached), so model.generate() — which uses the cache by default — fails. -Force use_cache=False during generation and compare. -""" - -from difflib import SequenceMatcher - -import torch -from data import DEFAULT_PROMPT, CosheSample7Dataset -from peft import PeftModel -from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig - -proc = AutoProcessor.from_pretrained("google/gemma-4-E2B-it") -proc.tokenizer.padding_side = "left" -bnb = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_use_double_quant=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - llm_int8_skip_modules=[ - "model.audio_tower", - "model.vision_tower", - "model.embed_audio", - "model.embed_vision", - "lm_head", - ], -) -m = AutoModelForMultimodalLM.from_pretrained( - "google/gemma-4-E2B-it", - quantization_config=bnb, - dtype=torch.bfloat16, - device_map="auto", - attn_implementation="eager", -) -pm = PeftModel.from_pretrained(m, "/train/out/e2b-overfit-v3") -pm.eval() - -ds = CosheSample7Dataset( - "/data/coshe-eval/sample7", max_seconds=30.0, target_max_chars=540 -) -print("===== generate(use_cache=False) =====", flush=True) -ratios = [] -for item in ds: - msgs = [ - { - "role": "user", - "content": [ - {"type": "text", "text": DEFAULT_PROMPT}, - {"type": "audio", "audio": item["audio"]}, - ], - } - ] - ptext = proc.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True) - inp = proc(text=[ptext], audio=[item["audio"]], return_tensors="pt").to(pm.device) - inlen = inp["input_ids"].shape[-1] - with torch.inference_mode(): - out = pm.generate(**inp, max_new_tokens=320, do_sample=False, use_cache=False) - gen = proc.decode(out[0][inlen:], skip_special_tokens=True).strip() - r = SequenceMatcher(None, gen, item["target"]).ratio() - ratios.append(r) - print(f"\n {item['name']} sim={r:.3f}", flush=True) - print(f" TGT: {item['target'][:220]}", flush=True) - print(f" GEN: {gen[:220]}", flush=True) -print(f"\n MEAN char-sim = {sum(ratios)/len(ratios):.3f}", flush=True) diff --git a/extras/asr-services/providers/gemma4/finetune/debug_inprocess.py b/extras/asr-services/providers/gemma4/finetune/debug_inprocess.py deleted file mode 100644 index ec19759b..00000000 --- a/extras/asr-services/providers/gemma4/finetune/debug_inprocess.py +++ /dev/null @@ -1,137 +0,0 @@ -"""Train in-process and eval WITHOUT save/reload, to localize the loss-0-but-wrong bug. - -If in-process teacher-forced loss is ~0 AND generation reproduces -> the bug is in -save/reload (adapter not persisted/loaded correctly). -If in-process loss is ~0 but generation still fails -> teacher-forcing/label issue. -If in-process loss is also high -> the Trainer's reported step loss was the illusion. -""" - -import torch -from data import DEFAULT_PROMPT, CosheSample7Dataset, Gemma4AudioCollator -from peft import LoraConfig, get_peft_model -from transformers import ( - AutoModelForMultimodalLM, - AutoProcessor, - BitsAndBytesConfig, - Trainer, - TrainingArguments, -) - -MODEL = "google/gemma-4-E2B-it" -LORA_TARGETS = ( - r".*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" -) - -processor = AutoProcessor.from_pretrained(MODEL) -bnb = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_use_double_quant=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - llm_int8_skip_modules=[ - "model.audio_tower", - "model.vision_tower", - "model.embed_audio", - "model.embed_vision", - "lm_head", - ], -) -model = AutoModelForMultimodalLM.from_pretrained( - MODEL, - quantization_config=bnb, - dtype=torch.bfloat16, - device_map="auto", - attn_implementation="eager", -) -model.config.use_cache = False -for p in model.parameters(): - p.requires_grad = False -model.gradient_checkpointing_enable( - gradient_checkpointing_kwargs={"use_reentrant": False} -) -model.enable_input_require_grads() -model = get_peft_model( - model, - LoraConfig( - r=16, - lora_alpha=32, - lora_dropout=0.0, - bias="none", - task_type="CAUSAL_LM", - target_modules=LORA_TARGETS, - ), -) -model.print_trainable_parameters() - -ds = CosheSample7Dataset( - "/data/coshe-eval/sample7", max_seconds=30.0, target_max_chars=540 -) -collator = Gemma4AudioCollator(processor) - -trainer = Trainer( - model=model, - args=TrainingArguments( - output_dir="/train/out/_inproc", - per_device_train_batch_size=1, - num_train_epochs=30, - learning_rate=2e-4, - lr_scheduler_type="constant", - bf16=True, - logging_steps=20, - save_strategy="no", - report_to=[], - remove_unused_columns=False, - gradient_checkpointing=False, - ), - train_dataset=ds, - data_collator=collator, -) -trainer.train() - -# ---- eval in the SAME process, no save/reload ---- -model.eval() -print("\n===== IN-PROCESS EVAL (no save/reload) =====", flush=True) -for item in ds: - b = collator([item]) - b = {k: v.to(model.device) for k, v in b.items()} - with torch.inference_mode(): - o = model(**b) - ng = model(**{k: v for k, v in b.items() if k != "labels"}) - lg = ng.logits[0] - lb = b["labels"][0] - p = lg[:-1].argmax(-1) - g = lb[1:] - s = g != -100 - n = int(s.sum()) - m = int((p[s] == g[s]).sum()) - print( - f" {item['name']:16s} model.loss={float(o.loss):.4f} argmax-match={m}/{n} " - f"({100*m/max(n,1):.1f}%)", - flush=True, - ) - -# generation on sample 0 using the SAME collator prompt path -model.config.use_cache = True -item = ds[0] -ptext = processor.apply_chat_template( - [ - { - "role": "user", - "content": [ - {"type": "text", "text": DEFAULT_PROMPT}, - {"type": "audio", "audio": item["audio"]}, - ], - } - ], - tokenize=False, - add_generation_prompt=True, -) -inp = processor(text=[ptext], audio=[item["audio"]], return_tensors="pt").to( - model.device -) -inlen = inp["input_ids"].shape[-1] -with torch.inference_mode(): - out = model.generate(**inp, max_new_tokens=256, do_sample=False) -gen = processor.decode(out[0][inlen:], skip_special_tokens=True) -print(f"\n TARGET[:300]: {item['target'][:300]}", flush=True) -print(f" GEN[:300] : {gen[:300]}", flush=True) diff --git a/extras/asr-services/providers/gemma4/finetune/debug_overfit.py b/extras/asr-services/providers/gemma4/finetune/debug_overfit.py deleted file mode 100644 index 9326adff..00000000 --- a/extras/asr-services/providers/gemma4/finetune/debug_overfit.py +++ /dev/null @@ -1,186 +0,0 @@ -"""Diagnose why teacher-forced loss -> 0 but generation doesn't reproduce targets. - -Checks two things on the trained v3 adapter (target_max_chars=540): - - 1. TOKEN-FORMAT PARITY: does the training collator's input construction - (apply_chat_template(tokenize=False) -> processor(text=str, audio=)) produce - the SAME prompt token ids as the inference path - (apply_chat_template(tokenize=True, return_dict=True))? A mismatch (e.g. a - duplicated <bos>) would mean we train on one format and generate on another. - - 2. TEACHER-FORCED REPRODUCTION: feed the full (prompt+target) exactly as in - training, take argmax of the logits at each SUPERVISED position, and measure - how many match the gold target token. If this is ~100%, the model truly - memorized and any failure is in the generation path; if it's low, loss never - really reached 0 on the hard (early) tokens. -""" - -import argparse - -import torch -from data import DEFAULT_PROMPT, CosheSample7Dataset, Gemma4AudioCollator -from peft import PeftModel -from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig - - -def parse_args(): - p = argparse.ArgumentParser() - p.add_argument("--model", default="google/gemma-4-E2B-it") - p.add_argument("--adapter", default="/train/out/e2b-overfit-v3") - p.add_argument("--data_dir", default="/data/coshe-eval/sample7") - p.add_argument("--target_max_chars", type=int, default=540) - return p.parse_args() - - -def main(): - args = parse_args() - processor = AutoProcessor.from_pretrained(args.model) - tok = processor.tokenizer - - ds = CosheSample7Dataset( - args.data_dir, max_seconds=30.0, target_max_chars=args.target_max_chars - ) - item = ds[0] - - # ---- 1. TOKEN-FORMAT PARITY ---- - print("===== 1. TOKEN-FORMAT PARITY (prompt only) =====", flush=True) - user_turn = { - "role": "user", - "content": [ - {"type": "text", "text": DEFAULT_PROMPT}, - {"type": "audio", "audio": item["audio"]}, - ], - } - # training-style: tokenize=False -> processor(text=str, audio=) - ptext = processor.apply_chat_template( - [user_turn], tokenize=False, add_generation_prompt=True - ) - train_path = processor(text=[ptext], audio=[item["audio"]], return_tensors="pt") - train_ids = train_path["input_ids"][0] - # inference-style: tokenize=True, return_dict=True - infer_path = processor.apply_chat_template( - [user_turn], - tokenize=True, - return_dict=True, - return_tensors="pt", - add_generation_prompt=True, - ) - infer_ids = infer_path["input_ids"][0] - - print( - f" train-path len={len(train_ids)} first8={train_ids[:8].tolist()}", - flush=True, - ) - print( - f" infer-path len={len(infer_ids)} first8={infer_ids[:8].tolist()}", - flush=True, - ) - bos = tok.bos_token_id - print( - f" bos_token_id={bos} train leading bos count={int((train_ids[:3]==bos).sum())} " - f"infer leading bos count={int((infer_ids[:3]==bos).sum())}", - flush=True, - ) - same = len(train_ids) == len(infer_ids) and bool((train_ids == infer_ids).all()) - print(f" >>> IDENTICAL: {same}", flush=True) - if not same: - # show first divergence - n = min(len(train_ids), len(infer_ids)) - diff = (train_ids[:n] != infer_ids[:n]).nonzero() - first = int(diff[0]) if len(diff) else n - print(f" first divergence at pos {first}", flush=True) - print( - f" train[{first}:{first+6}]={train_ids[first:first+6].tolist()} " - f"-> {tok.convert_ids_to_tokens(train_ids[first:first+6].tolist())}", - flush=True, - ) - print( - f" infer[{first}:{first+6}]={infer_ids[first:first+6].tolist()} " - f"-> {tok.convert_ids_to_tokens(infer_ids[first:first+6].tolist())}", - flush=True, - ) - - # ---- 2. TEACHER-FORCED REPRODUCTION ---- - print( - "\n===== 2. TEACHER-FORCED ARGMAX REPRODUCTION (v3 adapter) =====", flush=True - ) - bnb = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_use_double_quant=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - llm_int8_skip_modules=[ - "model.audio_tower", - "model.vision_tower", - "model.embed_audio", - "model.embed_vision", - "lm_head", - ], - ) - model = AutoModelForMultimodalLM.from_pretrained( - args.model, - quantization_config=bnb, - dtype=torch.bfloat16, - device_map="auto", - attn_implementation="eager", - ) - if args.adapter: - model = PeftModel.from_pretrained(model, args.adapter) - model.eval() - - collator = Gemma4AudioCollator(processor) - item0 = ds[0] - batch = collator([item0]) - batch = {k: v.to(model.device) for k, v in batch.items()} - with torch.inference_mode(): - out_with_labels = model(**batch) # model's OWN loss (its internal shift) - out = model(**{k: v for k, v in batch.items() if k != "labels"}) - logits = out.logits[0] - labels = batch["labels"][0] - # my manual shift: logits[t] predicts token t+1 - pred = logits[:-1].argmax(-1) - gold = labels[1:] - sup = gold != -100 - manual_loss = torch.nn.functional.cross_entropy( - logits[:-1][sup].float(), gold[sup], reduction="mean" - ) - # NO-shift variant: logits[t] vs labels[t] (in case model pre-shifts) - sup0 = labels != -100 - noshift_loss = torch.nn.functional.cross_entropy( - logits[sup0].float(), labels[sup0], reduction="mean" - ) - print( - f" model.forward(labels=...).loss = {float(out_with_labels.loss):.4f}", - flush=True, - ) - print(f" my manual SHIFTED loss = {float(manual_loss):.4f}", flush=True) - print(f" my manual NO-shift loss = {float(noshift_loss):.4f}", flush=True) - - # with vs without adapter (model's own loss) - if args.adapter: - with torch.inference_mode(), model.disable_adapter(): - base_loss = float(model(**batch).loss) - print( - f" model loss WITH adapter={float(out_with_labels.loss):.4f} " - f"WITHOUT adapter (base)={base_loss:.4f}", - flush=True, - ) - - print("\n per-sample argmax-match (shifted):", flush=True) - for item in ds: - b = collator([item]) - b = {k: v.to(model.device) for k, v in b.items()} - with torch.inference_mode(): - o = model(**{k: v for k, v in b.items() if k != "labels"}) - lg = o.logits[0] - lb = b["labels"][0] - p = lg[:-1].argmax(-1) - g = lb[1:] - s = g != -100 - n = int(s.sum()) - m = int((p[s] == g[s]).sum()) - print(f" {item['name']:16s} match={m}/{n} ({100*m/max(n,1):.1f}%)", flush=True) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/gemma4/finetune/debug_reload_loss.py b/extras/asr-services/providers/gemma4/finetune/debug_reload_loss.py deleted file mode 100644 index db431e46..00000000 --- a/extras/asr-services/providers/gemma4/finetune/debug_reload_loss.py +++ /dev/null @@ -1,51 +0,0 @@ -import torch -from data import CosheSample7Dataset, Gemma4AudioCollator -from peft import PeftModel -from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig - -proc = AutoProcessor.from_pretrained("google/gemma-4-E2B-it") -bnb = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_use_double_quant=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - llm_int8_skip_modules=[ - "model.audio_tower", - "model.vision_tower", - "model.embed_audio", - "model.embed_vision", - "lm_head", - ], -) -m = AutoModelForMultimodalLM.from_pretrained( - "google/gemma-4-E2B-it", - quantization_config=bnb, - dtype=torch.bfloat16, - device_map="auto", - attn_implementation="eager", -) -pm = PeftModel.from_pretrained(m, "/train/out/e2b-overfit-v3") -pm.eval() - -bnorm = sum(float(p.float().norm()) for n, p in pm.named_parameters() if "lora_B" in n) -print("reloaded lora_B sum norm:", round(bnorm, 3), flush=True) - -ds = CosheSample7Dataset( - "/data/coshe-eval/sample7", max_seconds=30.0, target_max_chars=540 -) -col = Gemma4AudioCollator(proc) - -b0 = col([ds[0]]) -b0 = {k: v.to(pm.device) for k, v in b0.items()} -for uc in [True, False]: - pm.config.use_cache = uc - if hasattr(pm, "base_model"): - pm.base_model.config.use_cache = uc - with torch.inference_mode(): - loss = float(pm(**b0).loss) - print(f" use_cache={uc}: reloaded-adapter loss = {loss:.4f}", flush=True) - -with torch.inference_mode(), pm.disable_adapter(): - b = col([ds[0]]) - b = {k: v.to(pm.device) for k, v in b.items()} - print(" base (adapter disabled) loss =", round(float(pm(**b).loss), 4), flush=True) diff --git a/extras/asr-services/providers/gemma4/finetune/stream_client_test.py b/extras/asr-services/providers/gemma4/finetune/stream_client_test.py deleted file mode 100644 index 5b1bd0a5..00000000 --- a/extras/asr-services/providers/gemma4/finetune/stream_client_test.py +++ /dev/null @@ -1,71 +0,0 @@ -"""Test the gemma4 /stream WebSocket endpoint end-to-end. - -Streams a clip's PCM (16-bit LE, 16 kHz mono) in chunks, prints interim -messages as they arrive, sends CloseStream, and prints the final result — -exactly the Chronicle stt_stream contract. -""" - -import argparse -import asyncio -import json - -import numpy as np -import soundfile as sf -import websockets - -SR = 16000 - - -async def main(): - p = argparse.ArgumentParser() - p.add_argument("--url", default="ws://localhost:8767/stream") - p.add_argument( - "--audio", - default="extras/test-audios/audio/Cheap-vs-Expensive-Rug-Tufting-for-the-First-Time-original.wav", - ) - p.add_argument("--seconds", type=float, default=30.0) - p.add_argument("--chunk-ms", type=int, default=500) - p.add_argument("--pace", type=float, default=0.05, help="sleep between chunks (s)") - args = p.parse_args() - - a, sr = sf.read(args.audio, dtype="float32", frames=int(args.seconds * 44100)) - if a.ndim > 1: - a = a.mean(1) - if sr != SR: - import librosa - - a = librosa.resample(a, orig_sr=sr, target_sr=SR) - pcm = (np.clip(a, -1, 1) * 32767).astype("<i2").tobytes() - chunk = int(SR * args.chunk_ms / 1000) * 2 - print( - f"streaming {len(pcm)/SR/2:.1f}s of audio in {args.chunk_ms}ms chunks", - flush=True, - ) - - async with websockets.connect(args.url, max_size=None) as ws: - interims = [] - - async def receiver(): - async for raw in ws: - msg = json.loads(raw) - if msg["type"] == "interim": - interims.append(msg["text"]) - print(f" [interim] {msg['text'][:120]}", flush=True) - elif msg["type"] == "final": - print(f"\n[FINAL] {msg['text']}", flush=True) - print( - f"[FINAL] segments={len(msg.get('segments') or [])}", flush=True - ) - return - - recv_task = asyncio.create_task(receiver()) - for i in range(0, len(pcm), chunk): - await ws.send(pcm[i : i + chunk]) - await asyncio.sleep(args.pace) - await ws.send(json.dumps({"type": "CloseStream"})) - await asyncio.wait_for(recv_task, timeout=180) - print(f"\n{len(interims)} interim message(s) received", flush=True) - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/extras/asr-services/providers/gemma4/finetune/test_cache57.py b/extras/asr-services/providers/gemma4/finetune/test_cache57.py deleted file mode 100644 index 3b316bfc..00000000 --- a/extras/asr-services/providers/gemma4/finetune/test_cache57.py +++ /dev/null @@ -1,90 +0,0 @@ -"""On the A100/transformers-5.7: is the use_cache bug still present? - -A correct implementation gives IDENTICAL logits/loss for a single full forward -regardless of use_cache. We compare use_cache True vs False on one batch with the -current adapter. If equal -> cache is fine -> fast cached generation/eval. If they -differ -> bug persists -> use_cache=False for eval. -""" - -import argparse -import time - -import torch -from data import DEFAULT_PROMPT, CosheParquetDataset, Gemma4AudioCollator -from peft import PeftModel -from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig - -p = argparse.ArgumentParser() -p.add_argument("--adapter", default="/home/gemma4ft/out/full/checkpoint-994") -p.add_argument("--parquet_glob", default="/home/coshe-data/data/eval-*.parquet") -p.add_argument("--cache_path", default="/home/gemma4ft/out/coshe_full_cache.pkl") -args = p.parse_args() - -proc = AutoProcessor.from_pretrained("google/gemma-4-E2B-it") -bnb = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_use_double_quant=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - llm_int8_skip_modules=[ - "model.audio_tower", - "model.vision_tower", - "model.embed_audio", - "model.embed_vision", - "lm_head", - ], -) -model = AutoModelForMultimodalLM.from_pretrained( - "google/gemma-4-E2B-it", - quantization_config=bnb, - dtype=torch.bfloat16, - device_map="auto", - attn_implementation="eager", -) -model = PeftModel.from_pretrained(model, args.adapter) -model.eval() - -ds = CosheParquetDataset( - args.parquet_glob, - max_seconds=30.0, - target_max_chars=0, - cache_path=args.cache_path, - limit=8, -) -col = Gemma4AudioCollator(proc) -b = col([ds[0]]) -b = {k: v.to(model.device) for k, v in b.items()} - -print("=== single-forward loss: use_cache True vs False ===", flush=True) -for uc in [False, True]: - with torch.inference_mode(): - loss = float(model(**b, use_cache=uc).loss) - print(f" use_cache={uc}: loss={loss:.4f}", flush=True) - -print("\n=== generation speed: cache vs no-cache (4 clips) ===", flush=True) -proc.tokenizer.padding_side = "left" -texts, audios = [], [] -for it in [ds[i] for i in range(4)]: - msgs = [ - { - "role": "user", - "content": [ - {"type": "text", "text": DEFAULT_PROMPT}, - {"type": "audio", "audio": it["audio"]}, - ], - } - ] - texts.append( - proc.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True) - ) - audios.append(it["audio"]) -inp = proc(text=texts, audio=audios, return_tensors="pt", padding=True).to(model.device) -for uc in [True, False]: - t = time.time() - with torch.inference_mode(): - out = model.generate(**inp, max_new_tokens=200, do_sample=False, use_cache=uc) - dt = time.time() - t - txt = proc.decode(out[0][inp["input_ids"].shape[-1] :], skip_special_tokens=True)[ - :120 - ] - print(f" use_cache={uc}: {dt:.1f}s for 4 clips | sample: {txt!r}", flush=True) diff --git a/extras/asr-services/providers/gemma4/finetune/test_longaudio.py b/extras/asr-services/providers/gemma4/finetune/test_longaudio.py deleted file mode 100644 index f4ef0ece..00000000 --- a/extras/asr-services/providers/gemma4/finetune/test_longaudio.py +++ /dev/null @@ -1,64 +0,0 @@ -"""Does gemma4 accept a full ~55s clip in one forward (vs the 30s window)? - -Prints token/feature counts and runs a forward at 30s and ~60s to see if it -errors or how VRAM scales. Decides whether full 1-min clips are usable for -training without chunking. -""" - -import torch -from data import DEFAULT_PROMPT, _load_wav_16k_mono -from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig - -WAV = "/data/coshe-eval/sample7/audio/audio_13.wav" - -proc = AutoProcessor.from_pretrained("google/gemma-4-E2B-it") -bnb = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_use_double_quant=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - llm_int8_skip_modules=[ - "model.audio_tower", - "model.vision_tower", - "model.embed_audio", - "model.embed_vision", - "lm_head", - ], -) -model = AutoModelForMultimodalLM.from_pretrained( - "google/gemma-4-E2B-it", - quantization_config=bnb, - dtype=torch.bfloat16, - device_map="auto", - attn_implementation="eager", -) -model.config.use_cache = False -model.eval() - -for secs in [30.0, 60.0]: - audio = _load_wav_16k_mono(WAV, secs) - msgs = [ - { - "role": "user", - "content": [ - {"type": "text", "text": DEFAULT_PROMPT}, - {"type": "audio", "audio": audio}, - ], - } - ] - text = proc.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True) - batch = proc(text=[text], audio=[audio], return_tensors="pt").to(model.device) - torch.cuda.reset_peak_memory_stats() - try: - with torch.inference_mode(): - out = model(**batch) - peak = torch.cuda.max_memory_allocated() / 1e9 - print( - f"audio={secs:.0f}s samples={len(audio)} " - f"input_ids={batch['input_ids'].shape[-1]} " - f"audio_frames={batch['input_features'].shape[1]} " - f"forward OK logits={tuple(out.logits.shape)} peakVRAM={peak:.1f}GB", - flush=True, - ) - except Exception as e: - print(f"audio={secs:.0f}s FAILED: {type(e).__name__}: {e}", flush=True) diff --git a/extras/asr-services/providers/gemma4/finetune/test_patch_loss.py b/extras/asr-services/providers/gemma4/finetune/test_patch_loss.py deleted file mode 100644 index 3867da7f..00000000 --- a/extras/asr-services/providers/gemma4/finetune/test_patch_loss.py +++ /dev/null @@ -1,131 +0,0 @@ -"""Instrument the gemma4 cache-mask logic to fix the use_cache prefill bug. - -Runs a single full forward (use_cache=True) on a memorized sample7 clip with the -s7-full adapter and prints what create_causal_mask_mapping sees (cache type, -get_seq_length, is_initialized, chosen is_first_iteration) + the resulting loss. -Goal: find the right prefill test so loss drops 4.49 -> ~0.1 with the cache on. -""" - -import torch -import transformers.models.gemma4.modeling_gemma4 as g4 -from data import DEFAULT_PROMPT, CosheSample7Dataset, Gemma4AudioCollator -from peft import PeftModel -from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig - -_orig = g4.create_causal_mask_mapping -_seen = {} - - -def _instrumented( - config, - inputs_embeds, - attention_mask, - past_key_values, - position_ids, - mm_token_type_ids=None, - pixel_values=None, - is_training=False, - is_first_iteration=None, - **kwargs, -): - pkv = past_key_values - info = { - "pkv_type": type(pkv).__name__ if pkv is not None else None, - "is_initialized": ( - getattr(pkv, "is_initialized", "n/a") if pkv is not None else None - ), - "get_seq_length": (pkv.get_seq_length() if pkv is not None else None), - "mm_token_type_ids_none": mm_token_type_ids is None, - "q_len": inputs_embeds.shape[1], - "pixel_values_none": pixel_values is None, - } - if not _seen: - _seen.update(info) - print("MASK CALL INFO:", info, flush=True) - return _orig( - config, - inputs_embeds, - attention_mask, - past_key_values, - position_ids, - mm_token_type_ids=mm_token_type_ids, - pixel_values=pixel_values, - is_training=is_training, - is_first_iteration=is_first_iteration, - **kwargs, - ) - - -g4.create_causal_mask_mapping = _instrumented - -proc = AutoProcessor.from_pretrained("google/gemma-4-E2B-it") -bnb = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_use_double_quant=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - llm_int8_skip_modules=[ - "model.audio_tower", - "model.vision_tower", - "model.embed_audio", - "model.embed_vision", - "lm_head", - ], -) -model = AutoModelForMultimodalLM.from_pretrained( - "google/gemma-4-E2B-it", - quantization_config=bnb, - dtype=torch.bfloat16, - device_map="auto", - attn_implementation="eager", -) -model = PeftModel.from_pretrained(model, "/train/out/s7-full") -model.eval() - -ds = CosheSample7Dataset( - "/data/coshe-eval/sample7", max_seconds=30.0, target_max_chars=0 -) -col = Gemma4AudioCollator(proc) -b = col([ds[0]]) -b = {k: v.to(model.device) for k, v in b.items()} - -for uc in [False, True]: - model.config.use_cache = uc - _seen.clear() - with torch.inference_mode(): - loss = float(model(**b, use_cache=uc).loss) - print(f"use_cache={uc}: loss={loss:.4f}", flush=True) - - -# now try forcing is_first_iteration=True always (prefill-style) to confirm theory -def _force_true( - config, - inputs_embeds, - attention_mask, - past_key_values, - position_ids, - mm_token_type_ids=None, - pixel_values=None, - is_training=False, - is_first_iteration=None, - **kwargs, -): - return _orig( - config, - inputs_embeds, - attention_mask, - past_key_values, - position_ids, - mm_token_type_ids=mm_token_type_ids, - pixel_values=pixel_values, - is_training=is_training, - is_first_iteration=True, - **kwargs, - ) - - -g4.create_causal_mask_mapping = _force_true -model.config.use_cache = True -with torch.inference_mode(): - loss = float(model(**b, use_cache=True).loss) -print(f"use_cache=True + force is_first_iteration=True: loss={loss:.4f}", flush=True) diff --git a/extras/asr-services/tests/capture_vibevoice_ground_truth.py b/extras/asr-services/tests/capture_vibevoice_ground_truth.py deleted file mode 100644 index 41ff338b..00000000 --- a/extras/asr-services/tests/capture_vibevoice_ground_truth.py +++ /dev/null @@ -1,168 +0,0 @@ -""" -Capture VibeVoice transcription output as ground truth for regression tests. - -Sends the 4-minute test audio to a running VibeVoice service and saves the -full response (segments, text, words) as a JSON fixture. - -To exercise batching/stitching on the 4-min file, restart the service with -small batch windows first: - - cd extras/asr-services - BATCH_THRESHOLD_SECONDS=60 BATCH_DURATION_SECONDS=90 BATCH_OVERLAP_SECONDS=15 \ - docker compose up vibevoice-asr -d - -Then run the capture: - - uv run python tests/capture_vibevoice_ground_truth.py - -To restore normal settings afterwards: - - docker compose up vibevoice-asr -d # env vars unset → uses config defaults - -The output is saved to tests/fixtures/vibevoice_4min_ground_truth.json. -Review the output manually before committing — it becomes the reference -for segment overlap and stitching regression tests. -""" - -import argparse -import json -import sys -from pathlib import Path - -import httpx - -DEFAULT_AUDIO = ( - Path(__file__).parent.parent.parent.parent - / "tests" - / "test_assets" - / "DIY_Experts_Glass_Blowing_16khz_mono_4min.wav" -) -DEFAULT_OUTPUT = Path(__file__).parent / "fixtures" / "vibevoice_4min_ground_truth.json" - - -def capture(service_url: str, audio_path: str, output_path: str) -> dict: - """Send audio to VibeVoice /transcribe and return the parsed result.""" - audio_path = Path(audio_path) - if not audio_path.exists(): - print(f"ERROR: Audio file not found: {audio_path}") - sys.exit(1) - - print(f"Sending {audio_path.name} to {service_url}/transcribe ...") - - with open(audio_path, "rb") as f: - # VibeVoice may return NDJSON for long audio — read the full - # response and extract the final "result" event. - with httpx.Client(timeout=httpx.Timeout(600.0)) as client: - resp = client.post( - f"{service_url}/transcribe", - files={"file": (audio_path.name, f, "audio/wav")}, - ) - resp.raise_for_status() - - content_type = resp.headers.get("content-type", "") - - if "application/x-ndjson" in content_type: - # Parse NDJSON — last "result" line is the transcription - data = None - for line in resp.text.strip().split("\n"): - line = line.strip() - if not line: - continue - event = json.loads(line) - if event.get("type") == "progress": - current = event.get("current", "?") - total = event.get("total", "?") - print(f" Progress: batch {current}/{total}") - elif event.get("type") == "result": - data = event - if data is None: - print("ERROR: NDJSON stream ended without a result event") - sys.exit(1) - else: - data = resp.json() - - # Summarize - segments = data.get("segments", []) - words = data.get("words", []) - text = data.get("text", "") - - print(f"\nResult summary:") - print(f" Text length: {len(text)} chars") - print(f" Segments: {len(segments)}") - print(f" Words: {len(words)}") - - if segments: - last_seg = segments[-1] - print(f" Duration: {last_seg.get('end', 0):.1f}s") - - # Check for overlaps - overlaps = [] - for i in range(len(segments) - 1): - overlap = segments[i]["end"] - segments[i + 1]["start"] - if overlap > 0.5: - overlaps.append( - ( - i, - i + 1, - overlap, - segments[i]["text"][:40], - segments[i + 1]["text"][:40], - ) - ) - - if overlaps: - print(f"\n WARNING: {len(overlaps)} segment overlaps detected!") - for idx_a, idx_b, dur, text_a, text_b in overlaps: - print( - f' [{idx_a}→{idx_b}] {dur:.2f}s overlap: "{text_a}" / "{text_b}"' - ) - else: - print(f" No segment overlaps detected") - - # Save - output = Path(output_path) - output.parent.mkdir(parents=True, exist_ok=True) - with open(output, "w") as f: - json.dump(data, f, indent=2) - - print(f"\nSaved to {output}") - return data - - -def main(): - parser = argparse.ArgumentParser( - description="Capture VibeVoice ground truth", - formatter_class=argparse.RawDescriptionHelpFormatter, - epilog=""" -To force batching on the 4-min test file, restart the service first: - - cd extras/asr-services - BATCH_THRESHOLD_SECONDS=60 BATCH_DURATION_SECONDS=90 BATCH_OVERLAP_SECONDS=15 \\ - docker compose up vibevoice-asr -d - -Then run this script. Restore defaults after with: - - docker compose up vibevoice-asr -d -""", - ) - parser.add_argument( - "--url", - default="http://localhost:8767", - help="VibeVoice service URL (default: http://localhost:8767)", - ) - parser.add_argument( - "--audio", - default=str(DEFAULT_AUDIO), - help="Path to test audio file", - ) - parser.add_argument( - "--output", - default=str(DEFAULT_OUTPUT), - help="Output JSON path", - ) - args = parser.parse_args() - capture(args.url, args.audio, args.output) - - -if __name__ == "__main__": - main() From 99c3410527e97f813ffe2562f4293d6e2f2dbfcf Mon Sep 17 00:00:00 2001 From: Ankush Malaker <43288948+AnkushMalaker@users.noreply.github.com> Date: Fri, 17 Jul 2026 14:12:47 +0000 Subject: [PATCH 06/16] config: enable diarization and memory executor settings --- config/config.yml.template | 1 + config/defaults.yml | 20 ++++++++++++++++++-- 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/config/config.yml.template b/config/config.yml.template index 5c2d1acf..fec21287 100644 --- a/config/config.yml.template +++ b/config/config.yml.template @@ -296,6 +296,7 @@ models: # insertions on Hindi speech (158% corpus WER vs 12.9% with hi on CoSHE-100) language: hi word_timestamps: 'true' + diarize: 'true' response: type: json extract: diff --git a/config/defaults.yml b/config/defaults.yml index 6d9f509d..c4d2bd17 100644 --- a/config/defaults.yml +++ b/config/defaults.yml @@ -376,6 +376,7 @@ models: # insertions on Hindi speech (158% corpus WER vs 12.9% with hi on CoSHE-100) language: hi word_timestamps: 'true' + diarize: 'true' response: type: json extract: @@ -554,6 +555,21 @@ memory: # Chronicle's agentic Markdown vault is the only memory provider. provider: chronicle timeout_seconds: 1200 + # How the memory agent executes vault writes: + # direct — built-in tool-calling loop via the configured LLM (metered API calls) + # codex — OpenAI Codex CLI running on the vault directory (uses a ChatGPT + # subscription; needs the codex binary + auth in CODEX_HOME — the + # wizard mounts ~/.codex into the containers) + agent_executor: direct + codex: + model: "gpt-5.6-terra" # faster/lower-cost Codex model for read-heavy vault work + reasoning_effort: low # keep routine ingestion economical; raise for quality evals + # codex --sandbox. workspace-write (writes confined to the vault dir) needs + # bubblewrap, which cannot run nested inside the rootless-podman containers — + # there the container is the isolation boundary and the wizard sets + # danger-full-access instead. + sandbox_mode: workspace-write + timeout_seconds: 900 # per-run wall clock before the run is failed extraction: enabled: true prompt: | @@ -576,7 +592,7 @@ speaker_recognition: hf_token: ${oc.env:HF_TOKEN,''} # Speaker identification threshold - similarity_threshold: 0.15 + similarity_threshold: 0.5 # Diarization chunking configuration (speaker service self-managed chunking) # Maximum audio duration (seconds) for single PyAnnote call @@ -707,7 +723,7 @@ backend: # "pyannote" always re-diarizes via the speaker service diarization: diarization_source: provider - similarity_threshold: 0.15 + similarity_threshold: 0.5 min_duration: 0.5 collar: 2.0 min_duration_off: 1.5 From e24da67798ae35689856a867811a38dc5b82bb7a Mon Sep 17 00:00:00 2001 From: Ankush Malaker <43288948+AnkushMalaker@users.noreply.github.com> Date: Fri, 17 Jul 2026 14:12:52 +0000 Subject: [PATCH 07/16] ci: add Python unit test coverage lanes --- .coveragerc | 22 ++++ .github/workflows/README.md | 15 +-- .github/workflows/python-tests.yml | 156 +++++++++++++++++++++++++++++ 3 files changed, 186 insertions(+), 7 deletions(-) create mode 100644 .coveragerc create mode 100644 .github/workflows/python-tests.yml diff --git a/.coveragerc b/.coveragerc new file mode 100644 index 00000000..5b30b17e --- /dev/null +++ b/.coveragerc @@ -0,0 +1,22 @@ +[run] +branch = True +relative_files = True +source = + config_manager + discovery + services + setup_utils + status + updates + wizard + +[report] +precision = 1 +show_missing = True +skip_empty = True + +[xml] +output = coverage-reports/coverage.xml + +[html] +directory = coverage-reports/html diff --git a/.github/workflows/README.md b/.github/workflows/README.md index bacc73ab..a9bbc1a9 100644 --- a/.github/workflows/README.md +++ b/.github/workflows/README.md @@ -6,11 +6,12 @@ Documentation for CI/CD workflows and test automation. Chronicle uses **three separate test workflows** to balance fast PR feedback with comprehensive testing: -| Workflow | Trigger | Test Coverage | API Keys | Purpose | +| Workflow | Trigger | Test selection | API Keys | Purpose | |----------|---------|---------------|----------|---------| -| `robot-tests.yml` | All PRs | ~70% (no-API tests) | ❌ Not required | Fast PR validation | -| `full-tests-with-api.yml` | Push to dev/main | 100% (full suite) | ✅ Required | Comprehensive validation | -| `pr-tests-with-api.yml` | PR label trigger | 100% (full suite) | ✅ Required | Pre-merge API testing | +| `python-tests.yml` | Relevant PRs, dev/main | Root, backend, and ASR pytest lanes | Not required | Unit tests and branch coverage reports | +| `robot-tests.yml` | All PRs | No-API Robot subset | Not required | Fast PR validation | +| `full-tests-with-api.yml` | Push to dev/main | Full Robot selection | Required | Comprehensive validation | +| `pr-tests-with-api.yml` | PR label trigger | Full Robot selection | Required | Pre-merge API testing | ## Workflow Details @@ -35,7 +36,7 @@ on: - **Makefile Target**: `make test-no-api OUTPUTDIR=results-no-api` - **Results**: `results-no-api/` - **Time**: ~10-15 minutes -- **Coverage**: ~70% of test suite +- **Selection**: Robot cases that do not require secrets, GPUs, or excluded slow/SDK environments **Benefits**: - Fast feedback on PRs @@ -134,7 +135,7 @@ if: contains(github.event.pull_request.labels.*.name, 'test-with-api-keys') **Normal PR Workflow**: 1. Push your branch 2. Create PR -3. `robot-tests.yml` runs automatically (~70% coverage) +3. `robot-tests.yml` runs the no-API Robot selection automatically 4. Fix any failures 5. Merge when tests pass @@ -142,7 +143,7 @@ if: contains(github.event.pull_request.labels.*.name, 'test-with-api-keys') 1. Push your branch 2. Create PR 3. Ask maintainer to add `test-with-api-keys` label -4. `pr-tests-with-api.yml` runs (100% coverage) +4. `pr-tests-with-api.yml` runs the full Robot selection 5. Fix any failures 6. Merge when tests pass diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml new file mode 100644 index 00000000..ad82e7de --- /dev/null +++ b/.github/workflows/python-tests.yml @@ -0,0 +1,156 @@ +name: Python Tests + +on: + pull_request: + paths: + - "*.py" + - "setup-requirements.txt" + - ".coveragerc" + - "tests/unit/**" + - "backends/advanced/src/**" + - "backends/advanced/tests/**" + - "backends/advanced/pyproject.toml" + - "backends/advanced/uv.lock" + - "extras/asr-services/common/**" + - "extras/asr-services/providers/**" + - "extras/asr-services/tests/**" + - "extras/asr-services/pyproject.toml" + - "extras/asr-services/uv.lock" + - ".github/workflows/python-tests.yml" + push: + branches: [dev, main] + workflow_dispatch: + +permissions: + contents: read + +jobs: + root-unit: + name: Root tooling unit tests + runs-on: ubuntu-latest + timeout-minutes: 10 + env: + PYTHONPATH: ${{ github.workspace }}/backends/advanced/src + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install uv + uses: astral-sh/setup-uv@v4 + with: + version: "latest" + + - name: Run unit tests with coverage + run: >- + uv run + --with-requirements setup-requirements.txt + --with pytest + --with pytest-cov + pytest tests/unit + --cov + --cov-config=.coveragerc + --cov-report=term-missing + --cov-report=xml + --cov-report=html + + - name: Upload coverage report + if: always() + uses: actions/upload-artifact@v4 + with: + name: root-unit-coverage + path: coverage-reports/ + if-no-files-found: warn + retention-days: 14 + + advanced-backend-unit: + name: Advanced backend unit tests + runs-on: ubuntu-latest + timeout-minutes: 15 + defaults: + run: + working-directory: backends/advanced + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install uv + uses: astral-sh/setup-uv@v4 + with: + version: "latest" + + - name: Install test dependencies + run: uv sync --locked --group test + + - name: Run unit tests with coverage + run: >- + uv run --group test pytest + --ignore=tests/test_audio_persistence_mongodb.py + --cov=advanced_omi_backend + --cov-report=term-missing + --cov-report=xml + --cov-report=html + + - name: Upload coverage report + if: always() + uses: actions/upload-artifact@v4 + with: + name: advanced-backend-unit-coverage + path: backends/advanced/coverage-reports/ + if-no-files-found: warn + retention-days: 14 + + asr-unit: + name: ASR unit tests + runs-on: ubuntu-latest + timeout-minutes: 10 + defaults: + run: + working-directory: extras/asr-services + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install uv + uses: astral-sh/setup-uv@v4 + with: + version: "latest" + + - name: Install test dependencies + run: uv sync --locked --group test + + - name: Run unit tests with coverage + run: >- + uv run --group test pytest + --ignore=tests/test_parakeet_service.py + --cov=common + --cov=providers + --cov-report=term-missing + --cov-report=xml + --cov-report=html + + - name: Upload coverage report + if: always() + uses: actions/upload-artifact@v4 + with: + name: asr-unit-coverage + path: extras/asr-services/coverage-reports/ + if-no-files-found: warn + retention-days: 14 From 34d965bddba0f6b9acfb6437646991a0f9f92ad7 Mon Sep 17 00:00:00 2001 From: Ankush Malaker <43288948+AnkushMalaker@users.noreply.github.com> Date: Fri, 17 Jul 2026 14:14:24 +0000 Subject: [PATCH 08/16] test: enforce canonical Robot Framework tags --- tests/Makefile | 29 +++-- tests/README.md | 23 ++-- tests/browser/browser_auth.robot | 1 + .../test_llm_custom_provider.robot | 1 + .../test_transcription_url.robot | 1 + tests/endpoints/client_queue_tests.robot | 20 ++-- tests/integration/sdk_tests.robot | 1 + tests/robot-tags.json | 21 ++++ tests/scripts/validate_robot_tags.py | 112 ++++++++++++++++++ tests/tags.md | 21 ++-- tests/test-requirements.txt | 1 + 11 files changed, 197 insertions(+), 34 deletions(-) create mode 100644 tests/robot-tags.json create mode 100644 tests/scripts/validate_robot_tags.py diff --git a/tests/Makefile b/tests/Makefile index 6ef60dbf..627e7fb3 100644 --- a/tests/Makefile +++ b/tests/Makefile @@ -1,7 +1,7 @@ # Chronicle Test Makefile # Shortcuts for running tests and managing test containers -.PHONY: help all clean configure \ +.PHONY: help all endpoints integration infra configuration clean configure validate-tags \ containers-start containers-stop containers-restart containers-rebuild \ containers-start-rebuild containers-clean containers-status containers-logs \ start stop restart rebuild start-rebuild status logs \ @@ -11,7 +11,7 @@ # Default output directory OUTPUTDIR ?= results -TEST_DIR = endpoints integration infrastructure asr +TEST_DIR = endpoints integration infrastructure asr configuration SERVICE ?= chronicle-backend-test # Test configuration file (set this to use different configs) @@ -47,10 +47,12 @@ help: @echo " make status - Show container status" @echo "" @echo "Running Tests:" + @echo " make validate-tags - Validate Robot tags without starting containers" @echo " make all - Run all tests (excludes slow/sdk/requires-gpu)" @echo " make endpoints - Run only endpoint tests" @echo " make integration - Run only integration tests" @echo " make infra - Run only infrastructure tests" + @echo " make configuration - Run container-independent configuration tests" @echo "" @echo "ASR Tests:" @echo " make test-asr - Run ASR protocol tests (no GPU required)" @@ -108,7 +110,7 @@ help: # Run all tests (excludes slow, sdk, and requires-gpu tests for faster feedback) # Creates a persistent fixture conversation that won't be deleted between suites -all: +all: validate-tags @echo "Running all tests (excluding slow, sdk, and requires-gpu tests)..." CREATE_FIXTURE=true uv run --with-requirements test-requirements.txt robot --outputdir $(OUTPUTDIR) \ --name "All Tests" \ @@ -120,7 +122,7 @@ all: $(TEST_DIR) # Run only endpoint tests -endpoints: +endpoints: validate-tags @echo "Running endpoint tests..." uv run --with-requirements test-requirements.txt robot --outputdir $(OUTPUTDIR) \ --name "Endpoint Tests" \ @@ -129,7 +131,7 @@ endpoints: endpoints # Run only integration tests -integration: +integration: validate-tags @echo "Running integration tests..." CREATE_FIXTURE=true uv run --with-requirements test-requirements.txt robot --outputdir $(OUTPUTDIR) \ --name "Integration Tests" \ @@ -138,7 +140,7 @@ integration: integration # Run only infrastructure tests -infra: +infra: validate-tags @echo "Running infrastructure tests..." uv run --with-requirements test-requirements.txt robot --outputdir $(OUTPUTDIR) \ --name "Infrastructure Tests" \ @@ -146,6 +148,15 @@ infra: --loglevel INFO:INFO \ infrastructure +# Run container-independent configuration tests +configuration: validate-tags + @echo "Running configuration tests..." + uv run --with-requirements test-requirements.txt robot --outputdir $(OUTPUTDIR) \ + --name "Configuration Tests" \ + --console verbose \ + --loglevel INFO:INFO \ + configuration + # Show current test configuration show-config: @echo "Current Test Configuration:" @@ -155,6 +166,10 @@ show-config: @echo " make test CONFIG=deepgram-openai.yml" @echo " make test CONFIG=/path/to/custom.yml" +# Validate suite metadata before spending time starting the composed environment. +validate-tags: + @uv run --with-requirements test-requirements.txt python3 scripts/validate_robot_tags.py + # Clean up test output files clean: @echo "Cleaning test outputs..." @@ -289,7 +304,7 @@ test-with-api-keys: # Run tests without API keys (CI mode) # Switches to mock-services.yml config and excludes requires-api-keys tests -test-no-api: +test-no-api: validate-tags @echo "🔄 Running tests without API keys (CI mode)..." @$(MAKE) containers-stop @TEST_CONFIG_FILE=/app/test-configs/mock-services.yml $(MAKE) containers-start diff --git a/tests/README.md b/tests/README.md index b199a2d1..523ffcb4 100644 --- a/tests/README.md +++ b/tests/README.md @@ -5,13 +5,19 @@ Start containers and run tests: ```bash cd tests -make test # Start containers + run all tests (excludes slow/sdk) +make test # Start containers + run tests (excludes slow/sdk/GPU) +``` + +Validate suite tags without starting containers: + +```bash +make validate-tags ``` Or step by step: ```bash make start # Start test containers -make test-all # Run all tests (excludes slow/sdk) +make all # Run all tests (excludes slow/sdk/GPU) make stop # Stop containers ``` @@ -22,9 +28,10 @@ make stop # Stop containers Run specific test suites: ```bash -make test-endpoints # API endpoint tests (~40 tests, fast) -make test-integration # End-to-end workflows (~15 tests, slower) -make test-infra # Infrastructure resilience tests (~5 tests) +make endpoints # API endpoint tests +make integration # End-to-end workflows +make infra # Infrastructure resilience tests +make configuration # Container-independent configuration tests ``` ### Special Test Categories @@ -201,7 +208,7 @@ cat tests/logs/2026-01-17_14-30-45/chronicle-backend-test.log Chronicle tests are separated into two execution paths: -### 1. No API Keys Required (~70% of tests) +### 1. No API Keys Required These tests run without external API dependencies: - Endpoint tests (CRUD operations, permissions) - Infrastructure tests (workers, queues, health checks) @@ -231,7 +238,7 @@ OPENAI_API_KEY=your-key-here **Faster iteration:** 1. Start containers once: `make start` -2. Run specific test suite: `make test-endpoints` +2. Run specific test suite: `make endpoints` 3. Keep containers running between test runs 4. Only rebuild when code changes: `make rebuild` @@ -252,7 +259,7 @@ uv run --with-requirements test-requirements.txt robot \ make rebuild # 3. Run specific test suite -make test-endpoints +make endpoints # 4. View logs if needed make logs SERVICE=chronicle-backend-test diff --git a/tests/browser/browser_auth.robot b/tests/browser/browser_auth.robot index d03c0017..42af3b9c 100644 --- a/tests/browser/browser_auth.robot +++ b/tests/browser/browser_auth.robot @@ -3,6 +3,7 @@ Documentation Browser Authentication Tests Library Browser Resource ../setup/setup_keywords.robot Suite Setup Suite Setup +Test Tags permissions e2e diff --git a/tests/configuration/test_llm_custom_provider.robot b/tests/configuration/test_llm_custom_provider.robot index 66d7fa81..7d4e5ee9 100644 --- a/tests/configuration/test_llm_custom_provider.robot +++ b/tests/configuration/test_llm_custom_provider.robot @@ -4,6 +4,7 @@ Library OperatingSystem Library Collections Library String Library ../libs/ConfigTestHelper.py +Test Tags infra *** Keywords *** Setup Temp Config diff --git a/tests/configuration/test_transcription_url.robot b/tests/configuration/test_transcription_url.robot index 618771df..ea5edb6b 100644 --- a/tests/configuration/test_transcription_url.robot +++ b/tests/configuration/test_transcription_url.robot @@ -2,6 +2,7 @@ Documentation Tests for Transcription Service URL Configuration Library Collections Library ../libs/ConfigTestHelper.py +Test Tags infra *** Test Cases *** Vibevoice Url Without Http Prefix diff --git a/tests/endpoints/client_queue_tests.robot b/tests/endpoints/client_queue_tests.robot index 03d21c4a..225fdffe 100644 --- a/tests/endpoints/client_queue_tests.robot +++ b/tests/endpoints/client_queue_tests.robot @@ -13,7 +13,7 @@ Suite Teardown Delete All Sessions Get Active Clients Test [Documentation] Test getting active client information - [Tags] client active positive + [Tags] infra Create API Session admin_session ${response}= GET On Session admin_session /api/clients/active @@ -31,7 +31,7 @@ Get Active Clients Test Get Queue Jobs Test [Documentation] Test getting queue jobs with pagination - [Tags] queue jobs positive + [Tags] queue Create API Session admin_session &{params}= Create Dictionary limit=20 offset=0 @@ -53,7 +53,7 @@ Get Queue Jobs Test Get Queue Jobs With Different Limits Test [Documentation] Test queue jobs pagination with different limits - [Tags] queue jobs pagination positive + [Tags] queue Get Anonymous Session anon_session Create API Session admin_session @@ -78,7 +78,7 @@ Get Queue Jobs With Different Limits Test Get Queue Statistics Test [Documentation] Test getting queue statistics - [Tags] queue statistics positive + [Tags] queue Get Anonymous Session anon_session Create API Session admin_session @@ -95,7 +95,7 @@ Get Queue Statistics Test Get Queue Health Test [Documentation] Test getting queue health status - [Tags] queue health positive + [Tags] queue health Get Anonymous Session anon_session Create API Session admin_session @@ -112,7 +112,7 @@ Get Queue Health Test Queue Jobs User Isolation Test [Documentation] Test that regular users only see their own queue jobs - [Tags] queue security isolation + [Tags] queue permissions Get Anonymous Session anon_session Create API Session admin_session @@ -140,7 +140,7 @@ Queue Jobs User Isolation Test Invalid Queue Parameters Test [Documentation] Test queue endpoints with invalid parameters - [Tags] queue negative validation + [Tags] queue Get Anonymous Session anon_session Create API Session admin_session @@ -162,7 +162,7 @@ Invalid Queue Parameters Test Unauthorized Client Access Test [Documentation] Test that client endpoints require authentication - [Tags] client security negative + [Tags] infra permissions ${session}= Get Anonymous Session session @@ -172,7 +172,7 @@ Unauthorized Client Access Test Unauthorized Queue Access Test [Documentation] Test that queue endpoints require authentication - [Tags] queue security negative + [Tags] queue permissions ${session}= Get Anonymous Session session # Try to access queue jobs without token @@ -186,7 +186,7 @@ Unauthorized Queue Access Test Client Manager Integration Test [Documentation] Test client manager functionality - [Tags] client manager integration + [Tags] infra Create API Session admin_session diff --git a/tests/integration/sdk_tests.robot b/tests/integration/sdk_tests.robot index 3f71247f..0d5d1780 100644 --- a/tests/integration/sdk_tests.robot +++ b/tests/integration/sdk_tests.robot @@ -16,6 +16,7 @@ Variables ../setup/test_env.py Suite Setup Suite Setup Suite Teardown Suite Teardown +Test Tags sdk *** Variables *** ${BACKEND_URL} http://localhost:8001 diff --git a/tests/robot-tags.json b/tests/robot-tags.json new file mode 100644 index 00000000..be784042 --- /dev/null +++ b/tests/robot-tags.json @@ -0,0 +1,21 @@ +{ + "business": [ + "permissions", + "conversation", + "memory", + "chat", + "queue", + "health", + "infra", + "audio-upload", + "audio-batch", + "audio-streaming", + "e2e" + ], + "execution": [ + "requires-api-keys", + "slow", + "sdk", + "requires-gpu" + ] +} diff --git a/tests/scripts/validate_robot_tags.py b/tests/scripts/validate_robot_tags.py new file mode 100644 index 00000000..6ea6899e --- /dev/null +++ b/tests/scripts/validate_robot_tags.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python3 +"""Validate Robot test tags against Chronicle's canonical allowlist.""" + +import argparse +import json +from pathlib import Path + +from robot.api.parsing import get_model +from robot.parsing.model.blocks import TestCase, TestCaseSection +from robot.parsing.model.statements import Tags, TestTags + +TEST_SUITE_DIRS = ( + "endpoints", + "integration", + "infrastructure", + "asr", + "browser", + "configuration", +) + + +def load_allowed_tags(config_path: Path) -> set[str]: + groups = json.loads(config_path.read_text()) + return {tag for tags in groups.values() for tag in tags} + + +def validate_file(path: Path, allowed_tags: set[str]) -> list[str]: + errors: list[str] = [] + model = get_model(path) + suite_tags: set[str] = set() + + for section in model.sections: + for node in getattr(section, "body", ()): + if isinstance(node, TestTags): + suite_tags.update(node.values) + + invalid_suite_tags = suite_tags - allowed_tags + if invalid_suite_tags: + errors.append( + f"{path}: invalid suite tags: {', '.join(sorted(invalid_suite_tags))}" + ) + + for section in model.sections: + if not isinstance(section, TestCaseSection): + continue + for test in section.body: + if not isinstance(test, TestCase): + continue + tag_statement = next( + (node for node in test.body if isinstance(node, Tags)), None + ) + test_tags = set(tag_statement.values) if tag_statement else set() + invalid_tags = test_tags - allowed_tags + location = f"{path}:{test.lineno} ({test.name})" + + if invalid_tags: + errors.append( + f"{location}: invalid tags: {', '.join(sorted(invalid_tags))}" + ) + if not suite_tags and not test_tags: + errors.append(f"{location}: no business or execution tag") + + return errors + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "paths", + nargs="*", + type=Path, + default=[Path(name) for name in TEST_SUITE_DIRS], + help="Robot files or suite directories (defaults to all test suite directories)", + ) + parser.add_argument( + "--config", + type=Path, + default=Path(__file__).resolve().parents[1] / "robot-tags.json", + help="Path to the canonical tag allowlist", + ) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + allowed_tags = load_allowed_tags(args.config) + files: set[Path] = set() + + for path in args.paths: + if path.is_dir(): + files.update(path.rglob("*.robot")) + elif path.suffix == ".robot" and path.is_file(): + files.add(path) + else: + print(f"Tag validation path does not exist or is not a Robot file: {path}") + return 2 + + errors = [ + error for path in sorted(files) for error in validate_file(path, allowed_tags) + ] + if errors: + print("Robot tag validation failed:") + for error in errors: + print(f"- {error}") + return 1 + + print(f"Validated {len(files)} Robot suites against {len(allowed_tags)} tags") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/tags.md b/tests/tags.md index 514dcc08..8024ce4d 100644 --- a/tests/tags.md +++ b/tests/tags.md @@ -4,7 +4,8 @@ This document defines the standard tags used across the Chronicle test suite. ## Simplified Tag Set -Chronicle uses a **minimal, focused tag set** for test organization. Only 15 tags are permitted. +Chronicle uses 11 business tags and four execution-requirement tags. The 15 permitted tags are +defined in [`robot-tags.json`](robot-tags.json), which is the canonical machine-readable allowlist. ## Tag Format @@ -102,7 +103,7 @@ Chronicle uses a **minimal, focused tag set** for test organization. Only 15 tag - Connection resilience tests - Heavy integration tests with multiple service restarts - Excluded from default `make test` runs for faster feedback -- Run explicitly with `make test-slow` or `make test-all-with-slow` +- Run explicitly with `make test-slow` or `make test-all-with-slow-and-sdk` **`sdk`** - Tests for unreleased SDK functionality - SDK integration tests @@ -197,7 +198,7 @@ Use 2-3 tags only when testing interactions between components: ## Prohibited Tags -**DO NOT create or use any tags other than the 14 approved tags above.** +**DO NOT create or use any tags other than the 15 approved tags above.** Commonly misused tags that should NOT be used: - ❌ `positive`, `negative` - Test outcome is in the results, not tags @@ -233,9 +234,10 @@ robot --include audio-upload --include audio-streaming tests/ ### Before Adding a New Tag **STOP!** Ask yourself: -1. Can I use one of the existing 11 tags? -2. Is this tag really necessary for test organization? -3. Have I checked with the team? +1. Can I use one of the 11 business tags? +2. Does the test genuinely need one of the four execution-requirement tags? +3. Is a new tag really necessary for test selection? +4. Have I checked with the team? **New tags require team approval and must be added to this document first.** @@ -243,9 +245,10 @@ robot --include audio-upload --include audio-streaming tests/ When updating tags across test files: 1. Update all affected test files -2. Update this document (tags.md) -3. Update TESTING_GUIDELINES.md if rules changed -4. Document the change in the commit message +2. Update `robot-tags.json` +3. Update this document (tags.md) +4. Update TESTING_GUIDELINES.md if rules changed +5. Document the change in the commit message ## Tag Statistics diff --git a/tests/test-requirements.txt b/tests/test-requirements.txt index 2ea12f65..ed431460 100644 --- a/tests/test-requirements.txt +++ b/tests/test-requirements.txt @@ -8,3 +8,4 @@ websockets pymongo omegaconf pyyaml +ruamel-yaml From 2327d4278ce6ea18207e66f7b1b4bedfa8449677 Mon Sep 17 00:00:00 2001 From: Ankush Malaker <43288948+AnkushMalaker@users.noreply.github.com> Date: Fri, 17 Jul 2026 14:14:54 +0000 Subject: [PATCH 09/16] fix(app): stop swallowing Deepgram key storage errors --- app/src/utils/storage.ts | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/app/src/utils/storage.ts b/app/src/utils/storage.ts index 9ef3e79d..6a0ca2b8 100644 --- a/app/src/utils/storage.ts +++ b/app/src/utils/storage.ts @@ -100,26 +100,20 @@ export const saveDeepgramApiKey = async (apiKey: string | null): Promise<void> = try { if (apiKey) { await AsyncStorage.setItem(DEEPGRAM_API_KEY_KEY, apiKey); - console.log('[Storage] Deepgram API Key saved.'); // Don't log the key itself for security } else { await AsyncStorage.removeItem(DEEPGRAM_API_KEY_KEY); - console.log('[Storage] Deepgram API Key removed.'); } } catch (error) { - console.error('[Storage] Error saving Deepgram API Key:', error); + throw error; } }; export const getDeepgramApiKey = async (): Promise<string | null> => { try { const apiKey = await AsyncStorage.getItem(DEEPGRAM_API_KEY_KEY); - if (apiKey) { - console.log('[Storage] Retrieved Deepgram API Key.'); - } return apiKey; } catch (error) { - console.error('[Storage] Error retrieving Deepgram API Key:', error); - return null; + throw error; } }; From 580d80e62a6512cdb1de0b23571c3967d2470f1f Mon Sep 17 00:00:00 2001 From: Ankush Malaker <43288948+AnkushMalaker@users.noreply.github.com> Date: Fri, 17 Jul 2026 14:15:10 +0000 Subject: [PATCH 10/16] docs: align contributor guidance and coverage ignores --- .gitignore | 4 ++++ AGENTS.md | 22 +++++++++++++++------- 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/.gitignore b/.gitignore index 66c21082..76d03b45 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,8 @@ **/__pycache__ +.coverage +.coverage.* +**/coverage-reports/ +**/htmlcov/ *.wav # Wake-word notification tones are bundled assets, not user audio — keep them tracked # so they're baked into the wakeword-service image (and served to all clients). diff --git a/AGENTS.md b/AGENTS.md index e07cf736..5b65b3e4 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -37,6 +37,8 @@ uv run --with-requirements setup-requirements.txt python wizard.py **Note on Convenience Scripts**: Chronicle provides wrapper scripts (`./wizard.sh`, `./start.sh`, `./restart.sh`, `./stop.sh`, `./status.sh`) that simplify the longer `uv run --with-requirements setup-requirements.txt python` commands. Use these for everyday operations. +**Temporary Tooling Rule**: When a Python tool or dependency is needed only for a one-off task, run it ephemerally with `uv run --with <package> ...`; do not install it into a project environment or add it to project dependencies. For a package-owned CLI, use `uvx --from <package> <entrypoint>`. This applies to browser automation, screenshots, data inspection, and other temporary utilities. + ### Setup Documentation For detailed setup instructions and troubleshooting, see: - **[@quickstart.md](quickstart.md)**: Beginner-friendly step-by-step setup guide @@ -99,13 +101,13 @@ make test # Start containers + run all tests # Or step by step make start # Start test containers (with health checks) -make test-all # Run all test suites +make all # Run all test suites make stop # Stop containers (preserves volumes) # Run specific test suites -make test-endpoints # API endpoint tests (~40 tests, fast) -make test-integration # End-to-end workflows (~15 tests, slower) -make test-infra # Infrastructure resilience (~5 tests) +make endpoints # API endpoint tests +make integration # End-to-end workflows +make infra # Infrastructure resilience # Quick iteration (reuse existing containers) make test-quick # Run tests without restarting containers @@ -611,13 +613,13 @@ For detailed technical documentation, see: Before writing any Robot Framework test: 1. **Read [@tests/TESTING_GUIDELINES.md](tests/TESTING_GUIDELINES.md)** for comprehensive testing patterns and standards -2. **Check [@tests/tags.md](tests/tags.md)** for approved tags - ONLY 11 tags are permitted +2. **Check [@tests/tags.md](tests/tags.md)** for approved tags - only the 11 business tags and 4 execution tags are permitted 3. **SCAN existing resource files** for keywords - NEVER write code that duplicates existing keywords 4. **Follow the Arrange-Act-Assert pattern** with inline verifications (not abstracted to keywords) Key Testing Rules: - **Check Existing Keywords FIRST**: Before writing ANY test code, scan relevant resource files (`websocket_keywords.robot`, `queue_keywords.robot`, `conversation_keywords.robot`, etc.) for existing keywords -- **Tags**: ONLY use the 11 approved tags from tags.md, tab-separated (e.g., `[Tags] infra audio-streaming`) +- **Tags**: ONLY use the 15 approved tags from tags.md, tab-separated (e.g., `[Tags] infra audio-streaming`) - **Verifications**: Write assertions directly in tests, not in resource keywords - **Keywords**: Only create keywords for reusable setup/action operations AFTER confirming no existing keyword exists - **Resources**: Always check existing resource files before creating new keywords or duplicating logic @@ -625,7 +627,7 @@ Key Testing Rules: **DO NOT:** - Write inline code without checking if a keyword already exists for that operation -- Create custom tags (use only the 11 approved tags) +- Create custom tags (use only the 15 approved tags) - Abstract verifications into keywords (keep them inline in tests) - Use space-separated tags (must be tab-separated) - Skip reading the guidelines before writing tests @@ -635,6 +637,12 @@ Check if the src/ is volume mounted. If not, do compose build so that code chang Check `docs/backend/` for up-to-date information on the advanced backend. All docker projects have .dockerignore following the exclude pattern. That means files need to be included for them to be visible to docker. The uv package manager is used for all python projects. Wherever you'd call `python3 main.py` you'd call `uv run python main.py` +For temporary Python-backed tooling that is not part of the repo dependencies, prefer `uv run --with <package> python ...` instead of installing packages into the project. Use `uvx --from <package> <entrypoint>` when invoking a package's own CLI. For browser scripts/screenshots, use `uv run --with playwright python - <<'PY'` and import `playwright.sync_api`; for the Playwright CLI use `uvx --from playwright playwright ...`. Do not add transient Playwright/npm packages to `package.json` just to drive a one-off check. + +**Compute-Intensive Workloads:** +- Chronicle is designed for heavy data and AI processing. Re-encoding or recomputation is acceptable when it improves correctness or output quality. +- Avoid wasteful repeated work: cache reusable artifacts, fingerprint model and configuration inputs, and reuse valid intermediate results. +- Prefer GPU acceleration whenever the workload and deployed service support it. **Container Engine (Docker or Podman):** - The project supports **both Docker and Podman**. The active engine is set by `container_engine` in `config/config.yml` (default `docker`); prefer the lifecycle scripts (`./start.sh`/`./stop.sh`/`./restart.sh`) which route through the selected engine. For one-off manual commands under Podman use `podman-compose` (not `docker compose`). See **[@docs/podman.md](docs/podman.md)**. From d60af3a25c2e34087234cb7ae4e75f63deda912b Mon Sep 17 00:00:00 2001 From: Ankush Malaker <43288948+AnkushMalaker@users.noreply.github.com> Date: Fri, 17 Jul 2026 14:15:41 +0000 Subject: [PATCH 11/16] docs: document testing coverage and data archives --- docs/README.md | 3 + docs/backend/data-archive.md | 132 +++++++++++++++++++++ docs/test-coverage-audit.md | 214 +++++++++++++++++++++++++++++++++++ docs/testing.md | 77 +++++++++++++ 4 files changed, 426 insertions(+) create mode 100644 docs/backend/data-archive.md create mode 100644 docs/test-coverage-audit.md create mode 100644 docs/testing.md diff --git a/docs/README.md b/docs/README.md index f281f0c1..28a282a0 100644 --- a/docs/README.md +++ b/docs/README.md @@ -7,6 +7,8 @@ day-to-day operation, and use [AGENTS.md](../AGENTS.md) for development conventi ## System - [Project overview](overview.md): components, deployment topology, and repository layout +- [Testing and coverage](testing.md): fast Python lanes, coverage reports, and integration tests +- [Test coverage audit](test-coverage-audit.md): current baseline, gaps, and cleanup plan - [Audio pipeline](audio-pipeline-architecture.md): session, transcription, and memory flow - [Initialization system](init-system.md): setup wizard and service orchestration - [Podman](podman.md): rootless containers, GPU access, and engine migration @@ -16,6 +18,7 @@ day-to-day operation, and use [AGENTS.md](../AGENTS.md) for development conventi - [Authentication](backend/auth.md): user identity, JWTs, and protected endpoints - [Memory system](backend/memories.md): agentic Markdown vault and retrieval +- [Data archive and memory rebuild](backend/data-archive.md): full export/import and clean vault reconstruction - [Plugin configuration](backend/plugin-configuration.md): configuration and secret boundaries - [Plugin development](backend/plugin-development-guide.md): plugin lifecycle and APIs diff --git a/docs/backend/data-archive.md b/docs/backend/data-archive.md new file mode 100644 index 00000000..e0058da0 --- /dev/null +++ b/docs/backend/data-archive.md @@ -0,0 +1,132 @@ +# Data Archive and Memory Rebuild + +Chronicle can export its durable data to a checksummed `.chronicle` archive, +restore that archive, and reconstruct the Markdown memory vault from the active +conversation transcripts. + +## What is archived + +The archive contains: + +- every MongoDB collection as a BSON stream, including users, conversations, + transcript versions, annotations, chat data, and `audio_chunks`; +- the original Opus bytes and timing metadata stored in each audio chunk; +- `data/conversation_docs/`, `data/memory_md/`, and legacy + `data/audio_chunks/` files; +- a versioned manifest with the document count, byte size, and SHA-256 digest + of every archive member. + +BSON is used instead of JSON so MongoDB ObjectIds, datetimes, binary audio, and +nested transcript data round-trip without conversion. Archives include password +hashes and personal data from the users collection. Store them as sensitive data. + +## Maintenance window + +Export is not a transactional MongoDB snapshot. For a consistent archive, stop +new device ingestion and the services that write data while running maintenance: + +```bash +cd backends/advanced +docker compose stop chronicle-backend workers annotation-cron +./chronicle-data.sh export +docker compose start chronicle-backend workers annotation-cron +``` + +`chronicle-data.sh` uses a one-off Compose container, so the backend service may +remain stopped. Set `CONTAINER_ENGINE=podman` or `COMPOSE_CMD=podman-compose` for +a Podman deployment. + +## Export and verify + +With no output argument, archives are written under `data/backups/`: + +```bash +./chronicle-data.sh export +./chronicle-data.sh export /app/data/backups/before-upgrade.chronicle +./chronicle-data.sh verify /app/data/backups/before-upgrade.chronicle +``` + +The exporter writes to a partial file and renames it only after completion. The +importer verifies the complete manifest and every checksum before changing MongoDB +or filesystem data. + +During import, Chronicle first groups conversations by chunk structure, then decodes +only matching candidates and fingerprints their PCM samples. This detects the same +clip even when its Opus container bytes differ. If identical audio occurs more than +once in the archive, only the earliest conversation by `created_at` is imported; its +transcript versions are retained and the later conversation plus related audio +chunks, annotations, waveforms, and other conversation-scoped records are skipped. +Merge imports also compare against audio already in MongoDB, where the existing +conversation wins. Every skipped duplicate is written to the application log and +printed by the CLI with both conversation IDs. + +## Restore an archive + +Merge restore upserts documents by MongoDB `_id` and overlays filesystem files: + +```bash +./chronicle-data.sh import /app/data/backups/before-upgrade.chronicle +``` + +Replace restore clears every collection and filesystem root represented by the +archive before restoring it: + +```bash +./chronicle-data.sh import /app/data/backups/before-upgrade.chronicle \ + --replace --force +``` + +Use `--database-only` to leave the current vault and legacy filesystem audio +untouched. + +## Restore and rebuild derived data + +`--rebuild-from` imports the durable MongoDB records, deliberately skips the +archived `memory_audit` collection and vault files, clears current derived memory +state for all users, and queues active transcripts from the selected stage: + +```bash +./chronicle-data.sh import /app/data/backups/before-upgrade.chronicle \ + --replace --rebuild-from memory --force + +./chronicle-data.sh import /app/data/backups/before-upgrade.chronicle \ + --replace --rebuild-from speakers --force +``` + +The `speakers` mode runs speaker recognition first and creates new active transcript +versions before starting memory reconstruction. It processes every non-deleted +conversation with an active transcript; `memory_excluded=true` conversations receive +speaker processing but remain excluded from memory. The `memory` mode starts directly +from the imported active transcripts. + +Transcript-only conversations that have no stored audio chunks are reported and +skipped by the speaker stage. They still participate in memory reconstruction from +their imported active transcript. + +Speaker and memory jobs are ordered chronologically within each user through RQ +dependencies. A failed speaker conversation is logged and skipped without blocking +later conversations; its memory is rebuilt from the unchanged active transcript. +Different users can rebuild in parallel. A copy of the old vault is written to +`data/backups/memory_vault_<timestamp>.tar.gz` before it is cleared; pass +`--no-vault-backup` to disable that copy. + +Start the workers after the import command queues the replay chains: + +```bash +docker compose start workers chronicle-backend annotation-cron +``` + +## Rebuild memory from current data + +The same clean rebuild can run without importing an archive: + +```bash +./chronicle-data.sh rebuild-memory --dry-run +./chronicle-data.sh rebuild-memory --force +./chronicle-data.sh rebuild-memory --rebuild-from speakers --force +./chronicle-data.sh rebuild-memory --user-id 507f1f77bcf86cd799439011 --force +``` + +The rebuild refuses to start when an existing queued, deferred, scheduled, or +running speaker or memory job targets any selected conversation. Syncthing +`.stfolder` and `.stignore` markers are retained while vault content is cleared. diff --git a/docs/test-coverage-audit.md b/docs/test-coverage-audit.md new file mode 100644 index 00000000..5bf183b5 --- /dev/null +++ b/docs/test-coverage-audit.md @@ -0,0 +1,214 @@ +# Test Coverage Audit + +Audit date: 2026-07-15 +Git baseline: `0e2eefa5` on `dev` +Scope: Git-tracked code is the reproducible baseline. Current untracked source and tests were +examined separately because they are not available to CI or a fresh contributor checkout. + +## Executive summary + +Chronicle has broad end-to-end scenario coverage, especially around backend endpoints, audio +streaming, queues, and conversations. It does not currently have a trustworthy repository-wide +code coverage number. + +The best available measurement is for the advanced Python backend: + +- The current working-tree pytest suite executed about **20.0% of tracked backend statements** + (`4,475 / 22,321`). This is an exposure measurement, not a passing baseline: the run had 127 + passes, 14 failures, and 33 errors after two collection failures were excluded. +- The smaller, Git-tracked, locally collectable unit subset exercised **8% of the statements + Coverage.py could discover in that run** and had 69 passes and 3 failures. Coverage discovery + omitted unimported namespace-package directories, so 8% is an upper bound rather than a full + backend denominator. +- The backend's most important orchestration code is the least directly tested area: workers had + **8.4%** measured coverage and controllers had **11.2%** in the broader working-tree run. +- All TypeScript/JavaScript products have **no configured test runner and no direct test files**: + the main web UI, mobile app, speaker-recognition UI, and vault graph UI. +- CI runs Robot tests and one speaker-recognition integration test, but it does **not run backend, + root-tooling, or ASR pytest suites**, and it publishes no line or branch coverage. + +The project therefore has meaningful workflow protection, but weak feedback on which branches and +failure paths are exercised. Robot's documented "70% coverage" means percentage of Robot tests +selected, not percentage of application code covered. + +## Inventory + +### Test suites + +| Surface | Current tests | Audit result | +| --- | --- | --- | +| Advanced backend | 7 tracked top-level pytest files; 12 more top-level files are currently untracked | Default collection fails; no pytest CI job | +| Cross-repo Robot suite | 232 tests across 34 suites | 118 endpoint, 64 integration, 28 ASR, 17 configuration, 4 infrastructure, 1 browser | +| Root setup/lifecycle tooling | 5 pytest files, about 80 test functions | Collection fails in 2 files; remaining run had 46 pass / 10 fail | +| ASR services | 4 pytest files, 63 collected cases | 49 pass / 9 skip / 2 fail / 3 environment errors | +| Speaker recognition | 1 large pytest integration scenario plus 2 ad hoc scripts | Secret-, model-, and Docker-dependent CI only | +| Frontends | No test files or test scripts | Build/lint only; no measured coverage | +| TTS and most plugins/extras | No direct automated tests | Some behavior is reached indirectly by Robot tests | + +There are 232 Robot cases, but the normal no-API lane selects only 166 after tag exclusions. The +suite contains 17 tests with unconditional `Skip`, including eight mobile placeholders and five SDK +placeholders. These should not be counted as implemented coverage. + +### Advanced backend line coverage + +This table comes from the broader current-working-tree pytest run. Failing tests still execute +lines, so the numbers indicate code reached by tests, not verified behavior. + +| Backend area | Statements | Covered | Coverage | +| --- | ---: | ---: | ---: | +| Models | 711 | 479 | 67.4% | +| Plugins framework | 453 | 136 | 30.0% | +| Routers | 3,598 | 943 | 26.2% | +| Utilities | 1,147 | 262 | 22.8% | +| Services | 5,490 | 1,207 | 22.0% | +| Package root | 2,838 | 625 | 22.0% | +| Observability | 246 | 51 | 20.7% | +| Clients | 260 | 47 | 18.1% | +| Controllers | 4,169 | 467 | 11.2% | +| Workers | 3,532 | 297 | 8.4% | + +Across the measured backend, 34 files had zero executed lines and 46 more were below 20%. The +largest high-risk gaps include: + +- `app_factory.py`: 289 statements, 0% +- `task_manager.py`: 192 statements, 0% +- memory agent edit/tool modules: 611 combined statements, 0% +- worker orchestrator modules: 465 combined statements, 0% +- `system_controller.py`: 1,103 statements, 10.0% +- `websocket_controller.py`: 712 statements, 9.3% +- `conversation_jobs.py`: 715 statements, 12.2% +- `transcription_jobs.py`: 561 statements, 8.9% +- `speaker_recognition_client.py`: 624 statements, 5.9% + +The root management tests measured 18% across five imported modules. `updates.py` reached 93%, but +`services.py`, `config_manager.py`, `discovery.py`, and `setup_utils.py` were between 10% and 20%. +This is not a whole-root percentage because two configured modules were never imported. + +## Suite health findings + +### 1. Test execution is not reproducible from a clean checkout + +- Advanced backend collection references deleted or changed APIs in `test_obsidian_service.py` and + `test_vad_analysis.py`. +- Root tests reference removed `merge_configs`, `read_config_yml`, and service lifecycle APIs. +- The ASR default suite starts Docker from an otherwise unit-oriented `pytest tests` command; Docker + absence becomes an error instead of an integration skip. Two independent ROCm Dockerfile checks + also currently fail. +- Robot dry-run found 16 failures in configuration suites because `ruamel-yaml` is imported but is + absent from `tests/test-requirements.txt`. +- The browser suite requires `rfbrowser init`, but that setup is absent from the test runner and CI. + +### 2. Test location does not consistently describe test type + +- `backends/advanced/tests/` mixes pure units, MongoDB integration tests, and manual graph-validation + scripts. The untracked graph scripts are automatically collected by pytest but expect command-line + arguments as fixtures. +- Root `tests/unit/` mixes root lifecycle scripts with advanced-backend configuration behavior. +- Robot `endpoints/` tests are full-stack HTTP contract/integration tests, not unit endpoint tests. +- Robot `integration/` includes genuine end-to-end flows, SDK contract tests, and placeholder mobile + scenarios. +- Hardware-, secret-, browser-, and container-dependent tests are separated mainly by tags, but tag + hygiene is not enforced. + +### 3. Tags and documentation have drifted + +- `tests/tags.md` says there are 15 approved tags, later says 14, and finally asks whether one of 11 + tags can be used. +- Ten tests in `client_queue_tests.robot` use prohibited tags such as `positive`, `negative`, + `security`, `client`, `jobs`, and `integration`. +- Eighteen Robot tests have no tags. +- Only four of nine SDK tests carry the `sdk` tag. The other five are placeholder skips that remain in + default selection. +- `tests/README.md` and `AGENTS.md` document nonexistent Make targets including `test-all`, + `test-endpoints`, `test-integration`, and `test-infra`; the actual targets are `all`, `endpoints`, + `integration`, and `infra`. + +### 4. CI enforces workflows, not code coverage + +The required OSS-friendly PR lane runs 166 no-secret Robot cases. That is valuable, but no workflow +runs the fast Python tests or any frontend tests. There is no `.coveragerc`, branch coverage setting, +coverage artifact, combined coverage job, patch coverage check, or minimum threshold. + +Path filters also mean changes to root lifecycle tooling, ASR, TTS, plugins, and frontends do not +trigger the main Robot workflow unless they touch its listed backend/test paths. + +## Proposed ownership model + +Tests should live with the component whose code they primarily validate. The repository-level +`tests/` directory should own cross-component acceptance behavior only. + +| Type | Location | Allowed dependencies | PR policy | +| --- | --- | --- | --- | +| Unit | `<component>/tests/unit/` | No network, Docker, database, secrets, or model downloads | Required; seconds | +| Component | `<component>/tests/component/` | In-process app plus fakes/in-memory stores | Required; minutes | +| Integration | `<component>/tests/integration/` | Real database/service containers | Required where CPU/no-secret | +| Contract | `tests/contract/` | Black-box API against composed Chronicle | Required no-secret subset | +| End to end | `tests/e2e/` | Multiple services and full user workflow | Required smoke subset; full scheduled | +| Environment | `tests/environment/{browser,gpu,hardware,secrets}/` | Explicit specialized runner | Optional/scheduled unless relevant | + +Use pytest markers matching execution requirements, not business domains: +`unit`, `component`, `integration`, `docker`, `gpu`, `external_api`, and `slow`. Keep Robot business +tags for selecting product behavior, but validate them against one canonical allowlist in CI. + +## Recommended rollout + +### Remediation started + +The first infrastructure tranche was implemented after this baseline was recorded: + +- Root tooling, advanced backend, and ASR now have branch-coverage configuration and a dedicated + Python CI workflow that uploads XML and HTML reports. +- Fast backend and ASR lanes explicitly exclude their MongoDB and Parakeet container integration + modules. Manual graph-validation scripts are no longer eligible for default pytest collection. +- Robot tags now have a machine-readable allowlist and validation command. Invalid legacy tags were + mapped to the existing taxonomy, and all SDK placeholders inherit the `sdk` execution tag. +- The 17 container-independent configuration tests are part of `make all`; they pass locally. +- The ASR fast lane now produces an **11.3% branch-coverage baseline** while retaining its two known + failing Dockerfile assertions for the parallel test-repair work. + +Coverage thresholds remain intentionally disabled until the known collection and assertion failures +are repaired. The new Python CI jobs will surface those failures rather than hiding them. + +### Phase 0: make the existing signal trustworthy + +1. Add separate CI jobs for root pytest, advanced-backend unit pytest, and ASR unit pytest. +2. Move or exclude manual graph-validation scripts so default pytest cannot collect them. +3. Split Docker/Mongo/GPU tests from unit commands and make missing prerequisites explicit skips. +4. Repair or remove tests for deleted APIs; do not retain compatibility code solely to satisfy stale + tests. +5. Fix Robot dependencies, browser setup/selection, placeholder selection, tags, Make targets, and + documentation. + +### Phase 1: establish coverage without blocking cleanup + +1. Configure line and branch coverage per Python component and publish XML/HTML artifacts. +2. Run the backend process under Coverage.py in Robot containers, save parallel data files, and + combine them with pytest coverage. Report pytest and Robot flags separately as well as combined. +3. Add Vitest and React Testing Library to both Vite UIs; add the Expo-supported Jest/React Native + Testing Library stack to the mobile app. +4. Record the baseline by component. Initially enforce only: green tests, no reduction in total + coverage, and about 80% changed-line coverage. Do not impose an arbitrary 80% repository-wide + target on legacy code. + +### Phase 2: cover failure-prone boundaries + +Prioritize tests around conversation lifecycle, transcription/audio persistence, worker retries and +idempotency, task cleanup, memory-agent edits, authentication/data isolation, and plugin failure +containment. For frontends, prioritize auth/token state, API error states, reconnect behavior, audio +session state, and destructive actions before snapshot-heavy presentation tests. + +### Phase 3: ratchet by subsystem + +Once the suites are green and stable, set modest per-component floors and raise them gradually. +Coverage should remain a discovery tool: mutation testing or focused fault injection on lifecycle +and worker code will reveal weak assertions that line coverage alone cannot detect. + +## Definition of done for the cleanup + +- One documented command runs every fast, no-secret test from a clean checkout. +- Every test has one owning component and one execution class. +- No placeholders are reported as passing coverage. +- All PRs run Python unit/component tests, frontend tests, and the no-secret contract smoke suite. +- Coverage is reported per component and combined for the advanced backend. +- Changed code cannot reduce coverage, while legacy coverage increases through a ratchet rather than + a one-time repository-wide threshold. diff --git a/docs/testing.md b/docs/testing.md new file mode 100644 index 00000000..8a72ad62 --- /dev/null +++ b/docs/testing.md @@ -0,0 +1,77 @@ +# Testing and coverage + +Chronicle separates fast Python tests from the composed Robot Framework suite. Test commands should +be run from the directory shown below so each component uses its own dependency and coverage +configuration. + +## Python tests + +### Root setup and lifecycle tooling + +From the repository root: + +```bash +PYTHONPATH=backends/advanced/src uv run \ + --with-requirements setup-requirements.txt \ + --with pytest \ + --with pytest-cov \ + pytest tests/unit \ + --cov \ + --cov-config=.coveragerc \ + --cov-report=term-missing \ + --cov-report=xml \ + --cov-report=html +``` + +### Advanced backend + +The default fast lane excludes the MongoDB integration module and manual scripts under +`tests/scripts/`: + +```bash +cd backends/advanced +uv sync --locked --group test +uv run --group test pytest \ + --ignore=tests/test_audio_persistence_mongodb.py \ + --cov=advanced_omi_backend \ + --cov-report=term-missing \ + --cov-report=xml \ + --cov-report=html +``` + +Run the MongoDB integration module explicitly when the test database is available: + +```bash +uv run --group test pytest tests/test_audio_persistence_mongodb.py +``` + +### ASR services + +The default fast lane excludes the Parakeet container integration module: + +```bash +cd extras/asr-services +uv sync --locked --group test +uv run --group test pytest \ + --ignore=tests/test_parakeet_service.py \ + --cov=common \ + --cov=providers \ + --cov-report=term-missing \ + --cov-report=xml \ + --cov-report=html +``` + +Run `tests/test_parakeet_service.py` explicitly on a machine prepared for its Docker and model +requirements. + +Each command writes XML and HTML output to the component's `coverage-reports/` directory. Branch +coverage is enabled. No minimum percentage is enforced while the baseline suites are being made +green; CI still fails when a test fails. + +## Integration and end-to-end tests + +The repository-level `tests/` directory owns composed Robot Framework behavior. Its current +commands and environment setup are documented in [`tests/README.md`](../tests/README.md). + +The coverage baseline and test-ownership cleanup plan are recorded in the +[test coverage audit](test-coverage-audit.md). From 3992cd2a93ad40f721c0c63797e89cb19dd74aab Mon Sep 17 00:00:00 2001 From: Ankush Malaker <43288948+AnkushMalaker@users.noreply.github.com> Date: Fri, 17 Jul 2026 14:16:15 +0000 Subject: [PATCH 12/16] chore: ignore generated ASR training outputs --- .gitignore | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.gitignore b/.gitignore index 76d03b45..db21425c 100644 --- a/.gitignore +++ b/.gitignore @@ -38,6 +38,7 @@ plugins/*/config.yml.backup example/* **/node_modules/* **/ollama-data/* +**/model_cache **/model_cache/* .vscode/* **/audio_chunks/* @@ -112,6 +113,8 @@ backends/advanced-backend/data/speaker_model_cache/ *.bin *.sqlite3 *checkpoints +extras/asr-services/providers/*/finetune/out/ +extras/asr-services/providers/*/finetune/split_adapters/ # k8s config From 9f781dfbb327480d55b99cdb504ef4909e9e40ed Mon Sep 17 00:00:00 2001 From: Ankush Malaker <43288948+AnkushMalaker@users.noreply.github.com> Date: Mon, 20 Jul 2026 09:58:37 +0000 Subject: [PATCH 13/16] feat: add dashboard screenshot skill --- skills/screenshots/SKILL.md | 77 ++++++++++++ .../screenshots/scripts/capture_dashboard.py | 111 ++++++++++++++++++ 2 files changed, 188 insertions(+) create mode 100644 skills/screenshots/SKILL.md create mode 100644 skills/screenshots/scripts/capture_dashboard.py diff --git a/skills/screenshots/SKILL.md b/skills/screenshots/SKILL.md new file mode 100644 index 00000000..5d93b1e2 --- /dev/null +++ b/skills/screenshots/SKILL.md @@ -0,0 +1,77 @@ +--- +name: screenshots +description: Captures consistent screenshots of every Chronicle web or Expo screen with Playwright, including authenticated routes and dynamic pages. Use when documenting UI, reviewing visual regressions, or asked to screenshot all pages. +--- + +# Screenshot skill + +Capture screenshots from a running app; do not add Playwright or other browser tooling to a project dependency file for a one-off capture. Use the repository rule: + +```bash +uv run --with playwright python - <<'PY' +from playwright.sync_api import sync_playwright +print("browser automation runs ephemerally") +PY +``` + +## Playwright with uv + +Use `uv run --with playwright` for Python browser automation. If Chromium is not already available, install it ephemerally with `uvx --from playwright playwright install chromium`; do not add Playwright to `package.json` or a Python project environment. Use the exact dashboard host allowed by Chronicle CORS—normally `http://localhost:5173`, not `127.0.0.1`. + +Use a 1440×1000 viewport, a fixed device scale factor, `full_page=True`, and wait for `networkidle` plus any page-specific loading indicator. Save images to `artifacts/screenshots/` unless the task requests a committed asset. For a LAN deployment with a self-signed development certificate, pass `--ignore-https-errors`; never use it for an untrusted public host. + +## Chronicle authentication + +The web dashboard is protected by a JWT. Prefer UI login when testing authentication. For screenshot-only work, an API login is faster: + +1. Read `ADMIN_EMAIL` and `ADMIN_PASSWORD` from `backends/advanced/.env` or the process environment without printing them. +2. `POST {backend}/auth/jwt/login` as form fields `username` and `password`. +3. Before navigating to a protected route, inject the returned `access_token` into local storage. The key is `{base-path-or-root}_token` (`root_token` for the normal `/` base path). +4. Load the route and verify it did not redirect to `/login`. The app verifies the token through `GET /users/me`. + +Never put a token, password, or authenticated screenshot with personal data in a commit. + +## Dashboard workflow + +1. Start the dashboard and backend using the repository’s normal commands. Confirm the actual URL and port. +2. Read `backends/advanced/webui/src/App.tsx` and build the route inventory from its `<Route>` entries. At present the inventory includes `/login`, `/`, `/live-record`, `/chat`, `/conversations`, `/conversations/:id`, `/memory-ledger`, `/users`, `/system`, `/system-errors`, `/settings`, `/upload`, `/queue`, `/plugins`, `/finetuning`, `/network`, `/data-audit`, `/speaker-enrollment`, and `/wakeword-lab`. +3. Use a real conversation ID for `/conversations/:id`; discover one through the UI or an authenticated API request. Never invent an ID and call that page complete. +4. Capture `/login` unauthenticated. Log in through the UI, then capture every protected route with the same viewport, theme, and browser state. +5. Wait for the page heading/content and loading indicators to settle. Use `full_page=True`, and save predictable names such as `01-login.png`, `02-conversations.png`, and `03-conversation-detail-<id>.png`. +6. Record failures separately. A redirect to login, error boundary, blank page, missing fixture, or unresolved loading state is a failed capture—not a screenshot of the requested page. + +Use the generic framework for any set of public or authenticated routes: + +```bash +uv run --with playwright python skills/screenshots/scripts/capture_dashboard.py \ + --base-url http://localhost:5173 \ + --authenticate \ + --route speaker-enrollment=/speaker-enrollment \ + --route system=/system +``` + +For public pages, omit `--authenticate`. Use `--backend-url`, `--env-file`, and `--output-dir` when the deployment differs. It rejects login redirects and blank pages rather than saving misleading screenshots. + +For a same-origin HTTPS proxy such as Chronicle's LAN deployment, point both URLs at the proxy and allow its development certificate: + +```bash +uv run --with playwright python skills/screenshots/scripts/capture_dashboard.py \ + --base-url https://192.168.0.110 \ + --backend-url https://192.168.0.110 \ + --ignore-https-errors \ + --authenticate \ + --route conversation=/conversations/<conversation-id> +``` + +## Expo screens + +The file-based screens are `app/app/index.tsx`, `app/app/diagnostics.tsx`, and `app/app/settings.tsx`. Capture them with Expo’s supported target (`npm run web`, simulator, or device); Playwright can drive the web target, but native-only Bluetooth, audio, share, and permission states must be supplied with fixtures or captured manually on a simulator/device. Include both light and dark themes when visual coverage is the goal. + +## Review checklist + +- Every route in `App.tsx` and every Expo screen has a result. +- Dynamic routes use valid fixture data. +- Authenticated and unauthenticated states are intentional. +- Screenshots use the same viewport, scale, and theme. +- No secrets, tokens, or personal data appear in committed images. +- Output artifacts are stored outside source directories and are ignored unless explicitly requested for commit. diff --git a/skills/screenshots/scripts/capture_dashboard.py b/skills/screenshots/scripts/capture_dashboard.py new file mode 100644 index 00000000..af4ca4b7 --- /dev/null +++ b/skills/screenshots/scripts/capture_dashboard.py @@ -0,0 +1,111 @@ +"""Capture arbitrary Chronicle dashboard routes with optional API authentication.""" + +from __future__ import annotations + +import argparse +import os +import re +from pathlib import Path +from urllib.parse import urlsplit + +from playwright.sync_api import TimeoutError as PlaywrightTimeoutError +from playwright.sync_api import sync_playwright + + +def env_values(path: Path) -> dict[str, str]: + values: dict[str, str] = {} + if not path.exists(): + return values + for line in path.read_text().splitlines(): + match = re.match(r"^([A-Za-z_][A-Za-z0-9_]*)=(.*)$", line.strip()) + if match: + values[match.group(1)] = match.group(2).strip().strip("'\"") + return values + + +def parse_route(value: str) -> tuple[str, str]: + try: + name, route = value.split("=", 1) + except ValueError as error: + raise argparse.ArgumentTypeError("routes must be NAME=/path") from error + if not name or not route.startswith("/"): + raise argparse.ArgumentTypeError("routes must be NAME=/path") + return name, route + + +def default_backend_url(base_url: str) -> str: + parsed = urlsplit(base_url) + return f"{parsed.scheme}://{parsed.hostname}:8000" + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--base-url", default="http://localhost:5173") + parser.add_argument("--backend-url") + parser.add_argument("--env-file", default="backends/advanced/.env") + parser.add_argument("--authenticate", action="store_true") + parser.add_argument( + "--ignore-https-errors", + action="store_true", + help="Accept a self-signed development certificate for the browser and login request", + ) + parser.add_argument("--output-dir", default="artifacts/screenshots") + parser.add_argument("--route", type=parse_route, action="append", required=True) + args = parser.parse_args() + + base_url = args.base_url.rstrip("/") + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + with sync_playwright() as playwright: + browser = playwright.chromium.launch() + page = browser.new_page( + viewport={"width": 1440, "height": 1000}, + device_scale_factor=1, + ignore_https_errors=args.ignore_https_errors, + ) + + if args.authenticate: + env = env_values(Path(args.env_file)) + email = os.getenv("ADMIN_EMAIL", env.get("ADMIN_EMAIL")) + password = os.getenv("ADMIN_PASSWORD", env.get("ADMIN_PASSWORD")) + if not email or not password: + raise RuntimeError("ADMIN_EMAIL and ADMIN_PASSWORD are required") + request = playwright.request.new_context( + base_url=args.backend_url or default_backend_url(base_url), + ignore_https_errors=args.ignore_https_errors, + ) + response = request.post( + "/auth/jwt/login", form={"username": email, "password": password} + ) + if not response.ok: + raise RuntimeError(f"login failed with HTTP {response.status}") + token = response.json()["access_token"] + base_path = urlsplit(base_url).path.strip("/") or "root" + page.add_init_script( + f"localStorage.setItem({base_path!r} + '_token', {token!r})" + ) + request.dispose() + + for name, route in args.route: + page.goto(f"{base_url}{route}", wait_until="domcontentloaded") + try: + page.wait_for_load_state("networkidle", timeout=10_000) + except PlaywrightTimeoutError: + # Dashboards with persistent SSE/HMR connections may never become idle. + pass + page.wait_for_timeout(1_000) + body = " ".join(page.locator("body").inner_text().split()) + if "/login" in page.url and route != "/login": + raise RuntimeError(f"{route} redirected to login") + if not body: + raise RuntimeError(f"{route} rendered a blank page") + path = output_dir / f"{name}.png" + page.screenshot(path=str(path), full_page=True) + print(f"saved {path} ({page.url})") + + browser.close() + + +if __name__ == "__main__": + main() From 0cf4708954ccda7f39c3476efdc85c1604a5d527 Mon Sep 17 00:00:00 2001 From: Ankush Malaker <43288948+AnkushMalaker@users.noreply.github.com> Date: Mon, 20 Jul 2026 10:06:24 +0000 Subject: [PATCH 14/16] chore: ignore experimental ASR finetuning workspaces --- .gitignore | 4 +- .../providers/gemma4/finetune/README.md | 128 ------- .../providers/gemma4/finetune/ab_empties.py | 97 ----- .../providers/gemma4/finetune/align_spike.py | 79 ---- .../gemma4/finetune/bench_coshe_12b.py | 342 ------------------ .../gemma4/finetune/build_windowed_dataset.py | 130 ------- .../providers/gemma4/finetune/data.py | 251 ------------- .../gemma4/finetune/data_interleave.py | 132 ------- .../gemma4/finetune/data_windowed.py | 57 --- .../providers/gemma4/finetune/data_wispr.py | 15 - .../gemma4/finetune/eval_interleave.py | 107 ------ .../gemma4/finetune/eval_test_split.py | 152 -------- .../providers/gemma4/finetune/eval_wer.py | 176 --------- .../gemma4/finetune/eval_windowed_stitch.py | 319 ---------------- .../gemma4/finetune/eval_wispr_stitch.py | 148 -------- .../gemma4/finetune/gemini_transcribe.py | 69 ---- .../gemma4/finetune/gemma4_cache_patch.py | 64 ---- .../providers/gemma4/finetune/infer.py | 169 --------- .../providers/gemma4/finetune/introspect.py | 162 --------- .../gemma4/finetune/poll_until_target.sh | 32 -- .../gemma4/finetune/run_split_pipeline.sh | 49 --- .../gemma4/finetune/smoke_interleave.py | 93 ----- .../providers/gemma4/finetune/split_status.py | 36 -- .../providers/gemma4/finetune/sweep_12b_v2.py | 254 ------------- .../providers/gemma4/finetune/train.py | 203 ----------- .../gemma4/finetune/train_interleave.py | 150 -------- .../gemma4/finetune/train_lora_split.py | 199 ---------- .../gemma4/finetune/train_lora_windowed.py | 160 -------- .../gemma4/finetune/train_until_wer.py | 302 ---------------- .../providers/gemma4/finetune/train_wispr.py | 168 --------- .../providers/gemma4/finetune/verify_fix.py | 53 --- .../providers/gemma4/finetune/vm_exec.py | 90 ----- .../providers/gemma4/finetune/window_coshe.py | 129 ------- .../gemma4/finetune/window_target.py | 25 -- .../providers/nemo/finetune/cut_segments.py | 94 ----- .../providers/nemo/finetune/eval_full.py | 120 ------ .../providers/nemo/finetune/fa_align.py | 114 ------ .../providers/nemo/finetune/make_manifest.py | 167 --------- .../providers/nemo/finetune/run_deepgram.py | 125 ------- .../nemo/finetune/segment_by_alignment.py | 206 ----------- .../providers/nemo/finetune/train_full.py | 154 -------- .../providers/nemo/finetune/train_overfit.py | 204 ----------- .../providers/nemo/finetune/train_split.py | 148 -------- .../providers/qwen3_asr/finetune/README.md | 94 ----- .../qwen3_asr/finetune/coshe_infer.py | 168 --------- .../qwen3_asr/finetune/introspect.py | 62 ---- .../qwen3_asr/finetune/make_manifest.py | 80 ---- .../qwen3_asr/finetune/make_manifest_full.py | 70 ---- .../qwen3_asr/finetune/qwen3_asr_sft.py | 334 ----------------- .../finetune/train_qwen3_until_wer.py | 314 ---------------- .../providers/qwen3_asr/finetune/verify.py | 76 ---- 51 files changed, 2 insertions(+), 7072 deletions(-) delete mode 100644 extras/asr-services/providers/gemma4/finetune/README.md delete mode 100644 extras/asr-services/providers/gemma4/finetune/ab_empties.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/align_spike.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/bench_coshe_12b.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/build_windowed_dataset.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/data.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/data_interleave.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/data_windowed.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/data_wispr.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/eval_interleave.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/eval_test_split.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/eval_wer.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/eval_windowed_stitch.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/eval_wispr_stitch.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/gemini_transcribe.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/gemma4_cache_patch.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/infer.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/introspect.py delete mode 100755 extras/asr-services/providers/gemma4/finetune/poll_until_target.sh delete mode 100644 extras/asr-services/providers/gemma4/finetune/run_split_pipeline.sh delete mode 100644 extras/asr-services/providers/gemma4/finetune/smoke_interleave.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/split_status.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/sweep_12b_v2.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/train.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/train_interleave.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/train_lora_split.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/train_lora_windowed.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/train_until_wer.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/train_wispr.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/verify_fix.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/vm_exec.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/window_coshe.py delete mode 100644 extras/asr-services/providers/gemma4/finetune/window_target.py delete mode 100644 extras/asr-services/providers/nemo/finetune/cut_segments.py delete mode 100644 extras/asr-services/providers/nemo/finetune/eval_full.py delete mode 100644 extras/asr-services/providers/nemo/finetune/fa_align.py delete mode 100644 extras/asr-services/providers/nemo/finetune/make_manifest.py delete mode 100644 extras/asr-services/providers/nemo/finetune/run_deepgram.py delete mode 100644 extras/asr-services/providers/nemo/finetune/segment_by_alignment.py delete mode 100644 extras/asr-services/providers/nemo/finetune/train_full.py delete mode 100644 extras/asr-services/providers/nemo/finetune/train_overfit.py delete mode 100644 extras/asr-services/providers/nemo/finetune/train_split.py delete mode 100644 extras/asr-services/providers/qwen3_asr/finetune/README.md delete mode 100644 extras/asr-services/providers/qwen3_asr/finetune/coshe_infer.py delete mode 100644 extras/asr-services/providers/qwen3_asr/finetune/introspect.py delete mode 100644 extras/asr-services/providers/qwen3_asr/finetune/make_manifest.py delete mode 100644 extras/asr-services/providers/qwen3_asr/finetune/make_manifest_full.py delete mode 100644 extras/asr-services/providers/qwen3_asr/finetune/qwen3_asr_sft.py delete mode 100644 extras/asr-services/providers/qwen3_asr/finetune/train_qwen3_until_wer.py delete mode 100644 extras/asr-services/providers/qwen3_asr/finetune/verify.py diff --git a/.gitignore b/.gitignore index db21425c..9c3b5469 100644 --- a/.gitignore +++ b/.gitignore @@ -113,8 +113,8 @@ backends/advanced-backend/data/speaker_model_cache/ *.bin *.sqlite3 *checkpoints -extras/asr-services/providers/*/finetune/out/ -extras/asr-services/providers/*/finetune/split_adapters/ +# Experimental provider-local model fine-tuning workspaces +extras/asr-services/providers/*/finetune/ # k8s config diff --git a/extras/asr-services/providers/gemma4/finetune/README.md b/extras/asr-services/providers/gemma4/finetune/README.md deleted file mode 100644 index 457b258e..00000000 --- a/extras/asr-services/providers/gemma4/finetune/README.md +++ /dev/null @@ -1,128 +0,0 @@ -# Gemma 4 audio QLoRA fine-tuning - -QLoRA (4-bit NF4) LoRA fine-tuning for `google/gemma-4-E2B-it` / `E4B-it` on an -audio→text ASR task. Built to validate the training+inference loop and the -single-4090 hardware setup by **overfitting** the 7-clip CoSHE-Eval `sample7` -subset (Hinglish, ~55s clips truncated to the model's 30s audio window). - -## Approach (grounded in the official recipes) - -Mirrors HuggingFace's official `fine_tune_gemma3n_on_audio.ipynb` data pipeline, -adapted to the **gemma4** family (this repo's model is `Gemma4ForConditionalGeneration`, -loaded via `AutoModelForMultimodalLM`, transformers >= 5.5): - -- **Chat-message training examples**: `{user: [audio, prompt]}` + `{assistant: target}`, - rendered with `apply_chat_template(tokenize=False)` then `processor(text=, audio=, padding=True)`. -- **Loss only on the transcription**: labels = `input_ids` with the prompt prefix - (user turn incl. expanded audio soft tokens + `<|turn>model\n`), pad, and all - multimodal special tokens masked to `-100`. (The official notebook leaves the - prompt unmasked; we mask it — the cleaner ASR objective the research flagged.) -- **QLoRA**: NF4 + double-quant + bf16 compute. LoRA (r=16, α=32) on the **text - decoder only** (`language_model.*.{q,k,v,o,gate,up,down}_proj`, regex-scoped so - it doesn't hit the identically-named audio-tower linears). Audio/vision towers - and multimodal embedders are **frozen and kept in bf16** (not quantized). - -### Two gemma4-specific gotchas (vs. the gemma3n recipes) - -1. **Don't 4-bit-quantize the audio/vision towers.** Their `Gemma4ClippableLinear` - calls `torch.finfo(weight.dtype)` for gradient clipping, which throws on - uint8-stored 4-bit weights. We pass `llm_int8_skip_modules=["model.audio_tower", - "model.vision_tower", "model.embed_audio", "model.embed_vision", "lm_head"]`. - Note the `model.` prefix is required — transformers' `should_convert_module` - matches skip patterns with `re.match` (anchored at start), so bare `audio_tower` - never matches `model.audio_tower....`. -2. **Don't use `prepare_model_for_kbit_training`.** It upcasts every non-4bit - param to fp32, making the text embeddings fp32 while the (bf16) audio tower - stays bf16 → Gemma4's multimodal `masked_scatter` merge errors on the dtype - clash. We do manual prep (freeze base + `gradient_checkpointing_enable(use_reentrant=False)` - + `enable_input_require_grads`), keeping everything bf16. - -## Environment - -Runs inside the `chronicle-asr-gemma4` image (transformers 5.5, torch cu126) with -`peft` + `bitsandbytes` added. A persistent container is the fastest iteration loop: - -```bash -cd extras/asr-services -docker run -d --name gemma4-train --gpus all \ - -v "$PWD/model_cache:/models" \ - -v "$PWD/providers/gemma4/finetune:/train" \ - -v "$PWD/../ml-experiments/data/coshe-eval:/data/coshe-eval" \ - -e HF_HOME=/models \ - chronicle-asr-gemma4:latest sleep infinity -docker exec gemma4-train uv pip install --python /app/.venv/bin/python \ - "peft>=0.17.0" "bitsandbytes>=0.46.1" -``` - -## Run - -```bash -# train (E2B overfit) -docker exec gemma4-train bash -c 'export PATH=/app/.venv/bin:$PATH; cd /train && \ - python train.py --model google/gemma-4-E2B-it --output_dir /train/out/e2b-overfit \ - --epochs 60 --lr 2e-4' - -# verify: transcribe the training clips with the adapter, compare to ground truth -docker exec gemma4-train bash -c 'export PATH=/app/.venv/bin:$PATH; cd /train && \ - python infer.py --model google/gemma-4-E2B-it --adapter /train/out/e2b-overfit' -``` - -Switch `--model google/gemma-4-E4B-it` for the larger model (also fits 24GB in 4-bit). - -## Files - -- `data.py` — dataset (manifest → 16k mono, ≤30s) + multimodal collator with label masking -- `train.py` — 4-bit load, LoRA, HF `Trainer` overfit loop -- `infer.py` — adapter inference + char-similarity vs. ground truth -- `introspect.py` — dumps processor keys / token ids / module names for a model - -## Results (E2B overfit smoke test) — IT OVERFITS - -The overfit works: with `r=16` attn+MLP LoRA on the text decoder, **5 of 7 clips -reproduce the ground-truth transcript exactly** (char-sim 1.000) and mean char-sim -is **0.869** (the two misses, 0.66 / 0.43, are perfect for the first ~150 chars -then drift in the back half — the part of the 540-char target that runs past the -audible 30s, i.e. correct behaviour). Final teacher-forced loss ~0.006–0.12. - -### The bug that hid it: gemma4 produces wrong logits with the KV cache on - -Until we found this, generation looked broken (char-sim ≈ base 0.07). Root cause: - -**`use_cache=True` makes gemma4's multimodal forward produce wrong audio-conditioned -logits (transformers 5.5.0).** Measured on a fully-overfit adapter, same weights, -same input: - -| forward | loss on a memorized clip | -|---------|--------------------------| -| `use_cache=False` (training path) | **0.10** | -| `use_cache=True` (generation default) | **4.49** (≈ base 4.84) | - -Training runs with `use_cache=False`, so loss collapsed correctly. But -`model.generate()` defaults to the cache, so every generation silently used the -broken path and reverted to near-base output. **Fix: generate with -`use_cache=False`** (`infer.py` sets `model.config.use_cache=False` and passes -`use_cache=False` to `generate`). Slower (no KV cache → full recompute per token), -but correct. The proper fix is upstream in the gemma4 modeling code. - -This was NOT exposure bias and NOT a capacity problem — the earlier v1/v2 "fixes" -(full transcript, lm_head+embed at r64) were all measured through the broken cached -path and so looked like failures/degeneration. - -## What this validated - -- **Overfit demonstrated end-to-end**: load 4-bit → LoRA → train → save → reload → - generate reproduces the training transcripts (mean char-sim 0.869, 5/7 exact). -- 4-bit QLoRA on a single 4090 works for gemma4 audio. -- VRAM: lean r16 ≈ 18–19 GB; r64 + lm_head/embed_tokens ≈ 23.7 GB (fits, barely). - bf16 audio/vision towers stay resident; activations gradient-checkpointed over - ~750 audio + text tokens. -- Speed: ~7 min / 60 epochs (r16). E4B (~5 GB in 4-bit) also fits with room. -- **Debugging note**: `debug_inprocess.py` (train+eval, no save/reload), - `debug_reload_loss.py` (use_cache toggle), `debug_gen_nocache.py` (no-cache - generation) are the scripts that localized the cache bug — keep for reference. - -## For a real fine-tune (next step) - -Need many ≤30s clips with transcripts aligned to the audible window (not full-clip -transcripts truncated by char count). Generation must use `use_cache=False` until -the upstream cache bug is fixed; budget for the slower decode. diff --git a/extras/asr-services/providers/gemma4/finetune/ab_empties.py b/extras/asr-services/providers/gemma4/finetune/ab_empties.py deleted file mode 100644 index 53b69530..00000000 --- a/extras/asr-services/providers/gemma4/finetune/ab_empties.py +++ /dev/null @@ -1,97 +0,0 @@ -"""A/B diagnostic for the 12B empty-output / bail behavior on CoSHE. - -Runs the first N deterministic clips under several decode configs, prints -GT vs HYP, and dumps the RAW (pre-strip) output for empties so we can see -exactly what the model emits (immediate EOS? thinking-only? refusal?). -""" - -import glob -import os -import re -import sys - -sys.path.insert(0, "/home/gemma4ft") -import pyarrow.parquet as pq -import torch -from bench_coshe_12b import SR, VERBATIM_PROMPT, decode_audio, select_names -from transformers import AutoModelForMultimodalLM, AutoProcessor - -MODEL = "google/gemma-4-12B-it" -DATA = "/home/coshe-data/data" -N = 25 - -proc = AutoProcessor.from_pretrained(MODEL) -model = AutoModelForMultimodalLM.from_pretrained( - MODEL, dtype=torch.bfloat16, device_map="auto" -).eval() -print("model loaded\n", flush=True) - -sel, shards = select_names(DATA, 500, "coshe-eval") -clips = [] -for s in shards: - t = pq.read_table(s, columns=["audio_file_name", "transcription", "audio"]) - for nm, tr, au in zip( - t.column(0).to_pylist(), t.column(1).to_pylist(), t.column(2).to_pylist() - ): - if nm in sel and au and au.get("bytes"): - clips.append((nm, decode_audio(au["bytes"]), tr)) - if len(clips) >= N: - break - if len(clips) >= N: - break - - -def strip_ch(raw): - parts = [] - for part in raw.split("<channel|>"): - parts.append(part.split("<|channel>")[0] if "<|channel>" in part else part) - t = "".join(parts) - t = re.sub(r"<\|?[a-z_]+\|?>", " ", t) - return re.sub(r"\s+", " ", t).strip() - - -def gen(audio, enable_thinking, min_new): - msgs = [ - { - "role": "user", - "content": [ - {"type": "text", "text": VERBATIM_PROMPT}, - {"type": "audio", "audio": audio}, - ], - } - ] - inp = proc.apply_chat_template( - msgs, - tokenize=True, - return_dict=True, - return_tensors="pt", - add_generation_prompt=True, - enable_thinking=enable_thinking, - ).to(model.device) - il = inp["input_ids"].shape[-1] - kw = dict(max_new_tokens=1024, do_sample=False, no_repeat_ngram_size=3) - if min_new: - kw["min_new_tokens"] = min_new - with torch.inference_mode(): - out = model.generate(**inp, **kw) - raw = proc.decode(out[0][il:], skip_special_tokens=False) - return raw, strip_ch(raw) - - -CONFIGS = [ - ("A: think=False min0 (current)", False, 0), - ("B: think=False min32 (force gen)", False, 32), - ("C: think=True min0 (real think)", True, 0), -] -for label, et, mn in CONFIGS: - empt = 0 - print(f"\n########## CONFIG {label} ##########", flush=True) - for nm, audio, gt in clips: - raw, hyp = gen(audio, et, mn) - if not hyp: - empt += 1 - print(f"[{nm} {len(audio)/SR:.0f}s] EMPTY RAW={raw[:140]!r}", flush=True) - else: - print(f"[{nm} {len(audio)/SR:.0f}s] GT : {gt[:80]}", flush=True) - print(f"{' '*len(nm)} HYP: {hyp[:80]}", flush=True) - print(f"==> {label}: empty={empt}/{len(clips)}", flush=True) diff --git a/extras/asr-services/providers/gemma4/finetune/align_spike.py b/extras/asr-services/providers/gemma4/finetune/align_spike.py deleted file mode 100644 index f36ebdd0..00000000 --- a/extras/asr-services/providers/gemma4/finetune/align_spike.py +++ /dev/null @@ -1,79 +0,0 @@ -"""Spike: forced-align a handful of CoSHE clips (Devanagari+Latin Hinglish) with the -MMS-based ctc-forced-aligner to check timestamp quality before building the full -windowed dataset. Prints per-word timestamps + sanity stats.""" - -import glob -import io -import sys -import tempfile -import wave - -import numpy as np -import pyarrow.parquet as pq -import soundfile as sf -import torch -from ctc_forced_aligner import ( - generate_emissions, - get_alignments, - get_spans, - load_alignment_model, - load_audio, - postprocess_results, - preprocess_text, -) - -N = int(sys.argv[1]) if len(sys.argv) > 1 else 5 -device = "cuda" if torch.cuda.is_available() else "cpu" -print(f"device={device}", flush=True) - -# grab N clips from first shard's first row group -f = sorted(glob.glob("/mnt/d/datasets/CoSHE-Eval/data/eval-*.parquet"))[0] -t = pq.ParquetFile(f).read_row_group( - 0, columns=["audio_file_name", "transcription", "audio"] -) -clips = [] -for i in range(min(N, t.num_rows)): - clips.append( - ( - t["audio_file_name"][i].as_py(), - t["transcription"][i].as_py(), - t["audio"][i].as_py()["bytes"], - ) - ) - -model, tokenizer = load_alignment_model( - device, dtype=torch.float16 if device == "cuda" else torch.float32 -) - -for name, text, ab in clips: - wav, sr = sf.read(io.BytesIO(ab)) - if wav.ndim > 1: - wav = wav.mean(1) - dur = len(wav) / sr - # write 16k mono temp for load_audio - with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tf: - sf.write(tf.name, wav, sr) - audio_waveform = load_audio(tf.name, model.dtype, model.device) - emissions, stride = generate_emissions(model, audio_waveform, batch_size=1) - tokens_starred, text_starred = preprocess_text(text, romanize=True, language="hin") - segments, scores, blank = get_alignments(emissions, tokens_starred, tokenizer) - spans = get_spans(tokens_starred, segments, blank) - word_ts = postprocess_results(text_starred, spans, stride, scores) - nwords_gt = len(text.split()) - monotonic = all( - word_ts[i]["end"] >= word_ts[i]["start"] for i in range(len(word_ts)) - ) and all( - word_ts[i + 1]["start"] >= word_ts[i]["start"] - 0.05 - for i in range(len(word_ts) - 1) - ) - cov = word_ts[-1]["end"] if word_ts else 0 - print( - f"\n=== {name} dur={dur:.1f}s gt_words={nwords_gt} aligned={len(word_ts)} " - f"last_end={cov:.1f}s coverage={cov/dur*100:.0f}% monotonic={monotonic}", - flush=True, - ) - for w in word_ts[:6] + (["..."] if len(word_ts) > 12 else []) + word_ts[-6:]: - if w == "...": - print(" ...") - continue - print(f" [{w['start']:6.2f}-{w['end']:6.2f}] {w['text']}") diff --git a/extras/asr-services/providers/gemma4/finetune/bench_coshe_12b.py b/extras/asr-services/providers/gemma4/finetune/bench_coshe_12b.py deleted file mode 100644 index 1b66b09f..00000000 --- a/extras/asr-services/providers/gemma4/finetune/bench_coshe_12b.py +++ /dev/null @@ -1,342 +0,0 @@ -"""Benchmark base google/gemma-4-12B-it on CoSHE-Eval (Hinglish). - -Faithful to the gemma4-asr service path that produced the E4B 26.23% number: - * same diarization prompt - * audio AFTER text, apply_chat_template + generate (model default sampling) - * parse_response -> strip "Speaker N:" labels -> join - * long clips (>30s) split into 30s windows, window texts concatenated - -Selection matches mlexp.runners.dataset.run_coshe(limit=500, seed="coshe-eval"): -all audio_file_names in shard order, shuffled with random.Random(seed), first N. - -Output JSONL (score_coshe schema): {audio_file_name, transcription, hyp, asr_seconds, duration_s} - - python bench_coshe_12b.py --data_dir /home/coshe-data/data --out /home/gemma4ft/out/coshe_12b.jsonl \ - --model google/gemma-4-12B-it --limit 500 --seed coshe-eval -""" - -import argparse -import glob -import io -import json -import os -import random -import re -import time - -import numpy as np -import pyarrow.parquet as pq -import soundfile as sf -import torch -from transformers import AutoModelForMultimodalLM, AutoProcessor - -SR = 16000 -# The gemma4-12B "unified" (encoder-free) model has NO hard 30s audio cap — its -# feature extractor encodes linearly at 25 tok/s with no truncation. CoSHE clips -# are all <=60s, so we feed each clip whole (window_seconds defaults large). -# Keep windowing logic for safety on hypothetical >window_seconds clips. - -# Setup A — HF gemma-4-12B-it model-card ASR prompt (NOT the E4B diarization -# prompt). "in its original language" generalizes the card's {LANGUAGE} -# placeholder for CoSHE's code-switched Hindi/English. Plain transcription. -HF_ASR_PROMPT = ( - "Transcribe the following speech segment in its original language.\n\n" - "Follow these specific instructions for formatting the answer:\n" - "* Only output the transcription, with no newlines.\n" - "* When transcribing numbers, write the digits, i.e. write 1.7 and not " - "one point seven, and write 3 instead of three." -) -# Setup B — strict verbatim transcription prompt (guards against the model -# translating Hinglish to English / adding commentary). -VERBATIM_PROMPT = ( - "You are an expert transcriptionist. Transcribe the provided Hinglish audio " - "verbatim exactly as it is spoken. Output ONLY the transcribed text. Do not " - "translate it to English, do not summarize, and do not add any conversational " - "pleasantries or commentary." -) -PROMPTS = {"hf_asr": HF_ASR_PROMPT, "verbatim": VERBATIM_PROMPT} -SPEAKER_LINE_RE = re.compile(r"^(Speaker \d+):\s*(.+)$", re.MULTILINE) - - -def parse_args(): - p = argparse.ArgumentParser() - p.add_argument("--model", default="google/gemma-4-12B-it") - p.add_argument("--data_dir", default="/home/coshe-data/data") - p.add_argument("--out", default="/home/gemma4ft/out/coshe_12b.jsonl") - p.add_argument("--limit", type=int, default=500) - p.add_argument("--seed", default="coshe-eval") - # CoSHE transcripts are long/dense (mean ~882 chars); 512 truncates and - # inflates WER (deletion trap). 1024 fits the ~60s clips without truncation. - p.add_argument("--max_new_tokens", type=int, default=1024) - p.add_argument( - "--window_seconds", - type=float, - default=600.0, - help="split clips longer than this; default large = whole-clip", - ) - # HF card recommends temp=1.0/top_p=0.95/top_k=64 (general). Lower temp for ASR. - p.add_argument("--temperature", type=float, default=1.0) - p.add_argument("--top_p", type=float, default=0.95) - p.add_argument("--top_k", type=int, default=64) - p.add_argument( - "--greedy", action="store_true", help="do_sample=False (overrides temp)" - ) - p.add_argument("--repetition_penalty", type=float, default=1.0) - p.add_argument( - "--no_repeat_ngram_size", - type=int, - default=0, - help="block repeated n-grams (e.g. 3) to stop greedy loops", - ) - p.add_argument( - "--min_new_tokens", - type=int, - default=0, - help="force >=N generated tokens (stops immediate-EOS empties)", - ) - p.add_argument("--prompt", choices=list(PROMPTS), default="hf_asr") - return p.parse_args() - - -def select_names(data_dir, limit, seed): - """Replicate run_coshe deterministic selection.""" - shards = sorted(glob.glob(os.path.join(data_dir, "eval-*.parquet"))) - names = [] - for s in shards: - names.extend( - pq.read_table(s, columns=["audio_file_name"]).column(0).to_pylist() - ) - if limit and limit < len(names): - random.Random(seed).shuffle(names) - return set(names[:limit]), shards - return set(names), shards - - -def load_done(out_path): - done = set() - if os.path.exists(out_path): - with open(out_path) as f: - for line in f: - try: - rec = json.loads(line) - except Exception: - continue - if "error" in rec: - continue - if rec.get("audio_file_name"): - done.add(rec["audio_file_name"]) - return done - - -def decode_audio(raw): - audio, sr = sf.read(io.BytesIO(raw), dtype="float32") - if audio.ndim > 1: - audio = audio.mean(axis=1) - if sr != SR: - import librosa - - audio = librosa.resample(audio, orig_sr=sr, target_sr=SR) - return np.ascontiguousarray(audio, dtype=np.float32) - - -def windows(audio, window_seconds): - """Whole clip if it fits; else non-overlapping windows of window_seconds.""" - w = int(window_seconds * SR) - if len(audio) <= w: - return [audio] - return [audio[i : i + w] for i in range(0, len(audio), w)] - - -class Model: - def __init__( - self, - model_id, - max_new_tokens, - window_seconds=600.0, - prompt=HF_ASR_PROMPT, - temperature=1.0, - top_p=0.95, - top_k=64, - greedy=False, - repetition_penalty=1.0, - no_repeat_ngram_size=0, - min_new_tokens=0, - ): - self.max_new_tokens = max_new_tokens - self.window_seconds = window_seconds - self.prompt = prompt - self.do_sample = not greedy - self.temperature = temperature - self.top_p = top_p - self.top_k = top_k - self.repetition_penalty = repetition_penalty - self.no_repeat_ngram_size = no_repeat_ngram_size - self.min_new_tokens = min_new_tokens - self.processor = AutoProcessor.from_pretrained(model_id) - self.model = AutoModelForMultimodalLM.from_pretrained( - model_id, dtype=torch.bfloat16, device_map="auto" - ) - self.model.eval() - - def _decode(self, outputs, input_len): - # Decode WITH special tokens so the channel markers survive, then strip - # the thinking channel exactly like the chat template's strip_thinking - # macro: drop everything inside <|channel>thought ... <channel|> blocks. - raw = self.processor.decode(outputs[0][input_len:], skip_special_tokens=False) - parts = [] - for part in raw.split("<channel|>"): - parts.append(part.split("<|channel>")[0] if "<|channel>" in part else part) - text = "".join(parts) - # remove any residual angle-bracket control tokens (<|...|>, <...|>, <eos>, etc.) - text = re.sub(r"<\|?[a-z_]+\|?>", " ", text) - text = re.sub(r"\s+", " ", text).strip() - return text - - def transcribe_window(self, audio): - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": self.prompt}, - {"type": "audio", "audio": audio}, - ], - } - ] - inputs = self.processor.apply_chat_template( - messages, - tokenize=True, - return_dict=True, - return_tensors="pt", - add_generation_prompt=True, - enable_thinking=False, - ).to(self.model.device) - input_len = inputs["input_ids"].shape[-1] - gen_kwargs = { - "max_new_tokens": self.max_new_tokens, - "repetition_penalty": self.repetition_penalty, - } - if self.no_repeat_ngram_size: - gen_kwargs["no_repeat_ngram_size"] = self.no_repeat_ngram_size - if self.min_new_tokens: - # the 12B-it often bails to immediate EOS (empty) on long clips; - # forcing a floor unlocks the transcription it was withholding. - gen_kwargs["min_new_tokens"] = self.min_new_tokens - if self.do_sample: - gen_kwargs.update( - do_sample=True, - temperature=self.temperature, - top_p=self.top_p, - top_k=self.top_k, - ) - else: - gen_kwargs.update(do_sample=False) # greedy / deterministic - with torch.inference_mode(): - outputs = self.model.generate(**inputs, **gen_kwargs) - raw_text = self._decode(outputs, input_len) - matches = list(SPEAKER_LINE_RE.finditer(raw_text)) - clean = ( - " ".join(m.group(2).strip() for m in matches if m.group(2).strip()) - if matches - else raw_text.strip() - ) - if clean == "[NO SPEECH]": - clean = "" - return clean - - def transcribe(self, audio): - parts = [self.transcribe_window(w) for w in windows(audio, self.window_seconds)] - return " ".join(p for p in parts if p).strip() - - -def main(): - args = parse_args() - selected, shards = select_names(args.data_dir, args.limit, args.seed) - done = load_done(args.out) - target = selected - done - print( - f"selected={len(selected)} done={len(selected & done)} todo={len(target)}", - flush=True, - ) - if not target: - print("nothing to do", flush=True) - return - - os.makedirs(os.path.dirname(args.out), exist_ok=True) - m = Model( - args.model, - args.max_new_tokens, - args.window_seconds, - prompt=PROMPTS[args.prompt], - temperature=args.temperature, - top_p=args.top_p, - top_k=args.top_k, - greedy=args.greedy, - repetition_penalty=args.repetition_penalty, - no_repeat_ngram_size=args.no_repeat_ngram_size, - min_new_tokens=args.min_new_tokens, - ) - print( - f"model loaded: {args.model} (prompt={args.prompt} thinking=False " - f"do_sample={not args.greedy} temp={args.temperature} " - f"top_p={args.top_p} top_k={args.top_k} rep_pen={args.repetition_penalty} " - f"no_repeat_ngram={args.no_repeat_ngram_size} min_new={args.min_new_tokens} " - f"max_new={args.max_new_tokens} window={args.window_seconds})", - flush=True, - ) - - n_ok = n_err = 0 - t0 = time.time() - with open(args.out, "a") as out_f: - for shard in shards: - if not target: - break - pf = pq.ParquetFile(shard) - for batch in pf.iter_batches(batch_size=8): - for row in batch.to_pylist(): - name = row["audio_file_name"] - if name not in target: - continue - target.discard(name) - try: - audio = decode_audio(row["audio"]["bytes"]) - dur = len(audio) / SR - ts = time.time() - hyp = m.transcribe(audio) - dt = time.time() - ts - rec = { - "audio_file_name": name, - "transcription": row["transcription"], - "hyp": hyp, - "asr_seconds": round(dt, 2), - "duration_s": round(dur, 2), - } - out_f.write(json.dumps(rec, ensure_ascii=False) + "\n") - out_f.flush() - n_ok += 1 - done_n = n_ok + n_err - print( - f" [{done_n}/{len(selected)}] {name} dur={dur:.0f}s " - f"gen={dt:.1f}s hyp[:80]={hyp[:80]!r}", - flush=True, - ) - except Exception as e: - n_err += 1 - out_f.write( - json.dumps( - {"audio_file_name": name, "error": str(e)}, - ensure_ascii=False, - ) - + "\n" - ) - out_f.flush() - print(f" ERR {name}: {e}", flush=True) - if not target: - break - print( - f"DONE ok={n_ok} err={n_err} elapsed={time.time()-t0:.0f}s out={args.out}", - flush=True, - ) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/gemma4/finetune/build_windowed_dataset.py b/extras/asr-services/providers/gemma4/finetune/build_windowed_dataset.py deleted file mode 100644 index 78d25a82..00000000 --- a/extras/asr-services/providers/gemma4/finetune/build_windowed_dataset.py +++ /dev/null @@ -1,130 +0,0 @@ -"""Materialize the windowed CoSHE training set: slice each train/val window's audio out -of the CoSHE parquet and write 16k-mono WAVs + a manifest, honoring the clip-level -20/10/70 split (windows inherit their clip's split, so no clip leaks across splits). - -Consumes windows.jsonl from window_coshe.py (honest <=30s GT-accurate windows). Test -clips are NOT sliced here — they are kept whole for windowed-stitch evaluation; we only -emit their clip list + per-window text. - -Out (under --out_dir): - audio/<clip>__w<idx>.wav 16k mono int16, one per train/val window - manifest.jsonl {name, clip, win_idx, split, audio, dur, text} - windowed_split.json {train:[names], val:[names], test_clips:[clips]} -""" - -import argparse -import glob -import io -import json -import wave -from collections import defaultdict -from pathlib import Path - -import numpy as np -import pyarrow.parquet as pq -import soundfile as sf - - -def safe(clip, idx): - return f"{clip.replace('.wav','')}__w{idx}" - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("--windows", default="/home/coshe_windowed/windows_v2.jsonl") - ap.add_argument("--split", default="/home/coshe_windowed/split_20_10_70.json") - ap.add_argument("--exclude", default="") - ap.add_argument("--parquet_glob", default="/home/coshe/data/eval-*.parquet") - ap.add_argument("--out_dir", default="/home/coshe_windowed") - ap.add_argument( - "--materialize", - default="train,val", - help="which splits get sliced WAVs (test kept whole for stitch eval)", - ) - args = ap.parse_args() - - drop = {x for x in args.exclude.split(",") if x} - mat = set(args.materialize.split(",")) - split = json.load(open(args.split)) - clip_split = {} - for s in ("train", "val", "test"): - for c in split[s]: - if c not in drop: - clip_split[c] = s - - wins_by_clip = defaultdict(list) - for line in open(args.windows): - r = json.loads(line) - if r["audio_file_name"] in drop: - continue - wins_by_clip[r["audio_file_name"]].append(r) - - out_dir = Path(args.out_dir) - audio_dir = out_dir / "audio" - audio_dir.mkdir(parents=True, exist_ok=True) - man = open(out_dir / "manifest.jsonl", "w") - split_out = { - "train": [], - "val": [], - "test_clips": sorted(c for c, s in clip_split.items() if s == "test"), - } - - want_clips = {c for c, s in clip_split.items() if s in mat} - written = 0 - for f in sorted(glob.glob(args.parquet_glob)): - t = pq.ParquetFile(f).read(columns=["audio_file_name", "audio"]) - names = t["audio_file_name"].to_pylist() - for ci, name in enumerate(names): - if name not in want_clips: - continue - s = clip_split[name] - wav, sr = sf.read(io.BytesIO(t["audio"][ci].as_py()["bytes"])) - if wav.ndim > 1: - wav = wav.mean(1) - if sr != 16000: - import librosa - - wav = librosa.resample( - wav.astype(np.float32), orig_sr=sr, target_sr=16000 - ) - sr = 16000 - pcm = (np.clip(wav, -1, 1) * 32767).astype(np.int16) - for w in wins_by_clip[name]: - nm = safe(name, w["win_idx"]) - a0, a1 = int(w["start"] * sr), int(w["end"] * sr) - seg = pcm[a0:a1] - path = audio_dir / f"{nm}.wav" - with wave.open(str(path), "wb") as wf: - wf.setnchannels(1) - wf.setsampwidth(2) - wf.setframerate(sr) - wf.writeframes(seg.tobytes()) - man.write( - json.dumps( - { - "name": nm, - "clip": name, - "win_idx": w["win_idx"], - "split": s, - "audio": str(path), - "dur": round(len(seg) / sr, 3), - "text": w["text"], - }, - ensure_ascii=False, - ) - + "\n" - ) - split_out[s].append(nm) - written += 1 - print(f" {Path(f).name}: written so far={written}", flush=True) - man.close() - json.dump(split_out, open(out_dir / "windowed_split.json", "w"), ensure_ascii=False) - print( - f"DONE windows materialized={written} " - f"train={len(split_out['train'])} val={len(split_out['val'])} " - f"test_clips={len(split_out['test_clips'])}" - ) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/gemma4/finetune/data.py b/extras/asr-services/providers/gemma4/finetune/data.py deleted file mode 100644 index 4af22542..00000000 --- a/extras/asr-services/providers/gemma4/finetune/data.py +++ /dev/null @@ -1,251 +0,0 @@ -"""Dataset + collator for Gemma 4 audio QLoRA fine-tuning on CoSHE-Eval sample7. - -The dataset yields raw {audio: float32 16k mono, target: transcript} items. -The collator does the multimodal preprocessing per batch (the HF gemma recipe -pattern): apply_chat_template(tokenize=False) -> processor(text, audio) -> -build labels. - -Label masking (loss only on the transcription): - * everything in the prompt prefix (user turn incl. expanded audio soft - tokens + the `<|turn>model\n` generation marker) -> -100 - * pad tokens and any multimodal special tokens -> -100 -Only the assistant transcription (+ its closing turn token) is supervised. -""" - -import json -import os -import wave -from pathlib import Path - -import numpy as np -import torch - -DEFAULT_PROMPT = ( - "Transcribe the following speech segment in its original language and identify different speakers. " - "Follow these specific instructions for formatting the answer:\n" - "* Label each speaker as Speaker 1, Speaker 2, etc.\n" - "* Format each turn as 'Speaker N: <text>' on its own line.\n" - "* Start a new line when the speaker changes.\n" - "* When transcribing numbers, write the digits, i.e. write 1.7 and not " - "one point seven, and write 3 instead of three.\n" - "* If the audio is silence or contains no speech, respond with exactly: [NO SPEECH]" -) - - -def _load_wav_16k_mono(path: str, max_seconds: float) -> np.ndarray: - with wave.open(path, "rb") as wf: - sr = wf.getframerate() - nch = wf.getnchannels() - raw = wf.readframes(wf.getnframes()) - audio = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0 - if nch > 1: - audio = audio.reshape(-1, nch).mean(axis=1) - if sr != 16000: - import librosa - - audio = librosa.resample(audio, orig_sr=sr, target_sr=16000) - if max_seconds: - audio = audio[: int(16000 * max_seconds)] - return np.ascontiguousarray(audio, dtype=np.float32) - - -class CosheSample7Dataset(torch.utils.data.Dataset): - """Reads coshe-eval/sample7/manifest.json -> {audio, target} items.""" - - def __init__( - self, data_dir: str, max_seconds: float = 30.0, target_max_chars: int = 0 - ): - d = Path(data_dir) - manifest = json.loads((d / "manifest.json").read_text()) - self.items = [] - for row in manifest: - wav = d / "audio" / row["audio_file_name"] - if not wav.exists(): - continue - target = row["transcription"].strip() - # When audio is truncated to the model's 30s window, the back half of - # the full-clip transcript is ungrounded. Truncating the target to ~the - # portion the model can actually hear makes the overfit reproducible at - # generation (CoSHE clips run ~18 chars/sec of speech). - if target_max_chars: - target = target[:target_max_chars] - self.items.append( - { - "audio": _load_wav_16k_mono(str(wav), max_seconds), - "target": target, - "name": row["audio_file_name"], - } - ) - if not self.items: - raise RuntimeError(f"No samples found under {data_dir}") - - def __len__(self): - return len(self.items) - - def __getitem__(self, idx): - return self.items[idx] - - -def _decode_audio_bytes_16k_mono(raw: bytes, max_seconds: float) -> np.ndarray: - import io - - import soundfile as sf - - audio, sr = sf.read(io.BytesIO(raw), dtype="float32") - if audio.ndim > 1: - audio = audio.mean(axis=1) - if sr != 16000: - import librosa - - audio = librosa.resample(audio, orig_sr=sr, target_sr=16000) - if max_seconds: - audio = audio[: int(16000 * max_seconds)] - return np.ascontiguousarray(audio, dtype=np.float32) - - -class CosheParquetDataset(torch.utils.data.Dataset): - """Full CoSHE-Eval: reads parquet shards (audio as embedded wav bytes). - - Pre-decodes everything into RAM once (≈3.8GB for the full 1985 clips at 30s) - so multi-epoch overfitting doesn't re-decode each epoch. `limit` caps samples - for quick recipe checks; `max_seconds` truncates to the model's 30s window. - """ - - def __init__( - self, - parquet_glob: str, - max_seconds: float = 30.0, - target_max_chars: int = 0, - limit: int = 0, - cache_path: str = "", - ): - import glob - import pickle - - import pyarrow.parquet as pq - - # Decoding+resampling all 1985 clips off the HDD is ~40 min; cache the - # decoded (audio, target) items so repeated training runs load instantly. - if cache_path and os.path.exists(cache_path): - with open(cache_path, "rb") as f: - self.items = pickle.load(f) - if target_max_chars: - for it in self.items: - it["target"] = it["target"][:target_max_chars] - if limit: - self.items = self.items[:limit] - return - - paths = sorted(glob.glob(parquet_glob)) - if not paths: - raise RuntimeError(f"No parquet shards match {parquet_glob}") - self.items = [] - for path in paths: - t = pq.read_table( - path, columns=["audio_file_name", "transcription", "audio"] - ) - names = t.column("audio_file_name").to_pylist() - trans = t.column("transcription").to_pylist() - audio_col = t.column("audio").to_pylist() # list of {bytes, path} - for name, tr, au in zip(names, trans, audio_col): - if au is None or au.get("bytes") is None or not tr: - continue - self.items.append( - { - "audio": _decode_audio_bytes_16k_mono(au["bytes"], max_seconds), - "target": tr.strip(), # full target; truncation applied below - "name": name, - } - ) - if limit and len(self.items) >= limit: - break - if limit and len(self.items) >= limit: - break - if not self.items: - raise RuntimeError(f"No samples loaded from {parquet_glob}") - if cache_path: - with open(cache_path, "wb") as f: - pickle.dump(self.items, f) - if target_max_chars: - for it in self.items: - it["target"] = it["target"][:target_max_chars] - - def __len__(self): - return len(self.items) - - def __getitem__(self, idx): - return self.items[idx] - - -class Gemma4AudioCollator: - def __init__( - self, processor, prompt: str = DEFAULT_PROMPT, mask_prompt: bool = True - ): - self.processor = processor - self.prompt = prompt - self.mask_prompt = mask_prompt - tok = processor.tokenizer - tok.padding_side = "right" - self.pad_id = tok.pad_token_id - # multimodal special tokens to exclude from the loss - self.special_ids = [ - t - for t in [ - getattr(tok, "audio_token_id", None), - getattr(tok, "image_token_id", None), - getattr(tok, "boi_token_id", None), - getattr(tok, "eoi_token_id", None), - getattr(tok, "boa_token_id", None), - getattr(tok, "eoa_token_id", None), - ] - if t is not None - ] - - def _user_turn(self, audio): - return { - "role": "user", - "content": [ - {"type": "text", "text": self.prompt}, - {"type": "audio", "audio": audio}, - ], - } - - def __call__(self, features): - audios = [f["audio"] for f in features] - texts_full = [] - for f in features: - msgs = [ - self._user_turn(f["audio"]), - { - "role": "assistant", - "content": [{"type": "text", "text": f["target"]}], - }, - ] - texts_full.append( - self.processor.apply_chat_template( - msgs, tokenize=False, add_generation_prompt=False - ) - ) - batch = self.processor( - text=texts_full, audio=audios, return_tensors="pt", padding=True - ) - - labels = batch["input_ids"].clone() - labels[labels == self.pad_id] = -100 - for tid in self.special_ids: - labels[labels == tid] = -100 - - if self.mask_prompt: - for i, f in enumerate(features): - ptext = self.processor.apply_chat_template( - [self._user_turn(f["audio"])], - tokenize=False, - add_generation_prompt=True, - ) - plen = self.processor( - text=[ptext], audio=[f["audio"]], return_tensors="pt" - )["input_ids"].shape[1] - labels[i, :plen] = -100 - - batch["labels"] = labels - return batch diff --git a/extras/asr-services/providers/gemma4/finetune/data_interleave.py b/extras/asr-services/providers/gemma4/finetune/data_interleave.py deleted file mode 100644 index 88a50cf3..00000000 --- a/extras/asr-services/providers/gemma4/finetune/data_interleave.py +++ /dev/null @@ -1,132 +0,0 @@ -"""Interleaved multi-chunk dataset + collator for the audio->formatted(pasted_text) task. - -Each example is ONE clip: its audio is fed as N <=28s chunks (multiple audio blocks in a -single user turn), the target is the WHOLE pasted_text, and the prompt is conditioned on the -app the dictation was going into (Wispr formats per-app). No text windowing — formatted text -isn't word-alignable, so the model attends to all chunks and emits the full formatted clip. - -Loss is on the target only (prompt prefix incl. all audio soft-tokens + pad + mm-special -tokens masked to -100), matching the single-audio Gemma4AudioCollator. -""" - -import json -import wave - -import numpy as np -import torch - - -def _load_wav_f32(path: str) -> np.ndarray: - with wave.open(path, "rb") as wf: - raw = wf.readframes(wf.getnframes()) - return np.ascontiguousarray( - np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0 - ) - - -def build_prompt(app: str) -> str: - """App-aware formatting instruction. Wispr formats dictation differently per target app.""" - where = f"the {app} app" if app else "the target app" - return ( - f"You are formatting voice dictation that will be typed into {where}. " - "The audio is provided as one or more consecutive segments of a single dictation. " - "Transcribe all of it and format the result the way it should appear when typed into " - f"{where} (punctuation, capitalization, digits as digits, app-appropriate style). " - "Output only the final formatted text, with no commentary." - ) - - -class ChunkInterleaveDataset(torch.utils.data.Dataset): - """One item per clip: {chunks:[float32 16k], target:str, app:str, clip:str}.""" - - def __init__(self, manifest_path: str, split: str): - self.items = [] - for line in open(manifest_path): - r = json.loads(line) - if r["split"] != split: - continue - self.items.append( - { - "chunks": [_load_wav_f32(p) for p in r["chunks"]], - "target": r["target"].strip(), - "app": r.get("app", ""), - "clip": r["clip"], - } - ) - if not self.items: - raise RuntimeError(f"no '{split}' rows in {manifest_path}") - - def __len__(self): - return len(self.items) - - def __getitem__(self, idx): - return self.items[idx] - - -class InterleaveCollator: - """Batches variable-#chunk examples. Uses the batch processor with audio as a per-example - list of arrays (nested); the chat template emits one audio placeholder per chunk.""" - - def __init__(self, processor, mask_prompt: bool = True): - self.processor = processor - self.mask_prompt = mask_prompt - tok = processor.tokenizer - tok.padding_side = "right" - self.pad_id = tok.pad_token_id - self.special_ids = [ - t - for t in [ - getattr(tok, "audio_token_id", None), - getattr(tok, "image_token_id", None), - getattr(tok, "boi_token_id", None), - getattr(tok, "eoi_token_id", None), - getattr(tok, "boa_token_id", None), - getattr(tok, "eoa_token_id", None), - ] - if t is not None - ] - - def _user_turn(self, chunks, app): - content = [{"type": "text", "text": build_prompt(app)}] - content += [{"type": "audio", "audio": a} for a in chunks] - return {"role": "user", "content": content} - - def __call__(self, features): - # processor wants a FLAT list of all audios across the batch, matched to the audio - # placeholders (one per chunk) in text order. - audios = [a for f in features for a in f["chunks"]] - texts_full = [] - for f in features: - msgs = [ - self._user_turn(f["chunks"], f["app"]), - { - "role": "assistant", - "content": [{"type": "text", "text": f["target"]}], - }, - ] - texts_full.append( - self.processor.apply_chat_template( - msgs, tokenize=False, add_generation_prompt=False - ) - ) - batch = self.processor( - text=texts_full, audio=audios, return_tensors="pt", padding=True - ) - - labels = batch["input_ids"].clone() - labels[labels == self.pad_id] = -100 - for tid in self.special_ids: - labels[labels == tid] = -100 - if self.mask_prompt: - for i, f in enumerate(features): - ptext = self.processor.apply_chat_template( - [self._user_turn(f["chunks"], f["app"])], - tokenize=False, - add_generation_prompt=True, - ) - plen = self.processor( - text=[ptext], audio=f["chunks"], return_tensors="pt" - )["input_ids"].shape[1] - labels[i, :plen] = -100 - batch["labels"] = labels - return batch diff --git a/extras/asr-services/providers/gemma4/finetune/data_windowed.py b/extras/asr-services/providers/gemma4/finetune/data_windowed.py deleted file mode 100644 index fc1ad496..00000000 --- a/extras/asr-services/providers/gemma4/finetune/data_windowed.py +++ /dev/null @@ -1,57 +0,0 @@ -"""Dataset over the windowed CoSHE manifest (build_windowed_dataset.py output). - -Each row is one honest <=30s window: a sliced 16k-mono WAV + its EXACT ground-truth text -(forced-alignment-derived, no first-30s truncation hack). Yields the same -{audio, target, name} shape the Gemma4AudioCollator expects, so training reuses the -existing collator unchanged. - -Plain-transcript targets: CoSHE has no speaker labels, so PLAIN_PROMPT asks only for a -verbatim transcript (digits as digits). Train and eval MUST use the same prompt. -""" - -import json -import wave - -import numpy as np -import torch - -PLAIN_PROMPT = ( - "Transcribe the following speech segment verbatim in its original language " - "(Hindi-English code-mixed). Write digits as digits (e.g. 3, not three). " - "Output only the transcript, no speaker labels." -) - - -def _load_wav_f32(path: str) -> np.ndarray: - with wave.open(path, "rb") as wf: - raw = wf.readframes(wf.getnframes()) - return np.ascontiguousarray( - np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0 - ) - - -class WindowedManifestDataset(torch.utils.data.Dataset): - """Loads the windows for one split (train/val). Audio is decoded eagerly into RAM - (train+val is ~1GB) so multi-epoch training doesn't re-read WAVs.""" - - def __init__(self, manifest_path: str, split: str): - self.items = [] - for line in open(manifest_path): - r = json.loads(line) - if r["split"] != split: - continue - self.items.append( - { - "audio": _load_wav_f32(r["audio"]), - "target": r["text"].strip(), - "name": r["name"], - } - ) - if not self.items: - raise RuntimeError(f"no '{split}' rows in {manifest_path}") - - def __len__(self): - return len(self.items) - - def __getitem__(self, idx): - return self.items[idx] diff --git a/extras/asr-services/providers/gemma4/finetune/data_wispr.py b/extras/asr-services/providers/gemma4/finetune/data_wispr.py deleted file mode 100644 index cfe1c395..00000000 --- a/extras/asr-services/providers/gemma4/finetune/data_wispr.py +++ /dev/null @@ -1,15 +0,0 @@ -"""Wispr dictation FT: shared English verbatim prompt + windowed dataset. - -The target is asr_text (Wispr's raw verbatim ASR, treated as ground truth) cut into honest -<=28s windows. English single-speaker dictation, so the prompt asks for a plain verbatim -English transcript (no speaker labels, no Hinglish wording). Train, eval, and the base -baseline MUST all use WISPR_PROMPT. -""" - -from data_windowed import WindowedManifestDataset # noqa: F401 (re-exported) - -WISPR_PROMPT = ( - "Transcribe the following speech segment verbatim in English. " - "Write digits as digits (e.g. 3, not three). " - "Only output the transcription text itself, with no commentary or explanation." -) diff --git a/extras/asr-services/providers/gemma4/finetune/eval_interleave.py b/extras/asr-services/providers/gemma4/finetune/eval_interleave.py deleted file mode 100644 index 17e23030..00000000 --- a/extras/asr-services/providers/gemma4/finetune/eval_interleave.py +++ /dev/null @@ -1,107 +0,0 @@ -"""Evaluate base or FT Gemma4-E2B on held-out clips for the interleave audio->pasted_text task. -Per test clip: feed all its <=28s chunks in one app-aware prompt (onestep, bf16, sdpa), -generate the formatted text in one shot, score against pasted_text (the target) with jiwer. -Same prompt/decode for base and adapter -> controlled delta. -""" - -import argparse -import json -import re -import statistics - -import jiwer -import torch -from data_interleave import _load_wav_f32, build_prompt -from peft import PeftModel -from transformers import AutoModelForMultimodalLM, AutoProcessor - -_B = re.compile(r"\[[^\]]*\]") -_A = re.compile(r"['’`]") -_P = re.compile(r"[^\w\s]") -_W = re.compile(r"\s+") - - -def norm(t): - t = _B.sub(" ", t).lower() - t = _A.sub("", t) - t = _P.sub(" ", t) - return _W.sub(" ", t).strip() - - -@torch.inference_mode() -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("--model", default="google/gemma-4-E2B-it") - ap.add_argument("--adapter", default="") - ap.add_argument("--manifest", default="/home/wispr_interleave/manifest.jsonl") - ap.add_argument("--out", required=True) - ap.add_argument("--max_new_tokens", type=int, default=512) - args = ap.parse_args() - - test = [ - json.loads(l) for l in open(args.manifest) if json.loads(l)["split"] == "test" - ] - proc = AutoProcessor.from_pretrained(args.model) - proc.tokenizer.padding_side = "left" - model = AutoModelForMultimodalLM.from_pretrained( - args.model, dtype=torch.bfloat16, device_map="auto" - ) # default attn (sdpa), bf16 - if args.adapter: - model = PeftModel.from_pretrained(model, args.adapter) - print(f"adapter: {args.adapter}", flush=True) - model.eval() - model.config.use_cache = True - - rows = [] - fout = open(args.out, "w") - for ci, r in enumerate(test): - chunks = [_load_wav_f32(p) for p in r["chunks"]] - content = [{"type": "text", "text": build_prompt(r.get("app", ""))}] - content += [{"type": "audio", "audio": a} for a in chunks] - inp = proc.apply_chat_template( - [{"role": "user", "content": content}], - tokenize=True, - return_dict=True, - return_tensors="pt", - add_generation_prompt=True, - enable_thinking=False, - ).to(model.device) - g = model.generate( - **inp, max_new_tokens=args.max_new_tokens, do_sample=False, use_cache=True - ) - hyp = proc.decode( - g[0][inp["input_ids"].shape[-1] :], skip_special_tokens=True - ).strip() - ref = r["target"] - rn, hn = norm(ref), norm(hyp) - wer = jiwer.wer(rn, hn) if hn else 1.0 - rows.append( - { - "clip": r["clip"], - "app": r.get("app", ""), - "wer": round(wer, 4), - "ref_words": len(rn.split()), - "n_chunks": len(chunks), - "transcription": ref, - "hyp": hyp, - } - ) - fout.write(json.dumps(rows[-1], ensure_ascii=False) + "\n") - fout.flush() - print(f"[{ci+1}/{len(test)}] {r['clip']} wer={wer:.3f}", flush=True) - fout.close() - - refs_n = [norm(r["transcription"]) for r in rows] - hyps_n = [norm(r["hyp"]) for r in rows] - summary = { - "adapter": args.adapter or "base", - "clips": len(rows), - "corpus_wer": round(jiwer.wer(refs_n, hyps_n), 4), - "median_wer": round(statistics.median(r["wer"] for r in rows), 4), - } - print("SUMMARY " + json.dumps(summary), flush=True) - json.dump(summary, open(args.out.replace(".jsonl", "_summary.json"), "w"), indent=2) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/gemma4/finetune/eval_test_split.py b/extras/asr-services/providers/gemma4/finetune/eval_test_split.py deleted file mode 100644 index 492b16ff..00000000 --- a/extras/asr-services/providers/gemma4/finetune/eval_test_split.py +++ /dev/null @@ -1,152 +0,0 @@ -"""Transcribe the held-out CoSHE test split with the base model and/or a LoRA adapter, -emitting mlexp-schema JSONL for scoring. - -Using the SAME harness for base and adapters means the base-vs-FT WER delta isolates the -adapter's effect (no service / audio-handling confound). Audio is the cached 30s window -(same as training); the same window + decode settings are applied to every model, so the -delta is controlled even though the absolute WER carries the 30s-truncation tail equally. -""" - -import argparse -import json -import os -import random -import time - -import torch -from data import DEFAULT_PROMPT, CosheParquetDataset -from peft import PeftModel -from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig -from window_target import apply_window_truncation - - -def load_done(path): - done = set() - if os.path.exists(path): - for line in open(path): - try: - done.add(json.loads(line)["audio_file_name"]) - except Exception: - pass - return done - - -@torch.inference_mode() -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("--model", default="google/gemma-4-E4B-it") - ap.add_argument("--adapter", default="") - ap.add_argument("--split_file", default="/home/gemma4ft/split_20_10_70.json") - ap.add_argument("--cache_path", default="/home/gemma4ft/out/coshe_full_cache.pkl") - ap.add_argument("--parquet_glob", default="/home/coshe-data/data/eval-*.parquet") - ap.add_argument("--out", required=True) - ap.add_argument("--batch_size", type=int, default=8) - ap.add_argument("--max_new_tokens", type=int, default=1024) - ap.add_argument("--limit", type=int, default=0, help="cap #clips (smoke test)") - ap.add_argument( - "--window_seconds", - type=float, - default=30.0, - help="proportionally truncate ref to this audio window (0=full)", - ) - ap.add_argument("--durations", default="/home/gemma4ft/durations.json") - args = ap.parse_args() - - proc = AutoProcessor.from_pretrained(args.model) - bnb = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_use_double_quant=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - llm_int8_skip_modules=[ - "model.audio_tower", - "model.vision_tower", - "model.embed_audio", - "model.embed_vision", - "lm_head", - ], - ) - model = AutoModelForMultimodalLM.from_pretrained( - args.model, - quantization_config=bnb, - dtype=torch.bfloat16, - device_map="auto", - attn_implementation="eager", - ) - if args.adapter: - model = PeftModel.from_pretrained(model, args.adapter) - print(f"loaded adapter: {args.adapter}", flush=True) - model.eval() - model.config.use_cache = True - proc.tokenizer.padding_side = "left" - - split = json.load(open(args.split_file)) - test_names = split["test"] - ds = CosheParquetDataset( - args.parquet_glob, max_seconds=30.0, cache_path=args.cache_path - ) - apply_window_truncation(list(ds), args.durations, args.window_seconds) - by = {it["name"]: it for it in ds} - items = [by[n] for n in test_names if n in by] - if args.limit and args.limit < len(items): - # seeded random subset so every model is evaluated on the SAME clips - items = sorted( - random.Random(0).sample(items, args.limit), key=lambda it: it["name"] - ) - done = load_done(args.out) - items = [it for it in items if it["name"] not in done] - print( - f"test split={len(test_names)} to_do={len(items)} already_done={len(done)}", - flush=True, - ) - - fout = open(args.out, "a") - t0 = time.time() - for s in range(0, len(items), args.batch_size): - batch = items[s : s + args.batch_size] - texts, audios = [], [] - for it in batch: - msgs = [ - { - "role": "user", - "content": [ - {"type": "text", "text": DEFAULT_PROMPT}, - {"type": "audio", "audio": it["audio"]}, - ], - } - ] - texts.append( - proc.apply_chat_template( - msgs, tokenize=False, add_generation_prompt=True - ) - ) - audios.append(it["audio"]) - inp = proc(text=texts, audio=audios, return_tensors="pt", padding=True).to( - model.device - ) - ts = time.time() - out = model.generate( - **inp, max_new_tokens=args.max_new_tokens, do_sample=False, use_cache=True - ) - dt = time.time() - ts - for i, it in enumerate(batch): - hyp = proc.decode( - out[i][inp["input_ids"].shape[-1] :], skip_special_tokens=True - ).strip() - rec = { - "audio_file_name": it["name"], - "transcription": it["target"], - "hyp": hyp, - "asr_seconds": round(dt / len(batch), 3), - "duration_s": round(len(it["audio"]) / 16000.0, 2), - "raw": hyp, - } - fout.write(json.dumps(rec, ensure_ascii=False) + "\n") - fout.flush() - print(f"[{s + len(batch)}/{len(items)}] {time.time() - t0:.0f}s", flush=True) - fout.close() - print(f"DONE out={args.out}", flush=True) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/gemma4/finetune/eval_wer.py b/extras/asr-services/providers/gemma4/finetune/eval_wer.py deleted file mode 100644 index 5558c288..00000000 --- a/extras/asr-services/providers/gemma4/finetune/eval_wer.py +++ /dev/null @@ -1,176 +0,0 @@ -"""Batched WER eval for a Gemma 4 audio LoRA adapter over the full CoSHE dataset. - -No-cache generation is slow (gemma4 cache bug), so we BATCH clips through -generate() with left padding to make evaluating ~2000 clips feasible. Reports -corpus WER (the goal metric) + mean per-clip WER. - - python eval_wer.py --adapter /train/out/full --parquet_glob '/coshe-full/data/eval-*.parquet' \ - --cache_path /train/out/coshe_full_cache.pkl --batch_size 16 [--limit N] -""" - -import argparse - -import jiwer -import torch -from data import DEFAULT_PROMPT, CosheParquetDataset, CosheSample7Dataset -from peft import PeftModel -from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig - -_WER_NORM = jiwer.Compose( - [ - jiwer.ToLowerCase(), - jiwer.RemovePunctuation(), - jiwer.RemoveMultipleSpaces(), - jiwer.Strip(), - jiwer.ReduceToListOfListOfWords(), - ] -) - - -def parse_args(): - p = argparse.ArgumentParser() - p.add_argument("--model", default="google/gemma-4-E2B-it") - p.add_argument("--adapter", default="") - p.add_argument("--parquet_glob", default="") - p.add_argument("--data_dir", default="/data/coshe-eval/sample7") - p.add_argument("--cache_path", default="") - p.add_argument("--limit", type=int, default=0) - p.add_argument("--max_seconds", type=float, default=30.0) - p.add_argument("--target_max_chars", type=int, default=0) - p.add_argument("--max_new_tokens", type=int, default=512) - p.add_argument("--batch_size", type=int, default=16) - p.add_argument( - "--use_cache", - action="store_true", - help="use KV cache (needs --patch for gemma4)", - ) - p.add_argument( - "--patch", action="store_true", help="apply gemma4 use_cache prefill bugfix" - ) - return p.parse_args() - - -def main(): - args = parse_args() - if args.patch: - import gemma4_cache_patch - - gemma4_cache_patch.apply() - print("applied gemma4 use_cache patch", flush=True) - proc = AutoProcessor.from_pretrained(args.model) - proc.tokenizer.padding_side = "left" # left padding for batched generation - - bnb = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_use_double_quant=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - llm_int8_skip_modules=[ - "model.audio_tower", - "model.vision_tower", - "model.embed_audio", - "model.embed_vision", - "lm_head", - ], - ) - model = AutoModelForMultimodalLM.from_pretrained( - args.model, - quantization_config=bnb, - dtype=torch.bfloat16, - device_map="auto", - attn_implementation="eager", - ) - if args.adapter: - model = PeftModel.from_pretrained(model, args.adapter) - model.eval() - # cached gemma4 multimodal forward is buggy unless --patch is applied - model.config.use_cache = args.use_cache - - if args.parquet_glob: - ds = CosheParquetDataset( - args.parquet_glob, - max_seconds=args.max_seconds, - target_max_chars=args.target_max_chars, - limit=args.limit, - cache_path=args.cache_path, - ) - else: - ds = CosheSample7Dataset( - args.data_dir, - max_seconds=args.max_seconds, - target_max_chars=args.target_max_chars, - ) - items = list(ds) - print(f"Evaluating {len(items)} clips, batch_size={args.batch_size}", flush=True) - - refs, hyps, wers = [], [], [] - for start in range(0, len(items), args.batch_size): - batch = items[start : start + args.batch_size] - texts, audios = [], [] - for it in batch: - msgs = [ - { - "role": "user", - "content": [ - {"type": "text", "text": DEFAULT_PROMPT}, - {"type": "audio", "audio": it["audio"]}, - ], - } - ] - texts.append( - proc.apply_chat_template( - msgs, tokenize=False, add_generation_prompt=True - ) - ) - audios.append(it["audio"]) - inp = proc(text=texts, audio=audios, return_tensors="pt", padding=True).to( - model.device - ) - in_len = inp["input_ids"].shape[-1] - with torch.inference_mode(): - out = model.generate( - **inp, - max_new_tokens=args.max_new_tokens, - do_sample=False, - use_cache=args.use_cache, - ) - for i, it in enumerate(batch): - gen = proc.decode(out[i][in_len:], skip_special_tokens=True).strip() - refs.append(it["target"]) - hyps.append(gen) - w = ( - jiwer.wer( - it["target"], - gen, - reference_transform=_WER_NORM, - hypothesis_transform=_WER_NORM, - ) - if it["target"].strip() - else 0.0 - ) - wers.append(w) - done = start + len(batch) - run_corpus = jiwer.wer( - refs, hyps, reference_transform=_WER_NORM, hypothesis_transform=_WER_NORM - ) - print( - f" {done}/{len(items)} running corpus WER={run_corpus*100:.2f}% " - f"batch mean WER={sum(wers[-len(batch):])/len(batch)*100:.2f}%", - flush=True, - ) - - corpus = jiwer.wer( - refs, hyps, reference_transform=_WER_NORM, hypothesis_transform=_WER_NORM - ) - mean = sum(wers) / len(wers) - n_exact = sum(1 for w in wers if w == 0.0) - print(f"\n==== FULL EVAL ====", flush=True) - print( - f"clips={len(wers)} exact(WER=0)={n_exact} " - f"mean per-clip WER={mean*100:.2f}% CORPUS WER={corpus*100:.2f}%", - flush=True, - ) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/gemma4/finetune/eval_windowed_stitch.py b/extras/asr-services/providers/gemma4/finetune/eval_windowed_stitch.py deleted file mode 100644 index 2eb5843d..00000000 --- a/extras/asr-services/providers/gemma4/finetune/eval_windowed_stitch.py +++ /dev/null @@ -1,319 +0,0 @@ -"""Evaluate a model (base or LoRA adapter) on the held-out CoSHE test clips via the -windowed-stitch path: cut each whole clip into its honest <=28s windows (the SAME -forced-alignment windows used for training, from windows_v2.jsonl), transcribe each window -with PLAIN_PROMPT, then concatenate window hyps in order -> a full-clip hypothesis scored -against the full GT transcript. - -Non-overlapping windows -> stitch is a plain concat (no overlap dedup). The same windowing -+ prompt + decode settings are used for base and every adapter, so the base->FT WER delta -is controlled. use_cache=True is safe (gemma4 cache bug fixed in transformers >=5.10). - -Out: mlexp-schema JSONL keyed by audio_file_name: {transcription (full GT), hyp (stitched), -n_wins, duration_s}. Resumable (skips clips already written). -""" - -import argparse -import glob -import io -import json -import os -import random -import time -from collections import defaultdict - -import numpy as np -import pyarrow.parquet as pq -import soundfile as sf -import torch -from data_windowed import PLAIN_PROMPT -from peft import PeftModel -from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig - -# Exact prod/benchmark prompt (interleave_probe.py PLAIN_PROMPT) that yields E2B ~10.6% -# median on CoSHE-100. gemma4 is prompt-conditioned; an elaborate prompt degrades base. -MINIMAL_PROMPT = ( - "Transcribe the following speech in its original language into text. " - "Only output the transcription text itself, with no commentary or explanation." -) -PROMPTS = {"plain": PLAIN_PROMPT, "minimal": MINIMAL_PROMPT} - - -def load_done(path): - done = set() - if os.path.exists(path): - for line in open(path): - try: - done.add(json.loads(line)["audio_file_name"]) - except Exception: - pass - return done - - -@torch.inference_mode() -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("--model", default="google/gemma-4-E2B-it") - ap.add_argument("--adapter", default="") - ap.add_argument("--windows", default="/home/coshe_windowed/windows_v2.jsonl") - ap.add_argument("--split", default="/home/coshe_windowed/windowed_split.json") - ap.add_argument("--parquet_glob", default="/home/coshe/data/eval-*.parquet") - ap.add_argument("--out", required=True) - ap.add_argument("--batch_size", type=int, default=8) - ap.add_argument("--max_new_tokens", type=int, default=256) - ap.add_argument( - "--no_repeat_ngram", - type=int, - default=0, - help="generate no_repeat_ngram_size; 3 kills greedy repetition-collapse", - ) - ap.add_argument("--repetition_penalty", type=float, default=1.0) - ap.add_argument( - "--limit", type=int, default=0, help="cap #test clips (seeded subset)" - ) - ap.add_argument( - "--quant", - default="bf16", - choices=["bf16", "4bit"], - help="bf16 = prod config (no quant); 4bit only for tiny GPUs", - ) - ap.add_argument( - "--prompt", - default="plain", - choices=["plain", "minimal"], - help="minimal = exact benchmark/prod prompt (E2B ~10.6%); plain = training prompt", - ) - ap.add_argument( - "--names_file", - default="", - help="JSON list of clip names to restrict eval to (∩ held-out)", - ) - ap.add_argument( - "--inject", - default="twostep", - choices=["twostep", "onestep"], - help="onestep = prod path apply_chat_template(tokenize=True,return_dict=True)", - ) - ap.add_argument( - "--fixed_win", - type=float, - default=0.0, - help="ignore fa windows; split each clip into fixed N-sec windows (benchmark style)", - ) - ap.add_argument( - "--attn", - default="", - help="attn_implementation; empty = model default (prod). 'eager' degrades gemma4 gen", - ) - args = ap.parse_args() - prompt_text = PROMPTS[args.prompt] - - test_clips = json.load(open(args.split))["test_clips"] - if args.names_file: - want = set(json.load(open(args.names_file))) - test_clips = [c for c in test_clips if c in want] - if args.limit and args.limit < len(test_clips): - test_clips = sorted(random.Random(0).sample(test_clips, args.limit)) - test_set = set(test_clips) - done = load_done(args.out) - todo = [c for c in test_clips if c not in done] - - wins_by_clip = defaultdict(list) - for line in open(args.windows): - r = json.loads(line) - if r["audio_file_name"] in test_set: - wins_by_clip[r["audio_file_name"]].append(r) - for c in wins_by_clip: - wins_by_clip[c].sort(key=lambda w: w["win_idx"]) - - # decode needed test clips from parquet (audio + full GT), slice window audio - gt, clip_audio = {}, {} - need = set(todo) - for f in sorted(glob.glob(args.parquet_glob)): - if not need: - break - t = pq.ParquetFile(f).read( - columns=["audio_file_name", "transcription", "audio"] - ) - names = t["audio_file_name"].to_pylist() - for ci, name in enumerate(names): - if name not in need: - continue - gt[name] = t["transcription"][ci].as_py().strip() - wav, sr = sf.read(io.BytesIO(t["audio"][ci].as_py()["bytes"])) - if wav.ndim > 1: - wav = wav.mean(1) - if sr != 16000: - import librosa - - wav = librosa.resample( - wav.astype(np.float32), orig_sr=sr, target_sr=16000 - ) - clip_audio[name] = np.ascontiguousarray(wav, dtype=np.float32) - need.discard(name) - print( - f"test clips total={len(test_clips)} todo={len(todo)} decoded={len(clip_audio)}", - flush=True, - ) - - proc = AutoProcessor.from_pretrained(args.model) - load_kw = dict(dtype=torch.bfloat16, device_map="auto") - if ( - args.attn - ): # benchmark/prod uses the MODEL DEFAULT (sdpa); eager degrades gemma4 gen - load_kw["attn_implementation"] = args.attn - if args.quant == "4bit": # QLoRA-training config; degrades small E2B at inference - load_kw["quantization_config"] = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_use_double_quant=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - llm_int8_skip_modules=[ - "model.audio_tower", - "model.vision_tower", - "model.embed_audio", - "model.embed_vision", - "lm_head", - ], - ) - model = AutoModelForMultimodalLM.from_pretrained(args.model, **load_kw) - if args.adapter: - model = PeftModel.from_pretrained(model, args.adapter) - print(f"loaded adapter: {args.adapter}", flush=True) - model.eval() - model.config.use_cache = True - proc.tokenizer.padding_side = "left" - - # optional: ignore fa windows, split each clip into fixed N-sec windows (benchmark style) - if args.fixed_win: - n = int(args.fixed_win * 16000) - wins_by_clip = { - c: [ - { - "win_idx": i, - "start": (i * n) / 16000, - "end": min(len(clip_audio[c]), (i + 1) * n) / 16000, - } - for i in range(max(1, (len(clip_audio[c]) + n - 1) // n)) - ] - for c in todo - } - - # flatten all windows of all todo clips into one work list, batch across clips - work = [] # (clip, win_idx, audio_slice) - for c in todo: - for w in wins_by_clip[c]: - a0, a1 = int(w["start"] * 16000), int(w["end"] * 16000) - work.append((c, w["win_idx"], clip_audio[c][a0:a1])) - print(f"windows to transcribe: {len(work)}", flush=True) - - gen_kw = dict( - max_new_tokens=args.max_new_tokens, - do_sample=False, - use_cache=True, - repetition_penalty=args.repetition_penalty, - ) - if args.no_repeat_ngram: - gen_kw["no_repeat_ngram_size"] = args.no_repeat_ngram - - hyp_parts = defaultdict(dict) # clip -> {win_idx: hyp} - fout = open(args.out, "a") - t0 = time.time() - bs = 1 if args.inject == "onestep" else args.batch_size - for s in range(0, len(work), bs): - batch = work[s : s + bs] - if ( - args.inject == "onestep" - ): # prod path: tokenize=True one-shot, per-window (bs=1) - import os as _os - import tempfile - import wave as _wave - - c, wi, au = batch[0] - # match the benchmark exactly: write a 16k PCM_16 WAV and pass the PATH (not the - # raw array) — the Gemma4 processor conditions differently on array vs file. - tf = tempfile.NamedTemporaryFile(suffix=".wav", delete=False) - with _wave.open(tf.name, "wb") as wf: - wf.setnchannels(1) - wf.setsampwidth(2) - wf.setframerate(16000) - wf.writeframes((np.clip(au, -1, 1) * 32767).astype(np.int16).tobytes()) - msgs = [ - { - "role": "user", - "content": [ - {"type": "text", "text": prompt_text}, - {"type": "audio", "audio": tf.name}, - ], - } - ] - inp = proc.apply_chat_template( - msgs, - tokenize=True, - return_dict=True, - return_tensors="pt", - add_generation_prompt=True, - enable_thinking=False, - ).to(model.device) - out = model.generate(**inp, **gen_kw) - _os.unlink(tf.name) - hyp_parts[c][wi] = proc.decode( - out[0][inp["input_ids"].shape[-1] :], skip_special_tokens=True - ).strip() - else: # legacy two-step path - texts, audios = [], [] - for _, _, au in batch: - msgs = [ - { - "role": "user", - "content": [ - {"type": "text", "text": prompt_text}, - {"type": "audio", "audio": au}, - ], - } - ] - texts.append( - proc.apply_chat_template( - msgs, - tokenize=False, - add_generation_prompt=True, - enable_thinking=False, - ) - ) - audios.append(au) - inp = proc(text=texts, audio=audios, return_tensors="pt", padding=True).to( - model.device - ) - out = model.generate(**inp, **gen_kw) - for i, (c, wi, _) in enumerate(batch): - hyp_parts[c][wi] = proc.decode( - out[i][inp["input_ids"].shape[-1] :], skip_special_tokens=True - ).strip() - if (s // bs) % 40 == 0: - print(f"[{s + len(batch)}/{len(work)}] {time.time() - t0:.0f}s", flush=True) - - # stitch + write per fully-completed clip - for c in todo: - parts = hyp_parts.get(c, {}) - if len(parts) != len(wins_by_clip[c]): - continue # incomplete (shouldn't happen); skip to retain resumability - stitched = " ".join(parts[i] for i in sorted(parts)).strip() - dur = round(len(clip_audio[c]) / 16000.0, 2) - fout.write( - json.dumps( - { - "audio_file_name": c, - "transcription": gt[c], - "hyp": stitched, - "n_wins": len(parts), - "duration_s": dur, - }, - ensure_ascii=False, - ) - + "\n" - ) - fout.close() - print(f"DONE out={args.out}", flush=True) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/gemma4/finetune/eval_wispr_stitch.py b/extras/asr-services/providers/gemma4/finetune/eval_wispr_stitch.py deleted file mode 100644 index a18d6781..00000000 --- a/extras/asr-services/providers/gemma4/finetune/eval_wispr_stitch.py +++ /dev/null @@ -1,148 +0,0 @@ -"""Evaluate base or LoRA-adapter Gemma4-E2B on the held-out Wispr test clips. - -Test clips are pre-sliced into the SAME fa windows used for training (manifest split=="test"). -Each window is transcribed with WISPR_PROMPT via the prod onestep path -(apply_chat_template(tokenize=True, return_dict=True), bf16, model-default attn=sdpa, -use_cache=True), then window hyps are concatenated per clip (non-overlapping -> plain concat) -into a full-clip hypothesis scored against asr_text (test_refs.jsonl) with jiwer. - -Same windowing + prompt + decode for base and every adapter -> the base->FT WER delta is -controlled. Run once with --adapter "" (base) and once with --adapter <dir> (FT). - -Usage: - python eval_wispr_stitch.py --manifest .../manifest.jsonl --refs .../test_refs.jsonl \ - --adapter "" --out base.jsonl -""" - -import argparse -import json -import re -import statistics -import time -from collections import defaultdict - -import jiwer -import torch -from data_wispr import WISPR_PROMPT -from peft import PeftModel -from transformers import AutoModelForMultimodalLM, AutoProcessor - -_BRACKET = re.compile(r"\[[^\]]*\]") -_APOS = re.compile(r"['’`]") -_PUNCT = re.compile(r"[^\w\s]") -_WS = re.compile(r"\s+") - - -def normalize(t: str) -> str: - t = _BRACKET.sub(" ", t).lower() - t = _APOS.sub("", t) - t = _PUNCT.sub(" ", t) - return _WS.sub(" ", t).strip() - - -@torch.inference_mode() -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("--model", default="google/gemma-4-E2B-it") - ap.add_argument("--adapter", default="") - ap.add_argument("--manifest", default="/home/wispr_windowed/manifest.jsonl") - ap.add_argument("--refs", default="/home/wispr_windowed/test_refs.jsonl") - ap.add_argument("--out", required=True) - ap.add_argument("--max_new_tokens", type=int, default=256) - args = ap.parse_args() - - # test windows grouped by clip - wins_by_clip = defaultdict(list) - for line in open(args.manifest): - r = json.loads(line) - if r["split"] == "test": - wins_by_clip[r["clip"]].append(r) - for c in wins_by_clip: - wins_by_clip[c].sort(key=lambda w: w["win_idx"]) - refs = { - json.loads(l)["audio_file_name"]: json.loads(l)["transcription"] - for l in open(args.refs) - } - - proc = AutoProcessor.from_pretrained(args.model) - proc.tokenizer.padding_side = "left" - model = AutoModelForMultimodalLM.from_pretrained( - args.model, dtype=torch.bfloat16, device_map="auto" - ) # default attn (sdpa), bf16 - if args.adapter: - model = PeftModel.from_pretrained(model, args.adapter) - print(f"adapter: {args.adapter}", flush=True) - model.eval() - model.config.use_cache = True - - gen_kw = dict(max_new_tokens=args.max_new_tokens, do_sample=False, use_cache=True) - t0 = time.time() - rows = [] - fout = open(args.out, "w") - for ci, (clip, wins) in enumerate(sorted(wins_by_clip.items())): - parts = [] - for w in wins: - msgs = [ - { - "role": "user", - "content": [ - {"type": "text", "text": WISPR_PROMPT}, - {"type": "audio", "audio": w["audio"]}, - ], - } - ] - inp = proc.apply_chat_template( - msgs, - tokenize=True, - return_dict=True, - return_tensors="pt", - add_generation_prompt=True, - enable_thinking=False, - ).to(model.device) - out = model.generate(**inp, **gen_kw) - parts.append( - proc.decode( - out[0][inp["input_ids"].shape[-1] :], skip_special_tokens=True - ).strip() - ) - hyp = " ".join(parts).strip() - ref = refs[clip] - rn, hn = normalize(ref), normalize(hyp) - wer = jiwer.wer(rn, hn) if hn else 1.0 - rows.append( - { - "audio_file_name": clip, - "wer": round(wer, 4), - "ref_words": len(rn.split()), - "n_wins": len(wins), - "transcription": ref, - "hyp": hyp, - } - ) - fout.write(json.dumps(rows[-1], ensure_ascii=False) + "\n") - fout.flush() - print( - f"[{ci+1}/{len(wins_by_clip)}] {clip} wer={wer:.3f} ({time.time()-t0:.0f}s)", - flush=True, - ) - fout.close() - - # corpus + median - tot_ref = sum(r["ref_words"] for r in rows) - refs_n = [normalize(refs[r["audio_file_name"]]) for r in rows] - hyps_n = [normalize(r["hyp"]) for r in rows] - corpus = jiwer.wer(refs_n, hyps_n) - median = statistics.median(r["wer"] for r in rows) - summary = { - "adapter": args.adapter or "base", - "clips": len(rows), - "corpus_wer": round(corpus, 4), - "median_wer": round(median, 4), - "total_ref_words": tot_ref, - } - print("SUMMARY " + json.dumps(summary), flush=True) - json.dump(summary, open(args.out.replace(".jsonl", "_summary.json"), "w"), indent=2) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/gemma4/finetune/gemini_transcribe.py b/extras/asr-services/providers/gemma4/finetune/gemini_transcribe.py deleted file mode 100644 index 4e633305..00000000 --- a/extras/asr-services/providers/gemma4/finetune/gemini_transcribe.py +++ /dev/null @@ -1,69 +0,0 @@ -"""Send an audio window to Gemini for a reference transcription.""" - -import argparse -import base64 -import io -import json -import os -import urllib.error -import urllib.request - -import soundfile as sf - -SR = 16000 - - -def main(): - p = argparse.ArgumentParser() - p.add_argument("--audio", default="/w/ankush-yap-fixed.wav") - p.add_argument("--offset", type=float, default=0.0) - p.add_argument("--seconds", type=float, default=60.0) - p.add_argument("--model", default="gemini-3.5-flash") - args = p.parse_args() - - a, sr = sf.read(args.audio, dtype="float32") - if a.ndim > 1: - a = a.mean(1) - if sr != SR: - import librosa - - a = librosa.resample(a, orig_sr=sr, target_sr=SR) - seg = a[int(args.offset * SR) : int((args.offset + args.seconds) * SR)] - buf = io.BytesIO() - sf.write(buf, seg, SR, format="WAV", subtype="PCM_16") - b64 = base64.b64encode(buf.getvalue()).decode() - print( - f"window {args.offset:.0f}-{args.offset+args.seconds:.0f}s, {len(b64)/1e6:.1f}MB b64", - flush=True, - ) - - key = os.environ["GOOGLE_API_KEY"] - url = f"https://generativelanguage.googleapis.com/v1beta/models/{args.model}:generateContent?key={key}" - body = { - "contents": [ - { - "parts": [ - { - "text": "Transcribe this audio verbatim. Output only the transcript text." - }, - {"inline_data": {"mime_type": "audio/wav", "data": b64}}, - ] - } - ] - } - req = urllib.request.Request( - url, - data=json.dumps(body).encode(), - headers={"Content-Type": "application/json"}, - ) - try: - resp = json.load(urllib.request.urlopen(req, timeout=120)) - cand = resp["candidates"][0]["content"]["parts"][0]["text"] - print("=== GEMINI TRANSCRIPT ===", flush=True) - print(cand, flush=True) - except urllib.error.HTTPError as e: - print("HTTP", e.code, e.read().decode()[:800], flush=True) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/gemma4/finetune/gemma4_cache_patch.py b/extras/asr-services/providers/gemma4/finetune/gemma4_cache_patch.py deleted file mode 100644 index 479a7253..00000000 --- a/extras/asr-services/providers/gemma4/finetune/gemma4_cache_patch.py +++ /dev/null @@ -1,64 +0,0 @@ -"""Monkeypatch for the gemma4 use_cache prefill bug (transformers 5.5.0). - -Bug: `create_causal_mask_mapping` only applies the audio/vision BIDIRECTIONAL -attention mask (`or_mask_function`) on the "first iteration" (prefill). It infers -prefill via: - - past_key_values is None or not past_key_values.is_initialized or pixel_values is not None - -But a freshly-created DynamicCache reports `is_initialized=True` *before any token -is cached*, so on the prefill forward (with use_cache=True / generate) this is -False → the audio soft tokens are masked purely causally instead of bidirectionally -→ the audio conditioning is wrong → logits collapse toward the base model -(measured: loss 4.49 with cache vs 0.10 without on a fully-overfit adapter). - -Fix: detect prefill by `get_seq_length() == 0` (no tokens cached yet) instead of -`is_initialized`. This makes generation correct AND fast (KV cache usable), turning -a ~100s/clip no-cache eval into normal cached generation. - -Usage: `import gemma4_cache_patch; gemma4_cache_patch.apply()` BEFORE loading the model. -""" - -import transformers.models.gemma4.modeling_gemma4 as _g4 - -_orig_create_causal_mask_mapping = _g4.create_causal_mask_mapping - - -def _patched_create_causal_mask_mapping( - config, - inputs_embeds, - attention_mask, - past_key_values, - position_ids, - mm_token_type_ids=None, - pixel_values=None, - is_training=False, - is_first_iteration=None, - **kwargs, -): - if is_first_iteration is None: - if past_key_values is None: - is_first_iteration = True - else: - try: - seqlen = past_key_values.get_seq_length() - except Exception: - seqlen = 0 - is_first_iteration = (seqlen == 0) or (pixel_values is not None) - return _orig_create_causal_mask_mapping( - config, - inputs_embeds, - attention_mask, - past_key_values, - position_ids, - mm_token_type_ids=mm_token_type_ids, - pixel_values=pixel_values, - is_training=is_training, - is_first_iteration=is_first_iteration, - **kwargs, - ) - - -def apply(): - _g4.create_causal_mask_mapping = _patched_create_causal_mask_mapping - return True diff --git a/extras/asr-services/providers/gemma4/finetune/infer.py b/extras/asr-services/providers/gemma4/finetune/infer.py deleted file mode 100644 index effd21c3..00000000 --- a/extras/asr-services/providers/gemma4/finetune/infer.py +++ /dev/null @@ -1,169 +0,0 @@ -"""Verify a Gemma 4 audio LoRA adapter by transcribing the training clips and -comparing to ground truth (overfit check). - -Loads the 4-bit base + adapter, generates on each sample7 clip, prints the -output next to the manifest target, and reports a crude char-level similarity -(SequenceMatcher ratio) so we can see the loss-near-zero overfit reflected in -the actual decoded text. - -Usage: - python infer.py --model google/gemma-4-E2B-it \ - --adapter /train/out/e2b-overfit --data_dir /data/coshe-eval/sample7 -""" - -import argparse -from difflib import SequenceMatcher - -import jiwer -import torch -from data import DEFAULT_PROMPT, CosheParquetDataset, CosheSample7Dataset -from peft import PeftModel -from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig - -# WER normalization: lowercase, strip punctuation, collapse whitespace. For an -# overfit, the model is trained on the target's exact script, so exact -# reproduction -> WER 0 (no romanization needed; raw mixed-script compares fine -# when hyp and ref use the same script, which they do after training). -_WER_NORM = jiwer.Compose( - [ - jiwer.ToLowerCase(), - jiwer.RemovePunctuation(), - jiwer.RemoveMultipleSpaces(), - jiwer.Strip(), - jiwer.ReduceToListOfListOfWords(), - ] -) - - -def compute_wer(ref: str, hyp: str) -> float: - if not ref.strip(): - return 0.0 if not hyp.strip() else 1.0 - return jiwer.wer( - ref, hyp, reference_transform=_WER_NORM, hypothesis_transform=_WER_NORM - ) - - -def parse_args(): - p = argparse.ArgumentParser() - p.add_argument("--model", default="google/gemma-4-E2B-it") - p.add_argument("--adapter", default="") - p.add_argument("--data_dir", default="/data/coshe-eval/sample7") - p.add_argument("--parquet_glob", default="") - p.add_argument("--limit", type=int, default=0) - p.add_argument("--cache_path", default="") - p.add_argument("--max_seconds", type=float, default=30.0) - p.add_argument("--target_max_chars", type=int, default=0) - p.add_argument("--max_new_tokens", type=int, default=512) - return p.parse_args() - - -def main(): - args = parse_args() - processor = AutoProcessor.from_pretrained(args.model) - processor.tokenizer.padding_side = "left" - - bnb = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_use_double_quant=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - llm_int8_skip_modules=[ - "model.audio_tower", - "model.vision_tower", - "model.embed_audio", - "model.embed_vision", - "lm_head", - ], - ) - model = AutoModelForMultimodalLM.from_pretrained( - args.model, - quantization_config=bnb, - dtype=torch.bfloat16, - device_map="auto", - attn_implementation="eager", - ) - tag = "BASE (no adapter)" - if args.adapter: - model = PeftModel.from_pretrained(model, args.adapter) - tag = f"ADAPTER {args.adapter}" - model.eval() - # CRITICAL: gemma4's multimodal forward produces WRONG audio-conditioned logits - # when the KV cache is enabled (transformers 5.5.0) — with use_cache=True a - # fully-overfit adapter scores loss ~4.5 (≈ base) instead of ~0.1, and - # generation reverts to near-base output. Generate with the cache OFF. - model.config.use_cache = False - - if args.parquet_glob: - dataset = CosheParquetDataset( - args.parquet_glob, - max_seconds=args.max_seconds, - target_max_chars=args.target_max_chars, - limit=args.limit, - cache_path=args.cache_path, - ) - else: - dataset = CosheSample7Dataset( - args.data_dir, - max_seconds=args.max_seconds, - target_max_chars=args.target_max_chars, - ) - print(f"\n===== INFERENCE: {tag} =====", flush=True) - ratios = [] - wers = [] - all_gens = [] - for item in dataset: - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": DEFAULT_PROMPT}, - {"type": "audio", "audio": item["audio"]}, - ], - } - ] - inputs = processor.apply_chat_template( - messages, - tokenize=True, - return_dict=True, - return_tensors="pt", - add_generation_prompt=True, - ).to(model.device) - in_len = inputs["input_ids"].shape[-1] - with torch.inference_mode(): - out = model.generate( - **inputs, - max_new_tokens=args.max_new_tokens, - do_sample=False, - use_cache=False, - ) - gen = processor.decode(out[0][in_len:], skip_special_tokens=True).strip() - target = item["target"] - ratio = SequenceMatcher(None, gen, target).ratio() - wer = compute_wer(target, gen) - ratios.append(ratio) - wers.append(wer) - all_gens.append(gen) - print( - f"\n--- {item['name']} (char-sim {ratio:.3f} WER {wer*100:.2f}%) ---", - flush=True, - ) - print(f" TARGET: {target[:240]}", flush=True) - print(f" OUTPUT: {gen[:240]}", flush=True) - - mean_wer = sum(wers) / len(wers) - # corpus WER (aggregate over all words), the headline metric for the goal - corpus_wer = jiwer.wer( - [d["target"] for d in dataset], - all_gens, - reference_transform=_WER_NORM, - hypothesis_transform=_WER_NORM, - ) - print(f"\nMean char-similarity: {sum(ratios)/len(ratios):.3f}", flush=True) - print( - f"Mean per-clip WER: {mean_wer*100:.2f}% Corpus WER: {corpus_wer*100:.2f}%", - flush=True, - ) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/gemma4/finetune/introspect.py b/extras/asr-services/providers/gemma4/finetune/introspect.py deleted file mode 100644 index 87255f76..00000000 --- a/extras/asr-services/providers/gemma4/finetune/introspect.py +++ /dev/null @@ -1,162 +0,0 @@ -"""Introspect a Gemma 4 E*B model + processor for QLoRA fine-tuning. - -Loads the model in 4-bit and dumps the facts we need to write the training -script correctly for THIS model (the gemma4 family), since published recipes -target gemma3n and key/module names can differ: - - 1. Processor output keys for a (text + audio) example, with shapes/dtypes. - 2. Audio/image special token ids exposed on the tokenizer/config. - 3. Module-name families for LoRA targets: text-decoder projections, the - audio projector, and embedding/lm_head modules. - -Run inside the gemma4 image (transformers >= 5.5). -""" - -import os -import wave -from collections import Counter - -import numpy as np -import torch -from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig - -MODEL_ID = os.getenv("ASR_MODEL", "google/gemma-4-E2B-it") -AUDIO = os.getenv("PROBE_AUDIO", "/data/coshe-eval/sample7/audio/audio_13.wav") -PROMPT = "Transcribe the following speech segment in its original language." - - -def load_wav_16k_mono(path: str, max_seconds: float = 30.0) -> np.ndarray: - with wave.open(path, "rb") as wf: - sr = wf.getframerate() - n = wf.getnframes() - raw = wf.readframes(n) - audio = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0 - # mono assumed (sample7 is mono). Resample to 16k via librosa. - if sr != 16000: - import librosa - - audio = librosa.resample(audio, orig_sr=sr, target_sr=16000) - audio = audio[: int(16000 * max_seconds)] - return audio - - -def main(): - print(f"=== MODEL: {MODEL_ID} ===", flush=True) - bnb = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_use_double_quant=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - ) - processor = AutoProcessor.from_pretrained(MODEL_ID) - model = AutoModelForMultimodalLM.from_pretrained( - MODEL_ID, - quantization_config=bnb, - dtype=torch.bfloat16, - device_map="auto", - ) - - tok = processor.tokenizer - print("\n=== SPECIAL TOKEN IDS ===", flush=True) - for attr in [ - "pad_token_id", - "eos_token_id", - "bos_token_id", - "audio_token_id", - "image_token_id", - "boi_token_id", - "eoi_token_id", - "boa_token_id", - "eoa_token_id", - ]: - print(f" tokenizer.{attr} = {getattr(tok, attr, '<none>')}", flush=True) - cfg = model.config - for attr in [ - "audio_token_id", - "image_token_id", - "boi_token_id", - "eoi_token_id", - "boa_token_id", - "eoa_token_id", - "eoa_token_index", - ]: - print(f" config.{attr} = {getattr(cfg, attr, '<none>')}", flush=True) - - print("\n=== PROCESSOR OUTPUT (text+audio, TRAINING form) ===", flush=True) - audio = load_wav_16k_mono(AUDIO) - print( - f" probe audio: {AUDIO} samples={len(audio)} ({len(audio)/16000:.1f}s)", - flush=True, - ) - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": PROMPT}, - {"type": "audio", "audio": audio}, - ], - }, - { - "role": "assistant", - "content": [{"type": "text", "text": "Speaker 1: hello world"}], - }, - ] - # tokenize=False then call processor with text+audio (the HF recipe pattern) - text = processor.apply_chat_template( - messages, tokenize=False, add_generation_prompt=False - ) - print(f" templated text (first 300 chars):\n {text[:300]!r}", flush=True) - batch = processor(text=[text], audio=[audio], return_tensors="pt", padding=True) - for k, v in batch.items(): - if hasattr(v, "shape"): - print(f" key {k:24s} shape={tuple(v.shape)} dtype={v.dtype}", flush=True) - else: - print(f" key {k:24s} type={type(v)}", flush=True) - - print("\n=== MODULE NAME FAMILIES (Linear leaf modules) ===", flush=True) - fam = Counter() - audio_like = set() - embed_like = set() - for name, mod in model.named_modules(): - leaf = name.split(".")[-1] - cls = type(mod).__name__ - if "Linear" in cls or "lora" in cls.lower(): - fam[leaf] += 1 - low = name.lower() - if any(t in low for t in ["audio", "speech", "conformer"]): - # record short module suffixes from the audio tower / projector - if "Linear" in cls or "Embedding" in cls or "proj" in low: - audio_like.add(f"{leaf} ({cls})") - if "embed" in low or leaf in ("lm_head",): - embed_like.add(f"{name} ({cls})") - print(" Linear leaf-name -> count:", flush=True) - for leaf, c in fam.most_common(): - print(f" {leaf:28s} x{c}", flush=True) - print("\n audio-tower / projector linear+embedding leaves (sample):", flush=True) - for s in sorted(audio_like)[:40]: - print(f" {s}", flush=True) - print("\n embedding / lm_head modules (full names, sample):", flush=True) - for s in sorted(embed_like)[:25]: - print(f" {s}", flush=True) - - print("\n=== top-level named children ===", flush=True) - for name, _ in model.named_children(): - print(f" {name}", flush=True) - # one level deeper into the LM + audio tower roots - for root in [ - "model", - "language_model", - "audio_tower", - "audio_model", - "multi_modal_projector", - ]: - sub = getattr(model, root, None) - if sub is not None: - kids = [n for n, _ in sub.named_children()] - print(f" {root}.children = {kids}", flush=True) - - print("\nDONE", flush=True) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/gemma4/finetune/poll_until_target.sh b/extras/asr-services/providers/gemma4/finetune/poll_until_target.sh deleted file mode 100755 index 09398adc..00000000 --- a/extras/asr-services/providers/gemma4/finetune/poll_until_target.sh +++ /dev/null @@ -1,32 +0,0 @@ -#!/bin/bash -# Poll the A100 training log via the Jupyter contents API (SSH-free) until the -# early-stopping loop reports it hit <2% WER ("TARGET REACHED" / "DONE"), or the -# training process dies, or a long timeout. Exits 0 with a summary so the agent is -# re-invoked to retrieve the adapter + pause the instance. -set -u -ENVBKP=/home/ankush/workspaces/friend-lite/.env.bkp -FT=/home/ankush/workspaces/friend-lite/extras/asr-services/providers/gemma4/finetune -JTOKEN=$(grep -hiE "^jarvislabs=" "$ENVBKP" | head -1 | cut -d= -f2- | tr -d '"'\'' \r') - -# Resolve host + jupyter token once. -URL=$(JLTOKEN="$JTOKEN" uv run --with "git+https://github.com/jarvislabsai/jlclient.git" python3 -c " -import os; from jlclient import jarvisclient; from jlclient.jarvisclient import * -jarvisclient.token=os.environ['JLTOKEN']; print(User.get_instances()[0].url)" 2>/dev/null | grep notebooks) -HOST=$(echo "$URL" | sed -E 's#https://([^/]+)/.*#\1#') -TOK=$(echo "$URL" | sed -E 's#.*token=##') -LOG="https://$HOST/api/contents/gemma4ft/out/train_full2.log?token=$TOK&type=file&format=text&content=1" -echo "polling host=$HOST" - -for i in $(seq 1 240); do # up to 240 * 4min = 16h - C=$(curl -s --max-time 90 "$LOG" 2>/dev/null) - # latest epoch-mean signal: last WER eval line and last loss/epoch - EVAL=$(echo "$C" | grep -aoE '\[epoch [0-9]+\][^\\]*(corpus WER = [0-9.]+%|skipping WER eval)' | tail -1) - TGT=$(echo "$C" | grep -aoE 'TARGET REACHED[^\\]*|DONE \(final adapter saved\)') - LOSSEP=$(echo "$C" | grep -aoE "'loss': '[0-9.]+', 'grad_norm'[^}]*'epoch': '[0-9.]+'" | tail -1 | sed -E "s/, 'grad_norm'[^,]*,/,/") - echo "[$(date +%H:%M)] iter=$i | $LOSSEP | eval: ${EVAL:-none}" - if [ -n "$TGT" ]; then echo "STOP_SIGNAL: $TGT"; exit 0; fi - # detect dead training: if process gone (no recent log growth is hard; rely on TARGET/DONE or manual) - sleep 240 -done -echo "poller timed out after 16h without TARGET" -exit 0 diff --git a/extras/asr-services/providers/gemma4/finetune/run_split_pipeline.sh b/extras/asr-services/providers/gemma4/finetune/run_split_pipeline.sh deleted file mode 100644 index 594a48c9..00000000 --- a/extras/asr-services/providers/gemma4/finetune/run_split_pipeline.sh +++ /dev/null @@ -1,49 +0,0 @@ -#!/bin/bash -# Drives the full split-LoRA generalization experiment on the A100, sequentially -# (single GPU). Each stage logs to its own file and appends a marker to a master -# progress log; a failed stage does not abort the rest. Window=30s "first-30s -# transcription" task throughout (proportional target truncation). -cd /home/gemma4ft || exit 1 -export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True -PROG=/home/gemma4ft/out/pipeline_progress.log -echo "PIPELINE_START $(date -u +%H:%M:%S)" > "$PROG" - -stage () { # stage <name> <logfile> <cmd...> - local name="$1" log="$2"; shift 2 - echo "STAGE_START $name $(date -u +%H:%M:%S)" >> "$PROG" - if "$@" > "$log" 2>&1; then echo "STAGE_OK $name $(date -u +%H:%M:%S)" >> "$PROG" - else echo "STAGE_FAILED $name rc=$? $(date -u +%H:%M:%S)" >> "$PROG"; fi -} - -COMMON="--lr 1e-5 --lora_r 16 --lora_alpha 32 --batch_size 2 --grad_accum 2 --patience 4 --epochs 40 --window_seconds 30" -EVAL_COMMON="--limit 700 --batch_size 8 --max_new_tokens 448 --window_seconds 30" - -# 1) decoder-only r16 -stage train_decoder out/split_decoder_r16.log \ - python3 train_lora_split.py --output_dir out/split_decoder_r16 $COMMON - -# 2) include-head r16 -stage train_head out/split_head_r16.log \ - python3 train_lora_split.py --output_dir out/split_head_r16 --include_head $COMMON - -# 3) base eval (no adapter) — the controlled baseline -stage eval_base out/eval_base.log \ - python3 eval_test_split.py --out out/base_test.jsonl $EVAL_COMMON - -# 4) ft decoder eval (guard: only if adapter saved) -if [ -f out/split_decoder_r16/adapter_model.safetensors ]; then - stage eval_ft_decoder out/eval_ft_decoder.log \ - python3 eval_test_split.py --adapter out/split_decoder_r16 --out out/ft_decoder_test.jsonl $EVAL_COMMON -else - echo "SKIP eval_ft_decoder (no adapter)" >> "$PROG" -fi - -# 5) ft head eval (guard) -if [ -f out/split_head_r16/adapter_model.safetensors ]; then - stage eval_ft_head out/eval_ft_head.log \ - python3 eval_test_split.py --adapter out/split_head_r16 --out out/ft_head_test.jsonl $EVAL_COMMON -else - echo "SKIP eval_ft_head (no adapter)" >> "$PROG" -fi - -echo "PIPELINE_DONE $(date -u +%H:%M:%S)" >> "$PROG" diff --git a/extras/asr-services/providers/gemma4/finetune/smoke_interleave.py b/extras/asr-services/providers/gemma4/finetune/smoke_interleave.py deleted file mode 100644 index c3671af7..00000000 --- a/extras/asr-services/providers/gemma4/finetune/smoke_interleave.py +++ /dev/null @@ -1,93 +0,0 @@ -"""Smoke test the multi-audio interleave collator + a forward pass BEFORE a full train. -Checks: processor accepts nested multi-audio, labels mask only the target, forward gives a -finite loss. Also runs a base-model generate on one multi-chunk clip to confirm interleave -inference works. -""" - -import sys - -import torch -from data_interleave import ChunkInterleaveDataset, InterleaveCollator, build_prompt -from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig - -MANIFEST = sys.argv[1] if len(sys.argv) > 1 else "/home/wispr_interleave/manifest.jsonl" -MODEL = "google/gemma-4-E2B-it" - -proc = AutoProcessor.from_pretrained(MODEL) -ds = ChunkInterleaveDataset(MANIFEST, "train") -print( - f"train clips={len(ds)}; chunk counts (first 10): {[len(ds[i]['chunks']) for i in range(min(10,len(ds)))]}" -) -# pick a 1-chunk and a multi-chunk example to stress the collator -idx_multi = next((i for i in range(len(ds)) if len(ds[i]["chunks"]) >= 2), 0) -idx_one = next((i for i in range(len(ds)) if len(ds[i]["chunks"]) == 1), 0) -feats = [ds[idx_one], ds[idx_multi]] -print( - f"batch: 1-chunk='{feats[0]['clip']}' ({len(feats[0]['chunks'])}), " - f"multi='{feats[1]['clip']}' ({len(feats[1]['chunks'])} chunks, app={feats[1]['app']})" -) - -coll = InterleaveCollator(proc) -batch = coll(feats) -print("batch keys:", list(batch.keys())) -print( - "input_ids", - tuple(batch["input_ids"].shape), - "| input_features", - tuple(batch["input_features"].shape) if "input_features" in batch else None, -) -sup = (batch["labels"] != -100).sum(dim=1).tolist() -print("supervised tokens/example (should ≈ target lengths):", sup) -# decode the supervised span of example 0 to confirm it's the target, not the prompt -lab = batch["labels"][0] -sup_ids = batch["input_ids"][0][lab != -100] -print("supervised decode[0]:", proc.decode(sup_ids, skip_special_tokens=True)[:120]) - -bnb = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_use_double_quant=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - llm_int8_skip_modules=[ - "model.audio_tower", - "model.vision_tower", - "model.embed_audio", - "model.embed_vision", - "lm_head", - ], -) -model = AutoModelForMultimodalLM.from_pretrained( - MODEL, - quantization_config=bnb, - dtype=torch.bfloat16, - device_map="auto", - attn_implementation="sdpa", -) -model.config.use_cache = False -batch = {k: v.to(model.device) for k, v in batch.items()} -with torch.no_grad(): - out = model(**batch) -print( - "FORWARD loss:", float(out.loss), "(finite:", torch.isfinite(out.loss).item(), ")" -) - -# interleave generate on the multi-chunk clip (base model) -model.config.use_cache = True -f = ds[idx_multi] -content = [{"type": "text", "text": build_prompt(f["app"])}] -content += [{"type": "audio", "audio": a} for a in f["chunks"]] -inp = proc.apply_chat_template( - [{"role": "user", "content": content}], - tokenize=True, - return_dict=True, - return_tensors="pt", - add_generation_prompt=True, - enable_thinking=False, -).to(model.device) -g = model.generate(**inp, max_new_tokens=200, do_sample=False, use_cache=True) -print( - "GEN (base, interleave):", - proc.decode(g[0][inp["input_ids"].shape[-1] :], skip_special_tokens=True)[:200], -) -print("TARGET:", f["target"][:200]) -print("SMOKE OK") diff --git a/extras/asr-services/providers/gemma4/finetune/split_status.py b/extras/asr-services/providers/gemma4/finetune/split_status.py deleted file mode 100644 index ea3e849f..00000000 --- a/extras/asr-services/providers/gemma4/finetune/split_status.py +++ /dev/null @@ -1,36 +0,0 @@ -"""Print a compact training status for a split-LoRA run: the val-loss series (from the -latest checkpoint's trainer_state.json), best eval_loss, current step, and any terminal -marker in the run's .log. Used by the progress monitor (the live stdout log is a single -\\r-joined line, so the structured state file is the reliable source).""" - -import glob -import json -import sys - -out = sys.argv[1].rstrip("/") -cks = sorted(glob.glob(out + "/checkpoint-*"), key=lambda p: int(p.split("-")[-1])) -evals, best, step = [], None, None -if cks: - s = json.load(open(cks[-1] + "/trainer_state.json")) - best = s.get("best_metric") - step = s.get("global_step") - for h in s["log_history"]: - if "eval_loss" in h: - evals.append((round(h["epoch"], 2), round(h["eval_loss"], 4))) -term = "" -try: - raw = open(out + ".log").read() - for m in [ - "DONE ", - "Traceback", - "CUDA out of memory", - "RuntimeError", - "OutOfMemory", - "Killed", - ]: - if m in raw: - term = m.strip() - break -except Exception: - pass -print(f"evals={len(evals)} step={step} best={best} series={evals[-6:]} TERM={term}") diff --git a/extras/asr-services/providers/gemma4/finetune/sweep_12b_v2.py b/extras/asr-services/providers/gemma4/finetune/sweep_12b_v2.py deleted file mode 100644 index e81666bd..00000000 --- a/extras/asr-services/providers/gemma4/finetune/sweep_12b_v2.py +++ /dev/null @@ -1,254 +0,0 @@ -"""Fast 12B ASR config sweep on a fixed CoSHE subset, via the REAL service path. - -Loads google/gemma-4-12B-it ONCE (Gemma4Transcriber, the production transcriber) -and runs a small fixed set of CoSHE-Eval clips through several configs, varying -prompt / windowing / decode strategy. Output is one JSONL per config in the -score_coshe schema so evaluate/score_coshe.py can score them directly. - -Unlike bench_coshe_12b.py (manual channel-strip decode, whole-clip 600s window) -this uses the transcriber's parse_response decode + 30s windowing + stitching — -i.e. exactly what the deployed gemma4-asr service does — so a winning config maps -straight to env knobs (GEMMA4_QUANT, TRANSCRIPTION_PROMPT, MAX_NEW_TOKENS, -BATCH_DURATION_SECONDS, BATCH_THRESHOLD_SECONDS). - -Run inside chronicle-asr-gemma4-smoke (transformers 5.10.1 + bnb), GPU, with -/models = HF cache, /coshe-data = CoSHE parquet, /data/coshe-eval mounted. -""" - -import argparse -import glob -import io -import json -import os -import tempfile -import time -import wave - -import numpy as np -import pyarrow.parquet as pq -import soundfile as sf - -SR = 16000 - -# 11 fixed samples chosen from the existing 500-sample run to span the -# distribution: catastrophic-tail (loops/empties), median, and easy clips. -SAMPLES = [ - "audio_1400.wav", - "audio_299.wav", - "audio_353.wav", - "audio_131.wav", - "audio_564.wav", - "audio_1404.wav", - "audio_243.wav", - "audio_666.wav", - "audio_1543.wav", - "audio_650.wav", - "audio_1938.wav", -] - -# Prompts under test ---------------------------------------------------------- -PROD_PROMPT = ( # current service default (DEFAULT_TRANSCRIPTION_PROMPT) - "Transcribe the following speech segment in its original language into text. " - "Only output the transcription text itself, with no commentary or explanation. " - "When transcribing numbers, write the digits, i.e. write 1.7 and not " - "one point seven, and write 3 instead of three." -) -HF_PROMPT = ( - "Transcribe the following speech segment in its original language.\n\n" - "Follow these specific instructions for formatting the answer:\n" - "* Only output the transcription, with no newlines.\n" - "* When transcribing numbers, write the digits, i.e. write 1.7 and not " - "one point seven, and write 3 instead of three." -) -VERBATIM_PROMPT = ( - "You are an expert transcriptionist. Transcribe the provided Hinglish audio " - "verbatim exactly as it is spoken. Output ONLY the transcribed text. Do not " - "translate it to English, do not summarize, and do not add any conversational " - "pleasantries or commentary." -) -HINGLISH_PROMPT = ( # explicit code-switch steer - "Transcribe the following Hindi-English code-mixed (Hinglish) speech verbatim. " - "Write each word in the language it is spoken. Only output the transcription " - "text itself, with no commentary. When transcribing numbers, write the digits." -) -# prod prompt WITHOUT the "write digits" clause (CoSHE refs spell numbers out, so -# the digit instruction inflates WER) and without code-switch wording. -NODIGIT_PROMPT = ( - "Transcribe the following speech segment in its original language into text. " - "Only output the transcription text itself, with no commentary or explanation." -) - -# Each config: (prompt, threshold_s, window_s, overlap_s, max_new, extra_gen) --- -WHOLE = 100000.0 -CONFIGS = { - # production defaults: 30s windows + overlap, greedy, 512 tokens - "prod": (PROD_PROMPT, 30.0, 30.0, 5.0, 512, {}), - # isolate windowing: production prompt but whole-clip (matches old bench geometry) - "prod_whole": (PROD_PROMPT, WHOLE, 30.0, 5.0, 1024, {}), - # anti-loop safety on the windowed path - "prod_norep3": (PROD_PROMPT, 30.0, 30.0, 5.0, 512, {"no_repeat_ngram_size": 3}), - # prompt variants on the windowed path - "hfprompt": (HF_PROMPT, 30.0, 30.0, 5.0, 512, {}), - "verbatim": (VERBATIM_PROMPT, 30.0, 30.0, 5.0, 512, {}), - "hinglish": (HINGLISH_PROMPT, 30.0, 30.0, 5.0, 512, {}), - # low-temp sampling on the windowed path - "sample_t03": ( - PROD_PROMPT, - 30.0, - 30.0, - 5.0, - 512, - {"do_sample": True, "temperature": 0.3, "top_p": 0.95, "top_k": 64}, - ), - # round 2: best prompts + anti-loop n-gram (tail-safe), windowed greedy - "hf_norep3": (HF_PROMPT, 30.0, 30.0, 5.0, 512, {"no_repeat_ngram_size": 3}), - "hinglish_norep3": ( - HINGLISH_PROMPT, - 30.0, - 30.0, - 5.0, - 512, - {"no_repeat_ngram_size": 3}, - ), - "nodigit_norep3": ( - NODIGIT_PROMPT, - 30.0, - 30.0, - 5.0, - 512, - {"no_repeat_ngram_size": 3}, - ), -} - - -def select_rows(data_dir, names): - """Pull the target rows from the parquet shards.""" - want = set(names) - found = {} - for s in sorted(glob.glob(os.path.join(data_dir, "eval-*.parquet"))): - pf = pq.ParquetFile(s) - for batch in pf.iter_batches(batch_size=64): - for row in batch.to_pylist(): - if row["audio_file_name"] in want: - found[row["audio_file_name"]] = row - if len(found) == len(want): - break - return found - - -def write_wav(raw_bytes, path): - audio, sr = sf.read(io.BytesIO(raw_bytes), dtype="float32") - if audio.ndim > 1: - audio = audio.mean(axis=1) - if sr != SR: - import librosa - - audio = librosa.resample(audio, orig_sr=sr, target_sr=SR) - pcm = np.clip(audio, -1.0, 1.0) - pcm16 = (pcm * 32767.0).astype(np.int16) - with wave.open(path, "wb") as wf: - wf.setnchannels(1) - wf.setsampwidth(2) - wf.setframerate(SR) - wf.writeframes(pcm16.tobytes()) - return len(audio) / SR - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("--data_dir", default="/coshe-data") - ap.add_argument("--out_dir", default="/data/coshe-eval/results/12b_sweep") - ap.add_argument("--model", default="google/gemma-4-12B-it") - ap.add_argument("--quant", default="4bit") - ap.add_argument("--configs", default="", help="comma list; default = all") - args = ap.parse_args() - - # Force the transcriber into the desired load mode BEFORE import/instantiation. - os.environ["GEMMA4_QUANT"] = args.quant - os.environ["GEMMA4_MTP"] = "0" # off: clean target distribution, faster load - os.environ["ASR_MODEL"] = args.model - os.environ["TORCH_DTYPE"] = "bfloat16" - - from providers.gemma4.transcriber import Gemma4Transcriber - - os.makedirs(args.out_dir, exist_ok=True) - - print(f"selecting {len(SAMPLES)} samples from {args.data_dir} ...", flush=True) - rows = select_rows(args.data_dir, SAMPLES) - missing = set(SAMPLES) - set(rows) - if missing: - print(f"WARNING missing samples: {missing}", flush=True) - - # Decode all target clips to temp WAVs once. - tmp = tempfile.mkdtemp(prefix="coshe_") - clips = {} # name -> (wav_path, ref_text, duration_s) - for name in SAMPLES: - if name not in rows: - continue - wav_path = os.path.join(tmp, name) - dur = write_wav(rows[name]["audio"]["bytes"], wav_path) - clips[name] = (wav_path, rows[name]["transcription"], dur) - print(f"decoded {len(clips)} clips", flush=True) - - # Load model once. - t = Gemma4Transcriber(model_id=args.model) - t0 = time.time() - t.load_model() - print(f"model loaded quant={args.quant} in {time.time()-t0:.0f}s", flush=True) - - # Monkeypatch _generate to merge per-config extra gen kwargs (anti-loop/sampling). - orig_generate = t._generate - t._extra_gen = {} - - def patched_generate(inputs, **kw): - kw.update(t._extra_gen) - return orig_generate(inputs, **kw) - - t._generate = patched_generate - - selected = args.configs.split(",") if args.configs else list(CONFIGS) - for cfg_name in selected: - prompt, thr, win, ov, max_new, extra = CONFIGS[cfg_name] - t.batch_threshold = thr - t.batch_duration = win - t.batch_overlap = ov - t.max_new_tokens = max_new - t._extra_gen = dict(extra) - out_path = os.path.join(args.out_dir, f"{cfg_name}.jsonl") - print( - f"\n=== CONFIG {cfg_name} (thr={thr} win={win} max_new={max_new} " - f"extra={extra}) -> {out_path}", - flush=True, - ) - with open(out_path, "w") as f: - for name in SAMPLES: - if name not in clips: - continue - wav_path, ref, dur = clips[name] - ts = time.time() - try: - res = t.transcribe(wav_path, prompt_override=prompt) - hyp = res.text - except Exception as e: - hyp = "" - print(f" ERR {name}: {e}", flush=True) - dt = time.time() - ts - rec = { - "audio_file_name": name, - "transcription": ref, - "hyp": hyp, - "asr_seconds": round(dt, 2), - "duration_s": round(dur, 2), - } - f.write(json.dumps(rec, ensure_ascii=False) + "\n") - f.flush() - print( - f" [{name}] dur={dur:.0f}s gen={dt:.1f}s " - f"hyp[:70]={hyp[:70]!r}", - flush=True, - ) - print("\nDONE", flush=True) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/gemma4/finetune/train.py b/extras/asr-services/providers/gemma4/finetune/train.py deleted file mode 100644 index a21b8dc8..00000000 --- a/extras/asr-services/providers/gemma4/finetune/train.py +++ /dev/null @@ -1,203 +0,0 @@ -"""QLoRA fine-tuning for Gemma 4 E*B on CoSHE-Eval sample7 (overfit smoke test). - -4-bit NF4 base (double quant, bf16 compute) + LoRA on the text decoder only -(audio tower + multimodal embedders frozen). Plain HF Trainer with the custom -multimodal collator in data.py. - -Usage (inside the gemma4 image venv): - python train.py \ - --model google/gemma-4-E2B-it \ - --data_dir /data/coshe-eval/sample7 \ - --output_dir /train/out/e2b-overfit \ - --epochs 40 --lr 2e-4 --batch_size 1 --grad_accum 1 -""" - -import argparse - -import torch -from data import CosheParquetDataset, CosheSample7Dataset, Gemma4AudioCollator -from peft import LoraConfig, get_peft_model -from transformers import ( - AutoModelForMultimodalLM, - AutoProcessor, - BitsAndBytesConfig, - Trainer, - TrainingArguments, -) - -# LoRA scoped to the text decoder (language_model) attention + MLP projections. -# Regex (PEFT re.fullmatch) keeps it off the identically-named audio-tower linears. -# NOTE: an experiment adding lm_head + embed_tokens here (rank 64) made free-running -# generation DEGENERATE into "Speaker 1: Speaker 1:..." loops — the blocker is target -# grounding, not output capacity, so we keep the lean attention+MLP set. -LORA_TARGETS = ( - r".*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" -) -# With --include_head, also adapt lm_head + input embeddings for more memorization -# capacity (safe now that generation uses use_cache=False; the earlier "degenerate" -# result with these was a cache-bug artifact, not a real failure). -LORA_TARGETS_HEAD = ( - r"(.*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" - r"|.*language_model\.embed_tokens|lm_head)" -) - - -def parse_args(): - p = argparse.ArgumentParser() - p.add_argument("--model", default="google/gemma-4-E2B-it") - p.add_argument("--data_dir", default="/data/coshe-eval/sample7") - p.add_argument( - "--parquet_glob", - default="", - help="if set, use full CoSHE parquet instead of sample7", - ) - p.add_argument("--limit", type=int, default=0, help="cap #samples (parquet only)") - p.add_argument( - "--cache_path", default="", help="pickle cache for decoded parquet audio" - ) - p.add_argument("--output_dir", default="/train/out/e2b-overfit") - p.add_argument("--epochs", type=float, default=40.0) - p.add_argument("--lr", type=float, default=2e-4) - p.add_argument("--batch_size", type=int, default=1) - p.add_argument("--grad_accum", type=int, default=1) - p.add_argument("--max_seconds", type=float, default=30.0) - p.add_argument("--target_max_chars", type=int, default=0) - p.add_argument("--lora_r", type=int, default=16) - p.add_argument("--lora_alpha", type=int, default=32) - p.add_argument("--lora_dropout", type=float, default=0.0) - p.add_argument( - "--optim", - default="adamw_torch", - help="e.g. adamw_8bit / paged_adamw_8bit to save VRAM", - ) - p.add_argument( - "--include_head", action="store_true", help="also LoRA lm_head + embed_tokens" - ) - p.add_argument( - "--resume", - action="store_true", - help="resume from latest checkpoint in output_dir", - ) - return p.parse_args() - - -def main(): - args = parse_args() - torch.manual_seed(0) - - print(f"Loading processor + 4-bit model: {args.model}", flush=True) - processor = AutoProcessor.from_pretrained(args.model) - - bnb = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_use_double_quant=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - # Keep the multimodal towers / embedders / lm_head in bf16. The audio - # tower's ClippableLinear calls torch.finfo(weight.dtype) for gradient - # clipping, which breaks on 4-bit (uint8-stored) weights. We only train - # LoRA on the text decoder anyway. - llm_int8_skip_modules=[ - "model.audio_tower", - "model.vision_tower", - "model.embed_audio", - "model.embed_vision", - "lm_head", - ], - ) - model = AutoModelForMultimodalLM.from_pretrained( - args.model, - quantization_config=bnb, - dtype=torch.bfloat16, - device_map="auto", - attn_implementation="eager", - ) - model.config.use_cache = False - - # Manual k-bit prep (NOT peft's prepare_model_for_kbit_training): that helper - # upcasts every non-4bit param to fp32, which makes the text embeddings fp32 - # while the (quant-skipped, bf16) audio tower stays bf16 -> Gemma4's - # multimodal masked_scatter merge then errors on the dtype mismatch. Keeping - # everything bf16 avoids that and matches the recipe's bf16-throughout advice. - for p in model.parameters(): - p.requires_grad = False - model.gradient_checkpointing_enable( - gradient_checkpointing_kwargs={"use_reentrant": False} - ) - model.enable_input_require_grads() - - lora = LoraConfig( - r=args.lora_r, - lora_alpha=args.lora_alpha, - lora_dropout=args.lora_dropout, - bias="none", - task_type="CAUSAL_LM", - target_modules=LORA_TARGETS_HEAD if args.include_head else LORA_TARGETS, - ) - model = get_peft_model(model, lora) - model.print_trainable_parameters() - - if args.parquet_glob: - dataset = CosheParquetDataset( - args.parquet_glob, - max_seconds=args.max_seconds, - target_max_chars=args.target_max_chars, - limit=args.limit, - cache_path=args.cache_path, - ) - else: - dataset = CosheSample7Dataset( - args.data_dir, - max_seconds=args.max_seconds, - target_max_chars=args.target_max_chars, - ) - print(f"Dataset: {len(dataset)} samples", flush=True) - collator = Gemma4AudioCollator(processor) - - targs = TrainingArguments( - output_dir=args.output_dir, - per_device_train_batch_size=args.batch_size, - gradient_accumulation_steps=args.grad_accum, - num_train_epochs=args.epochs, - learning_rate=args.lr, - lr_scheduler_type="constant", - warmup_steps=0, - bf16=True, - fp16=False, - logging_steps=1, - save_strategy="epoch", - save_total_limit=2, - optim=args.optim, - report_to=[], - remove_unused_columns=False, - dataloader_num_workers=0, - gradient_checkpointing=False, # already enabled via prepare_model_for_kbit_training - ) - - trainer = Trainer( - model=model, - args=targs, - train_dataset=dataset, - data_collator=collator, - ) - - # Resume from the latest checkpoint in output_dir if present and --resume set. - # /home is persisted on Jarvis Labs but the process dies on VM pause/restart, so - # resuming from the last per-epoch checkpoint avoids wasting progress. - import os as _os - - resume = False - if args.resume and _os.path.isdir(args.output_dir): - ckpts = [d for d in _os.listdir(args.output_dir) if d.startswith("checkpoint-")] - resume = len(ckpts) > 0 - print(f"Starting training... (resume={resume})", flush=True) - trainer.train(resume_from_checkpoint=resume) - - print(f"Saving adapter to {args.output_dir}", flush=True) - trainer.save_model(args.output_dir) - processor.save_pretrained(args.output_dir) - print("DONE", flush=True) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/gemma4/finetune/train_interleave.py b/extras/asr-services/providers/gemma4/finetune/train_interleave.py deleted file mode 100644 index 291bd39d..00000000 --- a/extras/asr-services/providers/gemma4/finetune/train_interleave.py +++ /dev/null @@ -1,150 +0,0 @@ -"""LoRA finetune Gemma4-E2B on the interleaved audio->pasted_text (app-aware formatting) task. -One example = one clip (N<=28s audio chunks in a single prompt) -> full formatted pasted_text. -Decoder-only LoRA, val-loss early stop. Mirrors train_wispr.py but uses the interleave -dataset/collator (multi-audio per example, formatted target). -""" - -import argparse -import json -import os - -import torch -from data_interleave import ChunkInterleaveDataset, InterleaveCollator -from peft import LoraConfig, get_peft_model -from transformers import ( - AutoModelForMultimodalLM, - AutoProcessor, - BitsAndBytesConfig, - EarlyStoppingCallback, - Trainer, - TrainingArguments, -) - -LORA_TARGETS_HEAD = ( - r"(.*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" - r"|.*language_model\.embed_tokens|lm_head)" -) -LORA_TARGETS = ( - r".*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" -) - - -def parse_args(): - p = argparse.ArgumentParser() - p.add_argument("--model", default="google/gemma-4-E2B-it") - p.add_argument("--manifest", default="/home/wispr_interleave/manifest.jsonl") - p.add_argument("--output_dir", required=True) - p.add_argument("--epochs", type=float, default=40.0) - p.add_argument("--lr", type=float, default=1e-5) - p.add_argument("--batch_size", type=int, default=1) - p.add_argument("--grad_accum", type=int, default=4) - p.add_argument("--eval_batch_size", type=int, default=1) - p.add_argument("--lora_r", type=int, default=16) - p.add_argument("--lora_alpha", type=int, default=32) - p.add_argument("--include_head", action="store_true") - p.add_argument("--optim", default="adamw_8bit") - p.add_argument("--patience", type=int, default=5) - p.add_argument("--attn", default="sdpa") - return p.parse_args() - - -def main(): - args = parse_args() - torch.manual_seed(0) - proc = AutoProcessor.from_pretrained(args.model) - bnb = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_use_double_quant=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - llm_int8_skip_modules=[ - "model.audio_tower", - "model.vision_tower", - "model.embed_audio", - "model.embed_vision", - "lm_head", - ], - ) - model = AutoModelForMultimodalLM.from_pretrained( - args.model, - quantization_config=bnb, - dtype=torch.bfloat16, - device_map="auto", - attn_implementation=args.attn, - ) - model.config.use_cache = False - for pm in model.parameters(): - pm.requires_grad = False - model.gradient_checkpointing_enable( - gradient_checkpointing_kwargs={"use_reentrant": False} - ) - model.enable_input_require_grads() - model = get_peft_model( - model, - LoraConfig( - r=args.lora_r, - lora_alpha=args.lora_alpha, - lora_dropout=0.0, - bias="none", - task_type="CAUSAL_LM", - target_modules=LORA_TARGETS_HEAD if args.include_head else LORA_TARGETS, - ), - ) - model.print_trainable_parameters() - - train_ds = ChunkInterleaveDataset(args.manifest, "train") - val_ds = ChunkInterleaveDataset(args.manifest, "val") - print(f"train clips={len(train_ds)} val clips={len(val_ds)}", flush=True) - collator = InterleaveCollator(proc) - - targs = TrainingArguments( - output_dir=args.output_dir, - per_device_train_batch_size=args.batch_size, - gradient_accumulation_steps=args.grad_accum, - per_device_eval_batch_size=args.eval_batch_size, - num_train_epochs=args.epochs, - learning_rate=args.lr, - lr_scheduler_type="constant", - warmup_steps=0, - bf16=True, - fp16=False, - logging_steps=10, - eval_strategy="epoch", - save_strategy="epoch", - save_total_limit=2, - load_best_model_at_end=True, - metric_for_best_model="eval_loss", - greater_is_better=False, - optim=args.optim, - report_to=[], - remove_unused_columns=False, - dataloader_num_workers=0, - gradient_checkpointing=False, - seed=0, - data_seed=0, - ) - trainer = Trainer( - model=model, - args=targs, - train_dataset=train_ds, - eval_dataset=val_ds, - data_collator=collator, - callbacks=[EarlyStoppingCallback(early_stopping_patience=args.patience)], - ) - print("Starting training...", flush=True) - trainer.train() - trainer.save_model(args.output_dir) - proc.save_pretrained(args.output_dir) - json.dump( - trainer.state.log_history, - open(os.path.join(args.output_dir, "log_history.json"), "w"), - ) - print( - f"DONE best_checkpoint={trainer.state.best_model_checkpoint} " - f"best_eval_loss={trainer.state.best_metric}", - flush=True, - ) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/gemma4/finetune/train_lora_split.py b/extras/asr-services/providers/gemma4/finetune/train_lora_split.py deleted file mode 100644 index 2a19bdf4..00000000 --- a/extras/asr-services/providers/gemma4/finetune/train_lora_split.py +++ /dev/null @@ -1,199 +0,0 @@ -"""Split-aware LoRA finetune of Gemma4-E4B on CoSHE — a *generalization* experiment. - -Unlike `train_until_wer.py` (which memorizes the full set to <2% WER), this trains on a -random 20% slice and selects the checkpoint with the lowest **validation loss** on a -held-out 10% slice. The remaining 70% (the test split) is never touched here — it is -scored separately for WER against the base-model baseline. - -Split comes from a JSON file ({"train":[names], "val":[names], "test":[names]}) so the -exact same partition is used here (training) and locally (scoring). Items are matched to -the decoded-audio cache by `name` (== audio_file_name). - -Strategy (per the user's spec): very low constant LR, LoRA adapter, early-stop when the -validation loss stops improving (HF EarlyStoppingCallback + load_best_model_at_end on -eval_loss). Two configs are supported via --include_head: - * decoder-only : LoRA on language_model q/k/v/o/gate/up/down (audio tower frozen) - * include_head : the above + LoRA on embed_tokens and lm_head -""" - -import argparse -import json -import os - -import torch -from data import CosheParquetDataset, Gemma4AudioCollator -from peft import LoraConfig, get_peft_model -from transformers import ( - AutoModelForMultimodalLM, - AutoProcessor, - BitsAndBytesConfig, - EarlyStoppingCallback, - Trainer, - TrainingArguments, -) -from window_target import apply_window_truncation - -LORA_TARGETS_HEAD = ( - r"(.*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" - r"|.*language_model\.embed_tokens|lm_head)" -) -LORA_TARGETS = ( - r".*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" -) - - -class ListDataset(torch.utils.data.Dataset): - def __init__(self, items): - self.items = items - - def __len__(self): - return len(self.items) - - def __getitem__(self, idx): - return self.items[idx] - - -def parse_args(): - p = argparse.ArgumentParser() - p.add_argument("--model", default="google/gemma-4-E4B-it") - p.add_argument("--parquet_glob", default="/home/coshe-data/data/eval-*.parquet") - p.add_argument("--cache_path", default="/home/gemma4ft/out/coshe_full_cache.pkl") - p.add_argument("--split_file", default="/home/gemma4ft/split_20_10_70.json") - p.add_argument("--output_dir", required=True) - p.add_argument("--epochs", type=float, default=40.0) - p.add_argument("--lr", type=float, default=1e-5) - p.add_argument("--batch_size", type=int, default=2) - p.add_argument("--grad_accum", type=int, default=2) - p.add_argument("--eval_batch_size", type=int, default=4) - p.add_argument("--lora_r", type=int, default=16) - p.add_argument("--lora_alpha", type=int, default=32) - p.add_argument("--include_head", action="store_true") - p.add_argument("--optim", default="adamw_8bit") - p.add_argument("--patience", type=int, default=4) - p.add_argument( - "--window_seconds", - type=float, - default=30.0, - help="proportionally truncate target to this audio window (0=full)", - ) - p.add_argument("--durations", default="/home/gemma4ft/durations.json") - return p.parse_args() - - -def main(): - args = parse_args() - torch.manual_seed(0) - proc = AutoProcessor.from_pretrained(args.model) - bnb = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_use_double_quant=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - llm_int8_skip_modules=[ - "model.audio_tower", - "model.vision_tower", - "model.embed_audio", - "model.embed_vision", - "lm_head", - ], - ) - model = AutoModelForMultimodalLM.from_pretrained( - args.model, - quantization_config=bnb, - dtype=torch.bfloat16, - device_map="auto", - attn_implementation="eager", - ) - model.config.use_cache = False - for pm in model.parameters(): - pm.requires_grad = False - model.gradient_checkpointing_enable( - gradient_checkpointing_kwargs={"use_reentrant": False} - ) - model.enable_input_require_grads() - model = get_peft_model( - model, - LoraConfig( - r=args.lora_r, - lora_alpha=args.lora_alpha, - lora_dropout=0.0, - bias="none", - task_type="CAUSAL_LM", - target_modules=LORA_TARGETS_HEAD if args.include_head else LORA_TARGETS, - ), - ) - model.print_trainable_parameters() - - split = json.load(open(args.split_file)) - ds = CosheParquetDataset( - args.parquet_glob, - max_seconds=30.0, - target_max_chars=0, - cache_path=args.cache_path, - ) - apply_window_truncation(list(ds), args.durations, args.window_seconds) - by = {it["name"]: it for it in ds} - train_items = [by[n] for n in split["train"] if n in by] - val_items = [by[n] for n in split["val"] if n in by] - missing_tr = len(split["train"]) - len(train_items) - missing_va = len(split["val"]) - len(val_items) - print( - f"train={len(train_items)} (missing {missing_tr}) " - f"val={len(val_items)} (missing {missing_va}) " - f"test(held out)={len(split['test'])}", - flush=True, - ) - collator = Gemma4AudioCollator(proc) - - targs = TrainingArguments( - output_dir=args.output_dir, - per_device_train_batch_size=args.batch_size, - gradient_accumulation_steps=args.grad_accum, - per_device_eval_batch_size=args.eval_batch_size, - num_train_epochs=args.epochs, - learning_rate=args.lr, - lr_scheduler_type="constant", - warmup_steps=0, - bf16=True, - fp16=False, - logging_steps=10, - eval_strategy="epoch", - save_strategy="epoch", - save_total_limit=2, - load_best_model_at_end=True, - metric_for_best_model="eval_loss", - greater_is_better=False, - optim=args.optim, - report_to=[], - remove_unused_columns=False, - dataloader_num_workers=0, - gradient_checkpointing=False, - seed=0, - data_seed=0, - ) - - trainer = Trainer( - model=model, - args=targs, - train_dataset=ListDataset(train_items), - eval_dataset=ListDataset(val_items), - data_collator=collator, - callbacks=[EarlyStoppingCallback(early_stopping_patience=args.patience)], - ) - - print("Starting training...", flush=True) - trainer.train() - trainer.save_model(args.output_dir) - proc.save_pretrained(args.output_dir) - json.dump( - trainer.state.log_history, - open(os.path.join(args.output_dir, "log_history.json"), "w"), - ) - # surface the chosen (best) checkpoint + its val loss for the writeup - best = getattr(trainer.state, "best_model_checkpoint", None) - best_loss = getattr(trainer.state, "best_metric", None) - print(f"DONE best_checkpoint={best} best_eval_loss={best_loss}", flush=True) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/gemma4/finetune/train_lora_windowed.py b/extras/asr-services/providers/gemma4/finetune/train_lora_windowed.py deleted file mode 100644 index 471dd8c3..00000000 --- a/extras/asr-services/providers/gemma4/finetune/train_lora_windowed.py +++ /dev/null @@ -1,160 +0,0 @@ -"""LoRA finetune of Gemma4-E2B on the HONEST windowed CoSHE set (20% train / 10% val), -val-loss early stopping. Successor to train_lora_split.py: instead of the "first-30s + -proportional char-truncation" hack, it trains on real <=30s (audio, GT-text) windows cut -at forced-alignment word timings (build_windowed_dataset.py). The held-out 70% test clips -are never touched here; they are scored separately via the windowed-stitch path. - -decoder-only vs --include_head as before. Same low constant LR + LoRA recipe. -""" - -import argparse -import json -import os - -import torch -from data import Gemma4AudioCollator -from data_windowed import PLAIN_PROMPT, WindowedManifestDataset -from peft import LoraConfig, get_peft_model -from transformers import ( - AutoModelForMultimodalLM, - AutoProcessor, - BitsAndBytesConfig, - EarlyStoppingCallback, - Trainer, - TrainingArguments, -) - -LORA_TARGETS_HEAD = ( - r"(.*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" - r"|.*language_model\.embed_tokens|lm_head)" -) -LORA_TARGETS = ( - r".*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" -) - - -def parse_args(): - p = argparse.ArgumentParser() - p.add_argument("--model", default="google/gemma-4-E2B-it") - p.add_argument("--manifest", default="/home/coshe_windowed/manifest.jsonl") - p.add_argument("--output_dir", required=True) - p.add_argument("--epochs", type=float, default=40.0) - p.add_argument("--lr", type=float, default=1e-5) - p.add_argument("--batch_size", type=int, default=2) - p.add_argument("--grad_accum", type=int, default=2) - p.add_argument("--eval_batch_size", type=int, default=4) - p.add_argument("--lora_r", type=int, default=16) - p.add_argument("--lora_alpha", type=int, default=32) - p.add_argument("--include_head", action="store_true") - p.add_argument("--optim", default="adamw_8bit") - p.add_argument("--patience", type=int, default=4) - p.add_argument( - "--attn", - default="sdpa", - help="sdpa matches prod/eval (gemma4 default); eager was the old buggy-at-eval choice", - ) - return p.parse_args() - - -def main(): - args = parse_args() - torch.manual_seed(0) - proc = AutoProcessor.from_pretrained(args.model) - bnb = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_use_double_quant=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - llm_int8_skip_modules=[ - "model.audio_tower", - "model.vision_tower", - "model.embed_audio", - "model.embed_vision", - "lm_head", - ], - ) - model = AutoModelForMultimodalLM.from_pretrained( - args.model, - quantization_config=bnb, - dtype=torch.bfloat16, - device_map="auto", - attn_implementation=args.attn, - ) - model.config.use_cache = False - for pm in model.parameters(): - pm.requires_grad = False - model.gradient_checkpointing_enable( - gradient_checkpointing_kwargs={"use_reentrant": False} - ) - model.enable_input_require_grads() - model = get_peft_model( - model, - LoraConfig( - r=args.lora_r, - lora_alpha=args.lora_alpha, - lora_dropout=0.0, - bias="none", - task_type="CAUSAL_LM", - target_modules=LORA_TARGETS_HEAD if args.include_head else LORA_TARGETS, - ), - ) - model.print_trainable_parameters() - - train_ds = WindowedManifestDataset(args.manifest, "train") - val_ds = WindowedManifestDataset(args.manifest, "val") - print(f"train windows={len(train_ds)} val windows={len(val_ds)}", flush=True) - collator = Gemma4AudioCollator(proc, prompt=PLAIN_PROMPT) - - targs = TrainingArguments( - output_dir=args.output_dir, - per_device_train_batch_size=args.batch_size, - gradient_accumulation_steps=args.grad_accum, - per_device_eval_batch_size=args.eval_batch_size, - num_train_epochs=args.epochs, - learning_rate=args.lr, - lr_scheduler_type="constant", - warmup_steps=0, - bf16=True, - fp16=False, - logging_steps=10, - eval_strategy="epoch", - save_strategy="epoch", - save_total_limit=2, - load_best_model_at_end=True, - metric_for_best_model="eval_loss", - greater_is_better=False, - optim=args.optim, - report_to=[], - remove_unused_columns=False, - dataloader_num_workers=0, - gradient_checkpointing=False, - seed=0, - data_seed=0, - ) - - trainer = Trainer( - model=model, - args=targs, - train_dataset=train_ds, - eval_dataset=val_ds, - data_collator=collator, - callbacks=[EarlyStoppingCallback(early_stopping_patience=args.patience)], - ) - - print("Starting training...", flush=True) - trainer.train() - trainer.save_model(args.output_dir) - proc.save_pretrained(args.output_dir) - json.dump( - trainer.state.log_history, - open(os.path.join(args.output_dir, "log_history.json"), "w"), - ) - print( - f"DONE best_checkpoint={trainer.state.best_model_checkpoint} " - f"best_eval_loss={trainer.state.best_metric}", - flush=True, - ) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/gemma4/finetune/train_until_wer.py b/extras/asr-services/providers/gemma4/finetune/train_until_wer.py deleted file mode 100644 index 6f4b93c1..00000000 --- a/extras/asr-services/providers/gemma4/finetune/train_until_wer.py +++ /dev/null @@ -1,302 +0,0 @@ -"""Overfit CoSHE, evaluating WER every N epochs and stopping when corpus WER < target. - -transformers 5.7 has the use_cache bug fixed, so eval uses fast cached generation. -Every `--eval_every` epochs: eval a fast subset; if subset WER < target, eval the -FULL set and stop only if that's also < target. Saves the adapter on stop. -""" - -import argparse -import time - -import jiwer -import torch -from data import DEFAULT_PROMPT, CosheParquetDataset, Gemma4AudioCollator -from peft import LoraConfig, PeftModel, get_peft_model -from transformers import ( - AutoModelForMultimodalLM, - AutoProcessor, - BitsAndBytesConfig, - Trainer, - TrainerCallback, - TrainingArguments, -) - -LORA_TARGETS_HEAD = ( - r"(.*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" - r"|.*language_model\.embed_tokens|lm_head)" -) -LORA_TARGETS = ( - r".*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" -) -_WER_NORM = jiwer.Compose( - [ - jiwer.ToLowerCase(), - jiwer.RemovePunctuation(), - jiwer.RemoveMultipleSpaces(), - jiwer.Strip(), - jiwer.ReduceToListOfListOfWords(), - ] -) - - -def corpus_wer(refs, hyps): - return jiwer.wer( - refs, hyps, reference_transform=_WER_NORM, hypothesis_transform=_WER_NORM - ) - - -@torch.inference_mode() -def eval_wer(model, proc, items, batch_size, max_new_tokens): - model.eval() - prev_uc = model.config.use_cache - model.config.use_cache = True - proc.tokenizer.padding_side = "left" - refs, hyps = [], [] - for s in range(0, len(items), batch_size): - batch = items[s : s + batch_size] - texts, audios = [], [] - for it in batch: - msgs = [ - { - "role": "user", - "content": [ - {"type": "text", "text": DEFAULT_PROMPT}, - {"type": "audio", "audio": it["audio"]}, - ], - } - ] - texts.append( - proc.apply_chat_template( - msgs, tokenize=False, add_generation_prompt=True - ) - ) - audios.append(it["audio"]) - inp = proc(text=texts, audio=audios, return_tensors="pt", padding=True).to( - model.device - ) - out = model.generate( - **inp, max_new_tokens=max_new_tokens, do_sample=False, use_cache=True - ) - for i, it in enumerate(batch): - hyps.append( - proc.decode( - out[i][inp["input_ids"].shape[-1] :], skip_special_tokens=True - ).strip() - ) - refs.append(it["target"]) - model.config.use_cache = prev_uc - model.train() - return corpus_wer(refs, hyps), refs, hyps - - -class WERStopCallback(TrainerCallback): - def __init__( - self, - proc, - subset, - full, - every, - target, - batch_size, - max_new_tokens, - out_dir, - loss_gate=0.15, - ): - self.proc, self.subset, self.full = proc, subset, full - self.every, self.target, self.bs, self.mnt = ( - every, - target, - batch_size, - max_new_tokens, - ) - self.out_dir = out_dir - self.loss_gate = loss_gate - - def _recent_loss(self, state): - # Mean of the last ~30 logged losses — the single last value is too noisy - # (per-batch variance 0.1–0.27) and was erratically skipping evals. - vals = [r["loss"] for r in state.log_history if "loss" in r] - if not vals: - return None - tail = vals[-30:] - return sum(tail) / len(tail) - - def on_epoch_end(self, args, state, control, model=None, **kwargs): - ep = int(round(state.epoch)) - if ep == 0 or ep % self.every != 0: - return control - # Free-running generation can't be good while teacher-forced loss is high - # (sample7: WER only collapsed once loss -> ~0). Skip the expensive WER eval - # until loss drops below the gate, to avoid wasting ~10min/eval early on. - rl = self._recent_loss(state) - if rl is not None and rl > self.loss_gate: - print( - f"[epoch {ep}] loss={rl:.3f} > {self.loss_gate}; skipping WER eval", - flush=True, - ) - return control - t = time.time() - wer, _, _ = eval_wer(model, self.proc, self.subset, self.bs, self.mnt) - print( - f"[epoch {ep}] subset({len(self.subset)}) corpus WER = {wer*100:.2f}% ({time.time()-t:.0f}s)", - flush=True, - ) - if wer < self.target: - fw, _, _ = eval_wer(model, self.proc, self.full, self.bs, self.mnt) - print( - f"[epoch {ep}] FULL({len(self.full)}) corpus WER = {fw*100:.2f}%", - flush=True, - ) - if fw < self.target: - print( - f"[epoch {ep}] TARGET REACHED (<{self.target*100:.0f}% on full). Stopping.", - flush=True, - ) - control.should_training_stop = True - return control - - -def parse_args(): - p = argparse.ArgumentParser() - p.add_argument("--model", default="google/gemma-4-E2B-it") - p.add_argument("--parquet_glob", default="/home/coshe-data/data/eval-*.parquet") - p.add_argument("--cache_path", default="/home/gemma4ft/out/coshe_full_cache.pkl") - p.add_argument("--output_dir", default="/home/gemma4ft/out/full") - p.add_argument("--epochs", type=float, default=150.0) - p.add_argument("--lr", type=float, default=2e-4) - p.add_argument("--batch_size", type=int, default=4) - p.add_argument("--lora_r", type=int, default=256) - p.add_argument("--lora_alpha", type=int, default=512) - p.add_argument("--optim", default="adamw_8bit") - p.add_argument("--include_head", action="store_true") - p.add_argument("--resume", action="store_true") - p.add_argument( - "--init_adapter", - default="", - help="load this adapter as trainable start (fresh optimizer)", - ) - p.add_argument("--eval_every", type=int, default=2) - p.add_argument("--wer_target", type=float, default=0.02) - p.add_argument("--eval_subset", type=int, default=300) - p.add_argument("--loss_gate", type=float, default=0.15) - p.add_argument("--eval_batch_size", type=int, default=24) - p.add_argument("--eval_max_new_tokens", type=int, default=512) - return p.parse_args() - - -def main(): - import os - - args = parse_args() - torch.manual_seed(0) - proc = AutoProcessor.from_pretrained(args.model) - bnb = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_use_double_quant=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - llm_int8_skip_modules=[ - "model.audio_tower", - "model.vision_tower", - "model.embed_audio", - "model.embed_vision", - "lm_head", - ], - ) - model = AutoModelForMultimodalLM.from_pretrained( - args.model, - quantization_config=bnb, - dtype=torch.bfloat16, - device_map="auto", - attn_implementation="eager", - ) - model.config.use_cache = False - for pm in model.parameters(): - pm.requires_grad = False - model.gradient_checkpointing_enable( - gradient_checkpointing_kwargs={"use_reentrant": False} - ) - model.enable_input_require_grads() - if args.init_adapter: - # Load a previously-trained adapter as the starting point but with a FRESH - # optimizer (so a new/higher LR takes effect) — used to accelerate a run - # whose loss descent has slowed, without losing learned weights. - print(f"Loading init adapter (trainable): {args.init_adapter}", flush=True) - model = PeftModel.from_pretrained(model, args.init_adapter, is_trainable=True) - else: - model = get_peft_model( - model, - LoraConfig( - r=args.lora_r, - lora_alpha=args.lora_alpha, - lora_dropout=0.0, - bias="none", - task_type="CAUSAL_LM", - target_modules=LORA_TARGETS_HEAD if args.include_head else LORA_TARGETS, - ), - ) - model.print_trainable_parameters() - - ds = CosheParquetDataset( - args.parquet_glob, - max_seconds=30.0, - target_max_chars=0, - cache_path=args.cache_path, - ) - items = list(ds) - subset = items[: args.eval_subset] - print(f"Dataset: {len(items)} samples; eval subset={len(subset)}", flush=True) - collator = Gemma4AudioCollator(proc) - - targs = TrainingArguments( - output_dir=args.output_dir, - per_device_train_batch_size=args.batch_size, - num_train_epochs=args.epochs, - learning_rate=args.lr, - lr_scheduler_type="constant", - warmup_steps=0, - bf16=True, - fp16=False, - logging_steps=10, - save_strategy="epoch", - save_total_limit=2, - optim=args.optim, - report_to=[], - remove_unused_columns=False, - dataloader_num_workers=0, - gradient_checkpointing=False, - ) - - cb = WERStopCallback( - proc, - subset, - items, - args.eval_every, - args.wer_target, - args.eval_batch_size, - args.eval_max_new_tokens, - args.output_dir, - loss_gate=args.loss_gate, - ) - trainer = Trainer( - model=model, - args=targs, - train_dataset=ds, - data_collator=collator, - callbacks=[cb], - ) - - resume = ( - args.resume - and os.path.isdir(args.output_dir) - and any(d.startswith("checkpoint-") for d in os.listdir(args.output_dir)) - ) - print(f"Starting training... (resume={resume})", flush=True) - trainer.train(resume_from_checkpoint=resume) - trainer.save_model(args.output_dir) - proc.save_pretrained(args.output_dir) - print("DONE (final adapter saved)", flush=True) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/gemma4/finetune/train_wispr.py b/extras/asr-services/providers/gemma4/finetune/train_wispr.py deleted file mode 100644 index ffb0257b..00000000 --- a/extras/asr-services/providers/gemma4/finetune/train_wispr.py +++ /dev/null @@ -1,168 +0,0 @@ -"""LoRA finetune of Gemma4-E2B on the Wispr dictation windows (60% train / 20% val), -val-loss early stopping. Mirrors train_lora_windowed.py but uses WISPR_PROMPT (English -verbatim) and the Wispr windowed manifest. Held-out 20% test clips are scored separately -by eval_wispr_stitch.py. decoder-only LoRA by default (--include_head optional). -""" - -import argparse -import json -import os - -import torch -from data import Gemma4AudioCollator -from data_wispr import WISPR_PROMPT, WindowedManifestDataset -from peft import LoraConfig, get_peft_model -from transformers import ( - AutoModelForMultimodalLM, - AutoProcessor, - BitsAndBytesConfig, - EarlyStoppingCallback, - Trainer, - TrainingArguments, -) - -LORA_TARGETS_HEAD = ( - r"(.*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" - r"|.*language_model\.embed_tokens|lm_head)" -) -LORA_TARGETS = ( - r".*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" -) -# Audio conformer linears live INSIDE Gemma4ClippableLinear wrappers, so target the inner -# `.linear` of each attention/FFN projection. Lets the model adapt its *hearing*, not just -# the text decoder (which a frozen encoder can't fix). -AUDIO_TARGETS = ( - r".*audio_tower.*\.(q_proj|k_proj|v_proj|post|relative_k_proj" - r"|ffw_layer_1|ffw_layer_2)\.linear" -) - - -def parse_args(): - p = argparse.ArgumentParser() - p.add_argument("--model", default="google/gemma-4-E2B-it") - p.add_argument("--manifest", default="/home/wispr_windowed/manifest.jsonl") - p.add_argument("--output_dir", required=True) - p.add_argument("--epochs", type=float, default=40.0) - p.add_argument("--lr", type=float, default=1e-5) - p.add_argument("--batch_size", type=int, default=2) - p.add_argument("--grad_accum", type=int, default=2) - p.add_argument("--eval_batch_size", type=int, default=4) - p.add_argument("--lora_r", type=int, default=16) - p.add_argument("--lora_alpha", type=int, default=32) - p.add_argument("--include_head", action="store_true") - p.add_argument( - "--audio_lora", - action="store_true", - help="also LoRA the audio conformer (adapt hearing, not just decoder)", - ) - p.add_argument("--optim", default="adamw_8bit") - p.add_argument("--patience", type=int, default=4) - p.add_argument("--attn", default="sdpa") - return p.parse_args() - - -def main(): - args = parse_args() - torch.manual_seed(0) - proc = AutoProcessor.from_pretrained(args.model) - bnb = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_use_double_quant=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - llm_int8_skip_modules=[ - "model.audio_tower", - "model.vision_tower", - "model.embed_audio", - "model.embed_vision", - "lm_head", - ], - ) - model = AutoModelForMultimodalLM.from_pretrained( - args.model, - quantization_config=bnb, - dtype=torch.bfloat16, - device_map="auto", - attn_implementation=args.attn, - ) - model.config.use_cache = False - for pm in model.parameters(): - pm.requires_grad = False - model.gradient_checkpointing_enable( - gradient_checkpointing_kwargs={"use_reentrant": False} - ) - model.enable_input_require_grads() - targets = LORA_TARGETS_HEAD if args.include_head else LORA_TARGETS - if args.audio_lora: - targets = f"({targets}|{AUDIO_TARGETS})" - model = get_peft_model( - model, - LoraConfig( - r=args.lora_r, - lora_alpha=args.lora_alpha, - lora_dropout=0.0, - bias="none", - task_type="CAUSAL_LM", - target_modules=targets, - ), - ) - model.print_trainable_parameters() - - train_ds = WindowedManifestDataset(args.manifest, "train") - val_ds = WindowedManifestDataset(args.manifest, "val") - print(f"train windows={len(train_ds)} val windows={len(val_ds)}", flush=True) - collator = Gemma4AudioCollator(proc, prompt=WISPR_PROMPT) - - targs = TrainingArguments( - output_dir=args.output_dir, - per_device_train_batch_size=args.batch_size, - gradient_accumulation_steps=args.grad_accum, - per_device_eval_batch_size=args.eval_batch_size, - num_train_epochs=args.epochs, - learning_rate=args.lr, - lr_scheduler_type="constant", - warmup_steps=0, - bf16=True, - fp16=False, - logging_steps=10, - eval_strategy="epoch", - save_strategy="epoch", - save_total_limit=2, - load_best_model_at_end=True, - metric_for_best_model="eval_loss", - greater_is_better=False, - optim=args.optim, - report_to=[], - remove_unused_columns=False, - dataloader_num_workers=0, - gradient_checkpointing=False, - seed=0, - data_seed=0, - ) - - trainer = Trainer( - model=model, - args=targs, - train_dataset=train_ds, - eval_dataset=val_ds, - data_collator=collator, - callbacks=[EarlyStoppingCallback(early_stopping_patience=args.patience)], - ) - - print("Starting training...", flush=True) - trainer.train() - trainer.save_model(args.output_dir) - proc.save_pretrained(args.output_dir) - json.dump( - trainer.state.log_history, - open(os.path.join(args.output_dir, "log_history.json"), "w"), - ) - print( - f"DONE best_checkpoint={trainer.state.best_model_checkpoint} " - f"best_eval_loss={trainer.state.best_metric}", - flush=True, - ) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/gemma4/finetune/verify_fix.py b/extras/asr-services/providers/gemma4/finetune/verify_fix.py deleted file mode 100644 index 2c6df283..00000000 --- a/extras/asr-services/providers/gemma4/finetune/verify_fix.py +++ /dev/null @@ -1,53 +0,0 @@ -"""Verify enable_thinking=False + no_repeat_ngram fixes the bad clips.""" - -import glob -import os -import sys - -sys.path.insert(0, "/home/gemma4ft") -import pyarrow.parquet as pq -from bench_coshe_12b import SR, VERBATIM_PROMPT, Model, decode_audio - -DATA = "/home/coshe-data/data" -TARGETS = [ - "audio_1400.wav", - "audio_299.wav", - "audio_160.wav", - "audio_1404.wav", - "audio_176.wav", - "audio_12.wav", -] - -want = {n: None for n in TARGETS} -for s in sorted(glob.glob(os.path.join(DATA, "eval-*.parquet"))): - t = pq.read_table(s, columns=["audio_file_name", "transcription", "audio"]) - for nm, tr, au in zip( - t.column(0).to_pylist(), t.column(1).to_pylist(), t.column(2).to_pylist() - ): - if nm in want and want[nm] is None and au and au.get("bytes"): - want[nm] = (decode_audio(au["bytes"]), tr) - if all(v is not None for v in want.values()): - break - -m = Model( - "google/gemma-4-12B-it", - max_new_tokens=1024, - prompt=VERBATIM_PROMPT, - greedy=True, - no_repeat_ngram_size=3, -) -print("model loaded (greedy, thinking=False, no_repeat_ngram=3)\n", flush=True) - -for nm in TARGETS: - if want[nm] is None: - print(f"{nm}: NOT FOUND") - continue - audio, gt = want[nm] - hyp = m.transcribe(audio) - print( - f"=== {nm} dur={len(audio)/SR:.0f}s hyp_words={len(hyp.split())} ===", - flush=True, - ) - print("GT :", gt[:200].replace(chr(10), " "), flush=True) - print("HYP:", hyp[:200].replace(chr(10), " ") if hyp else "(EMPTY)", flush=True) - print(flush=True) diff --git a/extras/asr-services/providers/gemma4/finetune/vm_exec.py b/extras/asr-services/providers/gemma4/finetune/vm_exec.py deleted file mode 100644 index 64576526..00000000 --- a/extras/asr-services/providers/gemma4/finetune/vm_exec.py +++ /dev/null @@ -1,90 +0,0 @@ -"""Run a shell command on the Jarvis VM via the Jupyter kernel API (SSH-free). - -SSH to the instance is currently throttled/blocked, but the Jupyter server is -reachable. This opens a kernel, executes a python snippet that shells out, streams -stdout back, and tears the kernel down. - -Usage: - JLTOKEN=... python vm_exec.py "<shell command>" -It resolves the instance URL+token via jlclient automatically. -""" - -import json -import os -import ssl -import sys -import time -import uuid - -import requests -from jlclient import jarvisclient -from jlclient.jarvisclient import * -from websocket import create_connection - -jarvisclient.token = os.environ["JLTOKEN"] -url = User.get_instances()[0].url # https://<host>/lab?token=<tok> -host = url.split("//")[1].split("/")[0] -tok = url.split("token=")[1] -base = f"https://{host}" - -cmd = sys.argv[1] -code = ( - "import subprocess;" - f"r=subprocess.run({cmd!r}, shell=True, capture_output=True, text=True);" - "print(r.stdout);" - "print(r.stderr) if r.stderr else None" -) - -r = requests.post(f"{base}/api/kernels", params={"token": tok}, timeout=30) -kid = r.json()["id"] -try: - ws = create_connection( - f"wss://{host}/api/kernels/{kid}/channels?token={tok}", - sslopt={"cert_reqs": ssl.CERT_NONE}, - timeout=30, - ) - msg_id = uuid.uuid4().hex - hdr = { - "msg_id": msg_id, - "username": "u", - "session": uuid.uuid4().hex, - "msg_type": "execute_request", - "version": "5.3", - } - ws.send( - json.dumps( - { - "header": hdr, - "parent_header": {}, - "metadata": {}, - "content": { - "code": code, - "silent": False, - "store_history": False, - "user_expressions": {}, - "allow_stdin": False, - "stop_on_error": True, - }, - "channel": "shell", - } - ) - ) - out = [] - t0 = time.time() - while time.time() - t0 < 600: - m = json.loads(ws.recv()) - if m.get("parent_header", {}).get("msg_id") != msg_id: - continue - mt = m.get("msg_type") - if mt == "stream": - out.append(m["content"]["text"]) - elif mt in ("execute_result", "display_data"): - out.append(str(m["content"]["data"].get("text/plain", ""))) - elif mt == "error": - out.append("\n".join(m["content"]["traceback"])) - elif mt == "status" and m["content"]["execution_state"] == "idle": - break - ws.close() - print("".join(out)) -finally: - requests.delete(f"{base}/api/kernels/{kid}", params={"token": tok}, timeout=30) diff --git a/extras/asr-services/providers/gemma4/finetune/window_coshe.py b/extras/asr-services/providers/gemma4/finetune/window_coshe.py deleted file mode 100644 index 3387233b..00000000 --- a/extras/asr-services/providers/gemma4/finetune/window_coshe.py +++ /dev/null @@ -1,129 +0,0 @@ -"""Window CoSHE clips into <=max-seconds segments with EXACT ground-truth targets. - -Input is `fa_words.jsonl` — a faithful forced alignment of the CoSHE ground-truth -transcript (per clip: the GT word sequence, each with start/end seconds). This is the -VibeVoice-ASR data recipe: a timestamped source gives word timings, and we map the exact -GT text onto that timeline. Because the alignment IS the GT (word sequence == GT.split()), -windowing is a lossless contiguous partition — concatenating window texts reproduces the -full transcript. - -Why window: Gemma4's audio encoder hears <=~30s (E2B hard-caps audio at ~31s / 786 tok), -and CoSHE clips run ~57s median. The prior finetune used a "first-30s + proportional -char-truncation" hack (no timestamps were available then); with real word timings we can -cut honest <=30s (audio, text) pairs and train/eval on the true windowed task. - -Cuts fall at WORD boundaries, preferring an inter-word silence (gap >= --min-gap) near the -window tail so boundaries land in pauses rather than mid-utterance. - -Out: windows.jsonl, one row per window: - {audio_file_name, win_idx, n_wins, start, end, dur, n_words, text, cut_gap} -""" - -import argparse -import json -import statistics as st -from pathlib import Path - - -def window_clip(words, max_s, min_s, min_gap): - """Partition `words` (list of {word,start,end}) into contiguous [i,j] index spans, - each spanning <= max_s seconds, preferring to end at a silence >= min_gap.""" - spans = [] - i, n = 0, len(words) - while i < n: - wstart = words[i]["start"] - j = i - while j + 1 < n and words[j + 1]["end"] - wstart <= max_s: - j += 1 - cut_gap = 0.0 - if j + 1 < n: # more words remain -> consider cutting at a pause - best_k, best_gap = j, -1.0 - for k in range(j, i, -1): - if words[k]["end"] - wstart < min_s: - break - gap = words[k + 1]["start"] - words[k]["end"] - if gap > best_gap: - best_gap, best_k = gap, k - if best_gap >= min_gap: - j, cut_gap = best_k, best_gap - spans.append((i, j, cut_gap)) - i = j + 1 - return spans - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("--fa", default="/home/ft/fa_words.jsonl") - ap.add_argument("--exclude", default="", help="comma-sep audio_file_names to drop") - ap.add_argument("--out", default="/home/coshe_windowed/windows.jsonl") - ap.add_argument("--max_seconds", type=float, default=28.0) - ap.add_argument("--min_seconds", type=float, default=5.0) - ap.add_argument("--min_gap", type=float, default=0.25) - ap.add_argument("--pad", type=float, default=0.1) - args = ap.parse_args() - - drop = {x for x in args.exclude.split(",") if x} - rows = [json.loads(l) for l in open(args.fa)] - Path(args.out).parent.mkdir(parents=True, exist_ok=True) - - out = open(args.out, "w") - n_clips = n_wins = n_dropped = 0 - durs, wins_per_clip, mismatches = [], [], 0 - for r in rows: - name = r["audio_file_name"] - if name in drop: - n_dropped += 1 - continue - words = r["words"] - if not words: - continue - spans = window_clip(words, args.max_seconds, args.min_seconds, args.min_gap) - n_clips += 1 - wins_per_clip.append(len(spans)) - # lossless check: contiguous partition reproduces the GT word sequence - flat = [words[k]["word"] for a, b, _ in spans for k in range(a, b + 1)] - if flat != [w["word"] for w in words]: - mismatches += 1 - for wi, (a, b, cut_gap) in enumerate(spans): - start = max(0.0, words[a]["start"] - args.pad) - end = words[b]["end"] + args.pad - text = " ".join(words[k]["word"] for k in range(a, b + 1)) - dur = round(end - start, 3) - durs.append(dur) - out.write( - json.dumps( - { - "audio_file_name": name, - "win_idx": wi, - "n_wins": len(spans), - "start": round(start, 3), - "end": round(end, 3), - "dur": dur, - "n_words": b - a + 1, - "text": text, - "cut_gap": round(cut_gap, 3), - }, - ensure_ascii=False, - ) - + "\n" - ) - n_wins += 1 - out.close() - - from collections import Counter - - c = Counter(wins_per_clip) - print(f"clips kept={n_clips} dropped(foreign)={n_dropped} windows={n_wins}") - print(f"partition-mismatch clips (should be 0): {mismatches}") - print(f"windows/clip: {dict(sorted(c.items()))}") - print( - f"window dur s: min {min(durs):.1f} med {st.median(durs):.1f} " - f"mean {st.mean(durs):.1f} max {max(durs):.1f} p95 {sorted(durs)[int(len(durs)*0.95)]:.1f}" - ) - over = sum(1 for d in durs if d > 30.0) - print(f"windows > 30s (should be ~0): {over}") - print(f"out -> {args.out}") - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/gemma4/finetune/window_target.py b/extras/asr-services/providers/gemma4/finetune/window_target.py deleted file mode 100644 index ec7a0744..00000000 --- a/extras/asr-services/providers/gemma4/finetune/window_target.py +++ /dev/null @@ -1,25 +0,0 @@ -"""Proportionally truncate transcripts to the audio window the model actually hears. - -CoSHE clips are mostly ~57s but Gemma4's audio encoder caps at 30s, and the dataset has -no word timestamps. So for a coherent "transcribe the first W seconds" task we truncate -each clip's transcript to the first `W/duration` fraction of its characters (≈ the words -spoken in the heard window, assuming roughly uniform speech rate). Clips <= W keep the -full transcript. Applied identically in training and eval so base/FT and ref/hyp stay -consistent. -""" - -import json - - -def apply_window_truncation(items, durations_path, window_seconds): - """Mutate `items` (list of {name,target,...}) in place; return it. No-op if - window_seconds is falsy.""" - if not window_seconds: - return items - dur = json.load(open(durations_path)) - for it in items: - d = dur.get(it["name"], 0.0) - if d and d > window_seconds: - n = max(1, int(len(it["target"]) * (window_seconds / d))) - it["target"] = it["target"][:n].rstrip() - return items diff --git a/extras/asr-services/providers/nemo/finetune/cut_segments.py b/extras/asr-services/providers/nemo/finetune/cut_segments.py deleted file mode 100644 index 26a3c1c8..00000000 --- a/extras/asr-services/providers/nemo/finetune/cut_segments.py +++ /dev/null @@ -1,94 +0,0 @@ -"""Slice source WAVs at the segment boundaries from segment_by_alignment.py and write a -NeMo training manifest of the short, correctly-targeted clips. - -Tool-agnostic: consumes segments.jsonl ({audio_file_name, seg_start, seg_end, text, -review}) + the full-clip WAV dir (from make_manifest.py --all). Emits one WAV per segment -and a NeMo manifest line per segment ({audio_filepath, text, duration, target_lang}). - -By default skips segments flagged review=true (low GT<->hyp match) — pass --keep-flagged -to include them. - -Usage: - python cut_segments.py --segments segments.jsonl --wav-dir /home/ft/data_full/wav \ - --out-dir /home/ft/data_seg --target-lang hi-IN -""" - -import argparse -import json -from pathlib import Path - -import soundfile as sf - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("--segments", required=True) - ap.add_argument( - "--wav-dir", - required=True, - help="dir of full-clip WAVs (stem = audio_file_name stem)", - ) - ap.add_argument("--out-dir", required=True) - ap.add_argument("--target-lang", default="hi-IN") - ap.add_argument("--min-seconds", type=float, default=0.5) - ap.add_argument("--keep-flagged", action="store_true") - args = ap.parse_args() - - out = Path(args.out_dir) - seg_wav_dir = out / "wav" - seg_wav_dir.mkdir(parents=True, exist_ok=True) - wav_dir = Path(args.wav_dir) - - # cache decoded source audio so we read each full clip once - cache: dict[str, tuple] = {} - n_seg = n_skip = 0 - with open(out / "train.json", "w") as mf: - for line in open(args.segments): - s = json.loads(line) - if s.get("review") and not args.keep_flagged: - n_skip += 1 - continue - dur = s["seg_end"] - s["seg_start"] - if dur < args.min_seconds or not s["text"].strip(): - n_skip += 1 - continue - name = s["audio_file_name"] - stem = Path(name).stem - if stem not in cache: - data, sr = sf.read(str(wav_dir / f"{stem}.wav"), dtype="float32") - if data.ndim > 1: - data = data.mean(axis=1) - cache[stem] = (data, sr) - data, sr = cache[stem] - a = max(0, int(s["seg_start"] * sr)) - b = min(len(data), int(s["seg_end"] * sr)) - if b - a < int(args.min_seconds * sr): - n_skip += 1 - continue - seg_name = f"{stem}_{int(s['seg_start']*1000):07d}_{int(s['seg_end']*1000):07d}.wav" - seg_path = seg_wav_dir / seg_name - sf.write(str(seg_path), data[a:b], sr, subtype="PCM_16") - mf.write( - json.dumps( - { - "audio_filepath": str(seg_path), - "text": s["text"], - "duration": round((b - a) / sr, 3), - "target_lang": args.target_lang, - "audio_file_name": seg_name, - "source_clip": name, - }, - ensure_ascii=False, - ) - + "\n" - ) - n_seg += 1 - print( - f"wrote {n_seg} segment WAVs + manifest, skipped {n_skip} " - f"-> {out / 'train.json'}", - flush=True, - ) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/nemo/finetune/eval_full.py b/extras/asr-services/providers/nemo/finetune/eval_full.py deleted file mode 100644 index ae1dc9cb..00000000 --- a/extras/asr-services/providers/nemo/finetune/eval_full.py +++ /dev/null @@ -1,120 +0,0 @@ -"""Evaluate a fine-tuned nemotron .nemo checkpoint on a CoSHE manifest. - -Loads a saved .nemo (restore_from), transcribes every clip in the manifest, and -writes one JSONL line per clip in the schema mlexp score_coshe.py consumes -({audio_file_name, transcription, hyp, asr_seconds, duration_s}). Score locally: - - uv run --with "jiwer,requests,tqdm" python3 \ - extras/ml-experiments/src/mlexp/evaluate/score_coshe.py \ - --result ft=ft.jsonl --result base=nemotron_auto_full.jsonl --out-dir report/ - -Resumable: clips already present in --out are skipped. - -Usage: - python eval_full.py --init /home/ft/out/warm/final.nemo \ - --manifest /home/ft/data_full/all.json --out /home/ft/results/ft_warm.jsonl -""" - -import argparse -import json -import time -from pathlib import Path - -import nemo.collections.asr as nemo_asr -import torch -from nemo.collections.asr.data.audio_to_text_lhotse_prompt_index import ( - LhotseSpeechToTextBpeDatasetWithPromptIndex, -) - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("--init", required=True, help="path to fine-tuned .nemo") - ap.add_argument("--manifest", required=True) - ap.add_argument("--out", required=True) - ap.add_argument("--target-lang", default="hi-IN") - ap.add_argument("--batch-size", type=int, default=1) - ap.add_argument( - "--max-symbols", - type=int, - default=None, - help="override RNN-T greedy max_symbols per frame (default cfg=10)", - ) - args = ap.parse_args() - - rows = [json.loads(l) for l in open(args.manifest)] - out_path = Path(args.out) - out_path.parent.mkdir(parents=True, exist_ok=True) - done = set() - if out_path.exists(): - for l in open(out_path): - try: - done.add(json.loads(l)["audio_file_name"]) - except Exception: - pass - todo = [r for r in rows if r["audio_file_name"] not in done] - print(f"{len(done)} done, {len(todo)} to eval", flush=True) - - # --init can be a local .nemo (fine-tuned) OR a HF model id (base, for comparison) - if Path(args.init).is_file(): - model = nemo_asr.models.ASRModel.restore_from(args.init) - else: - model = nemo_asr.models.ASRModel.from_pretrained(model_name=args.init) - model.eval() - - # RNN-T greedy decode caps at `max_symbols` tokens per encoder frame (default 10); - # too low truncates long / token-dense (Devanagari) segments. Bump to test. - if args.max_symbols: - from omegaconf import open_dict - - dec = model.cfg.decoding - with open_dict(dec): - dec.strategy = "greedy_batch" - dec.greedy.max_symbols = args.max_symbols - model.change_decoding_strategy(dec) - - # force the training target_lang prompt; num_workers=0 so this in-process - # override applies (worker subprocesses wouldn't see it). See train_overfit.py. - LhotseSpeechToTextBpeDatasetWithPromptIndex._get_prompt_index_for_cut = ( - lambda self, cut, _tl=args.target_lang: self._get_prompt_index(_tl) - ) - - t0 = time.time() - with open(out_path, "a") as out_f, torch.no_grad(): - for i in range(0, len(todo), args.batch_size): - batch = todo[i : i + args.batch_size] - wavs = [r["audio_filepath"] for r in batch] - t1 = time.time() - hyps = model.transcribe( - wavs, - batch_size=args.batch_size, - target_lang=args.target_lang, - num_workers=0, - verbose=False, - ) - dt = (time.time() - t1) / len(batch) - for r, h in zip(batch, hyps): - text = h.text if hasattr(h, "text") else str(h) - out_f.write( - json.dumps( - { - "audio_file_name": r["audio_file_name"], - "transcription": r["text"], - "hyp": text, - "asr_seconds": round(dt, 3), - "duration_s": r.get("duration", 0), - }, - ensure_ascii=False, - ) - + "\n" - ) - out_f.flush() - n = i + len(batch) - if n % 50 == 0 or n == len(todo): - el = time.time() - t0 - print(f" {n}/{len(todo)} ({el:.0f}s, {n/el:.2f} clip/s)", flush=True) - print(f"DONE -> {out_path}", flush=True) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/nemo/finetune/fa_align.py b/extras/asr-services/providers/nemo/finetune/fa_align.py deleted file mode 100644 index fd7c441d..00000000 --- a/extras/asr-services/providers/nemo/finetune/fa_align.py +++ /dev/null @@ -1,114 +0,0 @@ -"""FREE local forced alignment of CoSHE audio to its ground-truth transcript, on the -4090, via ctc-forced-aligner (MMS uroman CTC). Times the GROUND-TRUTH words directly — -no hyp<->GT reconciliation. - -IMPORTANT: the high-level get_word_stamps() does NOT romanize, so it silently drops -Devanagari (keeps only Latin words). We use the lower-level ONNX pipeline with -preprocess_text(romanize=True) so Hinglish (Devanagari + Latin) aligns fully. Needs -onnxruntime-gpu for CUDA. - -Emits the normalized JSONL segment_by_alignment.py consumes (words = GT words + times): - {"audio_file_name": "...", "words": [{"word": "<gt>", "start": 1.2, "end": 1.4}, ...]} - -Usage: - python fa_align.py --manifest overfit.json --out fa_words.jsonl --language hin -""" - -import argparse -import json -import os -import time -from pathlib import Path - -import onnxruntime -from ctc_forced_aligner import ( - MODEL_URL, - Tokenizer, - ensure_onnx_model, - generate_emissions, - get_alignments, - get_spans, - load_audio, - postprocess_results, - preprocess_text, -) - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("--manifest", required=True) - ap.add_argument("--out", required=True) - ap.add_argument( - "--language", - default="hin", - help="uroman iso (hin romanizes Devanagari; Latin passes through)", - ) - ap.add_argument("--batch-size", type=int, default=8) - ap.add_argument( - "--model-path", default=os.path.expanduser("~/.cache/ctc_fa/mms_fa.onnx") - ) - args = ap.parse_args() - - Path(args.model_path).parent.mkdir(parents=True, exist_ok=True) - ensure_onnx_model(args.model_path, MODEL_URL) - providers = onnxruntime.get_available_providers() - use = ( - ["CUDAExecutionProvider", "CPUExecutionProvider"] - if "CUDAExecutionProvider" in providers - else ["CPUExecutionProvider"] - ) - print(f"onnx providers: {providers} -> using {use}", flush=True) - session = onnxruntime.InferenceSession(args.model_path, providers=use) - tokenizer = Tokenizer() - - rows = [json.loads(l) for l in open(args.manifest)] - out_path = Path(args.out) - out_path.parent.mkdir(parents=True, exist_ok=True) - - t0 = time.time() - with open(out_path, "w") as out_f: - for r in rows: - audio = load_audio(r["audio_filepath"], ret_type="np") - emissions, stride = generate_emissions( - session, audio, batch_size=args.batch_size - ) - tokens_starred, text_starred = preprocess_text( - r["text"], - romanize=True, - language=args.language, - ) - segments, scores, blank = get_alignments( - emissions, tokens_starred, tokenizer - ) - spans = get_spans(tokens_starred, segments, blank) - word_ts = postprocess_results(text_starred, spans, stride, scores) - words = [ - { - "word": w["text"], - "start": round(w["start"], 3), - "end": round(w["end"], 3), - } - for w in word_ts - ] - out_f.write( - json.dumps( - { - "audio_file_name": r["audio_file_name"], - "words": words, - }, - ensure_ascii=False, - ) - + "\n" - ) - out_f.flush() - gt_n = len(r["text"].split()) - print( - f" {r['audio_file_name']}: {len(words)} aligned / {gt_n} GT words " - f"[{words[0]['start'] if words else 0:.1f}..{words[-1]['end'] if words else 0:.1f}s]", - flush=True, - ) - print(f"DONE {len(rows)} clips in {time.time()-t0:.0f}s -> {out_path}", flush=True) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/nemo/finetune/make_manifest.py b/extras/asr-services/providers/nemo/finetune/make_manifest.py deleted file mode 100644 index 9bd63c73..00000000 --- a/extras/asr-services/providers/nemo/finetune/make_manifest.py +++ /dev/null @@ -1,167 +0,0 @@ -"""Build NeMo manifests + 16 kHz mono WAVs from the CoSHE-Eval parquet shards. - -Nemotron-3.5-asr-streaming is ``EncDecRNNTBPEModelWithPrompt`` (NeMo 2.8.0rc0): a -cache-aware FastConformer-RNNT whose Lhotse dataloader reads a per-utterance -language from the manifest ``prompt_field`` (``target_lang``) and maps it through -the model's ``prompt_dictionary`` (hi-IN -> 6, en-US -> 0, ...). So every manifest -line needs ``audio_filepath``, ``text``, ``duration`` and ``target_lang``. - -This mirrors bench_coshe.decode_to_wav16k so train/eval audio is byte-identical to -the benchmark. It powers all three FT stages: - - Stage 1 (overfit smoke): --indices 0 1 (1-2 clips) - Stage 2 (full overfit): --all (1985 clips, one split) - Stage 3 (generalization): --split 0.2 --seed 0 (train/val/test jsonls) - -Usage: - python make_manifest.py --dataset /home/coshe/data --out-dir /home/ft/data \ - --target-lang hi-IN --indices 0 1 - python make_manifest.py --dataset /home/coshe/data --out-dir /home/ft/data \ - --target-lang hi-IN --split 0.2 --seed 0 -""" - -import argparse -import io -import json -import random -from pathlib import Path - -import librosa -import pyarrow.parquet as pq -import soundfile as sf - - -def decode_to_wav16k(audio_bytes: bytes, dst: str) -> float: - """Decode embedded audio bytes -> 16 kHz mono PCM16 WAV. Returns duration (s). - - Identical to bench_coshe.decode_to_wav16k so FT audio matches the benchmark. - """ - data, sr = sf.read(io.BytesIO(audio_bytes), dtype="float32", always_2d=False) - if data.ndim > 1: - data = data.mean(axis=1) - if sr != 16000: - data = librosa.resample(data, orig_sr=sr, target_sr=16000) - sr = 16000 - sf.write(dst, data, sr, subtype="PCM_16") - return len(data) / sr - - -def iter_rows(dataset: str): - """Yield (global_index, row_dict) over all eval-*.parquet shards in order.""" - shards = sorted(Path(dataset).glob("eval-*.parquet")) - if not shards: - raise SystemExit(f"No eval-*.parquet shards under {dataset}") - gi = 0 - for shard in shards: - pf = pq.ParquetFile(shard) - for batch in pf.iter_batches(batch_size=64): - for row in batch.to_pylist(): - yield gi, row - gi += 1 - - -def write_manifest(rows, wav_dir: Path, manifest_path: Path, target_lang: str) -> int: - wav_dir.mkdir(parents=True, exist_ok=True) - manifest_path.parent.mkdir(parents=True, exist_ok=True) - n = 0 - with open(manifest_path, "w") as mf: - for _gi, row in rows: - name = row["audio_file_name"] - wav_path = wav_dir / f"{Path(name).stem}.wav" - dur = decode_to_wav16k(row["audio"]["bytes"], str(wav_path)) - mf.write( - json.dumps( - { - "audio_filepath": str(wav_path), - "text": row["transcription"], - "duration": round(dur, 3), - "target_lang": target_lang, - "audio_file_name": name, - }, - ensure_ascii=False, - ) - + "\n" - ) - n += 1 - print(f" wrote {n} -> {manifest_path}", flush=True) - return n - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("--dataset", default="/home/coshe/data") - ap.add_argument("--out-dir", required=True) - ap.add_argument( - "--target-lang", - default="hi-IN", - help="prompt_dictionary key written to every manifest line", - ) - g = ap.add_mutually_exclusive_group(required=True) - g.add_argument("--indices", type=int, nargs="+", help="explicit global row indices") - g.add_argument("--names", nargs="+", help="explicit audio_file_name values") - g.add_argument("--all", action="store_true", help="every clip into one manifest") - g.add_argument( - "--split", type=float, help="held-out fraction, e.g. 0.2 (train=this frac)" - ) - ap.add_argument("--seed", type=int, default=0) - args = ap.parse_args() - - out = Path(args.out_dir) - wav_dir = out / "wav" - - if args.indices is not None: - want = set(args.indices) - rows = [(gi, r) for gi, r in iter_rows(args.dataset) if gi in want] - write_manifest(rows, wav_dir, out / "overfit.json", args.target_lang) - - elif args.names is not None: - want = set(args.names) - rows = [ - (gi, r) for gi, r in iter_rows(args.dataset) if r["audio_file_name"] in want - ] - got = {r["audio_file_name"] for _gi, r in rows} - missing = want - got - if missing: - raise SystemExit(f"names not found: {sorted(missing)}") - write_manifest(rows, wav_dir, out / "overfit.json", args.target_lang) - - elif args.all: - rows = list(iter_rows(args.dataset)) - write_manifest(rows, wav_dir, out / "all.json", args.target_lang) - - else: # --split: train = args.split fraction, then half of remainder = val, half = test - all_rows = list(iter_rows(args.dataset)) - idx = list(range(len(all_rows))) - random.Random(args.seed).shuffle(idx) - n = len(idx) - n_train = int(round(args.split * n)) - n_val = (n - n_train) // 2 - train_i = set(idx[:n_train]) - val_i = set(idx[n_train : n_train + n_val]) - test_i = set(idx[n_train + n_val :]) - print( - f"split seed={args.seed}: train={len(train_i)} val={len(val_i)} test={len(test_i)}", - flush=True, - ) - write_manifest( - [(gi, r) for gi, r in enumerate(all_rows) if gi in train_i], - wav_dir, - out / "train.json", - args.target_lang, - ) - write_manifest( - [(gi, r) for gi, r in enumerate(all_rows) if gi in val_i], - wav_dir, - out / "val.json", - args.target_lang, - ) - write_manifest( - [(gi, r) for gi, r in enumerate(all_rows) if gi in test_i], - wav_dir, - out / "test.json", - args.target_lang, - ) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/nemo/finetune/run_deepgram.py b/extras/asr-services/providers/nemo/finetune/run_deepgram.py deleted file mode 100644 index c1d2e146..00000000 --- a/extras/asr-services/providers/nemo/finetune/run_deepgram.py +++ /dev/null @@ -1,125 +0,0 @@ -"""Get word-level timestamps for CoSHE clips from Deepgram Nova-3 (multilingual), -emitting the normalized JSONL that segment_by_alignment.py consumes: - - {"audio_file_name": "...", "dg_text": "<deepgram hypothesis>", - "words": [{"word": "...", "start": 1.2, "end": 1.4}, ...]} - -Deepgram Nova-3 multilingual is the only option that advertises Hindi<->English -intra-utterance code-switching (CoSHE is exactly that) and returns precise word -timestamps. Already the project's transcription provider, so the key is in hand. -~$0.0092/min -> ~$16 for all 1985 clips (likely free under the $200 credit). - -Start small to de-risk alignment quality on real Hinglish before spending on all 1985: - DEEPGRAM_API_KEY=... python run_deepgram.py --manifest all.json \ - --out hyp_words.jsonl --limit 5 - -Then the full run (drop --limit). Resumable: clips already in --out are skipped. - -`dg_text` is kept so we also get a free Deepgram-vs-base WER on CoSHE later. -""" - -import argparse -import json -import os -import time -from pathlib import Path - -import requests - -DG_URL = "https://api.deepgram.com/v1/listen" - - -def transcribe(wav_path: str, key: str, model: str, language: str) -> dict: - """POST one WAV to Deepgram prerecorded; return its first alternative dict.""" - params = { - "model": model, - "language": language, - "punctuate": "true", - "smart_format": "false", - } - with open(wav_path, "rb") as f: - resp = requests.post( - DG_URL, - params=params, - headers={"Authorization": f"Token {key}", "Content-Type": "audio/wav"}, - data=f.read(), - timeout=120, - ) - resp.raise_for_status() - alts = resp.json()["results"]["channels"][0]["alternatives"] - return alts[0] if alts else {"transcript": "", "words": []} - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument( - "--manifest", required=True, help="all.json (audio_filepath + audio_file_name)" - ) - ap.add_argument("--out", required=True) - ap.add_argument("--model", default="nova-3") - ap.add_argument( - "--language", default="multi", help="'multi' = Nova-3 code-switching" - ) - ap.add_argument( - "--limit", type=int, default=None, help="only the first N clips (validation)" - ) - ap.add_argument("--names", nargs="+", help="only these audio_file_name values") - args = ap.parse_args() - - key = os.environ.get("DEEPGRAM_API_KEY") - if not key: - raise SystemExit("set DEEPGRAM_API_KEY") - - rows = [json.loads(l) for l in open(args.manifest)] - if args.names: - want = set(args.names) - rows = [r for r in rows if r["audio_file_name"] in want] - if args.limit: - rows = rows[: args.limit] - - out_path = Path(args.out) - out_path.parent.mkdir(parents=True, exist_ok=True) - done = set() - if out_path.exists(): - for l in open(out_path): - try: - done.add(json.loads(l)["audio_file_name"]) - except Exception: - pass - todo = [r for r in rows if r["audio_file_name"] not in done] - print( - f"{len(done)} done, {len(todo)} to transcribe (model={args.model} lang={args.language})", - flush=True, - ) - - t0 = time.time() - with open(out_path, "a") as out_f: - for i, r in enumerate(todo, 1): - try: - alt = transcribe(r["audio_filepath"], key, args.model, args.language) - except Exception as e: - print(f" ERR {r['audio_file_name']}: {str(e)[:160]}", flush=True) - continue - words = [ - {"word": w["word"], "start": w["start"], "end": w["end"]} - for w in alt.get("words", []) - ] - out_f.write( - json.dumps( - { - "audio_file_name": r["audio_file_name"], - "dg_text": alt.get("transcript", ""), - "words": words, - }, - ensure_ascii=False, - ) - + "\n" - ) - out_f.flush() - if i % 25 == 0 or i == len(todo): - print(f" {i}/{len(todo)} ({time.time()-t0:.0f}s)", flush=True) - print(f"DONE -> {out_path}", flush=True) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/nemo/finetune/segment_by_alignment.py b/extras/asr-services/providers/nemo/finetune/segment_by_alignment.py deleted file mode 100644 index be823047..00000000 --- a/extras/asr-services/providers/nemo/finetune/segment_by_alignment.py +++ /dev/null @@ -1,206 +0,0 @@ -"""Segment long CoSHE clips into <=N-second windows with CORRECT ground-truth targets, -using a timestamped ASR hypothesis as the timing source (forced-alignment substitute). - -Why: RNN-T loss is O(T*U*V) -> a 54s clip OOMs even at batch=1 (see -[[nemotron-coshe-finetune]]). The fix is short clips, but CoSHE has no word timestamps, -so we can't cut the transcript to match a shorter audio window. Solution: run a -timestamped ASR (Deepgram Nova-3 multilingual / Google Chirp), use its WORD TIMINGS to -find cut points, and map our exact GT text onto that timeline. - -This stage is TOOL-AGNOSTIC: it consumes a normalized per-clip word-timestamp JSONL: - {"audio_file_name": "...", "words": [{"word": "...", "start": 1.2, "end": 1.4}, ...]} -A thin adapter (see deepgram_words.py / chirp_words.py, TODO) converts each provider's -response into that shape. GT text comes from the CoSHE benchmark jsonl ("transcription"). - -Pipeline per clip: - 1. romanize GT words and hyp words (IndicXlit) so Devanagari-GT aligns to a hyp in - either script; - 2. sequence-align GT<->hyp (difflib) and assign each GT word an approx time (matched - words take the hyp time; unmatched GT words interpolate between matched anchors); - 3. greedily pack GT words into <=--max-seconds segments, preferring to cut at a hyp - PAUSE (word gap >= --min-gap) so boundaries fall in silence; - 4. emit one segment row per window: {audio_file_name, seg_start, seg_end, text, ...} - plus a confidence flag (match_ratio) so low-confidence segments can be reviewed. - -A later step (cut_segments.py, TODO) slices the WAVs at [seg_start, seg_end] and writes a -NeMo manifest. Nothing here is provider-specific or destructive. - -Usage: - python segment_by_alignment.py --timestamps hyp_words.jsonl \ - --gt nemotron_auto_full.jsonl --out segments.jsonl --max-seconds 18 --min-gap 0.3 -""" - -import argparse -import json -from difflib import SequenceMatcher -from pathlib import Path - -try: - from mlexp.utils.indicxlit import romanize as _romanize -except Exception: # script may run where mlexp/IndicXlit isn't importable - _romanize = None - - -def romanize_words(words: list[str]) -> list[str]: - """Lowercase romanized form per word for cross-script alignment. Falls back to a - plain lowercase when IndicXlit isn't available (Latin-only hyps still align).""" - if _romanize is None: - return [w.lower() for w in words] - # romanize the joined text once (IndicXlit is per-word internally + cached), re-split - rom = _romanize(" ".join(words)).lower().split() - # keep length aligned with input; fall back per-word if the join changed token count - if len(rom) == len(words): - return rom - return [_romanize(w).lower() for w in words] - - -def assign_times(gt_words, hyp_words): - """Return a list of (start, end) per GT word by aligning to the timed hyp words. - - Matched GT words inherit the hyp word's [start,end]; runs of unmatched GT words are - linearly interpolated between the surrounding matched anchors. Returns (times, - match_ratio) where match_ratio is the fraction of GT words that matched a hyp word. - """ - rom_gt = romanize_words([w["w"] for w in gt_words]) - rom_hyp = romanize_words([w["word"] for w in hyp_words]) - sm = SequenceMatcher(a=rom_gt, b=rom_hyp, autojunk=False) - - times = [None] * len(gt_words) - matched = 0 - for tag, i1, i2, j1, _j2 in sm.get_opcodes(): - if tag == "equal": - for k in range(i2 - i1): - h = hyp_words[j1 + k] - times[i1 + k] = (float(h["start"]), float(h["end"])) - matched += 1 - - # interpolate the None runs between anchors - n = len(times) - anchors = [(i, t) for i, t in enumerate(times) if t is not None] - if not anchors: - return None, 0.0 - # head/tail extrapolation: clamp to first/last anchor time - first_i, first_t = anchors[0] - last_i, last_t = anchors[-1] - for i in range(first_i): - times[i] = (first_t[0], first_t[0]) - for i in range(last_i + 1, n): - times[i] = (last_t[1], last_t[1]) - # interior gaps - for (ia, ta), (ib, tb) in zip(anchors, anchors[1:]): - gap = ib - ia - if gap <= 1: - continue - t0, t1 = ta[1], tb[0] - for k in range(1, gap): - frac = k / gap - t = t0 + (t1 - t0) * frac - times[ia + k] = (t, t) - return times, matched / max(1, len(gt_words)) - - -def pack_segments(gt_words, times, hyp_words, max_seconds, min_gap): - """Greedily group GT words into <=max_seconds windows, preferring cut points that - coincide with a hyp pause (a gap >= min_gap between consecutive hyp words).""" - # precompute pause times (end of a hyp word that is followed by a >=min_gap silence) - pauses = [] - for a, b in zip(hyp_words, hyp_words[1:]): - if float(b["start"]) - float(a["end"]) >= min_gap: - pauses.append((float(a["end"]) + float(b["start"])) / 2.0) - - segs = [] - cur_start_idx = 0 - seg_t0 = times[0][0] - for i in range(len(gt_words)): - cur_t1 = times[i][1] - if cur_t1 - seg_t0 >= max_seconds and i > cur_start_idx: - # find nearest pause at/before cur_t1 within this window, else cut here - window_pauses = [p for p in pauses if seg_t0 < p <= cur_t1] - cut_t = window_pauses[-1] if window_pauses else cur_t1 - segs.append((cur_start_idx, i, seg_t0, cut_t)) - cur_start_idx = i - seg_t0 = cut_t # contiguous: next segment starts where this one cut - # final segment - segs.append((cur_start_idx, len(gt_words), seg_t0, times[-1][1])) - - # merge a too-short final fragment (< 1/3 max) back into the previous segment - if len(segs) >= 2 and (segs[-1][3] - segs[-1][2]) < max_seconds / 3: - s_i, _e_i, t0, _t1 = segs[-2] - last = segs[-1] - segs[-2] = (s_i, last[1], t0, last[3]) - segs.pop() - return segs - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument( - "--timestamps", required=True, help="normalized word-timestamp JSONL" - ) - ap.add_argument( - "--gt", required=True, help="CoSHE jsonl with audio_file_name + transcription" - ) - ap.add_argument("--out", required=True) - ap.add_argument("--max-seconds", type=float, default=18.0) - ap.add_argument("--min-gap", type=float, default=0.3) - ap.add_argument( - "--min-match-ratio", - type=float, - default=0.4, - help="flag clips whose GT<->hyp match ratio is below this for review", - ) - args = ap.parse_args() - - gt = {} - for line in open(args.gt): - r = json.loads(line) - if "transcription" in r: - gt[r["audio_file_name"]] = r["transcription"] - - out_path = Path(args.out) - out_path.parent.mkdir(parents=True, exist_ok=True) - n_clip = n_seg = n_flag = 0 - with open(out_path, "w") as out_f: - for line in open(args.timestamps): - ts = json.loads(line) - name = ts["audio_file_name"] - if name not in gt or not ts.get("words"): - continue - gt_words = [{"w": w} for w in gt[name].split()] - hyp_words = ts["words"] - times, ratio = assign_times(gt_words, hyp_words) - if times is None: - continue - flagged = ratio < args.min_match_ratio - n_flag += flagged - for s_i, e_i, t0, t1 in pack_segments( - gt_words, times, hyp_words, args.max_seconds, args.min_gap - ): - text = " ".join(gt_words[k]["w"] for k in range(s_i, e_i)) - out_f.write( - json.dumps( - { - "audio_file_name": name, - "seg_start": round(t0, 3), - "seg_end": round(t1, 3), - "duration": round(t1 - t0, 3), - "text": text, - "match_ratio": round(ratio, 3), - "review": flagged, - }, - ensure_ascii=False, - ) - + "\n" - ) - n_seg += 1 - n_clip += 1 - print( - f"{n_clip} clips -> {n_seg} segments " - f"({n_seg/max(1,n_clip):.1f}/clip), {n_flag} clips flagged for review " - f"-> {out_path}", - flush=True, - ) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/nemo/finetune/train_full.py b/extras/asr-services/providers/nemo/finetune/train_full.py deleted file mode 100644 index a88931cc..00000000 --- a/extras/asr-services/providers/nemo/finetune/train_full.py +++ /dev/null @@ -1,154 +0,0 @@ -"""Stage 2 full-CoSHE overfit for nemotron-3.5-asr-streaming (NeMo RNNT-prompt). - -Memorize all 1985 CoSHE clips to <2% WER (capacity proof), via the warm->anneal -recipe proven on gemma4/qwen3 ([[gemma4_qlora_finetune]], [[qwen3_asr_coshe_overfit]]): - - warm: --model <hf id> --lr 2e-4 train until train-loss plateaus - anneal: --init <warm.nemo> --lr 3e-5 fresh optimizer -> loss collapses - -Full fine-tune (mode=full) already trains the RNNT joint's 13087-way output -projection, so it's the max-capacity setting (no separate "include_head" needed, -unlike the HF decoder-LM models). - -Runs full-length on correct full targets (CoSHE has no word timestamps, so audio -must NOT be truncated) — requires A100-80GB: one ~60s clip's RNNT joint is ~34 GB. - -Checkpoints every --save-every steps to --save-dir so the run is stop/resume/anneal --able. WER eval is a SEPARATE step (eval_full.py) — in-training transcribe competes -with the ~34 GB training footprint and risks OOM. - -Usage (A100-80GB, NeMo-main venv): - python train_full.py --manifest /home/ft/data_full/all.json \ - --lr 2e-4 --steps 200000 --batch-size 1 --save-dir /home/ft/out/warm - python train_full.py --init /home/ft/out/warm/step120000.nemo \ - --manifest /home/ft/data_full/all.json --lr 3e-5 --steps 20000 \ - --save-dir /home/ft/out/anneal -""" - -import argparse -import time -from pathlib import Path - -import lightning.pytorch as pl -import nemo.collections.asr as nemo_asr -import torch -from lightning.pytorch.callbacks import Callback -from omegaconf import open_dict - - -class PeriodicSave(Callback): - """Save a .nemo every `every` global steps (NeMo .nemo, not a Lightning ckpt).""" - - def __init__(self, save_dir: str, every: int): - self.save_dir = Path(save_dir) - self.every = every - self.save_dir.mkdir(parents=True, exist_ok=True) - - def on_train_batch_end(self, trainer, pl_module, *args, **kwargs): - step = trainer.global_step - if step > 0 and step % self.every == 0: - path = self.save_dir / f"step{step}.nemo" - pl_module.save_to(str(path)) - print(f"[ckpt] step {step} -> {path}", flush=True) - - -def build_train_cfg(model, manifest: str, batch_size: int, max_dur: float): - cfg = model.cfg.train_ds - with open_dict(cfg): - cfg.manifest_filepath = manifest - cfg.is_tarred = False - cfg.tarred_audio_filepaths = None - cfg.shard_manifests = False - cfg.use_lhotse = True - cfg.shuffle = True - cfg.batch_size = batch_size - cfg.batch_duration = None - cfg.bucketing_batch_size = None - cfg.num_buckets = 0 - cfg.max_duration = max_dur - cfg.min_duration = 0.1 - cfg.num_workers = 4 - cfg.prompt_field = "target_lang" - return cfg - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("--model", default="nvidia/nemotron-3.5-asr-streaming-0.6b") - ap.add_argument( - "--init", default=None, help="resume from a saved .nemo (fresh optimizer)" - ) - ap.add_argument("--manifest", required=True) - ap.add_argument("--steps", type=int, default=200000) - ap.add_argument("--lr", type=float, default=2e-4) - ap.add_argument("--batch-size", type=int, default=1) - ap.add_argument("--fused-batch-size", type=int, default=1) - ap.add_argument("--max-dur", type=float, default=65.0) - ap.add_argument("--mode", choices=["full", "decoder"], default="full") - ap.add_argument("--save-dir", required=True) - ap.add_argument("--save-every", type=int, default=20000) - args = ap.parse_args() - - torch.set_float32_matmul_precision("high") - - if args.init: - print(f"Resuming from {args.init} (fresh optimizer, lr={args.lr})", flush=True) - model = nemo_asr.models.ASRModel.restore_from(args.init) - else: - print(f"Loading base {args.model} (lr={args.lr})", flush=True) - model = nemo_asr.models.ASRModel.from_pretrained(model_name=args.model) - model.set_trainer(None) - - with open_dict(model.cfg.optim): - model.cfg.optim.name = "adamw" - model.cfg.optim.lr = args.lr - model.cfg.optim.weight_decay = 0.0 - model.cfg.optim.pop("sched", None) - - model.setup_training_data( - build_train_cfg(model, args.manifest, args.batch_size, args.max_dur) - ) - - if args.fused_batch_size > 0 and hasattr(model.joint, "set_fuse_loss_wer"): - model.joint.set_fuse_loss_wer(True, loss=model.loss, metric=model.wer) - model.joint.set_fused_batch_size(args.fused_batch_size) - print( - f"fused RNNT loss/wer enabled, fused_batch_size={args.fused_batch_size}", - flush=True, - ) - - if args.mode == "decoder": - for p in model.encoder.parameters(): - p.requires_grad = False - - tr = sum(p.numel() for p in model.parameters() if p.requires_grad) - tot = sum(p.numel() for p in model.parameters()) - print( - f"mode={args.mode} trainable={tr/1e6:.1f}M / {tot/1e6:.1f}M ({100*tr/tot:.1f}%)", - flush=True, - ) - - trainer = pl.Trainer( - max_steps=args.steps, - accelerator="gpu", - devices=1, - precision="bf16-mixed", - enable_checkpointing=False, - logger=False, - enable_progress_bar=True, - log_every_n_steps=50, - limit_val_batches=0, - num_sanity_val_steps=0, - callbacks=[PeriodicSave(args.save_dir, args.save_every)], - ) - model.set_trainer(trainer) - - t0 = time.time() - trainer.fit(model) - final = Path(args.save_dir) / "final.nemo" - model.save_to(str(final)) - print(f"train done in {time.time()-t0:.0f}s -> {final}", flush=True) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/nemo/finetune/train_overfit.py b/extras/asr-services/providers/nemo/finetune/train_overfit.py deleted file mode 100644 index 2c64b704..00000000 --- a/extras/asr-services/providers/nemo/finetune/train_overfit.py +++ /dev/null @@ -1,204 +0,0 @@ -"""Stage 1 overfit smoke test for nemotron-3.5-asr-streaming (NeMo RNNT-prompt). - -Goal: prove the hardware + data pipeline + NeMo training loop end-to-end by -memorizing 1-2 CoSHE clips to WER -> 0. Mirrors the gemma4/qwen3 `sample7` stage, -but the mechanics are NeMo-native (Lhotse manifest dataloader + Lightning), not -HF PEFT. - -Model: EncDecRNNTBPEModelWithPrompt (NeMo main). 24-layer Conformer encoder -(chunked_limited streaming) -> RNNT decoder+joint (13087 BPE). Per-utterance -language comes from the manifest `target_lang` field via the model's -prompt_dictionary (hi-IN -> 6). - -Tuning surface (--mode): - full : unfreeze everything (most reliable for a 1-2 clip overfit) - decoder : freeze encoder, train decoder + joint (+ prompt) only (lighter) - -Trains, then evaluates IN-PROCESS (no save/reload) so a WER->0 result can't be a -load-path artifact -- the exact trap that cost 3 runs on gemma4. - -Usage (on the GPU box, NeMo-main venv): - python train_overfit.py --model nvidia/nemotron-3.5-asr-streaming-0.6b \ - --manifest /home/ft/data/overfit.json --steps 400 --lr 1e-4 --mode full -""" - -import argparse -import json -import time - -import jiwer -import lightning.pytorch as pl # NeMo 2.x uses the `lightning` namespace -import nemo.collections.asr as nemo_asr -import torch -from nemo.collections.asr.data.audio_to_text_lhotse_prompt_index import ( - LhotseSpeechToTextBpeDatasetWithPromptIndex, -) -from omegaconf import open_dict - - -def build_train_cfg(model, manifest: str, batch_size: int, max_dur: float): - """Clone the model's train_ds cfg and point it at our tiny non-tarred manifest.""" - cfg = model.cfg.train_ds - with open_dict(cfg): - cfg.manifest_filepath = manifest - cfg.is_tarred = False - cfg.tarred_audio_filepaths = None - cfg.shard_manifests = False - cfg.use_lhotse = True - cfg.shuffle = True - cfg.batch_size = batch_size - # drop dynamic bucketing/batch_duration so a handful of clips form fixed batches - cfg.batch_duration = None - cfg.bucketing_batch_size = None - cfg.num_buckets = 0 - cfg.max_duration = max_dur # CoSHE clips run ~54s; default cap is 20 - cfg.min_duration = 0.1 - cfg.num_workers = 2 - cfg.prompt_field = "target_lang" # already the default, kept explicit - return cfg - - -def apply_freeze(model, mode: str): - if mode == "full": - for p in model.parameters(): - p.requires_grad = True - return - if mode == "decoder": - model.encoder.freeze() - for p in model.encoder.parameters(): - p.requires_grad = False - for mod in (model.decoder, model.joint): - for p in mod.parameters(): - p.requires_grad = True - return - raise SystemExit(f"unknown --mode {mode}") - - -def n_trainable(model) -> tuple[int, int]: - tot = sum(p.numel() for p in model.parameters()) - tr = sum(p.numel() for p in model.parameters() if p.requires_grad) - return tr, tot - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("--model", default="nvidia/nemotron-3.5-asr-streaming-0.6b") - ap.add_argument("--manifest", required=True) - ap.add_argument("--steps", type=int, default=400) - ap.add_argument("--lr", type=float, default=1e-4) - ap.add_argument("--batch-size", type=int, default=2) - ap.add_argument( - "--fused-batch-size", - type=int, - default=1, - help="RNNT joint fused loss/wer sub-batch; keeps the [B,T,U,V] " - "joint off-GPU on long CoSHE clips (0 disables fusion)", - ) - ap.add_argument("--max-dur", type=float, default=65.0) - ap.add_argument("--mode", choices=["full", "decoder"], default="full") - ap.add_argument("--target-lang", default="hi-IN") - ap.add_argument("--save", default=None, help="optional .nemo save path") - args = ap.parse_args() - - torch.set_float32_matmul_precision("high") - rows = [json.loads(l) for l in open(args.manifest)] - print( - f"Overfit on {len(rows)} clip(s): " f"{[r['audio_file_name'] for r in rows]}", - flush=True, - ) - - model = nemo_asr.models.ASRModel.from_pretrained(model_name=args.model) - model.set_trainer(None) - - # constant LR, no scheduler -> pure overfit signal (pretrain default lr=0.5/Noam). - # NeMo's setup_optimization does `optim['sched']['max_steps']=...` whenever a - # `sched` key is present, so the key must be REMOVED, not set to None. - with open_dict(model.cfg.optim): - model.cfg.optim.name = "adamw" - model.cfg.optim.lr = args.lr - model.cfg.optim.weight_decay = 0.0 - model.cfg.optim.pop("sched", None) - - model.setup_training_data( - build_train_cfg(model, args.manifest, args.batch_size, args.max_dur) - ) - - # RNN-T loss on long CoSHE clips (~59s -> T~737 frames, ~480 BPE tokens) would - # materialize a [B,T,U,13087] joint (~34 GB) and OOM a 40 GB GPU. Fuse loss+WER - # into the joint so it's computed in `fused_batch_size` sub-batches without ever - # holding the full tensor. - if args.fused_batch_size > 0 and hasattr(model.joint, "set_fuse_loss_wer"): - model.joint.set_fuse_loss_wer(True, loss=model.loss, metric=model.wer) - model.joint.set_fused_batch_size(args.fused_batch_size) - print( - f"fused RNNT loss/wer enabled, fused_batch_size={args.fused_batch_size}", - flush=True, - ) - - apply_freeze(model, args.mode) - tr, tot = n_trainable(model) - print( - f"mode={args.mode} trainable={tr/1e6:.1f}M / {tot/1e6:.1f}M ({100*tr/tot:.1f}%)", - flush=True, - ) - - trainer = pl.Trainer( - max_steps=args.steps, - accelerator="gpu", - devices=1, - precision="bf16-mixed", - enable_checkpointing=False, - logger=False, - enable_progress_bar=True, - log_every_n_steps=10, - limit_val_batches=0, - num_sanity_val_steps=0, - ) - model.set_trainer(trainer) - - t0 = time.time() - trainer.fit(model) - print(f"train done in {time.time()-t0:.0f}s", flush=True) - - # ---- in-process eval (no reload) ---- - model.eval() - wavs = [r["audio_filepath"] for r in rows] - refs = [r["text"] for r in rows] - - # transcribe() builds cuts with supervision.language=None. The prompt-index - # dataset's default "unified" mode reads that None language -> "Unknown prompt - # key: 'None'". Force every eval cut to our training target_lang index so the - # prompt matches training and no None lookup happens. - # num_workers=0 keeps the dataset in-process so this class override actually - # applies (worker subprocesses wouldn't see an in-process monkeypatch). - LhotseSpeechToTextBpeDatasetWithPromptIndex._get_prompt_index_for_cut = ( - lambda self, cut, _tl=args.target_lang: self._get_prompt_index(_tl) - ) - - with torch.no_grad(): - hyps = model.transcribe( - wavs, - batch_size=1, - target_lang=args.target_lang, - num_workers=0, - verbose=False, - ) - hyps = [h.text if hasattr(h, "text") else str(h) for h in hyps] - - for r, h in zip(rows, hyps): - w = jiwer.wer(r["text"], h) - print(f"\n[{r['audio_file_name']}] WER={w*100:.2f}%", flush=True) - print(f" REF: {r['text'][:160]}", flush=True) - print(f" HYP: {h[:160]}", flush=True) - corpus = jiwer.wer(refs, hyps) - print( - f"\n=== OVERFIT corpus WER = {corpus*100:.2f}% (target: ~0%) ===", flush=True - ) - - if args.save: - model.save_to(args.save) - print(f"saved -> {args.save}", flush=True) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/nemo/finetune/train_split.py b/extras/asr-services/providers/nemo/finetune/train_split.py deleted file mode 100644 index d0c53c0a..00000000 --- a/extras/asr-services/providers/nemo/finetune/train_split.py +++ /dev/null @@ -1,148 +0,0 @@ -"""Stage 3 generalization for nemotron-3.5-asr-streaming (NeMo RNNT-prompt). - -Fine-tune on a 20% split of the CoSHE SEGMENTS (leakage-free by source clip), with -val-loss early stopping, then eval the held-out test split vs base — does FT generalize? -Mirrors the gemma4 20% experiment ([[gemma4_coshe_20pct_generalization]]): very low LR, -EarlyStopping(patience), keep the best-val checkpoint. - -Unlike train_full.py (overfit, no val), this sets up validation_data + compute_eval_wer, -saves the best-val .nemo, and stops when val_wer stops improving. - -Usage: - python train_split.py --train s3_train.json --val s3_val.json \ - --lr 1e-5 --batch-size 8 --max-epochs 40 --patience 5 \ - --save-dir /home/ft/out/s3 --mode decoder -""" - -import argparse -import time -from pathlib import Path - -import lightning.pytorch as pl -import nemo.collections.asr as nemo_asr -import torch -from lightning.pytorch.callbacks import Callback, EarlyStopping -from omegaconf import open_dict - - -class BestNemoSaver(Callback): - """Save a .nemo whenever val_wer improves (NeMo .nemo, not a Lightning ckpt).""" - - def __init__(self, path: str, monitor: str = "val_wer"): - self.path = path - self.monitor = monitor - self.best = float("inf") - - def on_validation_end(self, trainer, pl_module): - v = trainer.callback_metrics.get(self.monitor) - if v is None: - return - v = float(v) - if v < self.best: - self.best = v - pl_module.save_to(self.path) - print(f"[best] {self.monitor}={v:.4f} -> saved {self.path}", flush=True) - - -def ds_cfg(model, manifest, batch_size, max_dur, shuffle): - cfg = model.cfg.train_ds - with open_dict(cfg): - cfg.manifest_filepath = manifest - cfg.is_tarred = False - cfg.tarred_audio_filepaths = None - cfg.shard_manifests = False - cfg.use_lhotse = True - cfg.shuffle = shuffle - cfg.batch_size = batch_size - cfg.batch_duration = None - cfg.bucketing_batch_size = None - cfg.num_buckets = 0 - cfg.max_duration = max_dur - cfg.min_duration = 0.1 - cfg.num_workers = 4 - cfg.prompt_field = "target_lang" - return cfg - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("--model", default="nvidia/nemotron-3.5-asr-streaming-0.6b") - ap.add_argument("--train", required=True) - ap.add_argument("--val", required=True) - ap.add_argument("--lr", type=float, default=1e-5) - ap.add_argument("--batch-size", type=int, default=8) - ap.add_argument("--fused-batch-size", type=int, default=1) - ap.add_argument("--max-dur", type=float, default=30.0) - ap.add_argument("--mode", choices=["full", "decoder"], default="decoder") - ap.add_argument("--max-epochs", type=int, default=40) - ap.add_argument("--patience", type=int, default=5) - ap.add_argument("--save-dir", required=True) - args = ap.parse_args() - - torch.set_float32_matmul_precision("high") - model = nemo_asr.models.ASRModel.from_pretrained(model_name=args.model) - model.set_trainer(None) - - # enable validation loss (off by default in the pretrain config) - with open_dict(model.cfg): - model.cfg.compute_eval_wer = True - with open_dict(model.cfg.optim): - model.cfg.optim.name = "adamw" - model.cfg.optim.lr = args.lr - model.cfg.optim.weight_decay = 0.0 - model.cfg.optim.pop("sched", None) - - model.setup_training_data( - ds_cfg(model, args.train, args.batch_size, args.max_dur, True) - ) - model.setup_validation_data( - ds_cfg(model, args.val, args.batch_size, args.max_dur, False) - ) - - if args.fused_batch_size > 0 and hasattr(model.joint, "set_fuse_loss_wer"): - model.joint.set_fuse_loss_wer(True, loss=model.loss, metric=model.wer) - model.joint.set_fused_batch_size(args.fused_batch_size) - - if args.mode == "decoder": - for p in model.encoder.parameters(): - p.requires_grad = False - - tr = sum(p.numel() for p in model.parameters() if p.requires_grad) - tot = sum(p.numel() for p in model.parameters()) - print( - f"mode={args.mode} lr={args.lr} trainable={tr/1e6:.1f}M/{tot/1e6:.1f}M", - flush=True, - ) - - save_dir = Path(args.save_dir) - save_dir.mkdir(parents=True, exist_ok=True) - best_path = str(save_dir / "best.nemo") - - trainer = pl.Trainer( - max_epochs=args.max_epochs, - accelerator="gpu", - devices=1, - precision="bf16-mixed", - enable_checkpointing=False, - logger=False, - enable_progress_bar=True, - log_every_n_steps=50, - num_sanity_val_steps=0, - check_val_every_n_epoch=1, - limit_val_batches=120, # ~960 val segs — fast, stable early-stop signal (full set decode/epoch is slow) - callbacks=[ - EarlyStopping( - monitor="val_wer", patience=args.patience, mode="min", verbose=True - ), - BestNemoSaver(best_path), - ], - ) - model.set_trainer(trainer) - - t0 = time.time() - trainer.fit(model) - print(f"train done in {time.time()-t0:.0f}s -> best {best_path}", flush=True) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/qwen3_asr/finetune/README.md b/extras/asr-services/providers/qwen3_asr/finetune/README.md deleted file mode 100644 index 46b5a58d..00000000 --- a/extras/asr-services/providers/qwen3_asr/finetune/README.md +++ /dev/null @@ -1,94 +0,0 @@ -# Qwen3-ASR LoRA fine-tuning on CoSHE (Hinglish) - -LoRA fine-tuning of `Qwen/Qwen3-ASR-0.6B` (AuT audio encoder + Qwen3 text decoder, built on -Qwen3-Omni) on the **CoSHE-Eval** Hinglish dataset (1985 clips, mixed Devanagari+Latin -code-switched transcripts). Mirrors the gemma4 finetune dir layout but uses Qwen's official SFT -pipeline + a gemma4-style anneal/early-stop loop. - -## Result — overfit goal achieved - -**Full-CoSHE corpus WER = 1.18%** (1355/1985 clips reproduced exactly) — overfit to <2% on all -1985 clips. The sample7 smoke test reproduces 7/7 clips exactly (0.00% WER). - -Qwen3-ASR is multimodal (audio→text) and natively produces **mixed Devanagari+Latin Hinglish** -(e.g. `complain कर लो technology evolve…`), with intra-sentence code-switching — confirmed -empirically on CoSHE. - -## Files - -- `coshe_infer.py` — transcribe CoSHE parquet shards with any Qwen3-ASR model → mlexp-schema JSONL - (used for the base-model + srota baselines). **Use `--max_new_tokens 2048`**: the 512/1024 - default truncates CoSHE's long transcripts and inflates WER (the deletion-heavy artifact seen on srota). -- `make_manifest.py` / `make_manifest_full.py` — build `train.jsonl` (sample7) / `train_full.jsonl` - (all 1985) in the official SFT schema: `{"audio": "/abs.wav", "text": "language None<asr_text>" + transcript}`. -- `introspect.py` — dump module names; LoRA targets the Qwen3 decoder `thinker.model.layers.*` - (q/k/v/o/gate/up/down), keeping off the frozen AuT `thinker.audio_tower`. -- `qwen3_asr_sft.py` — vendored official trainer (`QwenLM/Qwen3-ASR/finetuning`); reused for its - data/collator/`patch_outer_forward` (model.forward → `thinker.forward`) + label masking. -- `train_qwen3_until_wer.py` — LoRA training with gemma4-style anneal + loss-gated WER<target - early-stop (subset gate → full-1985 confirm). Reuses the sft collator. -- `verify.py` — load base + adapter, transcribe clips, report char-sim / exact / WER. - -## Recipe (what worked, and the gotchas) - -Runs on the Jarvis Labs A100 (torch 2.11/cu130; `pip install qwen-asr datasets jiwer peft`; note -qwen-asr pins transformers 4.57). Single GPU. - -1. **Smoke test (sample7, 7 clips):** `train_qwen3_until_wer.py --lora_r 16 --lr 2e-4 --epochs 50` - → loss → 0.0001, **7/7 exact** once eval uses `--eval_max_new_tokens 2048` (512 truncates → false - 12.6% WER). -2. **Full CoSHE (1985):** - - **Capacity matters.** Decoder-only LoRA (even r256) **plateaus at loss ~3.0** — it can't - memorize 1985 long transcripts with a frozen output layer. `--include_head` (LoRA also on - `lm_head` + `embed_tokens`, 239.8M params / 23.5%) breaks the wall → loss descends to ~0.08. - (`Qwen3ASRForConditionalGeneration` doesn't expose `get_input_embeddings`; the script patches - it to delegate to `thinker.model.embed_tokens` so PEFT can adapt the embeddings.) - - **Warm then anneal (the gemma4 lesson).** Constant `lr 2e-4` drives loss to ~0.08 but WER - **plateaus/bounces 11–22%** (memorizes some clips, too hot to settle the rest). Restart from - the plateaued checkpoint with a **fresh optimizer at `lr 3e-5`** (`--init_adapter <ckpt>`) → - loss collapses and **WER → 1.18% in a single anneal epoch.** - - **Eval OOM.** The in-training WER eval (`wrapper.transcribe`) competes with training memory; - use a small `--eval_chunk` (4–16) + `gc.collect()/empty_cache()` before eval, and - `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`. - -## Run - -```bash -# build full manifest (decodes 1985 clips → WAVs) -python3 make_manifest_full.py - -# warm phase -python3 train_qwen3_until_wer.py --model Qwen/Qwen3-ASR-0.6B \ - --train_file train_full.jsonl --output_dir out/coshe-full-head \ - --lr 2e-4 --batch_size 2 --lora_r 256 --lora_alpha 512 --include_head \ - --eval_every 2 --eval_subset 200 --eval_chunk 16 --wer_target 0.02 --eval_max_new_tokens 2048 - -# anneal from the plateaued checkpoint (fresh optimizer, lower LR) until WER < 2% -python3 train_qwen3_until_wer.py --model Qwen/Qwen3-ASR-0.6B \ - --train_file train_full.jsonl --output_dir out/coshe-anneal \ - --init_adapter out/coshe-full-head/checkpoint-<latest> \ - --lr 3e-5 --batch_size 2 --eval_every 1 --eval_subset 200 --eval_chunk 16 --wer_target 0.02 - -# verify -python3 verify.py --adapter out/coshe-anneal --train_file train_full.jsonl --max_new_tokens 2048 -``` - -## CoSHE WER comparison (IndicXlit-romanized, via mlexp `score_coshe`, 500 common clips) - -| model | corpus WER | note | -|-------|-----------|------| -| Qwen3-ASR-0.6B + CoSHE LoRA (this) | **1.18%** | overfit goal (training-callback jiwer, full 1985) | -| Qwen3-ASR-0.6B (base, @2048 tok) | **25.16%** | base 0.6B — competitive with the 4B gemma4 | -| gemma4-E4B (base) | 26.23% | 4B base model | -| srota = qwen3-asr-0.6b-hinglish (out-of-domain FT) | 48.85% | trained on HiACC+OpenSLR, not CoSHE | - -Findings: -- **The 1.18% is a memorization/overfit result** on the training set — it shows the model + LoRA - recipe can fully fit CoSHE, not a generalization number. -- **Base Qwen3-ASR-0.6B (25.2%) ≈ gemma4-E4B (26.2%)** on CoSHE despite being ~7× smaller — a strong - base model. (It has occasional greedy repetition-collapse: per-sample WER max 621%, so corpus WER - has a heavy tail; a repetition guard would help.) -- **srota (the community Hinglish finetune) is *worse* than base (48.8% vs 25.2%)** on CoSHE — the - out-of-domain SFT on HiACC/OpenSLR hurt generalization to this distribution (plus repetition-collapse - and, in the original 1024-token run, truncation). A cautionary data point: a same-language finetune - isn't automatically better on a different corpus. diff --git a/extras/asr-services/providers/qwen3_asr/finetune/coshe_infer.py b/extras/asr-services/providers/qwen3_asr/finetune/coshe_infer.py deleted file mode 100644 index 97682b8c..00000000 --- a/extras/asr-services/providers/qwen3_asr/finetune/coshe_infer.py +++ /dev/null @@ -1,168 +0,0 @@ -"""Transcribe the CoSHE-Eval dataset with a Qwen3-ASR model and emit mlexp-schema JSONL. - -Runs on the A100 (qwen-asr installed). Reads the ``eval-*.parquet`` shards, decodes each -row's audio bytes, batch-transcribes with ``Qwen3ASRModel`` (``language=None`` → language- -agnostic), and appends one JSONL line per clip in the exact schema mlexp's ``score_coshe`` -expects: - - {"audio_file_name", "transcription" (ground truth), "hyp", "asr_seconds", - "duration_s", "raw"} - -Resumable: rows whose ``audio_file_name`` already has a non-error line in the output are -skipped (mirrors ``mlexp.runners.dataset.load_done``). WER scoring (IndicXlit romanization + -jiwer) happens locally back in extras/ml-experiments, not here — this script only produces -hypotheses. -""" - -import argparse -import io -import json -import time -from glob import glob -from pathlib import Path - -import numpy as np -import pyarrow.parquet as pq -import soundfile as sf -from qwen_asr import Qwen3ASRModel - - -def load_done(out_path: Path) -> set[str]: - done: set[str] = set() - if not out_path.exists(): - return done - with open(out_path) as f: - for line in f: - try: - rec = json.loads(line) - except Exception: - continue - if "error" in rec: - continue - name = rec.get("audio_file_name") - if name: - done.add(name) - return done - - -def decode_audio(raw_bytes: bytes) -> tuple[np.ndarray, int]: - """WAV/audio bytes -> (float32 mono, sr). qwen_asr handles resampling internally.""" - data, sr = sf.read(io.BytesIO(raw_bytes), dtype="float32") - if data.ndim > 1: - data = data.mean(axis=1) - return np.ascontiguousarray(data), int(sr) - - -def main(): - p = argparse.ArgumentParser() - p.add_argument( - "--model", required=True, help="HF id or local path of the Qwen3-ASR model" - ) - p.add_argument("--out", required=True, help="output JSONL path") - p.add_argument("--parquet_glob", default="/home/coshe-data/data/eval-*.parquet") - p.add_argument( - "--limit", type=int, default=0, help="cap #clips (shard order); 0 = all" - ) - p.add_argument("--batch_size", type=int, default=16) - p.add_argument("--max_new_tokens", type=int, default=1024) - args = p.parse_args() - - out_path = Path(args.out) - out_path.parent.mkdir(parents=True, exist_ok=True) - done = load_done(out_path) - print(f"Model: {args.model}", flush=True) - print(f"Out: {out_path} (resuming, {len(done)} already done)", flush=True) - - import torch - - model = Qwen3ASRModel.from_pretrained( - args.model, - dtype=torch.bfloat16, - device_map="cuda:0", - max_new_tokens=args.max_new_tokens, - max_inference_batch_size=args.batch_size, - ) - - shards = sorted(glob(args.parquet_glob)) - n_ok = n_err = n_seen = 0 - t_start = time.time() - - def flush_batch(batch, out_f): - nonlocal n_ok, n_err - if not batch: - return - audios = [b["audio"] for b in batch] - try: - t0 = time.time() - results = model.transcribe(audios, language=[None] * len(audios)) - dt = (time.time() - t0) / len(audios) - except Exception as e: # batch-level failure -> record per-clip errors - for b in batch: - out_f.write( - json.dumps( - {"audio_file_name": b["name"], "error": str(e)}, - ensure_ascii=False, - ) - + "\n" - ) - n_err += 1 - out_f.flush() - return - for b, res in zip(batch, results): - rec = { - "audio_file_name": b["name"], - "transcription": b["gt"], - "hyp": res.text, - "asr_seconds": round(dt, 2), - "duration_s": b["duration_s"], - "raw": {"text": res.text, "language": getattr(res, "language", None)}, - } - out_f.write(json.dumps(rec, ensure_ascii=False) + "\n") - n_ok += 1 - out_f.flush() - - with open(out_path, "a") as out_f: - batch = [] - for shard in shards: - pf = pq.ParquetFile(shard) - for rb in pf.iter_batches(batch_size=64): - for row in rb.to_pylist(): - if args.limit and n_seen >= args.limit: - break - name = row["audio_file_name"] - n_seen += 1 - if name in done: - continue - audio, sr = decode_audio(row["audio"]["bytes"]) - batch.append( - { - "name": name, - "gt": row["transcription"], - "audio": (audio, sr), - "duration_s": round(len(audio) / sr, 2), - } - ) - if len(batch) >= args.batch_size: - flush_batch(batch, out_f) - batch = [] - if n_ok % 64 == 0: - el = time.time() - t_start - print( - f" done={n_ok} err={n_err} seen={n_seen} " - f"({el:.0f}s, {n_ok / max(el, 1):.2f} clip/s)", - flush=True, - ) - if args.limit and n_seen >= args.limit: - break - if args.limit and n_seen >= args.limit: - break - flush_batch(batch, out_f) - - print( - f"DONE: ok={n_ok} err={n_err} seen={n_seen} elapsed={time.time() - t_start:.0f}s", - flush=True, - ) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/qwen3_asr/finetune/introspect.py b/extras/asr-services/providers/qwen3_asr/finetune/introspect.py deleted file mode 100644 index 0c1eade4..00000000 --- a/extras/asr-services/providers/qwen3_asr/finetune/introspect.py +++ /dev/null @@ -1,62 +0,0 @@ -"""Dump Qwen3-ASR module structure to choose LoRA target modules (runs on the A100). - -Prints the top-level layout of ``model.thinker`` and the unique nn.Linear *leaf names*, -split into "audio/encoder" vs "decoder/text" so we can scope the LoRA target regex to the -Qwen3 text decoder and keep it off the AuT audio encoder (mirrors the gemma4 approach of -adapting the text decoder only). - -Loads on CPU (``device_map=None``, no .cuda()) so it can run while the GPU is busy. -""" - -import argparse -import re - -import torch -from qwen_asr import Qwen3ASRModel - - -def main(): - p = argparse.ArgumentParser() - p.add_argument("--model", default="Qwen/Qwen3-ASR-0.6B") - args = p.parse_args() - - wrapper = Qwen3ASRModel.from_pretrained( - args.model, dtype=torch.bfloat16, device_map=None - ) - model = wrapper.model - print("=== model class ===", model.__class__.__name__) - print("=== thinker children ===") - for n, _ in model.thinker.named_children(): - print(" ", n) - - # Collect every nn.Linear's full path; classify by which subtree it lives in. - audio_kw = ("audio", "encoder", "conv", "merger", "proj_audio") - lin_paths = [] - for name, mod in model.named_modules(): - if mod.__class__.__name__.endswith("Linear"): - lin_paths.append(name) - - def is_audio(path: str) -> bool: - return any(k in path.lower() for k in audio_kw) - - decoder = sorted({p.split(".")[-1] for p in lin_paths if not is_audio(p)}) - audio = sorted({p.split(".")[-1] for p in lin_paths if is_audio(p)}) - print("\n=== Linear leaf names: DECODER/text subtree ===", decoder) - print("=== Linear leaf names: AUDIO/encoder subtree ===", audio) - - # Show a few representative full paths so the regex can be anchored correctly. - print("\n=== sample decoder Linear paths ===") - for p_ in [x for x in lin_paths if not is_audio(x)][:8]: - print(" ", p_) - print("=== sample audio Linear paths ===") - for p_ in [x for x in lin_paths if is_audio(x)][:8]: - print(" ", p_) - - # Candidate regex (q/k/v/o/gate/up/down on the decoder layers, excluding audio). - cand = r"^(?!.*(audio|encoder)).*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)$" - n_match = sum(1 for p_ in lin_paths if re.search(cand, p_)) - print(f"\n=== candidate regex matches {n_match} Linear modules ===\n {cand}") - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/qwen3_asr/finetune/make_manifest.py b/extras/asr-services/providers/qwen3_asr/finetune/make_manifest.py deleted file mode 100644 index 8a826f73..00000000 --- a/extras/asr-services/providers/qwen3_asr/finetune/make_manifest.py +++ /dev/null @@ -1,80 +0,0 @@ -"""Build the sample7 overfit train.jsonl for Qwen3-ASR SFT (runs on the A100). - -Extracts the 7 CoSHE sample7 clips from the parquet shards by ``audio_file_name``, writes -each as a 16 kHz mono WAV, and emits a JSONL in the schema the official ``qwen3_asr_sft.py`` -expects: - - {"audio": "/abs/path.wav", "text": "language None<asr_text>" + <ground-truth transcript>} - -The ``language None`` prefix matches how we evaluate (``transcribe(language=None)``); WER is -computed on the parsed ``<asr_text>`` payload, so the language token itself is irrelevant. -""" - -import argparse -import io -import json -from glob import glob -from pathlib import Path - -import numpy as np -import pyarrow.parquet as pq -import soundfile as sf - -SAMPLE7 = [ - "audio_13.wav", - "audio_278.wav", - "audio_1320.wav", - "audio_58.wav", - "audio_1197.wav", - "audio_1149.wav", - "audio_1059.wav", -] - - -def main(): - p = argparse.ArgumentParser() - p.add_argument("--parquet_glob", default="/home/coshe-data/data/eval-*.parquet") - p.add_argument("--audio_dir", default="/home/qwen3ft/sample7_audio") - p.add_argument("--out", default="/home/qwen3ft/train.jsonl") - p.add_argument("--prefix", default="language None<asr_text>") - args = p.parse_args() - - audio_dir = Path(args.audio_dir) - audio_dir.mkdir(parents=True, exist_ok=True) - want = set(SAMPLE7) - found: dict[str, dict] = {} - - for shard in sorted(glob(args.parquet_glob)): - if len(found) == len(want): - break - pf = pq.ParquetFile(shard) - for rb in pf.iter_batches(batch_size=64): - for row in rb.to_pylist(): - name = row["audio_file_name"] - if name not in want or name in found: - continue - data, sr = sf.read(io.BytesIO(row["audio"]["bytes"]), dtype="float32") - if data.ndim > 1: - data = data.mean(axis=1) - wav_path = audio_dir / name - sf.write(str(wav_path), np.ascontiguousarray(data), sr) - found[name] = { - "audio": str(wav_path), - "text": args.prefix + row["transcription"], - } - if len(found) == len(want): - break - - missing = want - set(found) - if missing: - raise SystemExit(f"Missing clips not found in parquet: {sorted(missing)}") - - with open(args.out, "w") as f: - for name in SAMPLE7: - f.write(json.dumps(found[name], ensure_ascii=False) + "\n") - print(f"Wrote {len(found)} examples -> {args.out}") - print(f"Audio in {audio_dir}") - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/qwen3_asr/finetune/make_manifest_full.py b/extras/asr-services/providers/qwen3_asr/finetune/make_manifest_full.py deleted file mode 100644 index 0a2b38d8..00000000 --- a/extras/asr-services/providers/qwen3_asr/finetune/make_manifest_full.py +++ /dev/null @@ -1,70 +0,0 @@ -"""Build the FULL CoSHE train.jsonl for the <2% WER overfit (runs on the A100). - -Decodes every clip in the parquet shards to a 16 kHz mono WAV under ``--audio_dir`` and emits -a JSONL in the official ``qwen3_asr_sft.py`` schema: - - {"audio": "/abs/path.wav", "text": "language None<asr_text>" + <ground-truth transcript>} - -~1985 clips → a few GB of WAVs. Resumable: skips clips whose WAV already exists. Mirrors the -gemma4 CosheParquetDataset coverage (all shards, full transcripts). -""" - -import argparse -import io -import json -from glob import glob -from pathlib import Path - -import numpy as np -import pyarrow.parquet as pq -import soundfile as sf - - -def main(): - p = argparse.ArgumentParser() - p.add_argument("--parquet_glob", default="/home/coshe-data/data/eval-*.parquet") - p.add_argument("--audio_dir", default="/home/qwen3ft/coshe_audio") - p.add_argument("--out", default="/home/qwen3ft/train_full.jsonl") - p.add_argument("--prefix", default="language None<asr_text>") - p.add_argument("--limit", type=int, default=0, help="cap #clips (0 = all)") - args = p.parse_args() - - audio_dir = Path(args.audio_dir) - audio_dir.mkdir(parents=True, exist_ok=True) - n = 0 - with open(args.out, "w") as out_f: - for shard in sorted(glob(args.parquet_glob)): - pf = pq.ParquetFile(shard) - for rb in pf.iter_batches(batch_size=64): - for row in rb.to_pylist(): - if args.limit and n >= args.limit: - break - name = row["audio_file_name"] - wav_path = audio_dir / name - if not wav_path.exists(): - data, sr = sf.read( - io.BytesIO(row["audio"]["bytes"]), dtype="float32" - ) - if data.ndim > 1: - data = data.mean(axis=1) - sf.write(str(wav_path), np.ascontiguousarray(data), sr) - out_f.write( - json.dumps( - { - "audio": str(wav_path), - "text": args.prefix + row["transcription"], - }, - ensure_ascii=False, - ) - + "\n" - ) - n += 1 - if args.limit and n >= args.limit: - break - if args.limit and n >= args.limit: - break - print(f"Wrote {n} examples -> {args.out}; audio in {audio_dir}") - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/qwen3_asr/finetune/qwen3_asr_sft.py b/extras/asr-services/providers/qwen3_asr/finetune/qwen3_asr_sft.py deleted file mode 100644 index fc953131..00000000 --- a/extras/asr-services/providers/qwen3_asr/finetune/qwen3_asr_sft.py +++ /dev/null @@ -1,334 +0,0 @@ -# coding=utf-8 -# Copyright 2026 The Alibaba Qwen team. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import argparse -import os -import re -import shutil -from dataclasses import dataclass -from typing import Any, Dict, List, Optional - -import librosa -import torch -from datasets import load_dataset -from qwen_asr import Qwen3ASRModel -from transformers import GenerationConfig, Trainer, TrainerCallback, TrainingArguments - - -def patch_outer_forward(model): - cls = model.__class__ - if getattr(cls, "_forward_patched", False): - return - - if not hasattr(model, "thinker") or not hasattr(model.thinker, "forward"): - raise RuntimeError( - "Cannot patch forward: model has no `.thinker.forward`. " - "Your qwen3_asr model may be incompatible." - ) - - def forward( - self, - input_ids=None, - attention_mask=None, - input_features=None, - feature_attention_mask=None, - labels=None, - **kwargs, - ): - return self.thinker.forward( - input_ids=input_ids, - attention_mask=attention_mask, - input_features=input_features, - feature_attention_mask=feature_attention_mask, - labels=labels, - **kwargs, - ) - - cls.forward = forward - cls._forward_patched = True - - -_CKPT_RE = re.compile(r"^checkpoint-(\d+)$") - - -def find_latest_checkpoint(output_dir: str) -> Optional[str]: - if not output_dir or not os.path.isdir(output_dir): - return None - best_step = None - best_path = None - for name in os.listdir(output_dir): - m = _CKPT_RE.match(name) - if not m: - continue - step = int(m.group(1)) - path = os.path.join(output_dir, name) - if os.path.isdir(path) and (best_step is None or step > best_step): - best_step = step - best_path = path - return best_path - - -def load_audio(path: str, sr: int = 16000): - wav, _ = librosa.load(path, sr=sr, mono=True) - return wav - - -def build_prefix_messages(prompt: str, audio_array): - return [ - {"role": "system", "content": prompt or ""}, - {"role": "user", "content": [{"type": "audio", "audio": audio_array}]}, - ] - - -def make_preprocess_fn_prefix_only(processor): - def _preprocess(ex: Dict[str, Any]) -> Dict[str, Any]: - prompt = ex.get("prompt", "") - dummy_audio = None - prefix_msgs = build_prefix_messages(prompt, dummy_audio) - prefix_text = processor.apply_chat_template( - [prefix_msgs], add_generation_prompt=True, tokenize=False - )[0] - return { - "prompt": prompt, - "audio": ex["audio"], - "target": ex["text"], - "prefix_text": prefix_text, - } - - return _preprocess - - -@dataclass -class DataCollatorForQwen3ASRFinetuning: - processor: Any - sampling_rate: int = 16000 - - def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]: - audio_paths = [f["audio"] for f in features] - prefix_texts = [f["prefix_text"] for f in features] - targets = [f["target"] for f in features] - - eos = self.processor.tokenizer.eos_token or "" - full_texts = [pfx + tgt + eos for pfx, tgt in zip(prefix_texts, targets)] - audios = [load_audio(p, sr=self.sampling_rate) for p in audio_paths] - - full_inputs = self.processor( - text=full_texts, - audio=audios, - return_tensors="pt", - padding=True, - truncation=False, - ) - prefix_inputs = self.processor( - text=prefix_texts, - audio=audios, - return_tensors="pt", - padding=True, - truncation=False, - ) - - prefix_lens = prefix_inputs["attention_mask"].sum(dim=1).tolist() - labels = full_inputs["input_ids"].clone() - for i, pl in enumerate(prefix_lens): - labels[i, :pl] = -100 - - pad_id = self.processor.tokenizer.pad_token_id - if pad_id is not None: - labels[labels == pad_id] = -100 - - full_inputs["labels"] = labels - return full_inputs - - -class CastFloatInputsTrainer(Trainer): - def _prepare_inputs(self, inputs): - inputs = super()._prepare_inputs(inputs) - model_dtype = getattr(self.model, "dtype", None) - if model_dtype is not None: - for k, v in list(inputs.items()): - if torch.is_tensor(v) and v.is_floating_point(): - inputs[k] = v.to(dtype=model_dtype) - return inputs - - -def copy_required_hf_files_for_qwen_asr(src_dir: str, dst_dir: str): - os.makedirs(dst_dir, exist_ok=True) - required = [ - "config.json", - "generation_config.json", - "preprocessor_config.json", - "processor_config.json", - "tokenizer_config.json", - "tokenizer.json", - "special_tokens_map.json", - "chat_template.json", - "merges.txt", - "vocab.json", - ] - for fn in required: - src = os.path.join(src_dir, fn) - if os.path.exists(src): - shutil.copy2(src, os.path.join(dst_dir, fn)) - - -class MakeEveryCheckpointInferableCallback(TrainerCallback): - def __init__(self, base_model_path: str): - self.base_model_path = base_model_path - - def on_save(self, args: TrainingArguments, state, control, **kwargs): - if args.process_index != 0: - return control - - ckpt_dir = os.path.join(args.output_dir, f"checkpoint-{state.global_step}") - if not os.path.isdir(ckpt_dir): - ckpt_dir = kwargs.get("checkpoint", ckpt_dir) - - copy_required_hf_files_for_qwen_asr(self.base_model_path, ckpt_dir) - return control - - -def parse_args(): - p = argparse.ArgumentParser("Qwen3-ASR Finetuning") - - # Paths - p.add_argument("--model_path", type=str, default="Qwen/Qwen3-ASR-1.7B") - p.add_argument("--train_file", type=str, default="train.jsonl") - p.add_argument("--eval_file", type=str, default="") - p.add_argument("--output_dir", type=str, default="./qwen3-asr-finetuning-out") - - # Audio - p.add_argument("--sr", type=int, default=16000) - - # Train hyper-params - p.add_argument("--batch_size", type=int, default=32) - p.add_argument("--grad_acc", type=int, default=4) - p.add_argument("--lr", type=float, default=2e-5) - p.add_argument("--epochs", type=float, default=1) - p.add_argument("--log_steps", type=int, default=10) - p.add_argument("--lr_scheduler_type", type=str, default="linear") - p.add_argument("--warmup_ratio", type=float, default=0.02) - - # DataLoader - p.add_argument("--num_workers", type=int, default=4) - p.add_argument("--pin_memory", type=int, default=1) - p.add_argument("--persistent_workers", type=int, default=1) - p.add_argument("--prefetch_factor", type=int, default=2) - - # Save - p.add_argument("--save_strategy", type=str, default="steps") - p.add_argument("--save_steps", type=int, default=200) - p.add_argument("--save_total_limit", type=int, default=5) - - # Resume - p.add_argument("--resume_from", type=str, default="") - p.add_argument("--resume", type=int, default=0) - - return p.parse_args() - - -def main(): - args_cli = parse_args() - - if not args_cli.train_file: - raise ValueError( - "TRAIN_FILE is required (json/jsonl). Needs fields: audio, text, optional prompt" - ) - - use_bf16 = torch.cuda.is_available() and torch.cuda.get_device_capability(0)[0] >= 8 - asr_wrapper = Qwen3ASRModel.from_pretrained( - args_cli.model_path, - dtype=torch.bfloat16 if use_bf16 else torch.float16, - device_map=None, - ) - model = asr_wrapper.model - processor = asr_wrapper.processor - - patch_outer_forward(model) - model.generation_config = GenerationConfig.from_model_config(model.config) - - raw_ds = load_dataset( - "json", - data_files={ - "train": args_cli.train_file, - **({"validation": args_cli.eval_file} if args_cli.eval_file else {}), - }, - ) - ds = raw_ds.map(make_preprocess_fn_prefix_only(processor), num_proc=1) - - keep = {"prompt", "audio", "target", "prefix_text"} - for split in ds.keys(): - drop = [c for c in ds[split].column_names if c not in keep] - if drop: - ds[split] = ds[split].remove_columns(drop) - - collator = DataCollatorForQwen3ASRFinetuning( - processor=processor, sampling_rate=args_cli.sr - ) - - training_args = TrainingArguments( - output_dir=args_cli.output_dir, - per_device_train_batch_size=args_cli.batch_size, - gradient_accumulation_steps=args_cli.grad_acc, - learning_rate=args_cli.lr, - num_train_epochs=args_cli.epochs, - logging_steps=args_cli.log_steps, - lr_scheduler_type=args_cli.lr_scheduler_type, - warmup_ratio=args_cli.warmup_ratio, - dataloader_num_workers=args_cli.num_workers, - dataloader_pin_memory=(args_cli.pin_memory == 1), - dataloader_persistent_workers=(args_cli.persistent_workers == 1), - dataloader_prefetch_factor=( - args_cli.prefetch_factor if args_cli.num_workers > 0 else None - ), - save_strategy=args_cli.save_strategy, - save_steps=args_cli.save_steps, - save_total_limit=args_cli.save_total_limit, - save_safetensors=True, - eval_strategy="steps", - eval_steps=args_cli.save_steps, - do_eval=bool(args_cli.eval_file), - bf16=use_bf16, - fp16=not use_bf16, - ddp_find_unused_parameters=False, - remove_unused_columns=False, - report_to="none", - ) - - trainer = CastFloatInputsTrainer( - model=model, - args=training_args, - train_dataset=ds["train"], - eval_dataset=ds.get("validation", None), - data_collator=collator, - tokenizer=processor.tokenizer, - callbacks=[ - MakeEveryCheckpointInferableCallback(base_model_path=args_cli.model_path) - ], - ) - - resume_from = (args_cli.resume_from or "").strip() - if not resume_from and args_cli.resume == 1: - resume_from = find_latest_checkpoint(training_args.output_dir) or "" - - if resume_from: - if trainer.args.process_index == 0: - print(f"[resume] resume_from_checkpoint = {resume_from}") - trainer.train(resume_from_checkpoint=resume_from) - else: - trainer.train() - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/qwen3_asr/finetune/train_qwen3_until_wer.py b/extras/asr-services/providers/qwen3_asr/finetune/train_qwen3_until_wer.py deleted file mode 100644 index 7c5c89e8..00000000 --- a/extras/asr-services/providers/qwen3_asr/finetune/train_qwen3_until_wer.py +++ /dev/null @@ -1,314 +0,0 @@ -"""LoRA overfit of Qwen3-ASR on CoSHE sample7, with gemma4's anneal + WER<target early-stop. - -Reuses the official ``qwen3_asr_sft.py`` data/forward/label pipeline (prefix-masked target, -audio-path collator, ``patch_outer_forward`` → ``thinker.forward``) but: - - trains **LoRA on the Qwen3 text decoder only** (``thinker.model.layers.*`` q/k/v/o/gate/up/down), - AuT audio encoder (``thinker.audio_tower``) frozen — mirrors the gemma4 text-decoder-only recipe; - - runs `lr_scheduler_type="constant"` and supports a fresh-optimizer **anneal restart** via - ``--init_adapter`` (warm at high LR, then re-launch from the checkpoint at a lower LR to settle); - - a loss-gated ``WERStopCallback`` evaluates corpus WER on the 7 clips each ``--eval_every`` epochs - and stops at ``--wer_target`` (default 2%). transformers 4.57 has no use_cache bug, so eval uses - normal cached generation via ``wrapper.transcribe``. -""" - -import argparse -import time - -import jiwer -import torch -from peft import LoraConfig, PeftModel, get_peft_model -from qwen3_asr_sft import ( - DataCollatorForQwen3ASRFinetuning, - find_latest_checkpoint, - make_preprocess_fn_prefix_only, - patch_outer_forward, -) -from transformers import GenerationConfig, Trainer, TrainerCallback, TrainingArguments - -# LoRA on the Qwen3 text decoder only (excludes thinker.audio_tower.* and lm_head). -LORA_TARGETS = r".*thinker\.model\.layers\..*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" -# With --include_head, also adapt the output head + input embeddings for more memorization -# capacity on the full dataset (gemma4 needed this to drive full-CoSHE WER below 2%). -LORA_TARGETS_HEAD = ( - r"(.*thinker\.model\.layers\..*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)" - r"|.*thinker\.model\.embed_tokens|.*thinker\.lm_head)" -) - -_ASR_TAG = "<asr_text>" -_WER_NORM = jiwer.Compose( - [ - jiwer.ToLowerCase(), - jiwer.RemovePunctuation(), - jiwer.RemoveMultipleSpaces(), - jiwer.Strip(), - jiwer.ReduceToListOfListOfWords(), - ] -) - - -def strip_target(text: str) -> str: - """'language None<asr_text>foo' -> 'foo' (the payload we compare against).""" - return text.split(_ASR_TAG, 1)[1] if _ASR_TAG in text else text - - -def corpus_wer(refs, hyps): - return jiwer.wer( - refs, hyps, reference_transform=_WER_NORM, hypothesis_transform=_WER_NORM - ) - - -def char_sim(a: str, b: str) -> float: - from difflib import SequenceMatcher - - return SequenceMatcher(None, a, b).ratio() - - -@torch.inference_mode() -def eval_clips(wrapper, items, chunk=4): - """Transcribe clips (language=None), chunked, and return (wer, exact, refs, hyps).""" - refs = [strip_target(it["text"]) for it in items] - hyps = [] - for s in range(0, len(items), chunk): - paths = [it["audio"] for it in items[s : s + chunk]] - results = wrapper.transcribe(paths, language=[None] * len(paths)) - hyps.extend(r.text for r in results) - exact = sum(1 for r, h in zip(refs, hyps) if r.strip() == h.strip()) - return corpus_wer(refs, hyps), exact, refs, hyps - - -class WERStopCallback(TrainerCallback): - """Eval corpus WER each `every` epochs (loss-gated). Subset gate -> full confirm (gemma4 style).""" - - def __init__( - self, wrapper, items, every, target, loss_gate, eval_subset=0, chunk=4 - ): - self.wrapper, self.items = wrapper, items - self.every, self.target, self.loss_gate = every, target, loss_gate - self.chunk = chunk - self.subset = ( - items[:eval_subset] if eval_subset and eval_subset < len(items) else items - ) - - def _recent_loss(self, state): - vals = [r["loss"] for r in state.log_history if "loss" in r] - return sum(vals[-30:]) / len(vals[-30:]) if vals else None - - def on_epoch_end(self, args, state, control, model=None, **kwargs): - ep = int(round(state.epoch)) - if ep == 0 or ep % self.every != 0: - return control - rl = self._recent_loss(state) - if rl is not None and rl > self.loss_gate: - print( - f"[epoch {ep}] loss={rl:.3f} > gate {self.loss_gate}; skip WER eval", - flush=True, - ) - return control - was_training = model.training - model.eval() - import gc - - gc.collect() - torch.cuda.empty_cache() # free training caches before generation - t = time.time() - wer, exact, _, _ = eval_clips(self.wrapper, self.subset, chunk=self.chunk) - print( - f"[epoch {ep}] subset({len(self.subset)}) corpus WER = {wer*100:.2f}% " - f"exact={exact}/{len(self.subset)} ({time.time()-t:.0f}s)", - flush=True, - ) - if wer < self.target and len(self.subset) < len(self.items): - fw, fex, _, _ = eval_clips(self.wrapper, self.items, chunk=self.chunk) - print( - f"[epoch {ep}] FULL({len(self.items)}) corpus WER = {fw*100:.2f}% " - f"exact={fex}/{len(self.items)}", - flush=True, - ) - wer = fw - gc.collect() - torch.cuda.empty_cache() - if was_training: - model.train() - if wer < self.target: - print( - f"[epoch {ep}] TARGET REACHED (corpus WER < {self.target*100:.0f}%). Stopping.", - flush=True, - ) - control.should_training_stop = True - return control - - -def parse_args(): - p = argparse.ArgumentParser() - p.add_argument("--model", default="Qwen/Qwen3-ASR-0.6B") - p.add_argument("--train_file", default="/home/qwen3ft/train.jsonl") - p.add_argument("--output_dir", default="/home/qwen3ft/out/sample7-overfit") - p.add_argument("--epochs", type=float, default=150.0) - p.add_argument("--lr", type=float, default=2e-4) - p.add_argument("--batch_size", type=int, default=1) - p.add_argument("--grad_acc", type=int, default=1) - p.add_argument("--lora_r", type=int, default=16) - p.add_argument("--lora_alpha", type=int, default=32) - p.add_argument( - "--include_head", action="store_true", help="also LoRA lm_head + embed_tokens" - ) - p.add_argument("--optim", default="adamw_torch") - p.add_argument( - "--init_adapter", - default="", - help="load this adapter as trainable start (fresh optimizer) for LR anneal", - ) - p.add_argument("--resume", action="store_true") - p.add_argument("--eval_every", type=int, default=2) - p.add_argument( - "--eval_subset", type=int, default=0, help="WER-gate on first N clips; 0 = all" - ) - p.add_argument( - "--eval_chunk", - type=int, - default=4, - help="transcribe batch size during WER eval (small to avoid OOM atop training memory)", - ) - p.add_argument("--wer_target", type=float, default=0.02) - p.add_argument("--loss_gate", type=float, default=0.15) - p.add_argument( - "--eval_max_new_tokens", - type=int, - default=2048, - help="generation cap for WER eval; 512 truncates long CoSHE transcripts -> false high WER", - ) - p.add_argument("--sr", type=int, default=16000) - return p.parse_args() - - -def main(): - import json - import os - - args = parse_args() - torch.manual_seed(0) - - from qwen_asr import Qwen3ASRModel - - wrapper = Qwen3ASRModel.from_pretrained( - args.model, - dtype=torch.bfloat16, - device_map=None, - max_new_tokens=args.eval_max_new_tokens, - ) - model = wrapper.model - processor = wrapper.processor - patch_outer_forward(model) - # Qwen3ASRForConditionalGeneration doesn't expose (get|set)_input_embeddings at the top - # level; PEFT needs them when LoRA targets embed_tokens (--include_head). Delegate to thinker. - cls = type(model) - if not getattr(cls, "_emb_patched", False): - cls.get_input_embeddings = lambda self: self.thinker.model.embed_tokens - cls.set_input_embeddings = lambda self, v: setattr( - self.thinker.model, "embed_tokens", v - ) - cls._emb_patched = True - model.generation_config = GenerationConfig.from_model_config(model.config) - model.to("cuda") - model.config.use_cache = False - - for pm in model.parameters(): - pm.requires_grad = False - # No gradient checkpointing: a 0.6B + LoRA fits a 40GB A100 with full activations, and - # Qwen3ASRForConditionalGeneration doesn't expose get_input_embeddings at the top level - # (so enable_input_require_grads / grad-checkpointing error). LoRA on the decoder layers - # introduces the grad-requiring path on its own with the base frozen. - - if args.init_adapter: - print(f"Loading init adapter (trainable): {args.init_adapter}", flush=True) - model = PeftModel.from_pretrained(model, args.init_adapter, is_trainable=True) - else: - model = get_peft_model( - model, - LoraConfig( - r=args.lora_r, - lora_alpha=args.lora_alpha, - lora_dropout=0.0, - bias="none", - task_type="CAUSAL_LM", - target_modules=LORA_TARGETS_HEAD if args.include_head else LORA_TARGETS, - ), - ) - model.print_trainable_parameters() - - from datasets import load_dataset - - raw = load_dataset("json", data_files={"train": args.train_file}) - ds = raw.map(make_preprocess_fn_prefix_only(processor), num_proc=1) - keep = {"prompt", "audio", "target", "prefix_text"} - drop = [c for c in ds["train"].column_names if c not in keep] - if drop: - ds["train"] = ds["train"].remove_columns(drop) - collator = DataCollatorForQwen3ASRFinetuning( - processor=processor, sampling_rate=args.sr - ) - - # raw items (audio path + full target) for the WER callback - with open(args.train_file) as f: - items = [json.loads(line) for line in f if line.strip()] - - targs = TrainingArguments( - output_dir=args.output_dir, - per_device_train_batch_size=args.batch_size, - gradient_accumulation_steps=args.grad_acc, - num_train_epochs=args.epochs, - learning_rate=args.lr, - lr_scheduler_type="constant", - warmup_steps=0, - bf16=True, - fp16=False, - logging_steps=5, - save_strategy="epoch", - save_total_limit=2, - optim=args.optim, - report_to="none", - remove_unused_columns=False, - dataloader_num_workers=0, - gradient_checkpointing=False, - ) - - cb = WERStopCallback( - wrapper, - items, - args.eval_every, - args.wer_target, - args.loss_gate, - eval_subset=args.eval_subset, - chunk=args.eval_chunk, - ) - - class BF16CastTrainer(Trainer): - def _prepare_inputs(self, inputs): - inputs = super()._prepare_inputs(inputs) - for k, v in list(inputs.items()): - if torch.is_tensor(v) and v.is_floating_point(): - inputs[k] = v.to(dtype=torch.bfloat16) - return inputs - - trainer = BF16CastTrainer( - model=model, - args=targs, - train_dataset=ds["train"], - data_collator=collator, - tokenizer=processor.tokenizer, - callbacks=[cb], - ) - - resume = ( - args.resume - and os.path.isdir(args.output_dir) - and find_latest_checkpoint(args.output_dir) is not None - ) - print(f"Starting training... (resume={resume})", flush=True) - trainer.train(resume_from_checkpoint=resume) - trainer.save_model(args.output_dir) - print("DONE (adapter saved)", flush=True) - - -if __name__ == "__main__": - main() diff --git a/extras/asr-services/providers/qwen3_asr/finetune/verify.py b/extras/asr-services/providers/qwen3_asr/finetune/verify.py deleted file mode 100644 index 5757e7a2..00000000 --- a/extras/asr-services/providers/qwen3_asr/finetune/verify.py +++ /dev/null @@ -1,76 +0,0 @@ -"""Verify a Qwen3-ASR LoRA overfit: transcribe the sample7 clips and compare to ground truth. - -Loads the base model + the trained LoRA adapter (PeftModel), transcribes each train clip with -``language=None``, and prints hyp vs ground-truth with char-similarity + an exact-match count and -corpus WER (mirrors the gemma4 ``infer.py`` smoke check). -""" - -import argparse -import json -from difflib import SequenceMatcher - -import jiwer -import torch -from peft import PeftModel -from qwen_asr import Qwen3ASRModel - -_ASR_TAG = "<asr_text>" -_NORM = jiwer.Compose( - [ - jiwer.ToLowerCase(), - jiwer.RemovePunctuation(), - jiwer.RemoveMultipleSpaces(), - jiwer.Strip(), - jiwer.ReduceToListOfListOfWords(), - ] -) - - -def main(): - p = argparse.ArgumentParser() - p.add_argument("--model", default="Qwen/Qwen3-ASR-0.6B") - p.add_argument("--adapter", default="/home/qwen3ft/out/sample7-overfit") - p.add_argument("--train_file", default="/home/qwen3ft/train.jsonl") - p.add_argument("--max_new_tokens", type=int, default=2048) - args = p.parse_args() - - wrapper = Qwen3ASRModel.from_pretrained( - args.model, - dtype=torch.bfloat16, - device_map="cuda:0", - max_new_tokens=args.max_new_tokens, - ) - if args.adapter: - wrapper.model = PeftModel.from_pretrained(wrapper.model, args.adapter) - print(f"Loaded adapter: {args.adapter}", flush=True) - - with open(args.train_file) as f: - items = [json.loads(line) for line in f if line.strip()] - - refs = [it["text"].split(_ASR_TAG, 1)[-1] for it in items] - paths = [it["audio"] for it in items] - results = wrapper.transcribe(paths, language=[None] * len(paths)) - hyps = [r.text for r in results] - - exact = 0 - for it, ref, hyp in zip(items, refs, hyps): - cs = SequenceMatcher(None, ref.strip(), hyp.strip()).ratio() - ok = ref.strip() == hyp.strip() - exact += ok - name = it["audio"].split("/")[-1] - print( - f"\n=== {name} char-sim={cs:.3f} exact={ok} len(ref)={len(ref)} len(hyp)={len(hyp)} ===" - ) - print(f" REF: {ref[:200]}") - print(f" HYP: {hyp[:200]}") - print(f" REF tail: ...{ref[-120:]}") - print(f" HYP tail: ...{hyp[-120:]}") - - wer = jiwer.wer(refs, hyps, reference_transform=_NORM, hypothesis_transform=_NORM) - print( - f"\nSUMMARY: exact={exact}/{len(items)} corpus WER={wer*100:.2f}%", flush=True - ) - - -if __name__ == "__main__": - main() From 7a9d52e11322a58433e8d17471d8f5285022f44a Mon Sep 17 00:00:00 2001 From: Ankush Malaker <43288948+AnkushMalaker@users.noreply.github.com> Date: Mon, 20 Jul 2026 10:21:58 +0000 Subject: [PATCH 15/16] chore: remove one-off cleanup artifacts --- REPOSITORY_CLEANUP.md | 98 --- backends/advanced/pyproject.toml | 1 - .../scripts/audit_annotation_dataset.py | 630 ------------------ .../scripts/backfill_cluster_embeddings.py | 41 -- .../scripts/backfill_empty_segments.py | 90 --- .../scripts/collapse_device_registry.js | 86 --- .../scripts/evaluate_memory_executor.py | 186 ------ .../advanced/scripts/regen_titles_batch.py | 123 ---- .../scripts/reprocess_conversation.py | 96 --- .../advanced/scripts/scope_blank_segments.py | 68 -- .../advanced/scripts/settle_stuck_status.py | 92 --- .../tests/scripts/graph-validation/README.md | 57 -- .../scripts/graph-validation/sample_data.py | 190 ------ .../scripts/graph-validation/setup_schema.py | 139 ---- .../graph-validation/test_conversation_doc.py | 236 ------- .../graph-validation/test_entity_graph.py | 161 ----- .../graph-validation/test_fulltext_search.py | 215 ------ .../graph-validation/test_hybrid_search.py | 199 ------ .../graph-validation/test_vector_search.py | 174 ----- .../tests/scripts/graph-validation/utils.py | 292 -------- docs/README.md | 1 - docs/test-coverage-audit.md | 214 ------ docs/testing.md | 6 +- .../asr-services/scripts/validate_nemotron.py | 143 ---- .../scripts/backfill_segment_embeddings.py | 121 ---- 25 files changed, 1 insertion(+), 3658 deletions(-) delete mode 100644 REPOSITORY_CLEANUP.md delete mode 100644 backends/advanced/scripts/audit_annotation_dataset.py delete mode 100644 backends/advanced/scripts/backfill_cluster_embeddings.py delete mode 100644 backends/advanced/scripts/backfill_empty_segments.py delete mode 100644 backends/advanced/scripts/collapse_device_registry.js delete mode 100644 backends/advanced/scripts/evaluate_memory_executor.py delete mode 100644 backends/advanced/scripts/regen_titles_batch.py delete mode 100644 backends/advanced/scripts/reprocess_conversation.py delete mode 100644 backends/advanced/scripts/scope_blank_segments.py delete mode 100644 backends/advanced/scripts/settle_stuck_status.py delete mode 100644 backends/advanced/tests/scripts/graph-validation/README.md delete mode 100644 backends/advanced/tests/scripts/graph-validation/sample_data.py delete mode 100644 backends/advanced/tests/scripts/graph-validation/setup_schema.py delete mode 100644 backends/advanced/tests/scripts/graph-validation/test_conversation_doc.py delete mode 100644 backends/advanced/tests/scripts/graph-validation/test_entity_graph.py delete mode 100644 backends/advanced/tests/scripts/graph-validation/test_fulltext_search.py delete mode 100644 backends/advanced/tests/scripts/graph-validation/test_hybrid_search.py delete mode 100644 backends/advanced/tests/scripts/graph-validation/test_vector_search.py delete mode 100644 backends/advanced/tests/scripts/graph-validation/utils.py delete mode 100644 docs/test-coverage-audit.md delete mode 100644 extras/asr-services/scripts/validate_nemotron.py delete mode 100644 extras/speaker-recognition/scripts/backfill_segment_embeddings.py diff --git a/REPOSITORY_CLEANUP.md b/REPOSITORY_CLEANUP.md deleted file mode 100644 index fae8aaad..00000000 --- a/REPOSITORY_CLEANUP.md +++ /dev/null @@ -1,98 +0,0 @@ -# Repository cleanup inventory - -This document tracks the cleanup of the working tree started on 2026-07-17. - -## Initial state - -- Branch at inventory time: `dev` (matching `origin/dev`) -- Cleanup branch: `git-cleanup` -- Modified tracked files: 84 -- Untracked files: approximately 1,304 -- Untracked disk usage: approximately 25 GB -- Staged files at inventory time: none - -## Categories - -### Project work to review and commit - -- Advanced backend features and tests -- Advanced backend web UI changes -- ASR service source, configuration, and experiments that are suitable for the repository -- Speaker-recognition service and web UI changes -- CI and test infrastructure - -### Generated outputs likely needing repository ignore rules - -- Model checkpoints, adapters, and optimizer state -- Fine-tuning output directories -- Evaluation outputs and experiment logs -- Model and dataset caches -- Coverage reports - -These generated artifacts account for most of the untracked disk usage. Individual -source files inside experiment directories must be separated from generated outputs -before ignore rules are added. - -### Potentially personal or sensitive material - -- Root-level audio, screenshots, QR codes, transcripts, and notes -- `golden_data/` -- Conversation-vault evaluation artifacts under `artifacts/` -- Voice-cloning experiments and recordings -- Personal planning and research documents - -Do not move these without explicit confirmation. - -### Local-only files - -- Environment-file backups -- Caddy and TTS configuration backups -- Local logs, debug data, and model caches - -Use `.git/info/exclude` when a file is specific to this checkout and cannot be moved -under `untracked/`. Use repository `.gitignore` only for outputs routinely generated -for multiple contributors. - -## Privacy review - -The names `ankush`, `anushpa`, and `jahnvi` occur in both tracked and untracked -content. Tracked occurrences include examples, tests, documentation, and sample-vault -data. Untracked occurrences include evaluation artifacts, experiment logs, datasets, -research notes, and tests. The evaluation vault under `artifacts/` should be treated -as sensitive until reviewed. - -PII already present in Git history is out of scope. Current tracked and proposed new -content must be reviewed before the cleanup branch is finalized. - -## Rules for this cleanup - -- Confirm before moving any file. -- Do not delete or overwrite personal artifacts during categorization. -- Commit project work in coherent groups. -- Keep generated, local-only, and potentially sensitive material out of project - commits until explicitly classified. -- Do not add backward-compatibility code as part of cleanup. - -## Progress - -- [x] Initial read-only inventory -- [x] Create cleanup branch -- [x] Commit advanced backend project work (`4d52a0c1`) -- [x] Commit ASR service project work (`667f5716`) -- [x] Commit speaker-recognition project work (`43b25ae7`) -- [ ] Classify remaining untracked content -- [ ] Apply approved ignore and relocation decisions -- [ ] Clean PII from the current repository state - -## Verification notes - -- Repository pre-commit hooks passed for all three project commits. Black and isort - reformatted newly added Python files before the successful commits. -- The broad advanced-backend pytest run could not complete because existing MongoDB - persistence tests require unavailable local infrastructure. -- `backends/advanced/tests/test_vad_analysis.py` was deliberately left untracked. It - imports `silence_gaps_from_regions`, but that implementation is absent from the - working tree and Git history, so the test is currently incomplete. -- ASR checkpoints, optimizer state, adapter outputs, caches, and logs were not - committed. -- Speaker-recognition debug data and environment backups were not committed. diff --git a/backends/advanced/pyproject.toml b/backends/advanced/pyproject.toml index 91b6fcac..dec87b44 100644 --- a/backends/advanced/pyproject.toml +++ b/backends/advanced/pyproject.toml @@ -72,7 +72,6 @@ robotframework = "^6.1.1" [tool.pytest.ini_options] minversion = "8.0" testpaths = ["tests"] -norecursedirs = ["tests/scripts"] python_files = ["test_*.py", "*_test.py"] python_classes = ["Test*", "*Tests"] python_functions = ["test_*"] diff --git a/backends/advanced/scripts/audit_annotation_dataset.py b/backends/advanced/scripts/audit_annotation_dataset.py deleted file mode 100644 index a9f2b3e5..00000000 --- a/backends/advanced/scripts/audit_annotation_dataset.py +++ /dev/null @@ -1,630 +0,0 @@ -#!/usr/bin/env python3 -"""AI-assisted audio review of a Chronicle annotation dataset (read-only). - -Listens to every clip in an exported annotation dataset with an audio-capable -model (Gemini) and flags where the machine transcript disagrees with the audio. -The point is to *review a dataset once before sending it to human annotators* — so -the reviewer sees the likely-wrong spots instead of re-transcribing from scratch. - -This tool is deliberately standalone: it reads the export ZIP + writes a sidecar -report. It makes NO database, transcript, or memory changes. Staging the findings -as editor suggestions is a separate, later step (see ``--stage-suggestions`` in the -follow-up design — intentionally not implemented here until output quality is -reviewed). - -Why Gemini and not GPT: Gemini can directly analyse audio (transcription, -timing, diarization); OpenAI's text models can only judge transcript coherence, -they cannot listen. See https://ai.google.dev/gemini-api/docs/audio - -Usage ------ -Run on 3 clips first, inspect ``ai_audit.jsonl``, then run the whole set: - - uv run --with google-genai python scripts/audit_annotation_dataset.py \ - --dataset annotation_20260628_180119_adaf --limit 3 - - uv run --with google-genai python scripts/audit_annotation_dataset.py \ - --dataset annotation_20260628_180119_adaf - -The run is resumable: clips already present in the output file are skipped, so a -crash or Ctrl-C mid-run costs nothing. Pass ``--overwrite`` to start fresh. - -Requires ``GOOGLE_API_KEY`` (read from the environment, else from -``extras/ml-experiments/.env`` or ``backends/advanced/.env``). -""" - -from __future__ import annotations - -import argparse -import io -import json -import os -import sys -import tempfile -import time -import zipfile -from pathlib import Path -from typing import Any, Optional - -# --------------------------------------------------------------------------- -# Paths / config -# --------------------------------------------------------------------------- - -# scripts/ -> backends/advanced -> repo root -BACKEND_DIR = Path(__file__).resolve().parent.parent -REPO_ROOT = BACKEND_DIR.parent.parent -DEFAULT_EXPORTS_DIR = BACKEND_DIR / "data" / "exports" - -DEFAULT_MODEL = "gemini-3.1-pro-preview" -OUTPUT_NAME = "ai_audit.jsonl" - -# Headroom so a dense clip's thinking + JSON output can't hit the cap and truncate -# the reply into invalid JSON (findings-only output is otherwise small). -MAX_OUTPUT_TOKENS = 65536 -UPLOAD_PROCESSING_TIMEOUT_S = 180 -REQUEST_TIMEOUT_MS = 300_000 -# Retry a clip whose reply doesn't parse — the pro model is occasionally -# nondeterministic about emitting strictly valid JSON. -MAX_ATTEMPTS = 3 - - -ISSUE_TYPES = [ - "mistranscription", # wrong word(s) — the transcript says X, the audio says Y - "missing_speech", # audible speech absent from the transcript - "hallucinated_text", # transcript text with no corresponding audio - "wrong_speaker", # speech attributed to the wrong speaker - "punctuation", # punctuation/casing only; wording is correct - "script_mix", # Hindi/English written in the wrong script (romanized vs Devanagari) - "other", -] - -AUDIT_PROMPT = """\ -You are auditing an automatic speech-recognition (ASR) transcript against its \ -source audio. This transcript will be used to fine-tune an STT model, so \ -verbatim accuracy matters more than readability. - -You are given: -1. The audio clip. -2. The current machine transcript, split into numbered segments. - -Listen to the ENTIRE clip and compare it to each segment. Report every place the \ -transcript disagrees with what is actually said. Focus on content errors that \ -would teach the model wrong things: wrong words, missed speech, invented text, \ -and code-switch script errors. Ignore trivial stylistic choices. - -Rules for suggested_text: -* Give the corrected VERBATIM text for that segment (what is actually spoken). -* Keep Hindi in Devanagari and English in Roman/Latin script; do not transliterate. -* Preserve filler words, false starts, repetitions exactly as spoken. -* If a segment is already correct, do NOT include it in findings. - -Current transcript segments: -{segments_block} - -Respond with a single JSON object of exactly this shape (no markdown, no prose \ -outside the JSON): -{{ - "transcript_quality": one of ["good", "minor_issues", "major_issues", "unusable"], - "overall_notes": "one or two sentences on the clip's overall transcript quality", - "findings": [ - {{ - "segment_index": <int, the segment number above>, - "original_text": "<the segment's current text>", - "suggested_text": "<corrected verbatim text>", - "issue_type": one of {issue_types}, - "confidence": <float 0.0-1.0, how sure you are the transcript is wrong>, - "explanation": "<short reason, e.g. 'audio says X not Y'>" - }} - ] -}} -If the whole transcript is faithful, return "findings": []. -""" - - -def _load_api_key() -> str: - key = os.environ.get("GOOGLE_API_KEY") - if key: - return key - # Fall back to the two .env files that may hold it. - try: - from dotenv import dotenv_values - except ImportError: - dotenv_values = None - if dotenv_values is not None: - for env_path in ( - REPO_ROOT / "extras" / "ml-experiments" / ".env", - BACKEND_DIR / ".env", - ): - if env_path.exists(): - val = dotenv_values(env_path).get("GOOGLE_API_KEY") - if val: - return val - print( - "ERROR: GOOGLE_API_KEY not found in environment or " - "extras/ml-experiments/.env / backends/advanced/.env", - file=sys.stderr, - ) - sys.exit(2) - - -def _resolve_dataset_zip(args: argparse.Namespace) -> Path: - if args.dataset_path: - p = Path(args.dataset_path) - if p.is_dir(): - p = p / "dataset.zip" - if not p.exists(): - print(f"ERROR: dataset ZIP not found: {p}", file=sys.stderr) - sys.exit(2) - return p - exports_dir = Path(args.exports_dir) if args.exports_dir else DEFAULT_EXPORTS_DIR - zip_path = exports_dir / args.dataset / "dataset.zip" - if not zip_path.exists(): - print( - f"ERROR: dataset ZIP not found: {zip_path}\n" - f" (looked under exports dir {exports_dir})", - file=sys.stderr, - ) - sys.exit(2) - return zip_path - - -def _read_manifest(zf: zipfile.ZipFile) -> list[dict[str, Any]]: - records = [] - for line in zf.read("manifest.jsonl").decode("utf-8").splitlines(): - if line.strip(): - records.append(json.loads(line)) - return records - - -def _segments_block(record: dict[str, Any]) -> str: - """Numbered segment listing for the prompt. Falls back to the whole text.""" - segments = record.get("segments") or [] - if not segments: - return f"[0] (Unknown Speaker) {record.get('text', '').strip()}" - lines = [] - for i, seg in enumerate(segments): - speaker = seg.get("identified_as") or seg.get("speaker") or "Unknown Speaker" - start = seg.get("start", 0.0) - end = seg.get("end", 0.0) - text = (seg.get("text") or "").strip() - lines.append(f"[{i}] ({speaker} | {start:.1f}-{end:.1f}s) {text}") - return "\n".join(lines) - - -def _parse_json_response(raw: str) -> dict[str, Any]: - """Parse a model JSON reply, tolerating ```json fences and stray prose.""" - text = raw.strip() - if text.startswith("```"): - text = text.split("```", 2)[1] - if text.lstrip().startswith("json"): - text = text.lstrip()[4:] - text = text.strip() - # Grab the outermost {...} if the model wrapped it in prose. - if not text.startswith("{"): - start = text.find("{") - end = text.rfind("}") - if start != -1 and end != -1: - text = text[start : end + 1] - return json.loads(text) - - -def _normalize_findings( - parsed: dict[str, Any], record: dict[str, Any] -) -> dict[str, Any]: - segments = record.get("segments") or [] - findings = [] - for raw in parsed.get("findings") or []: - if not isinstance(raw, dict): - continue - try: - seg_index = int(raw.get("segment_index", -1)) - except (TypeError, ValueError): - seg_index = -1 - issue = str(raw.get("issue_type", "other")) - if issue not in ISSUE_TYPES: - issue = "other" - try: - confidence = float(raw.get("confidence", 0.0)) - except (TypeError, ValueError): - confidence = 0.0 - confidence = max(0.0, min(1.0, confidence)) - original = raw.get("original_text") - if original is None and 0 <= seg_index < len(segments): - original = segments[seg_index].get("text", "") - findings.append( - { - "segment_index": seg_index, - "original_text": original or "", - "suggested_text": raw.get("suggested_text") or "", - "issue_type": issue, - "confidence": round(confidence, 3), - "explanation": (raw.get("explanation") or "").strip(), - } - ) - quality = str(parsed.get("transcript_quality", "unknown")) - if quality not in {"good", "minor_issues", "major_issues", "unusable"}: - quality = "unknown" - return { - "transcript_quality": quality, - "overall_notes": (parsed.get("overall_notes") or "").strip(), - "findings": findings, - } - - -def _audit_clip( - client, - types, - *, - model: str, - record: dict[str, Any], - audio_bytes: bytes, - error_dir: Optional[Path] = None, -) -> dict[str, Any]: - """Upload one clip + transcript to Gemini and return the normalized audit.""" - prompt = AUDIT_PROMPT.format( - segments_block=_segments_block(record), - issue_types=json.dumps(ISSUE_TYPES), - ) - - with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: - tmp.write(audio_bytes) - tmp_path = tmp.name - - uploaded = None - try: - uploaded = client.files.upload(file=tmp_path) - deadline = time.time() + UPLOAD_PROCESSING_TIMEOUT_S - while uploaded.state == "PROCESSING": - if time.time() > deadline: - raise TimeoutError("Gemini file stuck in PROCESSING") - time.sleep(2) - uploaded = client.files.get(name=uploaded.name) - if uploaded.state == "FAILED": - raise RuntimeError("Gemini rejected the audio upload") - - last_error: Exception | None = None - last_raw = "" - for _attempt in range(MAX_ATTEMPTS): - resp = client.models.generate_content( - model=model, - contents=[uploaded, prompt], - config=types.GenerateContentConfig( - max_output_tokens=MAX_OUTPUT_TOKENS, - response_mime_type="application/json", - ), - ) - last_raw = resp.text or "" - finish = str(resp.candidates[0].finish_reason) if resp.candidates else "?" - try: - parsed = _parse_json_response(last_raw) - except json.JSONDecodeError as exc: - last_error = RuntimeError( - f"unparseable JSON (finish_reason={finish}): {exc}" - ) - continue - audit = _normalize_findings(parsed, record) - usage = resp.usage_metadata - audit["tokens"] = { - "input": getattr(usage, "prompt_token_count", 0) or 0, - "output": getattr(usage, "candidates_token_count", 0) or 0, - "thinking": getattr(usage, "thoughts_token_count", 0) or 0, - } - return audit - - # All attempts failed to parse — persist the last raw reply for inspection - # rather than throwing it away, then surface the error. - if error_dir is not None: - error_dir.mkdir(parents=True, exist_ok=True) - (error_dir / f"{record['clip_id']}.raw.txt").write_text(last_raw) - raise last_error or RuntimeError("audit failed") - finally: - if uploaded is not None: - try: - client.files.delete(name=uploaded.name) - except Exception: - pass - try: - os.unlink(tmp_path) - except OSError: - pass - - -async def _stage_suggestions(args: argparse.Namespace) -> int: - """Read ai_audit.jsonl and stage its findings as PENDING model suggestions. - - Findings are matched to the *imported* conversations (the memory-excluded - copies created by Upload → Annotation workspace) via ``external_source_id`` - computed by the very same parser the importer used, so segment indices and - ids line up exactly. Nothing is applied to a transcript — each finding becomes - a suggestion the editor renders for accept/edit/reject. - """ - from beanie import init_beanie - - from advanced_omi_backend.database import db - from advanced_omi_backend.models.annotation import ( - Annotation, - AnnotationSource, - AnnotationStatus, - AnnotationType, - ) - from advanced_omi_backend.models.conversation import Conversation - from advanced_omi_backend.models.user import User - from advanced_omi_backend.utils.annotation_import import parse_annotation_dataset - - zip_path = _resolve_dataset_zip(args) - out_path = Path(args.out) if args.out else zip_path.parent / OUTPUT_NAME - if not out_path.exists(): - print( - f"ERROR: no audit report at {out_path}. Run the audit (without " - f"--stage-suggestions) first.", - file=sys.stderr, - ) - return 2 - - audits: dict[str, dict[str, Any]] = {} - for line in out_path.read_text().splitlines(): - if line.strip(): - row = json.loads(line) - audits[row["clip_id"]] = row - - # Same parser the import endpoint used → identical dataset_id + clip_ids. - dataset = parse_annotation_dataset(zip_path.read_bytes()) - - await init_beanie(database=db, document_models=[User, Conversation, Annotation]) - - staged = 0 - skipped_existing = 0 - skipped_conf = 0 - skipped_bounds = 0 - missing_conv = 0 - touched: set[str] = set() - - for clip in dataset.clips: - row = audits.get(clip.clip_id) - if not row: - continue - external_source_id = f"{dataset.dataset_id}:{clip.clip_id}" - conv = await Conversation.find_one( - Conversation.external_source_type == "annotation_dataset", - Conversation.external_source_id == external_source_id, - ) - if not conv: - missing_conv += 1 - continue - - active = conv.active_transcript - segments = active.segments if active and active.segments else [] - seg_count = len(segments) - - if args.replace: - await Annotation.find( - Annotation.conversation_id == conv.conversation_id, - Annotation.annotation_type == AnnotationType.TRANSCRIPT, - Annotation.source == AnnotationSource.MODEL_SUGGESTION, - Annotation.status == AnnotationStatus.PENDING, - ).delete() - - for finding in row.get("findings", []): - if float(finding.get("confidence", 0.0)) < args.min_confidence: - skipped_conf += 1 - continue - seg_idx = finding.get("segment_index") - if seg_idx is None or seg_idx < 0 or seg_idx >= seg_count: - skipped_bounds += 1 - continue - - if not args.replace: - existing = await Annotation.find_one( - Annotation.conversation_id == conv.conversation_id, - Annotation.segment_index == seg_idx, - Annotation.source == AnnotationSource.MODEL_SUGGESTION, - Annotation.status == AnnotationStatus.PENDING, - ) - if existing: - skipped_existing += 1 - continue - - annotation = Annotation( - annotation_type=AnnotationType.TRANSCRIPT, - user_id=conv.user_id, - conversation_id=conv.conversation_id, - segment_index=seg_idx, - # Truthful "before": the segment's current text, matching the - # convention used when the editor creates transcript annotations. - original_text=segments[seg_idx].text, - corrected_text=finding.get("suggested_text", ""), - source=AnnotationSource.MODEL_SUGGESTION, - status=AnnotationStatus.PENDING, - ) - await annotation.save() - staged += 1 - touched.add(conv.conversation_id) - - print( - f"Dataset: {dataset.dataset_id}\n" - f"Staged {staged} suggestion(s) across {len(touched)} conversation(s).\n" - f" skipped: {skipped_existing} already-pending, " - f"{skipped_conf} below --min-confidence={args.min_confidence}, " - f"{skipped_bounds} out-of-bounds segment_index\n" - f" {missing_conv} clip(s) had no imported conversation " - f"(import the dataset via Upload → Annotation workspace first)\n" - f"\nReview them at: https://localhost/data-audit?dataset={dataset.dataset_id}" - ) - return 0 - - -def main() -> int: - ap = argparse.ArgumentParser(description=__doc__) - src = ap.add_mutually_exclusive_group(required=True) - src.add_argument( - "--dataset", - help="Export id under data/exports/, e.g. annotation_20260628_180119_adaf", - ) - src.add_argument( - "--dataset-path", - help="Explicit path to a dataset.zip (or a directory containing one)", - ) - ap.add_argument("--model", default=DEFAULT_MODEL, help=f"default: {DEFAULT_MODEL}") - ap.add_argument( - "--limit", - type=int, - default=None, - help="Only audit the first N not-yet-audited clips (use 3 for a first look)", - ) - ap.add_argument( - "--out", - default=None, - help=f"Output path (default: {OUTPUT_NAME} beside the dataset)", - ) - ap.add_argument( - "--overwrite", - action="store_true", - help="Ignore existing output and re-audit every clip", - ) - ap.add_argument( - "--exports-dir", default=None, help="Override data/exports location" - ) - ap.add_argument( - "--stage-suggestions", - action="store_true", - help=( - "Instead of auditing, read an existing ai_audit.jsonl and stage its " - "findings as PENDING model suggestions on the imported conversations " - "so they appear in the transcript editor. Requires the dataset to have " - "been imported (Upload → Annotation workspace) first." - ), - ) - ap.add_argument( - "--min-confidence", - type=float, - default=0.0, - help="When staging, only stage findings at or above this confidence (0.0-1.0)", - ) - ap.add_argument( - "--replace", - action="store_true", - help="When staging, clear existing pending model suggestions on each " - "conversation first (else duplicate segments are skipped)", - ) - args = ap.parse_args() - - if args.stage_suggestions: - import asyncio - - return asyncio.run(_stage_suggestions(args)) - - api_key = _load_api_key() - zip_path = _resolve_dataset_zip(args) - out_path = Path(args.out) if args.out else zip_path.parent / OUTPUT_NAME - - # Resume: collect clip_ids already in the output file. - done: set[str] = set() - if out_path.exists() and not args.overwrite: - for line in out_path.read_text().splitlines(): - if line.strip(): - try: - done.add(json.loads(line)["clip_id"]) - except (json.JSONDecodeError, KeyError): - pass - elif args.overwrite and out_path.exists(): - out_path.unlink() - - archive_bytes = zip_path.read_bytes() - with zipfile.ZipFile(io.BytesIO(archive_bytes)) as zf: - records = _read_manifest(zf) - audio_cache = { - r["clip_id"]: zf.read(r["audio_path"]) - for r in records - if r.get("audio_path") in zf.namelist() - } - - pending = [r for r in records if r["clip_id"] not in done] - if args.limit is not None: - pending = pending[: args.limit] - - print( - f"Dataset: {zip_path}\n" - f"Model: {args.model}\n" - f"Clips: {len(records)} total, {len(done)} already audited, " - f"{len(pending)} to do this run\n" - f"Output: {out_path}\n", - flush=True, - ) - if not pending: - print("Nothing to do. (Use --overwrite to re-audit.)") - return 0 - - from google import genai - from google.genai import types - - client = genai.Client( - api_key=api_key, - http_options=types.HttpOptions(timeout=REQUEST_TIMEOUT_MS), - ) - - ok = 0 - failed = 0 - quality_counts: dict[str, int] = {} - total_findings = 0 - t_start = time.time() - - with out_path.open("a", encoding="utf-8") as out_f: - for i, record in enumerate(pending, start=1): - clip_id = record["clip_id"] - audio_bytes = audio_cache.get(clip_id) - label = f"[{i}/{len(pending)}] {clip_id}" - if not audio_bytes: - print(f"{label} — SKIP (audio missing in ZIP)", flush=True) - failed += 1 - continue - t0 = time.time() - try: - audit = _audit_clip( - client, - types, - model=args.model, - record=record, - audio_bytes=audio_bytes, - error_dir=out_path.parent / "ai_audit_errors", - ) - except Exception as exc: # noqa: BLE001 - report + continue - print(f"{label} — ERROR: {exc}", flush=True) - failed += 1 - continue - - row = { - "clip_id": clip_id, - "conversation_id": record.get("conversation_id"), - "conversation_title": record.get("conversation_title"), - "audio_path": record.get("audio_path"), - "duration_seconds": record.get("duration_seconds"), - "model": args.model, - **audit, - } - out_f.write(json.dumps(row, ensure_ascii=False) + "\n") - out_f.flush() - - ok += 1 - n_find = len(audit["findings"]) - total_findings += n_find - quality_counts[audit["transcript_quality"]] = ( - quality_counts.get(audit["transcript_quality"], 0) + 1 - ) - print( - f"{label} — {audit['transcript_quality']}, " - f"{n_find} finding(s), {time.time() - t0:.1f}s", - flush=True, - ) - - print( - f"\nDone in {time.time() - t_start:.0f}s. " - f"audited={ok} failed={failed} total_findings={total_findings}\n" - f"quality: {quality_counts}\n" - f"Report: {out_path}", - flush=True, - ) - return 0 if failed == 0 else 1 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/backends/advanced/scripts/backfill_cluster_embeddings.py b/backends/advanced/scripts/backfill_cluster_embeddings.py deleted file mode 100644 index 77dac1e1..00000000 --- a/backends/advanced/scripts/backfill_cluster_embeddings.py +++ /dev/null @@ -1,41 +0,0 @@ -"""One-time backfill of per-cluster speaker centroids for existing conversations. - -Embeds one centroid per diarized speaker from each conversation's already-stored segment -boundaries (no re-diarization) and saves it to -``TranscriptVersion.metadata["cluster_centroids"]``. This lets the "reprocess impact" -finder re-identify past conversations against the current gallery with pure vector math. - -GPU-bound (runs the wespeaker embedder via the speaker service). Idempotent: skips -conversations that already have stored centroids unless you flip only_missing. - -Run inside the backend container: - python3 /app/scripts/backfill_cluster_embeddings.py # all missing - python3 /app/scripts/backfill_cluster_embeddings.py 5 # first 5 (smoke test) -""" - -import asyncio -import sys - -from beanie import init_beanie - -from advanced_omi_backend.controllers.drift_controller import ( - backfill_cluster_embeddings, -) -from advanced_omi_backend.database import db -from advanced_omi_backend.models.audio_chunk import AudioChunkDocument -from advanced_omi_backend.models.conversation import Conversation -from advanced_omi_backend.models.user import User - - -async def main() -> None: - limit = int(sys.argv[1]) if len(sys.argv) > 1 else None - # AudioChunkDocument is needed because backfill reconstructs conversation audio. - await init_beanie( - database=db, document_models=[User, Conversation, AudioChunkDocument] - ) - result = await backfill_cluster_embeddings(limit=limit, only_missing=True) - print(f"Backfill result: {result}") - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/backends/advanced/scripts/backfill_empty_segments.py b/backends/advanced/scripts/backfill_empty_segments.py deleted file mode 100644 index 3c64cdd3..00000000 --- a/backends/advanced/scripts/backfill_empty_segments.py +++ /dev/null @@ -1,90 +0,0 @@ -"""Backfill renderable segments for transcript versions that have text but no segments. - -The streaming transcription path used to store a transcript version with text + words -but an EMPTY segments list when the provider didn't diarize (and speaker recognition was -off/unavailable). The web UI only renders the segment list, so those conversations showed -"No transcript segments available" despite having a transcript. - -This script finds every transcript version where `transcript` has text but `segments` is -empty, and builds a single fallback "Speaker 0" segment from the version's words (or the -full text). It is idempotent — versions that already have segments are left untouched. - -Run inside the backend/worker container: - uv run python3 scripts/backfill_empty_segments.py # dry run (default) - uv run python3 scripts/backfill_empty_segments.py --apply # write changes -""" - -import argparse -import asyncio -import sys - -from beanie import init_beanie - -from advanced_omi_backend.database import db -from advanced_omi_backend.models.conversation import Conversation -from advanced_omi_backend.models.user import User - - -def _build_fallback_segment(version) -> "Conversation.SpeakerSegment": - """Build a single full-span fallback segment from a version's words/text.""" - words = version.words or [] - start = words[0].start if words else 0.0 - end = words[-1].end if words else 0.0 - return Conversation.SpeakerSegment( - speaker="Speaker 0", - start=start, - end=end, - text=version.transcript or "", - words=list(words), - ) - - -async def main(apply: bool) -> None: - await init_beanie(database=db, document_models=[User, Conversation]) - - scanned = 0 - fixed_versions = 0 - fixed_convs = 0 - - async for conv in Conversation.find_all(): - scanned += 1 - changed = False - for version in conv.transcript_versions: - has_text = bool((version.transcript or "").strip()) - if has_text and not version.segments: - seg = _build_fallback_segment(version) - version.segments = [seg] - # Mark provenance so it's distinguishable from real diarization - version.diarization_source = version.diarization_source or None - if isinstance(version.metadata, dict): - version.metadata.setdefault( - "segments_created_by", "backfill_fallback" - ) - fixed_versions += 1 - changed = True - print( - f" conv={conv.conversation_id[:12]} version={version.version_id} " - f"-> 1 fallback segment ({len(seg.text)} chars, {len(seg.words)} words)" - ) - if changed: - fixed_convs += 1 - if apply: - await conv.save() - - print( - f"\nScanned {scanned} conversations. " - f"{'Fixed' if apply else 'Would fix'} {fixed_versions} versions " - f"across {fixed_convs} conversations." - ) - if not apply and fixed_versions: - print("Dry run — re-run with --apply to write changes.") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - "--apply", action="store_true", help="Write changes (default: dry run)" - ) - args = parser.parse_args() - asyncio.run(main(args.apply)) - sys.exit(0) diff --git a/backends/advanced/scripts/collapse_device_registry.js b/backends/advanced/scripts/collapse_device_registry.js deleted file mode 100644 index 418b48b7..00000000 --- a/backends/advanced/scripts/collapse_device_registry.js +++ /dev/null @@ -1,86 +0,0 @@ -// One-time migration: collapse the bloated `registered_clients` map. -// -// Before stable client_ids existed, every reconnect of one device minted a new -// counter-suffixed client_id (havpe, havpe-2, … havpe-286), so the registry -// accumulated hundreds of rows for a handful of physical devices. This script -// rewrites each user's `registered_clients` keyed by the STABLE client_id -// (user_suffix + sanitized device_name) — one row per device — merging the -// counter duplicates and preserving any user-set friendly name + the earliest -// first_seen / latest last_seen. -// -// IDEMPOTENT: re-running yields the same result. Run it AFTER the stable-id -// backend code is deployed (otherwise the old counter immediately re-pollutes). -// -// docker exec <mongo> mongosh chronicle --file /path/collapse_device_registry.js -// DRY_RUN: set env DRYRUN=1 (default) to preview; DRYRUN=0 to apply. - -// mongosh exposes process.env -const APPLY = (typeof process !== "undefined" && process.env && process.env.DRYRUN === "0"); - -// Mirror backend generate_client_id sanitization: -// lowercase, keep [a-z0-9-], first 10 chars. -function sanitizeDevice(d) { - return String(d || "") - .toLowerCase() - .split("") - .filter((c) => /[a-z0-9-]/.test(c)) - .join("") - .slice(0, 10); -} - -function stableClientId(userIdHex, deviceName) { - return userIdHex.slice(-6) + "-" + sanitizeDevice(deviceName); -} - -let usersTouched = 0; -let totalBefore = 0; -let totalAfter = 0; - -db.users.find({ "registered_clients": { $exists: true, $ne: {} } }).forEach((u) => { - const idHex = u._id.toString(); - const rc = u.registered_clients || {}; - const keys = Object.keys(rc); - if (keys.length === 0) return; - - const collapsed = {}; - keys.forEach((k) => { - const e = rc[k] || {}; - // Without a device_name we can't derive a stable id — keep the row as-is. - const stableId = e.device_name ? stableClientId(idHex, e.device_name) : k; - - const cur = collapsed[stableId]; - if (!cur) { - collapsed[stableId] = { - client_id: stableId, - device_name: e.device_name || null, - name: e.name || e.device_name || stableId, - first_seen: e.first_seen || e.last_seen || new Date(), - last_seen: e.last_seen || e.first_seen || new Date(), - }; - } else { - // Merge duplicates: keep a real user-set name, widen the time range. - if ((!cur.name || cur.name === cur.device_name) && e.name && e.name !== e.device_name) { - cur.name = e.name; - } - if (e.first_seen && e.first_seen < cur.first_seen) cur.first_seen = e.first_seen; - if (e.last_seen && e.last_seen > cur.last_seen) cur.last_seen = e.last_seen; - } - }); - - const before = keys.length; - const after = Object.keys(collapsed).length; - totalBefore += before; - totalAfter += after; - if (after === before) return; // nothing to collapse for this user - - usersTouched += 1; - print(`${u.email}: ${before} -> ${after} [${Object.keys(collapsed).join(", ")}]`); - - if (APPLY) { - db.users.updateOne({ _id: u._id }, { $set: { registered_clients: collapsed } }); - } -}); - -print(""); -print(`${APPLY ? "APPLIED" : "DRY RUN (set DRYRUN=0 to apply)"}: ` + - `users touched=${usersTouched}, entries ${totalBefore} -> ${totalAfter}`); diff --git a/backends/advanced/scripts/evaluate_memory_executor.py b/backends/advanced/scripts/evaluate_memory_executor.py deleted file mode 100644 index e32912a1..00000000 --- a/backends/advanced/scripts/evaluate_memory_executor.py +++ /dev/null @@ -1,186 +0,0 @@ -#!/usr/bin/env python3 -"""Replay a transcript canary set into a fresh vault with one memory executor. - -This is intentionally separate from the rebuild path: it never touches Mongo, queues, -the live vault, or the audit ledger. It is a small reproducible quality experiment for -comparing the direct tool-calling agent with the Codex CLI agent. -""" - -from __future__ import annotations - -import argparse -import asyncio -import contextlib -import json -import sys -import time -from pathlib import Path - -DEFAULT_PREFIXES = ( - "5c0b2333", - "99e2a020", - "7620c7b4", - "ea282e40", - "a4ed37ac", - "19f5a281", - "d0de4521", - "fd3f7f7d", -) - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--executor", choices=("codex", "direct"), required=True) - parser.add_argument("--dataset", type=Path, required=True) - parser.add_argument("--output", type=Path, required=True) - parser.add_argument("--model", help="Codex model override recorded in the manifest") - parser.add_argument( - "--reasoning-effort", - choices=("none", "minimal", "low", "medium", "high", "xhigh", "max", "ultra"), - help="Codex reasoning-effort override recorded in the manifest", - ) - parser.add_argument( - "--conversation", - action="append", - dest="prefixes", - help="conversation id or unique prefix; repeatable (defaults to audited 8)", - ) - parser.add_argument( - "--keep-output", - action="store_true", - help="reuse the output vault; default refuses a non-empty destination", - ) - return parser.parse_args() - - -def _load_rows(dataset: Path, prefixes: tuple[str, ...]) -> list[dict]: - rows = [ - json.loads(line) for line in dataset.read_text().splitlines() if line.strip() - ] - selected = [] - for prefix in prefixes: - matches = [row for row in rows if row["conversation_id"].startswith(prefix)] - if len(matches) != 1: - raise SystemExit( - f"{prefix!r} matched {len(matches)} conversations (expected 1)" - ) - selected.append(matches[0]) - return selected - - -def _prepare_output(output: Path, keep: bool) -> None: - if output.exists() and any(output.iterdir()) and not keep: - raise SystemExit( - f"refusing non-empty output {output}; remove it or pass --keep-output" - ) - output.mkdir(parents=True, exist_ok=True) - - -@contextlib.contextmanager -def _isolated_vault_lock(*_args, **_kwargs): - """No Redis dependency: this process is the sole writer of a throwaway vault.""" - yield - - -async def _run(args: argparse.Namespace) -> int: - # The script lives at backends/advanced/scripts; make src imports work when invoked - # directly without requiring an editable install. - backend = Path(__file__).resolve().parents[1] - sys.path.insert(0, str(backend / "src")) - - from advanced_omi_backend.services.memory import vault_lock - from advanced_omi_backend.services.memory.agent import codex_agent - from advanced_omi_backend.services.memory.agent.codex_agent import CodexMemoryAgent - from advanced_omi_backend.services.memory.agent.memory_agent import MemoryAgent - from advanced_omi_backend.services.memory.vault_scaffold import seed_vault_scaffold - - # Production uses Redis to serialize writers. This evaluator owns a unique fresh - # directory and runs sequentially, so requiring the application stack would add no - # safety and would make the experiment needlessly hard to reproduce. - vault_lock.vault_run_lock = _isolated_vault_lock - vault_lock.vault_note_lock = _isolated_vault_lock - configured = codex_agent._codex_settings() if args.executor == "codex" else {} - if args.executor == "codex" and (args.model or args.reasoning_effort): - configured.update( - { - key: value - for key, value in { - "model": args.model, - "reasoning_effort": args.reasoning_effort, - }.items() - if value - } - ) - codex_agent._codex_settings = lambda: configured - - # Codex receives the vault as both subprocess cwd and --cd. Resolve once so a - # caller's relative path is not interpreted relative to itself by the CLI. - args.output = args.output.resolve() - args.dataset = args.dataset.resolve() - prefixes = tuple(args.prefixes or DEFAULT_PREFIXES) - rows = _load_rows(args.dataset, prefixes) - _prepare_output(args.output, args.keep_output) - seed_vault_scaffold(args.output) - agent_type = CodexMemoryAgent if args.executor == "codex" else MemoryAgent - - manifest = { - "executor": args.executor, - "model": args.model - or (configured.get("model") if args.executor == "codex" else None), - "reasoning_effort": args.reasoning_effort - or (configured.get("reasoning_effort") if args.executor == "codex" else None), - "dataset": str(args.dataset.resolve()), - "output": str(args.output.resolve()), - "started_at_epoch": time.time(), - "runs": [], - } - manifest_path = args.output / "evaluation-manifest.json" - failures = 0 - for index, row in enumerate(rows, 1): - conversation_id = row["conversation_id"] - print( - f"[{index}/{len(rows)}] {args.executor}: {conversation_id} ({row['n_chars']} chars)", - flush=True, - ) - started = time.perf_counter() - result = await agent_type(args.output).run( - row["transcript"], - conversation_id, - date=row.get("created_at"), - duration_minutes=(row.get("duration_s") or 0) / 60, - title=row.get("title"), - ) - note = args.output / "Conversations" / f"{conversation_id}.md" - ok = note.is_file() and not result.truncated and not result.stalled - failures += int(not ok) - manifest["runs"].append( - { - "conversation_id": conversation_id, - "ok": ok, - "elapsed_seconds": round(time.perf_counter() - started, 3), - "rounds": result.rounds, - "tool_calls": result.tool_calls, - "touched": result.touched, - "removed": result.removed, - "errors": result.errors, - "truncated": result.truncated, - "stalled": result.stalled, - "summary": result.summary, - } - ) - manifest_path.write_text(json.dumps(manifest, indent=2) + "\n") - - manifest["finished_at_epoch"] = time.time() - manifest["failures"] = failures - manifest_path.write_text(json.dumps(manifest, indent=2) + "\n") - print(f"vault: {args.output.resolve()}\nmanifest: {manifest_path.resolve()}") - return int(bool(failures)) - - -def main() -> int: - args = _parse_args() - return asyncio.run(_run(args)) - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/backends/advanced/scripts/regen_titles_batch.py b/backends/advanced/scripts/regen_titles_batch.py deleted file mode 100644 index 6e8d597d..00000000 --- a/backends/advanced/scripts/regen_titles_batch.py +++ /dev/null @@ -1,123 +0,0 @@ -"""Batch-regenerate titles/summaries for completed conversations stuck with a -placeholder title (e.g. "Recording...", "Audio Recording (Transcription Failed)"). - -These conversations have a real transcript but never got an LLM title (completed via -a path that skipped generate_title_summary, or the title job failed). This enqueues -generate_title_summary_job (reads the existing transcript — NO re-transcription) in -batches of BATCH_SIZE, polling until each title changes off the placeholder. - -Run inside the backend/worker container: - python3 /app/regen_titles_batch.py # dry run (lists targets) - python3 /app/regen_titles_batch.py --apply # enqueue + poll -""" - -import argparse -import asyncio -import sys -import time - -from beanie import init_beanie - -from advanced_omi_backend.database import db -from advanced_omi_backend.models.conversation import Conversation -from advanced_omi_backend.models.user import User - -BATCH_SIZE = 5 -POLL_SECS = 20 -MAX_WAIT_SECS = 8 * 60 - -PLACEHOLDER_TITLES = ("Reprocessing...", "Recording...", "Transcribing...") - - -def _is_placeholder_title(title: str) -> bool: - return "Audio Recording (" in title or title in PLACEHOLDER_TITLES - - -def _is_placeholder_conv(c: "Conversation") -> bool: - t = c.title or "" - return bool(t) and _is_placeholder_title(t) - - -async def _targets() -> list: - out = [] - async for c in Conversation.find_all(): - if ( - not c.deleted - and c.processing_status == Conversation.ConversationStatus.COMPLETED.value - and _is_placeholder_conv(c) - ): - out.append(c.conversation_id) - return out - - -async def _title_of(cid: str): - c = await Conversation.find_one(Conversation.conversation_id == cid) - return (c.title or "") if c else "" - - -async def main(apply: bool) -> None: - await init_beanie(database=db, document_models=[User, Conversation]) - ids = await _targets() - print(f"=== {len(ids)} completed convs with placeholder titles ===", flush=True) - for cid in ids: - print(f" {cid[:12]} title={await _title_of(cid)!r}", flush=True) - if not apply: - print("Dry run — re-run with --apply to enqueue.", flush=True) - return - - # Import the queue lazily (module-level side effects: redis/queue setup). - from advanced_omi_backend.controllers.queue_controller import default_queue - from advanced_omi_backend.workers.conversation_jobs import ( - generate_title_summary_job, - ) - - done = {} - nbatches = (len(ids) + BATCH_SIZE - 1) // BATCH_SIZE - for b in range(nbatches): - batch = ids[b * BATCH_SIZE : (b + 1) * BATCH_SIZE] - print(f"--- BATCH {b+1}/{nbatches}: {[c[:8] for c in batch]} ---", flush=True) - for cid in batch: - default_queue.enqueue( - generate_title_summary_job, - cid, - job_id=f"regen_title_{cid[:12]}", - job_timeout=300, - description=f"Regenerate title/summary for {cid[:8]}", - ) - print(f" enqueued {cid[:8]}", flush=True) - - deadline = time.time() + MAX_WAIT_SECS - while True: - await asyncio.sleep(POLL_SECS) - pending = [] - for cid in batch: - t = await _title_of(cid) - if (not t) or _is_placeholder_title(t): - pending.append(cid) - else: - done[cid] = t - print( - f" poll: {len(batch)-len(pending)}/{len(batch)} retitled; " - f"pending={[c[:8] for c in pending]}", - flush=True, - ) - if not pending or time.time() > deadline: - if pending: - print( - f" batch {b+1}: TIMEOUT, pending={[c[:8] for c in pending]}", - flush=True, - ) - break - - print("=== FINAL ===", flush=True) - for cid in ids: - print(f" {cid[:12]} -> {done.get(cid, '(unchanged)')!r}", flush=True) - print(f"Retitled {len(done)}/{len(ids)}.", flush=True) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--apply", action="store_true") - args = parser.parse_args() - asyncio.run(main(args.apply)) - sys.exit(0) diff --git a/backends/advanced/scripts/reprocess_conversation.py b/backends/advanced/scripts/reprocess_conversation.py deleted file mode 100644 index 8a39185d..00000000 --- a/backends/advanced/scripts/reprocess_conversation.py +++ /dev/null @@ -1,96 +0,0 @@ -#!/usr/bin/env python3 -"""Reprocess a conversation's transcript via the backend API. - -This triggers POST /api/conversations/{id}/reprocess-transcript, which re-runs -batch transcription using the *currently configured* default batch STT provider -(config/config.yml -> defaults.stt) and then the full post-conversation chain -(speaker recognition -> memory -> title/summary -> dispatch complete). - -There is no per-reprocess provider override in the backend, so the "config" is -whatever defaults.stt points at. To reprocess with smallest.ai Pulse (hi), set - defaults.stt: stt-smallest -in config/config.yml and restart backend + workers, then run this script. - -Usage: - uv run python3 scripts/reprocess_conversation.py <conversation_id> [<conversation_id> ...] - -Env (optional, defaults shown): - BACKEND_URL=http://localhost:8000 - ADMIN_EMAIL / ADMIN_PASSWORD (else read from backends/advanced/.env) -""" - -import json -import os -import sys -import urllib.error -import urllib.parse -import urllib.request -from pathlib import Path - -BACKEND_URL = os.environ.get("BACKEND_URL", "http://localhost:8000").rstrip("/") -ENV_FILE = Path(__file__).resolve().parent.parent / ".env" - - -def _read_env_file(path: Path) -> dict: - values = {} - if not path.exists(): - return values - for line in path.read_text().splitlines(): - line = line.strip() - if not line or line.startswith("#") or "=" not in line: - continue - key, _, val = line.partition("=") - values[key.strip()] = val.strip().strip("'").strip('"') - return values - - -def _creds() -> tuple[str, str]: - env = _read_env_file(ENV_FILE) - email = os.environ.get("ADMIN_EMAIL") or env.get("ADMIN_EMAIL") - password = os.environ.get("ADMIN_PASSWORD") or env.get("ADMIN_PASSWORD") - if not email or not password: - sys.exit("ERROR: ADMIN_EMAIL / ADMIN_PASSWORD not found (env or .env).") - return email, password - - -def login(email: str, password: str) -> str: - data = urllib.parse.urlencode({"username": email, "password": password}).encode() - req = urllib.request.Request( - f"{BACKEND_URL}/auth/jwt/login", - data=data, - headers={"Content-Type": "application/x-www-form-urlencoded"}, - method="POST", - ) - with urllib.request.urlopen(req, timeout=30) as resp: - token = json.loads(resp.read())["access_token"] - print(f"[auth] logged in as {email}") - return token - - -def reprocess(conversation_id: str, token: str) -> None: - req = urllib.request.Request( - f"{BACKEND_URL}/api/conversations/{conversation_id}/reprocess-transcript", - data=b"", - headers={"Authorization": f"Bearer {token}"}, - method="POST", - ) - try: - with urllib.request.urlopen(req, timeout=60) as resp: - body = json.loads(resp.read()) - print(f"[ok] {conversation_id}: {json.dumps(body, indent=2)}") - except urllib.error.HTTPError as e: - detail = e.read().decode(errors="replace") - print(f"[FAIL] {conversation_id}: HTTP {e.code} -> {detail}") - - -def main() -> None: - ids = sys.argv[1:] - if not ids: - sys.exit(__doc__) - token = login(*_creds()) - for conversation_id in ids: - reprocess(conversation_id, token) - - -if __name__ == "__main__": - main() diff --git a/backends/advanced/scripts/scope_blank_segments.py b/backends/advanced/scripts/scope_blank_segments.py deleted file mode 100644 index 0bbd63bf..00000000 --- a/backends/advanced/scripts/scope_blank_segments.py +++ /dev/null @@ -1,68 +0,0 @@ -"""Read-only scope: find conversations whose ACTIVE transcript version has words/text -but an empty segments list (so the WebUI renders blank). Breaks them down by -diarization_source to distinguish pyannote-wiped from provider/never-diarized. - -Run inside the backend/worker container (read-only, no writes): - python3 /app/scope_blank_segments.py -""" - -import asyncio -import sys -from collections import Counter - -from beanie import init_beanie - -from advanced_omi_backend.database import db -from advanced_omi_backend.models.conversation import Conversation -from advanced_omi_backend.models.user import User - - -async def main() -> None: - await init_beanie(database=db, document_models=[User, Conversation]) - - by_diar = Counter() - rows = [] - async for c in Conversation.find_all(): - if c.deleted: - continue - av = c.active_transcript - if not av: - continue - has_text = bool((av.transcript or "").strip()) - words = av.words or [] - segs = av.segments or [] - if (has_text or words) and not segs: - diar = av.diarization_source or "(none)" - by_diar[diar] += 1 - word_spk = sorted({w.speaker for w in words if w.speaker is not None}) - rows.append( - ( - c.conversation_id, - c.processing_status, - c.client_id, - round(c.audio_total_duration or 0, 1), - len(words), - word_spk, - diar, - av.provider, - (c.title or "")[:40], - ) - ) - - print( - f"=== {len(rows)} non-deleted convs: active version has words/text but 0 segments ===" - ) - print(f"by diarization_source: {dict(by_diar)}\n") - rows.sort(key=lambda r: -r[3]) - print( - f"{'conv':10} {'status':10} {'client':20} {'dur':>7} {'wrds':>4} {'wordspk':12} {'diar':10} {'provider':10} title" - ) - for cid, st, client, dur, nw, wspk, diar, prov, title in rows: - print( - f"{cid[:8]:10} {str(st):10} {str(client):20} {dur:>7} {nw:>4} {str(wspk):12} {str(diar):10} {str(prov):10} {title!r}" - ) - - -if __name__ == "__main__": - asyncio.run(main()) - sys.exit(0) diff --git a/backends/advanced/scripts/settle_stuck_status.py b/backends/advanced/scripts/settle_stuck_status.py deleted file mode 100644 index d0870e07..00000000 --- a/backends/advanced/scripts/settle_stuck_status.py +++ /dev/null @@ -1,92 +0,0 @@ -"""Settle conversations stuck in a non-terminal processing_status. - -The status reconciler (services/status_reconciler.py) only scans non-deleted -conversations, and the no-speech/dead-end paths used to leave processing_status at -"active" (with a stale "Reprocessing..."/"Audio Recording (...)" title) — see the -fixes in mark_conversation_deleted and reprocess_speakers. That left a backlog of -conversations (mostly soft-deleted) stuck "active"/None/legacy "transcription_failed". - -This one-off applies the SAME fact-derived logic (Conversation.apply_status, the single -owner of the field) to EVERY conversation — including deleted ones — and clears stale -placeholder titles. Idempotent: terminal, correctly-titled conversations are untouched. - -Run inside the backend/worker container: - uv run python3 scripts/settle_stuck_status.py # dry run (default) - uv run python3 scripts/settle_stuck_status.py --apply # write changes -""" - -import argparse -import asyncio -import sys - -from beanie import init_beanie - -from advanced_omi_backend.database import db -from advanced_omi_backend.models.conversation import Conversation -from advanced_omi_backend.models.user import User - -TERMINAL = { - Conversation.ConversationStatus.COMPLETED.value, - Conversation.ConversationStatus.FAILED.value, -} -PLACEHOLDER_TITLES = ("Reprocessing...", "Recording...", "Transcribing...") - - -def _is_placeholder_title(title: str) -> bool: - # "Audio Recording (Processing...)" / "(Batch Transcription...)" / "(Transcription - # Failed)" — require the parenthetical so real LLM titles like "Audio Recording - # Issues" are NOT matched. - return "Audio Recording (" in title or title in PLACEHOLDER_TITLES - - -async def main(apply: bool) -> None: - await init_beanie(database=db, document_models=[User, Conversation]) - - scanned = settled = titled = 0 - async for conv in Conversation.find_all(): - scanned += 1 - changed = False - - # Only touch non-terminal conversations (active / None / legacy strings) — - # leave already-settled completed/failed ones (and their titles) alone. - if conv.processing_status not in TERMINAL: - # Settle status from facts. apply_status: has transcript -> completed; - # settled & none -> failed; else active. - before = conv.processing_status - if conv.apply_status(settled=True): - settled += 1 - changed = True - print( - f" conv={conv.conversation_id[:12]} deleted={conv.deleted} " - f"status {before!r} -> {conv.processing_status!r}" - + (f" stage={conv.failure_stage}" if conv.failure_stage else "") - ) - - # Clear stale in-flight placeholder titles on these dead-end convs. - title = conv.title or "" - if title and _is_placeholder_title(title): - conv.title = None - titled += 1 - changed = True - print(f" conv={conv.conversation_id[:12]} cleared title {title!r}") - - if changed and apply: - await conv.save() - - verb = "Settled" if apply else "Would settle" - print( - f"\nScanned {scanned} conversations. {verb} {settled} statuses, " - f"cleared {titled} placeholder titles." - ) - if not apply and (settled or titled): - print("Dry run — re-run with --apply to write changes.") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - "--apply", action="store_true", help="Write changes (default: dry run)" - ) - args = parser.parse_args() - asyncio.run(main(args.apply)) - sys.exit(0) diff --git a/backends/advanced/tests/scripts/graph-validation/README.md b/backends/advanced/tests/scripts/graph-validation/README.md deleted file mode 100644 index b8475dd0..00000000 --- a/backends/advanced/tests/scripts/graph-validation/README.md +++ /dev/null @@ -1,57 +0,0 @@ -# FalkorDB Graph Memory Validation - -Standalone scripts to validate using FalkorDB as the unified memory store (replacing Qdrant) before modifying Chronicle's production code. - -## Prerequisites - -- FalkorDB running (`docker compose up falkordb -d`) -- OpenAI API key in `backends/advanced/.env` -- Python deps: `falkordb`, `openai`, `python-dotenv` (all in existing pyproject.toml) - -## Quick Start - -```bash -cd backends/advanced - -# 1. Start FalkorDB -docker compose up falkordb -d - -# 2. Create schema (indexes + constraints) -uv run python tests/scripts/graph-validation/setup_schema.py - -# 3. Insert sample data (generates real embeddings) -uv run python tests/scripts/graph-validation/sample_data.py - -# 4. Run tests -uv run python tests/scripts/graph-validation/test_vector_search.py -uv run python tests/scripts/graph-validation/test_fulltext_search.py -uv run python tests/scripts/graph-validation/test_hybrid_search.py -uv run python tests/scripts/graph-validation/test_entity_graph.py -uv run python tests/scripts/graph-validation/test_conversation_doc.py - -# 5. Cleanup -uv run python tests/scripts/graph-validation/setup_schema.py --cleanup -``` - -## Connection - -FalkorDB defaults: -- **Host**: `localhost` -- **Port**: `6381` -- **Graph name**: `chronicle` - -Override via environment variables: `FALKORDB_HOST`, `FALKORDB_PORT`, `FALKORDB_GRAPH`. - -## What Each Test Validates - -| Script | Tests | -|--------|-------| -| `test_vector_search.py` | Semantic similarity, score ordering, user_id scoping | -| `test_fulltext_search.py` | BM25 keyword search, multi-term AND, domain terms, empty results | -| `test_hybrid_search.py` | Vector+BM25 merge, recency bias, exact keyword boost | -| `test_entity_graph.py` | Entity traversal, cross-entity queries, entity listing | -| `test_conversation_doc.py` | LLM doc generation, section parsing, entity extraction reliability | - -## Test Labels (ConvDoc/ConvChunk/ConvEntity) - -All test data uses `Conv*` prefixed labels to avoid conflicting with existing schema. diff --git a/backends/advanced/tests/scripts/graph-validation/sample_data.py b/backends/advanced/tests/scripts/graph-validation/sample_data.py deleted file mode 100644 index 8157f8b0..00000000 --- a/backends/advanced/tests/scripts/graph-validation/sample_data.py +++ /dev/null @@ -1,190 +0,0 @@ -"""Insert sample conversation data into FalkorDB for validation. - -Reads sample .md files, chunks them by ### headers, generates real embeddings, -and stores as ConvDoc/ConvChunk/ConvEntity nodes with relationships. - -Usage: - uv run python tests/scripts/graph-validation/sample_data.py -""" - -import uuid -from pathlib import Path - -from utils import ( - generate_embeddings_sync, - get_graph, - parse_frontmatter, - parse_people_section, - print_fail, - print_header, - print_pass, - split_on_headers, -) - -SAMPLE_DIR = Path(__file__).parent / "sample_conversations" -TEST_USER_ID = "test_user_001" -OTHER_USER_ID = "other_user_002" # For user-scoping tests - - -def insert_conversation(graph, md_path: Path, user_id: str): - """Parse a conversation .md and insert into FalkorDB.""" - content = md_path.read_text(encoding="utf-8") - frontmatter = parse_frontmatter(content) - chunks = split_on_headers(content) - people = parse_people_section(content) - - conv_id = frontmatter.get("conversation_id", md_path.stem) - date = frontmatter.get("date", "2026-03-15T00:00:00") - speakers = frontmatter.get("speakers", "") - - # Extract title from first ## header - import re - - title_match = re.search(r"^##\s+(.+)$", content, re.MULTILINE) - title = title_match.group(1) if title_match else md_path.stem - - print(f"\n Processing: {title}") - print(f" Chunks: {len(chunks)}, People: {len(people)}") - - # Generate embeddings for all chunks - chunk_texts = [f"{c['section_title']}: {c['text']}" for c in chunks] - if not chunk_texts: - print_fail(f"No chunks extracted from {md_path.name}") - return - - print(f" Generating embeddings for {len(chunk_texts)} chunks...") - embeddings = generate_embeddings_sync(chunk_texts) - print(f" Embeddings generated ({len(embeddings[0])} dimensions)") - - # Create ConvDoc node - graph.query( - """ - MERGE (d:ConvDoc {conversation_id: $conv_id}) - SET d.title = $title, - d.date = $date, - d.user_id = $user_id, - d.speakers = $speakers, - d.file_path = $file_path - """, - params={ - "conv_id": conv_id, - "title": title, - "date": date, - "user_id": user_id, - "speakers": speakers, - "file_path": str(md_path.relative_to(SAMPLE_DIR.parent)), - }, - ) - - # Create ConvChunk nodes with embeddings and link to ConvDoc - for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)): - chunk_id = f"{conv_id}_chunk_{i:03d}" - graph.query( - """ - MATCH (d:ConvDoc {conversation_id: $conv_id}) - MERGE (c:ConvChunk {id: $chunk_id}) - SET c.text = $text, - c.section_title = $section_title, - c.embedding = vecf32($embedding), - c.user_id = $user_id, - c.created_at = $date, - c.chunk_index = $chunk_index - MERGE (d)-[:HAS_CHUNK]->(c) - """, - params={ - "conv_id": conv_id, - "chunk_id": chunk_id, - "text": chunk["text"], - "section_title": chunk["section_title"], - "embedding": embedding, - "user_id": user_id, - "date": date, - "chunk_index": i, - }, - ) - - # Create ConvEntity nodes from People section and link to chunks - for person in people: - entity_id = f"{user_id}_{person['name'].lower().replace(' ', '_')}" - graph.query( - """ - MERGE (e:ConvEntity {id: $entity_id}) - SET e.name = $name, - e.description = $description, - e.type = 'person', - e.user_id = $user_id - """, - params={ - "entity_id": entity_id, - "name": person["name"], - "description": person["description"], - "user_id": user_id, - }, - ) - - # Link entity to all chunks that mention their name - for i, chunk in enumerate(chunks): - if person["name"].lower() in chunk["text"].lower(): - chunk_id = f"{conv_id}_chunk_{i:03d}" - graph.query( - """ - MATCH (c:ConvChunk {id: $chunk_id}) - MATCH (e:ConvEntity {id: $entity_id}) - MERGE (c)-[:MENTIONS]->(e) - """, - params={ - "chunk_id": chunk_id, - "entity_id": entity_id, - }, - ) - - print_pass(f"Inserted: {title} ({len(chunks)} chunks, {len(people)} entities)") - - -def main(): - print_header("Inserting Sample Data") - - graph = get_graph() - - # Insert conversations for the primary test user - for md_file in sorted(SAMPLE_DIR.glob("*.md")): - insert_conversation(graph, md_file, TEST_USER_ID) - - # Insert one conversation for a different user (for scoping tests) - dentist_file = SAMPLE_DIR / "dentist_appointment.md" - if dentist_file.exists(): - print(f"\n Inserting duplicate for other user (scoping test)...") - # We need to modify the conv_id to avoid UNIQUE constraint clash - content = dentist_file.read_text(encoding="utf-8") - # Write a temp copy with different conv_id - temp_path = SAMPLE_DIR / "_other_user_dentist.md" - temp_path.write_text( - content.replace( - "conversation_id: conv_002", "conversation_id: conv_002_other" - ) - ) - insert_conversation(graph, temp_path, OTHER_USER_ID) - temp_path.unlink() - - # Verify counts - doc_result = graph.query("MATCH (d:ConvDoc) RETURN count(d) AS c") - doc_count = doc_result.result_set[0][0] - - chunk_result = graph.query("MATCH (c:ConvChunk) RETURN count(c) AS c") - chunk_count = chunk_result.result_set[0][0] - - entity_result = graph.query("MATCH (e:ConvEntity) RETURN count(e) AS c") - entity_count = entity_result.result_set[0][0] - - rel_result = graph.query("MATCH ()-[r:HAS_CHUNK|MENTIONS]->() RETURN count(r) AS c") - rel_count = rel_result.result_set[0][0] - - print_header("Summary") - print_pass(f"ConvDoc nodes: {doc_count}") - print_pass(f"ConvChunk nodes: {chunk_count}") - print_pass(f"ConvEntity nodes: {entity_count}") - print_pass(f"Relationships: {rel_count}") - - -if __name__ == "__main__": - main() diff --git a/backends/advanced/tests/scripts/graph-validation/setup_schema.py b/backends/advanced/tests/scripts/graph-validation/setup_schema.py deleted file mode 100644 index 7159baa1..00000000 --- a/backends/advanced/tests/scripts/graph-validation/setup_schema.py +++ /dev/null @@ -1,139 +0,0 @@ -"""Create FalkorDB schema for conversation memory validation. - -Creates constraints, vector index, and full-text index using ConvDoc/ConvChunk -labels to avoid conflicting with existing schema. - -Usage: - uv run python tests/scripts/graph-validation/setup_schema.py - uv run python tests/scripts/graph-validation/setup_schema.py --cleanup -""" - -import sys - -from utils import EMBEDDING_DIMENSIONS, get_graph, print_fail, print_header, print_pass - - -def setup(graph): - print_header("Creating Schema") - - # Constraints — FalkorDB doesn't support IF NOT EXISTS, use try/except - try: - graph.query( - "CREATE CONSTRAINT ON (d:ConvDoc) ASSERT d.conversation_id IS UNIQUE" - ) - print_pass("Constraint: ConvDoc.conversation_id UNIQUE") - except Exception as e: - if "already exists" in str(e).lower() or "already indexed" in str(e).lower(): - print_pass("Constraint: ConvDoc.conversation_id UNIQUE (already exists)") - else: - print_fail(f"Constraint ConvDoc.conversation_id: {e}") - - try: - graph.query("CREATE CONSTRAINT ON (c:ConvChunk) ASSERT c.id IS UNIQUE") - print_pass("Constraint: ConvChunk.id UNIQUE") - except Exception as e: - if "already exists" in str(e).lower() or "already indexed" in str(e).lower(): - print_pass("Constraint: ConvChunk.id UNIQUE (already exists)") - else: - print_fail(f"Constraint ConvChunk.id: {e}") - - try: - graph.query("CREATE CONSTRAINT ON (e:ConvEntity) ASSERT e.id IS UNIQUE") - print_pass("Constraint: ConvEntity.id UNIQUE") - except Exception as e: - if "already exists" in str(e).lower() or "already indexed" in str(e).lower(): - print_pass("Constraint: ConvEntity.id UNIQUE (already exists)") - else: - print_fail(f"Constraint ConvEntity.id: {e}") - - # Vector index - try: - graph.query( - f"CREATE VECTOR INDEX FOR (c:ConvChunk) ON (c.embedding) " - f"OPTIONS {{dimension: {EMBEDDING_DIMENSIONS}, similarityFunction: 'cosine'}}" - ) - print_pass( - f"Vector index: ConvChunk.embedding ({EMBEDDING_DIMENSIONS}d, cosine)" - ) - except Exception as e: - if "already exists" in str(e).lower() or "already indexed" in str(e).lower(): - print_pass( - f"Vector index: ConvChunk.embedding ({EMBEDDING_DIMENSIONS}d, cosine) (already exists)" - ) - else: - print_fail(f"Vector index: {e}") - - # Full-text index - try: - graph.query( - "CREATE FULLTEXT INDEX FOR (c:ConvChunk) ON (c.text, c.section_title)" - ) - print_pass("Full-text index: ConvChunk (text + section_title)") - except Exception as e: - if "already exists" in str(e).lower() or "already indexed" in str(e).lower(): - print_pass( - "Full-text index: ConvChunk (text + section_title) (already exists)" - ) - else: - print_fail(f"Full-text index: {e}") - - # Verify indexes exist - try: - result = graph.query("CALL db.indexes()") - idx_names = [] - for row in result.result_set: - # Each row is a list; index name/label varies by FalkorDB version - idx_names.append(str(row)) - print_pass(f"Indexes present: {len(result.result_set)} index(es) found") - except Exception as e: - print_fail(f"Could not verify indexes: {e}") - - print("\nSchema setup complete.") - - -def cleanup(graph): - print_header("Cleaning Up Test Data") - - # Delete all test nodes and relationships - result = graph.query( - "MATCH (n) WHERE n:ConvDoc OR n:ConvChunk OR n:ConvEntity " - "DETACH DELETE n RETURN count(n) AS deleted" - ) - if result.result_set: - deleted = result.result_set[0][0] - print_pass(f"Deleted {deleted} test nodes") - else: - print_pass("No test nodes to delete") - - # Drop indexes — FalkorDB doesn't support IF EXISTS for DROP, use try/except - for idx_query in [ - "DROP VECTOR INDEX ON :ConvChunk(embedding)", - "DROP FULLTEXT INDEX ON :ConvChunk(text, section_title)", - ]: - try: - graph.query(idx_query) - print_pass(f"Dropped index: {idx_query}") - except Exception as e: - print_fail(f"Failed to drop index ({idx_query}): {e}") - - # Drop constraints - for constraint_query in [ - "DROP CONSTRAINT ON (d:ConvDoc) ASSERT d.conversation_id IS UNIQUE", - "DROP CONSTRAINT ON (c:ConvChunk) ASSERT c.id IS UNIQUE", - "DROP CONSTRAINT ON (e:ConvEntity) ASSERT e.id IS UNIQUE", - ]: - try: - graph.query(constraint_query) - print_pass(f"Dropped constraint: {constraint_query}") - except Exception as e: - print_fail(f"Failed to drop constraint ({constraint_query}): {e}") - - print("\nCleanup complete.") - - -if __name__ == "__main__": - graph = get_graph() - if "--cleanup" in sys.argv: - cleanup(graph) - else: - setup(graph) diff --git a/backends/advanced/tests/scripts/graph-validation/test_conversation_doc.py b/backends/advanced/tests/scripts/graph-validation/test_conversation_doc.py deleted file mode 100644 index 7832e828..00000000 --- a/backends/advanced/tests/scripts/graph-validation/test_conversation_doc.py +++ /dev/null @@ -1,236 +0,0 @@ -"""Validate LLM conversation document generation and markdown parsing. - -Tests: -1. LLM produces a well-structured .md from a sample transcript -2. Parsed sections match expected structure (Summary, Key Facts, People, Action Items) -3. Entity extraction from People section is reliable -4. Header-based chunking produces correct splits - -This directly tests the review's concern about entity parsing fragility. - -Usage: - uv run python tests/scripts/graph-validation/test_conversation_doc.py -""" - -import asyncio -import json - -from openai import AsyncOpenAI -from utils import ( - OPENAI_API_KEY, - OPENAI_BASE_URL, - parse_action_items, - parse_frontmatter, - parse_people_section, - print_fail, - print_header, - print_pass, - split_on_headers, -) - -SAMPLE_TRANSCRIPT = """Speaker 0: Hey John, thanks for meeting. I wanted to go over the Q3 timeline. -John: Sure, so the main concern is the backend migration. It might block our September 15th deadline. -Speaker 0: Right. Sarah mentioned she wants to descope the auth rewrite. What do you think? -John: I agree with Sarah actually. The auth rewrite is nice to have but not critical for Q3. Let's push it to Q4. -Speaker 0: Makes sense. I'll schedule a meeting with Sarah to align on scope. Oh and John, do you prefer morning or afternoon standups? -John: Morning for sure. I'm much more productive before noon. -Speaker 0: Got it. One more thing - we need the budget review ready by Friday. Can you pull the numbers? -John: Yeah I'll have the draft ready by Thursday. Mike from finance wants to review it too. -Speaker 0: Perfect. I'll also check with the DevOps team about the migration timeline. Thanks John. -John: Sounds good, talk later.""" - -GENERATE_CONVERSATION_DOC_PROMPT = """\ -You are generating a structured conversation document from a transcript. - -Given a transcript with speaker labels, produce a markdown document with this EXACT structure: - ---- -conversation_id: {conversation_id} -date: {date} -speakers: [{speakers}] -duration_minutes: {duration} ---- - -## {Title - descriptive, 3-8 words} - -### Summary -{2-3 sentence summary of what was discussed} - -### Key Facts -{Bulleted list of specific facts, decisions, numbers, dates mentioned} - -### People -{Bulleted list in format: - Name (role/relationship, context)} -Include ALL named individuals — speakers, people mentioned, people referenced. -If a speaker is identified by name (e.g., "John" not "Speaker 0"), they MUST appear here. -Do not include unnamed roles or generic labels like "Speaker 0". - -### Action Items -{Bulleted list in format: - [ ] Action item description} -Use [x] for items already completed in the conversation. - -Rules: -- Every ### section MUST be present, even if empty (use "- None" for empty sections) -- People section: ONLY named individuals, format MUST be "- Name (description)" -- Key Facts: be specific — include dates, numbers, names -- Action Items: use checkbox format [ ] or [x] -- Do NOT add any sections beyond the four listed above -""" - - -async def generate_doc(transcript: str) -> str: - """Call LLM to generate a conversation document.""" - client = AsyncOpenAI(api_key=OPENAI_API_KEY, base_url=OPENAI_BASE_URL) - - response = await client.chat.completions.create( - model="gpt-4o-mini", - messages=[ - {"role": "system", "content": GENERATE_CONVERSATION_DOC_PROMPT}, - { - "role": "user", - "content": f"Transcript:\n{transcript}\n\nMetadata:\n- conversation_id: test_conv_llm\n- date: 2026-03-18T10:00:00\n- speakers: Speaker 0, John\n- duration: 8 minutes", - }, - ], - temperature=0.2, - ) - return response.choices[0].message.content.strip() - - -def test_structure(doc: str): - """Verify the generated doc has all required sections.""" - print_header("Test: Document Structure") - - required_sections = ["Summary", "Key Facts", "People", "Action Items"] - chunks = split_on_headers(doc) - found_sections = [c["section_title"] for c in chunks] - - print(f" Found sections: {found_sections}") - - for section in required_sections: - if section in found_sections: - print_pass(f"Section '{section}' present") - else: - print_fail(f"Section '{section}' MISSING") - - # Check frontmatter - fm = parse_frontmatter(doc) - if fm.get("conversation_id"): - print_pass(f"Frontmatter has conversation_id: {fm['conversation_id']}") - else: - print_fail("Frontmatter missing conversation_id") - - if fm.get("date"): - print_pass(f"Frontmatter has date: {fm['date']}") - else: - print_fail("Frontmatter missing date") - - return chunks - - -def test_entity_parsing(doc: str): - """Test that People section can be reliably parsed.""" - print_header("Test: Entity Parsing from People Section") - - people = parse_people_section(doc) - print(f" Parsed entities: {len(people)}") - for p in people: - print(f" - {p['name']} ({p['description']})") - - # Expected entities from the transcript - expected_names = {"John", "Sarah", "Mike"} - - found_names = {p["name"] for p in people} - for name in expected_names: - if name in found_names: - print_pass(f"Entity '{name}' extracted") - else: - print_fail(f"Entity '{name}' NOT extracted (found: {found_names})") - - # Check that descriptions are non-empty - has_descriptions = sum(1 for p in people if p["description"]) - if has_descriptions == len(people): - print_pass(f"All {len(people)} entities have descriptions") - elif has_descriptions > 0: - print_pass(f"{has_descriptions}/{len(people)} entities have descriptions") - else: - print_fail("No entities have descriptions") - - # Check for unnamed entries (these should NOT appear) - unnamed = [ - p for p in people if p["name"].lower().startswith(("speaker", "the ", "a ")) - ] - if not unnamed: - print_pass("No unnamed/generic entities (good)") - else: - print_fail(f"Found unnamed entities: {[p['name'] for p in unnamed]}") - - return people - - -def test_action_items(doc: str): - """Test Action Items parsing.""" - print_header("Test: Action Items Parsing") - - items = parse_action_items(doc) - print(f" Parsed items: {len(items)}") - for item in items: - status = "[x]" if item["done"] else "[ ]" - print(f" {status} {item['text']}") - - if items: - print_pass(f"Extracted {len(items)} action items") - else: - print_fail("No action items extracted") - - # Should have at least 2 action items from the transcript - if len(items) >= 2: - print_pass(f"At least 2 action items found (got {len(items)})") - else: - print_fail(f"Expected at least 2 action items, got {len(items)}") - - -def test_chunking(doc: str): - """Test header-based chunking produces correct splits.""" - print_header("Test: Header-Based Chunking") - - chunks = split_on_headers(doc) - print(f" Chunks: {len(chunks)}") - for c in chunks: - text_preview = c["text"][:80].replace("\n", " ") - print(f" [{c['section_title']}] {text_preview}...") - - # Should have 4 chunks (Summary, Key Facts, People, Action Items) - if len(chunks) >= 4: - print_pass(f"Got {len(chunks)} chunks (expected ≥4)") - else: - print_fail(f"Got {len(chunks)} chunks (expected ≥4)") - - # Each chunk should have non-empty text - empty_chunks = [c for c in chunks if not c["text"].strip()] - if not empty_chunks: - print_pass("All chunks have non-empty text") - else: - print_fail(f"{len(empty_chunks)} chunks have empty text") - - -async def main(): - print_header("Generating Conversation Document via LLM") - print(f" Transcript length: {len(SAMPLE_TRANSCRIPT)} chars") - print(f" Generating...") - - doc = await generate_doc(SAMPLE_TRANSCRIPT) - - print(f" Generated document: {len(doc)} chars") - print(f"\n{'─'*60}") - print(doc) - print(f"{'─'*60}\n") - - # Run all tests - test_structure(doc) - test_entity_parsing(doc) - test_action_items(doc) - test_chunking(doc) - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/backends/advanced/tests/scripts/graph-validation/test_entity_graph.py b/backends/advanced/tests/scripts/graph-validation/test_entity_graph.py deleted file mode 100644 index f80fedb3..00000000 --- a/backends/advanced/tests/scripts/graph-validation/test_entity_graph.py +++ /dev/null @@ -1,161 +0,0 @@ -"""Validate entity-based graph traversal search. - -Tests: -1. Find chunks mentioning a specific person via graph traversal -2. Cross-entity query (conversations mentioning both X and Y) -3. Entity listing per user - -Usage: - uv run python tests/scripts/graph-validation/test_entity_graph.py -""" - -from utils import get_graph, print_fail, print_header, print_pass, print_results_table - -TEST_USER = "test_user_001" - - -def _parse_results(result): - """Parse FalkorDB result_set into list of dicts using header names.""" - headers = [h[1] for h in result.header] - rows = [] - for row in result.result_set: - rows.append({headers[i]: row[i] for i in range(len(headers))}) - return rows - - -def test_entity_search(graph): - """Search for chunks mentioning 'John' via graph traversal.""" - print_header("Test: Entity Graph Traversal -- 'John'") - - result = graph.query( - """ - MATCH (e:ConvEntity {user_id: $user_id}) - WHERE toLower(e.name) CONTAINS toLower($entity_name) - WITH e - MATCH (chunk:ConvChunk)-[:MENTIONS]->(e) - MATCH (doc:ConvDoc)-[:HAS_CHUNK]->(chunk) - RETURN e.name AS entity, chunk.id AS chunk_id, - chunk.section_title AS section, chunk.text AS text, - doc.title AS title, doc.date AS date - ORDER BY doc.date DESC - """, - params={"user_id": TEST_USER, "entity_name": "John"}, - ) - results = _parse_results(result) - - print(f" Entity: 'John'") - print(f" Results: {len(results)}") - print_results_table(results, ["entity", "title", "section"]) - - if results: - print_pass(f"Found {len(results)} chunks mentioning John") - # All should be from Q3 meeting - titles = {r["title"] for r in results} - if all("Q3" in t for t in titles): - print_pass("All John mentions are in Q3 meeting (expected)") - else: - print_fail(f"Unexpected titles: {titles}") - else: - print_fail("No results for entity 'John'") - - -def test_entity_search_dr_chen(graph): - """Search for chunks mentioning 'Dr. Chen'.""" - print_header("Test: Entity Graph Traversal -- 'Dr. Chen'") - - result = graph.query( - """ - MATCH (e:ConvEntity {user_id: $user_id}) - WHERE toLower(e.name) CONTAINS toLower($entity_name) - WITH e - MATCH (chunk:ConvChunk)-[:MENTIONS]->(e) - MATCH (doc:ConvDoc)-[:HAS_CHUNK]->(chunk) - RETURN e.name AS entity, chunk.section_title AS section, - doc.title AS title, doc.date AS date - ORDER BY doc.date DESC - """, - params={"user_id": TEST_USER, "entity_name": "Dr. Chen"}, - ) - results = _parse_results(result) - - print(f" Entity: 'Dr. Chen'") - print(f" Results: {len(results)}") - print_results_table(results, ["entity", "title", "section"]) - - if results: - print_pass(f"Found {len(results)} chunks mentioning Dr. Chen") - if all("Dentist" in r.get("title", "") for r in results): - print_pass("All Dr. Chen mentions are in dentist conversation (expected)") - else: - print_fail("No results for entity 'Dr. Chen'") - - -def test_cross_entity_query(graph): - """Find conversations that mention BOTH John AND Sarah.""" - print_header("Test: Cross-Entity Query (John AND Sarah)") - - result = graph.query( - """ - MATCH (e1:ConvEntity {user_id: $user_id}) - WHERE toLower(e1.name) CONTAINS 'john' - MATCH (e2:ConvEntity {user_id: $user_id}) - WHERE toLower(e2.name) CONTAINS 'sarah' - MATCH (doc:ConvDoc)-[:HAS_CHUNK]->(c1:ConvChunk)-[:MENTIONS]->(e1) - MATCH (doc)-[:HAS_CHUNK]->(c2:ConvChunk)-[:MENTIONS]->(e2) - RETURN DISTINCT doc.conversation_id AS conv_id, doc.title AS title, - collect(DISTINCT e1.name) AS entity1, - collect(DISTINCT e2.name) AS entity2 - """, - params={"user_id": TEST_USER}, - ) - results = _parse_results(result) - - print(f" Cross-entity: John AND Sarah") - print(f" Results: {len(results)}") - print_results_table(results, ["title", "entity1", "entity2"]) - - if results: - print_pass(f"Found {len(results)} conversations mentioning both John and Sarah") - else: - print_fail("No conversations found with both John and Sarah") - - -def test_list_entities(graph): - """List all entities for a user with mention counts.""" - print_header("Test: List All Entities") - - result = graph.query( - """ - MATCH (e:ConvEntity {user_id: $user_id}) - OPTIONAL MATCH (chunk:ConvChunk)-[:MENTIONS]->(e) - RETURN e.name AS name, e.type AS type, e.description AS description, - count(chunk) AS mention_count - ORDER BY mention_count DESC - """, - params={"user_id": TEST_USER}, - ) - results = _parse_results(result) - - print(f" User: {TEST_USER}") - print(f" Entities: {len(results)}") - print_results_table(results, ["name", "type", "description", "mention_count"]) - - if results: - print_pass(f"Found {len(results)} entities for user") - # Should have at least John, Sarah, Dr. Chen, etc. - names = {r["name"] for r in results} - for expected in ["John", "Sarah", "Dr. Chen"]: - if expected in names: - print_pass(f"Entity '{expected}' found") - else: - print_fail(f"Entity '{expected}' not found in {names}") - else: - print_fail("No entities found") - - -if __name__ == "__main__": - graph = get_graph() - test_entity_search(graph) - test_entity_search_dr_chen(graph) - test_cross_entity_query(graph) - test_list_entities(graph) diff --git a/backends/advanced/tests/scripts/graph-validation/test_fulltext_search.py b/backends/advanced/tests/scripts/graph-validation/test_fulltext_search.py deleted file mode 100644 index 20d1ae86..00000000 --- a/backends/advanced/tests/scripts/graph-validation/test_fulltext_search.py +++ /dev/null @@ -1,215 +0,0 @@ -"""Validate FalkorDB full-text (BM25) search. - -Tests: -1. Single keyword search returns matching chunks -2. Multi-term AND search narrows results -3. BM25 scores are non-zero and ranked correctly -4. User scoping works with full-text results - -Usage: - uv run python tests/scripts/graph-validation/test_fulltext_search.py -""" - -from utils import get_graph, print_fail, print_header, print_pass, print_results_table - -TEST_USER = "test_user_001" - - -def _parse_results(result): - """Parse FalkorDB result_set into list of dicts using header names.""" - headers = [h[1] for h in result.header] - rows = [] - for row in result.result_set: - rows.append({headers[i]: row[i] for i in range(len(headers))}) - return rows - - -def test_single_keyword(graph): - """Search for 'deadline' -- should return Q3 meeting chunks.""" - print_header("Test: Single Keyword Full-Text Search") - - result = graph.query( - """ - CALL db.idx.fulltext.queryNodes('ConvChunk', $search_term) - YIELD node AS chunk, score - WHERE chunk.user_id = $user_id - MATCH (doc:ConvDoc)-[:HAS_CHUNK]->(chunk) - RETURN chunk.id AS chunk_id, chunk.section_title AS section, - chunk.text AS text, doc.title AS title, score - ORDER BY score DESC - """, - params={"search_term": "deadline", "user_id": TEST_USER}, - ) - results = _parse_results(result) - - print(f" Query: 'deadline'") - print(f" Results: {len(results)}") - print_results_table(results, ["score", "title", "section"]) - - if results: - print_pass(f"Full-text search returned {len(results)} results") - # Check that scores are positive - if all(r["score"] > 0 for r in results): - print_pass("All BM25 scores are positive") - else: - print_fail("Some scores are zero or negative") - else: - print_fail("No results for 'deadline'") - - return results - - -def test_multi_term_search(graph): - """Search for 'Q3 AND deadline' -- should narrow to Q3 meeting.""" - print_header("Test: Multi-Term AND Search") - - result = graph.query( - """ - CALL db.idx.fulltext.queryNodes('ConvChunk', $search_term) - YIELD node AS chunk, score - WHERE chunk.user_id = $user_id - MATCH (doc:ConvDoc)-[:HAS_CHUNK]->(chunk) - RETURN chunk.id AS chunk_id, chunk.section_title AS section, - doc.title AS title, score - ORDER BY score DESC - """, - params={"search_term": "Q3 AND deadline", "user_id": TEST_USER}, - ) - results = _parse_results(result) - - print(f" Query: 'Q3 AND deadline'") - print(f" Results: {len(results)}") - print_results_table(results, ["score", "title", "section"]) - - if results: - # All results should be from Q3 meeting - q3_results = [r for r in results if "Q3" in r.get("title", "")] - if len(q3_results) == len(results): - print_pass("All multi-term results are from Q3 meeting") - else: - non_q3 = [r["title"] for r in results if "Q3" not in r.get("title", "")] - print_fail(f"Some results are not Q3: {non_q3}") - else: - print_fail("No results for 'Q3 AND deadline'") - - -def test_dental_terms(graph): - """Search for dental-specific terms -- should return dentist chunks.""" - print_header("Test: Domain-Specific Terms") - - result = graph.query( - """ - CALL db.idx.fulltext.queryNodes('ConvChunk', $search_term) - YIELD node AS chunk, score - WHERE chunk.user_id = $user_id - MATCH (doc:ConvDoc)-[:HAS_CHUNK]->(chunk) - RETURN chunk.id AS chunk_id, chunk.section_title AS section, - doc.title AS title, score - ORDER BY score DESC - """, - params={"search_term": "crown molar", "user_id": TEST_USER}, - ) - results = _parse_results(result) - - print(f" Query: 'crown molar'") - print(f" Results: {len(results)}") - print_results_table(results, ["score", "title", "section"]) - - if results and "Dentist" in results[0].get("title", ""): - print_pass("Top result is dentist conversation") - elif results: - print_fail(f"Top result: {results[0].get('title')}") - else: - print_fail("No results for 'crown molar'") - - -def test_no_results(graph): - """Search for completely unrelated term -- should return empty.""" - print_header("Test: No Results for Unrelated Query") - - result = graph.query( - """ - CALL db.idx.fulltext.queryNodes('ConvChunk', $search_term) - YIELD node AS chunk, score - WHERE chunk.user_id = $user_id - RETURN chunk.id AS chunk_id, score - """, - params={"search_term": "quantum entanglement photon", "user_id": TEST_USER}, - ) - results = _parse_results(result) - - print(f" Query: 'quantum entanglement photon'") - print(f" Results: {len(results)}") - - if len(results) == 0: - print_pass("No results for unrelated query (expected)") - else: - print_fail(f"Got {len(results)} unexpected results") - - -def test_user_scoping_fulltext(graph): - """Verify full-text results are scoped to user.""" - print_header("Test: Full-Text User Scoping") - - # Search as test user -- should find all 3 conversations - result = graph.query( - """ - CALL db.idx.fulltext.queryNodes('ConvChunk', $search_term) - YIELD node AS chunk, score - WHERE chunk.user_id = $user_id - RETURN chunk.id AS chunk_id, chunk.user_id AS uid - """, - params={ - "search_term": "appointment OR deadline OR grocery", - "user_id": TEST_USER, - }, - ) - test_results = _parse_results(result) - - # Search as other user -- should find only their dentist copy - result = graph.query( - """ - CALL db.idx.fulltext.queryNodes('ConvChunk', $search_term) - YIELD node AS chunk, score - WHERE chunk.user_id = $user_id - RETURN chunk.id AS chunk_id, chunk.user_id AS uid - """, - params={ - "search_term": "appointment OR deadline OR grocery", - "user_id": "other_user_002", - }, - ) - other_results = _parse_results(result) - - print(f" TEST_USER results: {len(test_results)}") - print(f" OTHER_USER results: {len(other_results)}") - - if len(test_results) > len(other_results): - print_pass("TEST_USER has more results (expected -- more conversations)") - else: - print_pass( - f"Result counts: TEST={len(test_results)}, OTHER={len(other_results)}" - ) - - # Verify no cross-user contamination - test_uids = {r["uid"] for r in test_results} - other_uids = {r["uid"] for r in other_results} - - if test_uids <= {TEST_USER}: - print_pass("TEST_USER full-text results scoped correctly") - else: - print_fail(f"Leaked user_ids in TEST results: {test_uids}") - - if other_uids <= {"other_user_002"}: - print_pass("OTHER_USER full-text results scoped correctly") - else: - print_fail(f"Leaked user_ids in OTHER results: {other_uids}") - - -if __name__ == "__main__": - graph = get_graph() - test_single_keyword(graph) - test_multi_term_search(graph) - test_dental_terms(graph) - test_no_results(graph) - test_user_scoping_fulltext(graph) diff --git a/backends/advanced/tests/scripts/graph-validation/test_hybrid_search.py b/backends/advanced/tests/scripts/graph-validation/test_hybrid_search.py deleted file mode 100644 index 0105335e..00000000 --- a/backends/advanced/tests/scripts/graph-validation/test_hybrid_search.py +++ /dev/null @@ -1,199 +0,0 @@ -"""Validate hybrid search: vector + full-text + recency bias. - -Tests: -1. Hybrid merge produces better ranking than either alone -2. Recency bias correctly boosts recent conversations -3. Score components are traceable - -Usage: - uv run python tests/scripts/graph-validation/test_hybrid_search.py -""" - -from utils import ( - compute_hybrid_scores, - generate_embeddings_sync, - get_graph, - print_fail, - print_header, - print_pass, - print_results_table, -) - -TEST_USER = "test_user_001" - - -def _parse_results(result): - """Parse FalkorDB result_set into list of dicts using header names.""" - headers = [h[1] for h in result.header] - rows = [] - for row in result.result_set: - rows.append({headers[i]: row[i] for i in range(len(headers))}) - return rows - - -def vector_search(graph, query: str, user_id: str, limit: int = 20): - """Run vector search and return results.""" - embedding = generate_embeddings_sync([query])[0] - - result = graph.query( - """ - CALL db.idx.vector.queryNodes('ConvChunk', 'embedding', $limit, vecf32($vector)) - YIELD node AS chunk, score - WHERE chunk.user_id = $user_id - MATCH (doc:ConvDoc)-[:HAS_CHUNK]->(chunk) - RETURN chunk.id AS chunk_id, chunk.section_title AS section, - doc.title AS title, doc.date AS date, score - ORDER BY score DESC - """, - params={"vector": embedding, "limit": limit, "user_id": user_id}, - ) - return _parse_results(result) - - -def fulltext_search(graph, query: str, user_id: str): - """Run full-text search and return results.""" - result = graph.query( - """ - CALL db.idx.fulltext.queryNodes('ConvChunk', $search_term) - YIELD node AS chunk, score - WHERE chunk.user_id = $user_id - MATCH (doc:ConvDoc)-[:HAS_CHUNK]->(chunk) - RETURN chunk.id AS chunk_id, chunk.section_title AS section, - doc.title AS title, doc.date AS date, score - ORDER BY score DESC - """, - params={"search_term": query, "user_id": user_id}, - ) - return _parse_results(result) - - -def test_hybrid_merge(graph): - """Test that hybrid merge combines vector + full-text results.""" - print_header("Test: Hybrid Merge (Vector + Full-Text)") - - query = "project deadline September" - - vec_results = vector_search(graph, query, TEST_USER) - ft_results = fulltext_search(graph, query, TEST_USER) - - print(f" Query: '{query}'") - print(f" Vector results: {len(vec_results)}") - print(f" Full-text results: {len(ft_results)}") - - # Merge - merged = compute_hybrid_scores(vec_results, ft_results) - - print(f"\n Hybrid results (top 5):") - print_results_table( - merged[:5], - ["final_score", "relevance_score", "recency_score", "title", "section"], - ) - - if merged: - print_pass(f"Hybrid merge produced {len(merged)} results") - - # Check that both sources contribute - has_vector = any(r["vector_score"] > 0 for r in merged) - has_text = any(r["text_score"] > 0 for r in merged) - if has_vector: - print_pass("Vector scores contributing to results") - else: - print_fail("No vector scores in merged results") - if has_text: - print_pass("Full-text scores contributing to results") - else: - print_fail( - "No full-text scores in merged results (might be expected if query terms not exact)" - ) - else: - print_fail("Hybrid merge returned no results") - - -def test_recency_bias(graph): - """Test that recency bias boosts recent conversations.""" - print_header("Test: Recency Bias") - - # Use a generic query that matches multiple conversations - query = "appointment schedule plans" - - vec_results = vector_search(graph, query, TEST_USER) - ft_results = fulltext_search(graph, query, TEST_USER) - - # Compare with and without recency - no_recency = compute_hybrid_scores( - vec_results, ft_results, recency_half_life_days=99999, recency_floor=1.0 - ) - with_recency = compute_hybrid_scores( - vec_results, ft_results, recency_half_life_days=30, recency_floor=0.5 - ) - - print(f" Query: '{query}'") - print(f"\n WITHOUT recency bias:") - print_results_table( - no_recency[:5], - ["final_score", "recency_score", "title", "date"], - ) - print(f"\n WITH recency bias (30-day half-life):") - print_results_table( - with_recency[:5], - ["final_score", "recency_score", "title", "date"], - ) - - if with_recency: - # Most recent conversation (grocery, 2026-03-16) should rank higher with recency - recency_scores = {r["title"]: r["recency_score"] for r in with_recency} - if recency_scores: - print_pass("Recency scores computed for all results") - - # Check that newer dates have higher recency scores - dates_and_recency = [(r["date"], r["recency_score"]) for r in with_recency] - print(f"\n Date -> Recency mapping:") - for d, rc in sorted(set(dates_and_recency)): - print(f" {d} -> {rc:.4f}") - else: - print_fail("No recency scores found") - else: - print_fail("No results to test recency on") - - -def test_exact_keyword_boost(graph): - """Test that exact keyword match (BM25) boosts relevance.""" - print_header("Test: Exact Keyword Boost") - - # Search for "Dr. Chen" -- should strongly prefer dentist via BM25 - query_semantic = "doctor dental health" - query_exact = "Dr. Chen" - - vec_results = vector_search(graph, query_semantic, TEST_USER) - ft_results_exact = fulltext_search(graph, query_exact, TEST_USER) - - # Hybrid with exact keyword - merged = compute_hybrid_scores(vec_results, ft_results_exact) - - print(f" Semantic query: '{query_semantic}'") - print(f" Exact keyword: '{query_exact}'") - print(f"\n Hybrid results:") - print_results_table( - merged[:5], ["final_score", "vector_score", "text_score", "title", "section"] - ) - - if merged: - # Chunks with exact "Dr. Chen" match should have text_score > 0 - exact_matches = [r for r in merged if r["text_score"] > 0] - if exact_matches: - print_pass(f"{len(exact_matches)} results boosted by exact keyword match") - # These should all be from the dentist conversation - dentist_exact = [ - r for r in exact_matches if "Dentist" in r.get("title", "") - ] - if dentist_exact: - print_pass("Exact keyword matches are from dentist conversation") - else: - print_fail("No results had text_score > 0 for exact keyword") - - -if __name__ == "__main__": - graph = get_graph() - test_hybrid_merge(graph) - test_recency_bias(graph) - test_exact_keyword_boost(graph) diff --git a/backends/advanced/tests/scripts/graph-validation/test_vector_search.py b/backends/advanced/tests/scripts/graph-validation/test_vector_search.py deleted file mode 100644 index a2bec151..00000000 --- a/backends/advanced/tests/scripts/graph-validation/test_vector_search.py +++ /dev/null @@ -1,174 +0,0 @@ -"""Validate vector index search with user scoping. - -Tests: -1. Semantic search finds relevant chunks (project deadline -> Q3 meeting) -2. User scoping prevents cross-user data leakage -3. Score ordering is correct - -Usage: - uv run python tests/scripts/graph-validation/test_vector_search.py -""" - -from utils import ( - generate_embeddings_sync, - get_graph, - print_fail, - print_header, - print_pass, - print_results_table, -) - -TEST_USER = "test_user_001" -OTHER_USER = "other_user_002" - - -def _parse_vector_results(result, fields): - """Parse FalkorDB result_set into list of dicts using header names.""" - headers = [h[1] for h in result.header] - rows = [] - for row in result.result_set: - rows.append({headers[i]: row[i] for i in range(len(headers))}) - return rows - - -def test_basic_vector_search(graph): - """Search for project deadline -- should return Q3 meeting chunks.""" - print_header("Test: Basic Vector Search") - - query = "project deadline and timeline" - embedding = generate_embeddings_sync([query])[0] - - result = graph.query( - """ - CALL db.idx.vector.queryNodes('ConvChunk', 'embedding', $limit, vecf32($vector)) - YIELD node AS chunk, score - WHERE chunk.user_id = $user_id - MATCH (doc:ConvDoc)-[:HAS_CHUNK]->(chunk) - RETURN chunk.id AS chunk_id, chunk.section_title AS section, - chunk.text AS text, doc.title AS title, score - ORDER BY score DESC - """, - params={"vector": embedding, "limit": 10, "user_id": TEST_USER}, - ) - results = _parse_vector_results(result, []) - - print(f" Query: '{query}'") - print(f" Results: {len(results)}") - print_results_table(results, ["score", "title", "section"]) - - # Verify: top result should be from Q3 meeting - if results and "Q3" in results[0].get("title", ""): - print_pass("Top result is from Q3 meeting (expected)") - elif results: - print_fail(f"Top result is '{results[0].get('title')}', expected Q3 meeting") - else: - print_fail("No results returned") - - # Verify scores are in descending order - scores = [r["score"] for r in results] - if scores == sorted(scores, reverse=True): - print_pass("Scores in descending order") - else: - print_fail("Scores NOT in descending order") - - return results - - -def test_semantic_relevance(graph): - """Search for health/medical topic -- should return dentist chunks.""" - print_header("Test: Semantic Relevance") - - query = "dental health and medical procedures" - embedding = generate_embeddings_sync([query])[0] - - result = graph.query( - """ - CALL db.idx.vector.queryNodes('ConvChunk', 'embedding', $limit, vecf32($vector)) - YIELD node AS chunk, score - WHERE chunk.user_id = $user_id - MATCH (doc:ConvDoc)-[:HAS_CHUNK]->(chunk) - RETURN chunk.id AS chunk_id, chunk.section_title AS section, - doc.title AS title, score - ORDER BY score DESC - """, - params={"vector": embedding, "limit": 5, "user_id": TEST_USER}, - ) - results = _parse_vector_results(result, []) - - print(f" Query: '{query}'") - print_results_table(results, ["score", "title", "section"]) - - if results and "Dentist" in results[0].get("title", ""): - print_pass("Top result is dentist conversation (expected)") - elif results: - print_fail(f"Top result is '{results[0].get('title')}', expected Dentist") - else: - print_fail("No results returned") - - -def test_user_scoping(graph): - """Verify user_id filtering prevents cross-user results.""" - print_header("Test: User Scoping") - - query = "dentist crown procedure" - embedding = generate_embeddings_sync([query])[0] - - # Search as TEST_USER - result = graph.query( - """ - CALL db.idx.vector.queryNodes('ConvChunk', 'embedding', $limit, vecf32($vector)) - YIELD node AS chunk, score - WHERE chunk.user_id = $user_id - RETURN chunk.id AS chunk_id, chunk.user_id AS user_id, score - ORDER BY score DESC - """, - params={"vector": embedding, "limit": 20, "user_id": TEST_USER}, - ) - test_user_results = _parse_vector_results(result, []) - - # Search as OTHER_USER - result = graph.query( - """ - CALL db.idx.vector.queryNodes('ConvChunk', 'embedding', $limit, vecf32($vector)) - YIELD node AS chunk, score - WHERE chunk.user_id = $user_id - RETURN chunk.id AS chunk_id, chunk.user_id AS user_id, score - ORDER BY score DESC - """, - params={"vector": embedding, "limit": 20, "user_id": OTHER_USER}, - ) - other_user_results = _parse_vector_results(result, []) - - print(f" TEST_USER results: {len(test_user_results)}") - print(f" OTHER_USER results: {len(other_user_results)}") - - # Verify no cross-user leakage - test_user_ids = {r["user_id"] for r in test_user_results} - other_user_ids = {r["user_id"] for r in other_user_results} - - if test_user_ids <= {TEST_USER}: - print_pass("TEST_USER results contain only TEST_USER data") - else: - print_fail(f"TEST_USER results leaked: {test_user_ids}") - - if other_user_ids <= {OTHER_USER}: - print_pass("OTHER_USER results contain only OTHER_USER data") - else: - print_fail(f"OTHER_USER results leaked: {other_user_ids}") - - # Verify different result counts (other user has fewer conversations) - if len(test_user_results) > len(other_user_results): - print_pass( - f"TEST_USER has more results ({len(test_user_results)} vs {len(other_user_results)})" - ) - else: - print_pass( - f"Result counts: TEST={len(test_user_results)}, OTHER={len(other_user_results)}" - ) - - -if __name__ == "__main__": - graph = get_graph() - test_basic_vector_search(graph) - test_semantic_relevance(graph) - test_user_scoping(graph) diff --git a/backends/advanced/tests/scripts/graph-validation/utils.py b/backends/advanced/tests/scripts/graph-validation/utils.py deleted file mode 100644 index 0d5a8d1a..00000000 --- a/backends/advanced/tests/scripts/graph-validation/utils.py +++ /dev/null @@ -1,292 +0,0 @@ -"""Shared utilities for FalkorDB graph memory validation scripts.""" - -import asyncio -import math -import os -import re -import sys -from datetime import datetime, timezone -from pathlib import Path -from typing import Dict, List, Optional, Tuple - -from dotenv import load_dotenv -from falkordb import FalkorDB - -# Load .env from backends/advanced/ -# __file__ is .../backends/advanced/tests/scripts/graph-validation/utils.py -# parents: [0]=graph-validation, [1]=scripts, [2]=tests, [3]=advanced -_env_path = Path(__file__).resolve().parents[3] / ".env" -load_dotenv(_env_path) - -# --- FalkorDB Connection --- - -_falkordb_host = os.getenv("FALKORDB_HOST", "localhost") -if _falkordb_host == "falkordb": - _falkordb_host = "localhost" -FALKORDB_PORT = int(os.getenv("FALKORDB_PORT", "6381")) -FALKORDB_GRAPH = os.getenv("FALKORDB_GRAPH", "chronicle") - -# --- OpenAI Embeddings --- - -OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "") -OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1") -EMBEDDING_MODEL = "text-embedding-3-small" -EMBEDDING_DIMENSIONS = 1536 - - -def get_graph(host=None, port=None, graph_name=None): - """Create a FalkorDB graph connected to the local instance.""" - host = host or _falkordb_host - port = port or FALKORDB_PORT - graph_name = graph_name or FALKORDB_GRAPH - db = FalkorDB(host=host, port=port) - return db.select_graph(graph_name) - - -async def generate_embeddings(texts: List[str]) -> List[List[float]]: - """Generate embeddings using OpenAI API.""" - from openai import AsyncOpenAI - - client = AsyncOpenAI(api_key=OPENAI_API_KEY, base_url=OPENAI_BASE_URL) - response = await client.embeddings.create(model=EMBEDDING_MODEL, input=texts) - return [data.embedding for data in response.data] - - -def generate_embeddings_sync(texts: List[str]) -> List[List[float]]: - """Synchronous wrapper for embedding generation.""" - return asyncio.run(generate_embeddings(texts)) - - -# --- Markdown Chunking --- - - -def split_on_headers(markdown: str) -> List[Dict[str, str]]: - """Split markdown into chunks by ### headers. - - Returns list of dicts with 'section_title' and 'text' keys. - Frontmatter (---...---) is stripped and returned as metadata. - """ - # Strip frontmatter - fm_match = re.match(r"^---\s*\n(.*?)\n---\s*\n", markdown, re.DOTALL) - if fm_match: - content = markdown[fm_match.end() :] - else: - content = markdown - - chunks = [] - # Split on ### headers (h3) - parts = re.split(r"^(###\s+.+)$", content, flags=re.MULTILINE) - - # parts alternates: [pre-header text, header, body, header, body, ...] - # First element is text before any header - i = 0 - if parts[0].strip(): - chunks.append({"section_title": "Introduction", "text": parts[0].strip()}) - i = 1 - else: - i = 1 - - while i < len(parts): - if parts[i].startswith("###"): - title = parts[i].replace("###", "").strip() - body = parts[i + 1].strip() if i + 1 < len(parts) else "" - if body: - chunks.append({"section_title": title, "text": body}) - i += 2 - else: - i += 1 - - return chunks - - -def parse_frontmatter(markdown: str) -> Dict[str, str]: - """Extract YAML frontmatter as a simple dict.""" - fm_match = re.match(r"^---\s*\n(.*?)\n---\s*\n", markdown, re.DOTALL) - if not fm_match: - return {} - result = {} - for line in fm_match.group(1).splitlines(): - if ":" in line: - key, _, value = line.partition(":") - result[key.strip()] = value.strip().strip('"').strip("'") - return result - - -# --- Entity Parsing --- - - -def parse_people_section(markdown: str) -> List[Dict[str, str]]: - """Parse the ### People section to extract entity names and descriptions. - - Expected format: - ### People - - John (coworker, project lead) - - Dr. Sarah Chen (dentist, referred by Jane) - - "The IT guy" (unnamed, fixed the printer) - - Returns list of dicts with 'name' and 'description'. - """ - # Find the People section - match = re.search( - r"^###\s+People\s*\n(.*?)(?=^###|\Z)", markdown, re.MULTILINE | re.DOTALL - ) - if not match: - return [] - - people = [] - for line in match.group(1).strip().splitlines(): - line = line.strip() - if not line.startswith("-"): - continue - line = line[1:].strip() - - # Try to parse "Name (description)" format - paren_match = re.match(r'^["""]?(.+?)["""]?\s*\((.+?)\)\s*$', line) - if paren_match: - people.append( - { - "name": paren_match.group(1).strip(), - "description": paren_match.group(2).strip(), - } - ) - else: - # Just a name with no description - people.append({"name": line.strip('"').strip("'"), "description": ""}) - - return people - - -def parse_action_items(markdown: str) -> List[Dict[str, str]]: - """Parse ### Action Items section. - - Expected format: - ### Action Items - - [ ] Send Q3 report to John by Friday - - [x] Book dentist follow-up - - Returns list of dicts with 'text' and 'done' (bool). - """ - match = re.search( - r"^###\s+Action Items\s*\n(.*?)(?=^###|\Z)", markdown, re.MULTILINE | re.DOTALL - ) - if not match: - return [] - - items = [] - for line in match.group(1).strip().splitlines(): - line = line.strip() - if not line.startswith("-"): - continue - line = line[1:].strip() - - done = False - if line.startswith("[x]") or line.startswith("[X]"): - done = True - line = line[3:].strip() - elif line.startswith("[ ]"): - line = line[3:].strip() - - if line: - items.append({"text": line, "done": done}) - - return items - - -# --- Hybrid Search Scoring --- - - -def compute_hybrid_scores( - vector_results: List[Dict], - fulltext_results: List[Dict], - vector_weight: float = 0.7, - text_weight: float = 0.3, - recency_half_life_days: float = 30.0, - recency_floor: float = 0.5, -) -> List[Dict]: - """Merge vector and full-text results with recency bias. - - Each result dict must have: 'chunk_id', 'score', 'date' (ISO string or datetime). - Additional fields are preserved. - """ - now = datetime.now(timezone.utc) - merged: Dict[str, Dict] = {} - - for r in vector_results: - cid = r["chunk_id"] - merged[cid] = {**r, "vector_score": r["score"], "text_score": 0.0} - - for r in fulltext_results: - cid = r["chunk_id"] - if cid in merged: - merged[cid]["text_score"] = r["score"] - else: - merged[cid] = {**r, "vector_score": 0.0, "text_score": r["score"]} - - results = [] - for entry in merged.values(): - # Parse date - d = entry.get("date") - if isinstance(d, str): - d = datetime.fromisoformat(d.replace("Z", "+00:00")) - if d is None: - d = now - - if d.tzinfo is None: - d = d.replace(tzinfo=timezone.utc) - - age_days = (now - d).total_seconds() / 86400.0 - - relevance = ( - vector_weight * entry["vector_score"] + text_weight * entry["text_score"] - ) - recency = max( - recency_floor, math.exp(-0.693 * age_days / recency_half_life_days) - ) - entry["relevance_score"] = relevance - entry["recency_score"] = recency - entry["final_score"] = relevance * recency - results.append(entry) - - results.sort(key=lambda x: x["final_score"], reverse=True) - return results - - -# --- Pretty Printing --- - -GREEN = "\033[92m" -RED = "\033[91m" -YELLOW = "\033[93m" -BOLD = "\033[1m" -RESET = "\033[0m" - - -def print_pass(msg: str): - print(f" {GREEN}PASS{RESET} {msg}") - - -def print_fail(msg: str): - print(f" {RED}FAIL{RESET} {msg}") - - -def print_header(msg: str): - print(f"\n{BOLD}{'='*60}{RESET}") - print(f"{BOLD}{msg}{RESET}") - print(f"{BOLD}{'='*60}{RESET}") - - -def print_results_table(results: List[Dict], fields: List[str], max_text_len: int = 60): - """Print results as a simple table.""" - if not results: - print(" (no results)") - return - for i, r in enumerate(results): - parts = [] - for f in fields: - val = r.get(f, "") - if isinstance(val, float): - parts.append(f"{f}={val:.4f}") - elif isinstance(val, str) and len(val) > max_text_len: - parts.append(f"{f}={val[:max_text_len]}...") - else: - parts.append(f"{f}={val}") - print(f" [{i+1}] {', '.join(parts)}") diff --git a/docs/README.md b/docs/README.md index 28a282a0..45c1d88f 100644 --- a/docs/README.md +++ b/docs/README.md @@ -8,7 +8,6 @@ day-to-day operation, and use [AGENTS.md](../AGENTS.md) for development conventi - [Project overview](overview.md): components, deployment topology, and repository layout - [Testing and coverage](testing.md): fast Python lanes, coverage reports, and integration tests -- [Test coverage audit](test-coverage-audit.md): current baseline, gaps, and cleanup plan - [Audio pipeline](audio-pipeline-architecture.md): session, transcription, and memory flow - [Initialization system](init-system.md): setup wizard and service orchestration - [Podman](podman.md): rootless containers, GPU access, and engine migration diff --git a/docs/test-coverage-audit.md b/docs/test-coverage-audit.md deleted file mode 100644 index 5bf183b5..00000000 --- a/docs/test-coverage-audit.md +++ /dev/null @@ -1,214 +0,0 @@ -# Test Coverage Audit - -Audit date: 2026-07-15 -Git baseline: `0e2eefa5` on `dev` -Scope: Git-tracked code is the reproducible baseline. Current untracked source and tests were -examined separately because they are not available to CI or a fresh contributor checkout. - -## Executive summary - -Chronicle has broad end-to-end scenario coverage, especially around backend endpoints, audio -streaming, queues, and conversations. It does not currently have a trustworthy repository-wide -code coverage number. - -The best available measurement is for the advanced Python backend: - -- The current working-tree pytest suite executed about **20.0% of tracked backend statements** - (`4,475 / 22,321`). This is an exposure measurement, not a passing baseline: the run had 127 - passes, 14 failures, and 33 errors after two collection failures were excluded. -- The smaller, Git-tracked, locally collectable unit subset exercised **8% of the statements - Coverage.py could discover in that run** and had 69 passes and 3 failures. Coverage discovery - omitted unimported namespace-package directories, so 8% is an upper bound rather than a full - backend denominator. -- The backend's most important orchestration code is the least directly tested area: workers had - **8.4%** measured coverage and controllers had **11.2%** in the broader working-tree run. -- All TypeScript/JavaScript products have **no configured test runner and no direct test files**: - the main web UI, mobile app, speaker-recognition UI, and vault graph UI. -- CI runs Robot tests and one speaker-recognition integration test, but it does **not run backend, - root-tooling, or ASR pytest suites**, and it publishes no line or branch coverage. - -The project therefore has meaningful workflow protection, but weak feedback on which branches and -failure paths are exercised. Robot's documented "70% coverage" means percentage of Robot tests -selected, not percentage of application code covered. - -## Inventory - -### Test suites - -| Surface | Current tests | Audit result | -| --- | --- | --- | -| Advanced backend | 7 tracked top-level pytest files; 12 more top-level files are currently untracked | Default collection fails; no pytest CI job | -| Cross-repo Robot suite | 232 tests across 34 suites | 118 endpoint, 64 integration, 28 ASR, 17 configuration, 4 infrastructure, 1 browser | -| Root setup/lifecycle tooling | 5 pytest files, about 80 test functions | Collection fails in 2 files; remaining run had 46 pass / 10 fail | -| ASR services | 4 pytest files, 63 collected cases | 49 pass / 9 skip / 2 fail / 3 environment errors | -| Speaker recognition | 1 large pytest integration scenario plus 2 ad hoc scripts | Secret-, model-, and Docker-dependent CI only | -| Frontends | No test files or test scripts | Build/lint only; no measured coverage | -| TTS and most plugins/extras | No direct automated tests | Some behavior is reached indirectly by Robot tests | - -There are 232 Robot cases, but the normal no-API lane selects only 166 after tag exclusions. The -suite contains 17 tests with unconditional `Skip`, including eight mobile placeholders and five SDK -placeholders. These should not be counted as implemented coverage. - -### Advanced backend line coverage - -This table comes from the broader current-working-tree pytest run. Failing tests still execute -lines, so the numbers indicate code reached by tests, not verified behavior. - -| Backend area | Statements | Covered | Coverage | -| --- | ---: | ---: | ---: | -| Models | 711 | 479 | 67.4% | -| Plugins framework | 453 | 136 | 30.0% | -| Routers | 3,598 | 943 | 26.2% | -| Utilities | 1,147 | 262 | 22.8% | -| Services | 5,490 | 1,207 | 22.0% | -| Package root | 2,838 | 625 | 22.0% | -| Observability | 246 | 51 | 20.7% | -| Clients | 260 | 47 | 18.1% | -| Controllers | 4,169 | 467 | 11.2% | -| Workers | 3,532 | 297 | 8.4% | - -Across the measured backend, 34 files had zero executed lines and 46 more were below 20%. The -largest high-risk gaps include: - -- `app_factory.py`: 289 statements, 0% -- `task_manager.py`: 192 statements, 0% -- memory agent edit/tool modules: 611 combined statements, 0% -- worker orchestrator modules: 465 combined statements, 0% -- `system_controller.py`: 1,103 statements, 10.0% -- `websocket_controller.py`: 712 statements, 9.3% -- `conversation_jobs.py`: 715 statements, 12.2% -- `transcription_jobs.py`: 561 statements, 8.9% -- `speaker_recognition_client.py`: 624 statements, 5.9% - -The root management tests measured 18% across five imported modules. `updates.py` reached 93%, but -`services.py`, `config_manager.py`, `discovery.py`, and `setup_utils.py` were between 10% and 20%. -This is not a whole-root percentage because two configured modules were never imported. - -## Suite health findings - -### 1. Test execution is not reproducible from a clean checkout - -- Advanced backend collection references deleted or changed APIs in `test_obsidian_service.py` and - `test_vad_analysis.py`. -- Root tests reference removed `merge_configs`, `read_config_yml`, and service lifecycle APIs. -- The ASR default suite starts Docker from an otherwise unit-oriented `pytest tests` command; Docker - absence becomes an error instead of an integration skip. Two independent ROCm Dockerfile checks - also currently fail. -- Robot dry-run found 16 failures in configuration suites because `ruamel-yaml` is imported but is - absent from `tests/test-requirements.txt`. -- The browser suite requires `rfbrowser init`, but that setup is absent from the test runner and CI. - -### 2. Test location does not consistently describe test type - -- `backends/advanced/tests/` mixes pure units, MongoDB integration tests, and manual graph-validation - scripts. The untracked graph scripts are automatically collected by pytest but expect command-line - arguments as fixtures. -- Root `tests/unit/` mixes root lifecycle scripts with advanced-backend configuration behavior. -- Robot `endpoints/` tests are full-stack HTTP contract/integration tests, not unit endpoint tests. -- Robot `integration/` includes genuine end-to-end flows, SDK contract tests, and placeholder mobile - scenarios. -- Hardware-, secret-, browser-, and container-dependent tests are separated mainly by tags, but tag - hygiene is not enforced. - -### 3. Tags and documentation have drifted - -- `tests/tags.md` says there are 15 approved tags, later says 14, and finally asks whether one of 11 - tags can be used. -- Ten tests in `client_queue_tests.robot` use prohibited tags such as `positive`, `negative`, - `security`, `client`, `jobs`, and `integration`. -- Eighteen Robot tests have no tags. -- Only four of nine SDK tests carry the `sdk` tag. The other five are placeholder skips that remain in - default selection. -- `tests/README.md` and `AGENTS.md` document nonexistent Make targets including `test-all`, - `test-endpoints`, `test-integration`, and `test-infra`; the actual targets are `all`, `endpoints`, - `integration`, and `infra`. - -### 4. CI enforces workflows, not code coverage - -The required OSS-friendly PR lane runs 166 no-secret Robot cases. That is valuable, but no workflow -runs the fast Python tests or any frontend tests. There is no `.coveragerc`, branch coverage setting, -coverage artifact, combined coverage job, patch coverage check, or minimum threshold. - -Path filters also mean changes to root lifecycle tooling, ASR, TTS, plugins, and frontends do not -trigger the main Robot workflow unless they touch its listed backend/test paths. - -## Proposed ownership model - -Tests should live with the component whose code they primarily validate. The repository-level -`tests/` directory should own cross-component acceptance behavior only. - -| Type | Location | Allowed dependencies | PR policy | -| --- | --- | --- | --- | -| Unit | `<component>/tests/unit/` | No network, Docker, database, secrets, or model downloads | Required; seconds | -| Component | `<component>/tests/component/` | In-process app plus fakes/in-memory stores | Required; minutes | -| Integration | `<component>/tests/integration/` | Real database/service containers | Required where CPU/no-secret | -| Contract | `tests/contract/` | Black-box API against composed Chronicle | Required no-secret subset | -| End to end | `tests/e2e/` | Multiple services and full user workflow | Required smoke subset; full scheduled | -| Environment | `tests/environment/{browser,gpu,hardware,secrets}/` | Explicit specialized runner | Optional/scheduled unless relevant | - -Use pytest markers matching execution requirements, not business domains: -`unit`, `component`, `integration`, `docker`, `gpu`, `external_api`, and `slow`. Keep Robot business -tags for selecting product behavior, but validate them against one canonical allowlist in CI. - -## Recommended rollout - -### Remediation started - -The first infrastructure tranche was implemented after this baseline was recorded: - -- Root tooling, advanced backend, and ASR now have branch-coverage configuration and a dedicated - Python CI workflow that uploads XML and HTML reports. -- Fast backend and ASR lanes explicitly exclude their MongoDB and Parakeet container integration - modules. Manual graph-validation scripts are no longer eligible for default pytest collection. -- Robot tags now have a machine-readable allowlist and validation command. Invalid legacy tags were - mapped to the existing taxonomy, and all SDK placeholders inherit the `sdk` execution tag. -- The 17 container-independent configuration tests are part of `make all`; they pass locally. -- The ASR fast lane now produces an **11.3% branch-coverage baseline** while retaining its two known - failing Dockerfile assertions for the parallel test-repair work. - -Coverage thresholds remain intentionally disabled until the known collection and assertion failures -are repaired. The new Python CI jobs will surface those failures rather than hiding them. - -### Phase 0: make the existing signal trustworthy - -1. Add separate CI jobs for root pytest, advanced-backend unit pytest, and ASR unit pytest. -2. Move or exclude manual graph-validation scripts so default pytest cannot collect them. -3. Split Docker/Mongo/GPU tests from unit commands and make missing prerequisites explicit skips. -4. Repair or remove tests for deleted APIs; do not retain compatibility code solely to satisfy stale - tests. -5. Fix Robot dependencies, browser setup/selection, placeholder selection, tags, Make targets, and - documentation. - -### Phase 1: establish coverage without blocking cleanup - -1. Configure line and branch coverage per Python component and publish XML/HTML artifacts. -2. Run the backend process under Coverage.py in Robot containers, save parallel data files, and - combine them with pytest coverage. Report pytest and Robot flags separately as well as combined. -3. Add Vitest and React Testing Library to both Vite UIs; add the Expo-supported Jest/React Native - Testing Library stack to the mobile app. -4. Record the baseline by component. Initially enforce only: green tests, no reduction in total - coverage, and about 80% changed-line coverage. Do not impose an arbitrary 80% repository-wide - target on legacy code. - -### Phase 2: cover failure-prone boundaries - -Prioritize tests around conversation lifecycle, transcription/audio persistence, worker retries and -idempotency, task cleanup, memory-agent edits, authentication/data isolation, and plugin failure -containment. For frontends, prioritize auth/token state, API error states, reconnect behavior, audio -session state, and destructive actions before snapshot-heavy presentation tests. - -### Phase 3: ratchet by subsystem - -Once the suites are green and stable, set modest per-component floors and raise them gradually. -Coverage should remain a discovery tool: mutation testing or focused fault injection on lifecycle -and worker code will reveal weak assertions that line coverage alone cannot detect. - -## Definition of done for the cleanup - -- One documented command runs every fast, no-secret test from a clean checkout. -- Every test has one owning component and one execution class. -- No placeholders are reported as passing coverage. -- All PRs run Python unit/component tests, frontend tests, and the no-secret contract smoke suite. -- Coverage is reported per component and combined for the advanced backend. -- Changed code cannot reduce coverage, while legacy coverage increases through a ratchet rather than - a one-time repository-wide threshold. diff --git a/docs/testing.md b/docs/testing.md index 8a72ad62..1258ae60 100644 --- a/docs/testing.md +++ b/docs/testing.md @@ -25,8 +25,7 @@ PYTHONPATH=backends/advanced/src uv run \ ### Advanced backend -The default fast lane excludes the MongoDB integration module and manual scripts under -`tests/scripts/`: +The default fast lane excludes the MongoDB integration module: ```bash cd backends/advanced @@ -72,6 +71,3 @@ green; CI still fails when a test fails. The repository-level `tests/` directory owns composed Robot Framework behavior. Its current commands and environment setup are documented in [`tests/README.md`](../tests/README.md). - -The coverage baseline and test-ownership cleanup plan are recorded in the -[test coverage audit](test-coverage-audit.md). diff --git a/extras/asr-services/scripts/validate_nemotron.py b/extras/asr-services/scripts/validate_nemotron.py deleted file mode 100644 index f4cf09fa..00000000 --- a/extras/asr-services/scripts/validate_nemotron.py +++ /dev/null @@ -1,143 +0,0 @@ -"""Validate the Nemotron streaming ASR service on a WAV file. - -Exercises both endpoints the Chronicle backend uses: - * POST /transcribe — offline full-file decode - * WS /stream — cache-aware streaming (binary PCM in, interim/final out), - driven exactly like backend RegistryStreamingProvider: - raw PCM frames, then {"type":"CloseStream"} to finalize. - -Usage: - uv run python scripts/validate_nemotron.py <wav_path> [--port 8771] -""" - -import argparse -import asyncio -import json -import sys -import time -import wave - -import httpx -import websockets - -CHUNK_SECONDS = 0.25 # mirrors AudioStreamProducer's 0.25s chunks - - -def read_pcm(path: str): - with wave.open(path, "rb") as wf: - assert wf.getnchannels() == 1, "expected mono" - assert wf.getsampwidth() == 2, "expected 16-bit" - rate = wf.getframerate() - pcm = wf.readframes(wf.getnframes()) - return pcm, rate - - -async def test_offline(base_url: str, path: str): - print(f"\n=== /transcribe (offline) ===") - t0 = time.time() - with open(path, "rb") as f: - files = {"file": ("audio.wav", f, "audio/wav")} - async with httpx.AsyncClient(timeout=300) as client: - resp = await client.post(f"{base_url}/transcribe", files=files) - resp.raise_for_status() - data = resp.json() - dt = time.time() - t0 - print(f" time: {dt:.1f}s") - print(f" words: {len(data.get('words') or [])}") - print(f" text: {data.get('text', '')!r}") - return data.get("text", "") - - -async def test_streaming(ws_url: str, pcm: bytes, rate: int): - print(f"\n=== /stream (cache-aware streaming) ===") - chunk_bytes = int(rate * 2 * CHUNK_SECONDS) - interims = [] - final = None - t0 = time.time() - first_interim_at = None - # ping_interval=None: a client that continuously pushes audio doesn't rely on - # library-level keepalive (GIL-heavy RNNT decode can starve the server loop - # past a 20s ping window while it drains its backlog after CloseStream). - async with websockets.connect(ws_url, max_size=None, ping_interval=None) as ws: - - async def sender(): - for i in range(0, len(pcm), chunk_bytes): - await ws.send(pcm[i : i + chunk_bytes]) - await asyncio.sleep(CHUNK_SECONDS) # real-time pacing - await ws.send(json.dumps({"type": "CloseStream"})) - - send_task = asyncio.create_task(sender()) - try: - while True: - msg = await asyncio.wait_for(ws.recv(), timeout=120) - data = json.loads(msg) - if data.get("type") == "interim": - if first_interim_at is None: - first_interim_at = time.time() - t0 - interims.append(data.get("text", "")) - elif data.get("type") == "final": - final = data - break - finally: - send_task.cancel() - dt = time.time() - t0 - print(f" interims received: {len(interims)}") - if first_interim_at is not None: - print(f" first interim at: {first_interim_at:.2f}s into stream") - if interims: - print(f" sample interims:") - for t in interims[: min(3, len(interims))]: - print(f" - {t!r}") - print(f" - (last) {interims[-1]!r}") - print(f" final words: {len(final.get('words') or []) if final else 0}") - print(f" final text: {final.get('text', '') if final else None!r}") - print(f" total stream time: {dt:.1f}s") - return (final or {}).get("text", ""), interims - - -async def main(): - ap = argparse.ArgumentParser() - ap.add_argument("wav") - ap.add_argument("--port", type=int, default=8772) - ap.add_argument("--host", default="localhost") - args = ap.parse_args() - - base_url = f"http://{args.host}:{args.port}" - ws_url = f"ws://{args.host}:{args.port}/stream" - - # Wait for health - print(f"Waiting for {base_url}/health ...") - async with httpx.AsyncClient(timeout=10) as client: - for _ in range(120): - try: - r = await client.get(f"{base_url}/health") - if r.status_code == 200 and r.json().get("status") == "healthy": - print(f" healthy: {r.json()}") - break - except Exception: - pass - await asyncio.sleep(5) - else: - print("Service did not become healthy in time") - sys.exit(1) - - pcm, rate = read_pcm(args.wav) - print(f"audio: {len(pcm)} bytes @ {rate}Hz ({len(pcm)/(rate*2):.1f}s)") - - offline_text = await test_offline(base_url, args.wav) - stream_text, interims = await test_streaming(ws_url, pcm, rate) - - print("\n=== SUMMARY ===") - print(f"streaming produced {len(interims)} interim updates") - print(f"offline final text len: {len(offline_text)}") - print(f"streaming final text len: {len(stream_text)}") - ok = bool(stream_text.strip()) and len(interims) > 1 - print( - f"RESULT: {'PASS' if ok else 'FAIL'} " - f"(needs non-empty streaming final + multiple interims)" - ) - sys.exit(0 if ok else 2) - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/extras/speaker-recognition/scripts/backfill_segment_embeddings.py b/extras/speaker-recognition/scripts/backfill_segment_embeddings.py deleted file mode 100644 index 2afc4e97..00000000 --- a/extras/speaker-recognition/scripts/backfill_segment_embeddings.py +++ /dev/null @@ -1,121 +0,0 @@ -"""Backfill per-clip embeddings (SpeakerAudioSegment rows) from enrollment audio. - -Historically the service stored only a single averaged centroid per speaker -(`speakers.embedding_data`); the `speaker_audio_segments.embedding` column was never -populated. The enrollment-health audit (and a future multi-vector gallery) need one -embedding per enrolled clip, so this one-time importer walks the on-disk enrollment -audio, embeds each clip with the wespeaker model, and inserts a SpeakerAudioSegment row. - -Idempotent: skips any (speaker_id, filename) that already has a row, so it's safe to -re-run after new enrollments. New enrollments persist their own segments going forward -(see enrollment.py:save_segment_record), so this is only for pre-existing data. - -Run inside the speaker-service container: - podman exec speaker-recognition_speaker-service-gpu_1 \ - python3 /app/scripts/backfill_segment_embeddings.py -""" - -import glob -import json -import os -import sys -from pathlib import Path - -import numpy as np -import torch - -# Make the package importable when run as a plain script inside the container. -sys.path.insert(0, "/app/src") - -from pyannote.audio import Audio # noqa: E402 -from pyannote.audio.pipelines.speaker_verification import ( # noqa: E402 - PretrainedSpeakerEmbedding, -) - -from simple_speaker_recognition.database import get_db_session # noqa: E402 -from simple_speaker_recognition.database.models import ( # noqa: E402 - Speaker, - SpeakerAudioSegment, -) - -DATA_DIR = Path(os.getenv("SPEAKER_DATA_DIR", "/app/data")) -ENROLL_DIR = DATA_DIR / "enrollment_audio" -MIN_SAMPLES = 400 # wespeaker fbank window (25 ms @ 16 kHz) - - -def main() -> None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - embedder = PretrainedSpeakerEmbedding( - "pyannote/wespeaker-voxceleb-resnet34-LM", device=device - ) - loader = Audio(sample_rate=16000, mono="downmix") - - def embed(path: str): - wav, _ = loader(path) - wav = wav.unsqueeze(0) - if wav.shape[-1] < MIN_SAMPLES: - wav = torch.nn.functional.pad(wav, (0, MIN_SAMPLES - wav.shape[-1])) - with torch.inference_mode(): - e = embedder(wav.to(device)) - if isinstance(e, torch.Tensor): - e = e.cpu().numpy() - e = np.asarray(e).reshape(-1) - n = np.linalg.norm(e) - return (e / n) if np.isfinite(n) and n > 0 else None - - session = get_db_session() - try: - speaker_ids = {s.id for s in session.query(Speaker).all()} - # Existing (speaker_id, filename) pairs for idempotency. - existing = { - (row.speaker_id, os.path.basename(row.audio_file_path)) - for row in session.query(SpeakerAudioSegment).all() - } - - added = skipped = orphan = 0 - for user_dir in sorted(glob.glob(str(ENROLL_DIR / "*"))): - for sdir in sorted(glob.glob(f"{user_dir}/*")): - speaker_id = os.path.basename(sdir) - if speaker_id not in speaker_ids: - orphan += 1 - continue - for wav in sorted(glob.glob(f"{sdir}/*.wav")): - fname = os.path.basename(wav) - if (speaker_id, fname) in existing: - skipped += 1 - continue - try: - dur = float(loader.get_duration(wav)) - vec = embed(wav) - except Exception as e: # noqa: BLE001 - print(f" SKIP {wav}: {e}") - continue - if vec is None: - print(f" NULL embedding {wav}") - continue - rel = str(Path(wav).resolve().relative_to(ENROLL_DIR.resolve())) - session.add( - SpeakerAudioSegment( - speaker_id=speaker_id, - audio_file_path=rel, - original_file_path=fname, - start_time=0.0, - end_time=dur, - duration_seconds=dur, - embedding=json.dumps(vec.astype(np.float32).tolist()), - ) - ) - added += 1 - session.commit() - - session.commit() - print( - f"Backfill done: added={added} skipped(existing)={skipped} " - f"orphan_dirs(no speaker row)={orphan}" - ) - finally: - session.close() - - -if __name__ == "__main__": - main() From a7c84c12a445a140816d125f2205399de1bb9e52 Mon Sep 17 00:00:00 2001 From: Ankush Malaker <43288948+AnkushMalaker@users.noreply.github.com> Date: Mon, 20 Jul 2026 10:28:15 +0000 Subject: [PATCH 16/16] chore: archive local scripts and plans --- CONNECTION_RESILIENCE_AND_UX_PLAN.md | 294 ----------------- .../scripts/delete_all_conversations_api.py | 131 -------- backends/advanced/scripts/laptop_client.py | 275 ---------------- extras/ml-training/README.md | 27 +- .../whisper-adapter-finetuning/README.md | 202 ------------ .../evaluate_sneeze_model.py | 123 -------- .../prepare_sneeze_data.py | 80 ----- .../sneeze_data.jsonl | 23 -- .../train_sneeze.py | 295 ------------------ .../scripts/enroll_from_backup.py | 181 ----------- .../scripts/simple_enroll.py | 65 ---- .../scripts/test_diarize.py | 65 ---- .../scripts/test_workflow.py | 220 ------------- 13 files changed, 2 insertions(+), 1979 deletions(-) delete mode 100644 CONNECTION_RESILIENCE_AND_UX_PLAN.md delete mode 100755 backends/advanced/scripts/delete_all_conversations_api.py delete mode 100644 backends/advanced/scripts/laptop_client.py delete mode 100644 extras/ml-training/whisper-adapter-finetuning/README.md delete mode 100644 extras/ml-training/whisper-adapter-finetuning/evaluate_sneeze_model.py delete mode 100644 extras/ml-training/whisper-adapter-finetuning/prepare_sneeze_data.py delete mode 100644 extras/ml-training/whisper-adapter-finetuning/sneeze_data.jsonl delete mode 100644 extras/ml-training/whisper-adapter-finetuning/train_sneeze.py delete mode 100755 extras/speaker-recognition/scripts/enroll_from_backup.py delete mode 100644 extras/speaker-recognition/scripts/simple_enroll.py delete mode 100644 extras/speaker-recognition/scripts/test_diarize.py delete mode 100755 extras/speaker-recognition/scripts/test_workflow.py diff --git a/CONNECTION_RESILIENCE_AND_UX_PLAN.md b/CONNECTION_RESILIENCE_AND_UX_PLAN.md deleted file mode 100644 index 0731b748..00000000 --- a/CONNECTION_RESILIENCE_AND_UX_PLAN.md +++ /dev/null @@ -1,294 +0,0 @@ -# Connection Resilience & UX — Unified Plan - -> **Implementation status (branch `fix/optimizations-2`): IMPLEMENTED.** -> Done: Phase 0 (backend pong), Track 1 (1A BLE device-forget, 1B central auth -> manager + SecureStore + logout/forget split, 1C persistent WS retry, 1D zombie -> ping/pong, 1E foreground/relaunch reconnect, 1F Connection Doctor probe ladder -> + Tailscale-aware messaging), Track 2 (QR JSON bundle parse, SM URL/token -> storage, `serviceManager.ts` proxy+direct client, `NetworkOverview` screen with -> start/stop/restart, webui QR bundle), Track 3 (wearable + HAVPE relay reconnect). -> App `tsc --noEmit` clean; all edited Python `py_compile`s. -> **Two deliberate deferrals:** (a) 1F "Settings home" *relocation* of auth/URL into -> a separate route + compact chip — pure navigation reorg, the diagnosis/QR value -> shipped in-place; (b) Track 2 SM **token** in the QR — the SM token is a -> server-side secret not exposed to the browser; minting it into a QR is a security -> decision left to the user. The QR carries `backendUrl` + `serviceManagerUrl`; the -> backend-proxy transport works token-free whenever the backend is up. - -This is the **authoritative, consolidated** plan. It merges two prior efforts: - -- `RECONNECT_FIXES_PLAN.md` — the **mechanical reconnection** layer (8 gaps: never - permanently give up, never destroy saved state on one failure, detect zombie sockets, - recover on foreground/relaunch). Spans mobile app + wearable client + HAVPE relay. -- `MOBILE_APP_CONNECTION_UX_PLAN.md` — the **auth + diagnosis + network-control** layer - (silent token refresh so "logged in means logged in", honest actionable failures / - Connection Doctor, QR-first config, service-manager control). App-only (+ one webui change). - -Both originals are now superseded by this file and may be deleted once this is approved. - ---- - -## Why a merge (the collision) - -The two plans operate at different layers but **rewrite the same app code** — the -token-refresh / reconnect path in `app/src/hooks/useAudioStreamer.ts` (`attemptReLogin`, -`attemptReconnect`, `onclose`, NetInfo). - -- The UX plan's **Phase 1 central auth manager** is the natural *home* for re-login. -- The reconnect plan's **A2/A3** call re-login *during* a reconnect. - -If we land the mechanical reconnect first against today's inline `attemptReLogin`, Phase 1 -then tears it out — pure churn. So the merged order is: - -> **Build the central auth manager first**, then layer WS-reconnect on top of it. Anything -> that does *not* touch auth — BLE device-forget (A1), the backend `pong` reply, and the -> entire device-client track (B/C) — is independent and can proceed in parallel. - -Guiding principles (carried from both plans): -- **Persistent retry** — never permanently give up; never delete saved state on a single - failure; always keep a recovery path alive. -- **Set-once, then invisible** — configure backend URL + creds once (ideally by QR); if the - app says you're connected, actions work; failures are honest and actionable. - ---- - -## Tracks (what can run in parallel) - -| Track | Content | Depends on | Parallelizable? | -|-------|---------|-----------|-----------------| -| **0 · Backend pong** | `websocket_controller.py` reply `{type:'pong'}` (2 sites) | nothing | ✅ anytime | -| **1 · App auth core** | Auth manager → mechanical reconnect (A1–A4) → Connection Doctor | Track 0 for A2 | ⛓ internally ordered | -| **2 · App network control** | QR bundle → SM client → Network Overview → remote start | Track 1 auth manager | after Track 1 | -| **3 · Device clients** | Wearable (B1/B2/B3) + HAVPE relay (C1/C2) reconnect | Track 0 (uses pong) optional | ✅ fully independent of app | - -Track 3 has **no app dependency** and can be implemented start-to-finish alongside Track 1. - ---- - -## Phase 0 — Backend `pong` reply *(standalone, unblocks A2)* - -**File:** `backends/advanced/src/advanced_omi_backend/controllers/websocket_controller.py` -(`:1719-1721` and the second ping site `:1772`) - -Today the backend receives app-level `{type:'ping'}` and just logs it — no reply, so the -client can't tell a half-open socket from a healthy idle one. - -**Fix:** at *both* ping sites, reply `{type:'pong'}` over the same socket (mirror how other -control replies are sent). No client behavior depends on it until A2 ships, so this is safe -to land immediately. - -**Verify:** `cd tests && make test-quick` (WS control-message regression) + a manual -`wscat` ping to `/ws` confirming the `{type:"pong"}` reply. - ---- - -## Track 1 — App connection core (ordered) - -### 1A. BLE device-forget on a single launch-reconnect failure *(independent — can land first)* -**File:** `app/src/hooks/useAutoReconnect.ts:183-198` -**Problem:** the launch auto-reconnect `catch` (`:189-192`) calls -`saveLastConnectedDeviceId(null)` + `setLastKnownDeviceId(null)` — one failed attempt -(device briefly out of range at startup) permanently forgets the device. The continuous -backoff retry loop (`:90-168`) only fires on a `connected→disconnected` transition, so it -never covers a launch attempt that never connected. - -**Fix:** -- Extract the disconnect-retry machinery (`:103-167`) into a - `scheduleRetry(deviceId, isQuickFailure)` callback (backoff refs, countdown timer, - `connectToDevice` call). -- In the launch `catch`, **stop wiping the device**; call `scheduleRetry(lastKnownDeviceId, true)`. - The device id survives; retries continue with capped backoff (10s→5min). -- `handleCancelAutoReconnect` (`:207-216`) and the explicit Disconnect button - (`index.tsx:373-377,391-396`) remain the *only* paths that clear the saved device. - -> Pure BLE, no auth dependency. This is the highest-value mechanical fix and can ship before -> the auth manager. (This was the in-flight edit; re-do it cleanly here.) - -### 1B. Central auth / connection manager *(foundation for everything below)* -Fixes UX issues A1–A3 (creds disappear on logout, forced relogin while "logged in", -plaintext password). - -- **New** `app/src/contexts/AuthContext.tsx` (or `services/auth.ts`): single source of truth - for `{ backendUrl, email, token, status }`. -- **Silent refresh everywhere:** lift `useAudioStreamer.ts::attemptReLogin` into the manager - as a `fetchAuthed()` wrapper that retries once on 401/expiry using stored creds, then - surfaces a real failure only if re-login fails. *All* calls (health, future overview) go - through it. JWTs are 1h, so this is what makes "logged in" actually mean logged in. -- **Honest status:** "Logged in as X" reflects token validity / refreshing, not mere presence. -- **Split logout from forget:** `clearAuthData()` → - - **Log out** = clear token only; keep email (+ optionally password) → re-login is one tap. - - **Forget account** = clear everything (today's behavior, made explicit). -- **SecureStore:** move password (+ token) to `expo-secure-store`; keep email in AsyncStorage. - -**Touchpoints:** `app/src/utils/storage.ts` (SecureStore + split clear fns), -`app/src/components/AuthSection.tsx` (consume manager; logout vs forget), -`app/src/hooks/useAudioStreamer.ts` (relocate `attemptReLogin`). - -### 1C. Persistent WS retry — remove permanent give-up *(A3; consumes 1B)* -**File:** `app/src/hooks/useAudioStreamer.ts:255-262` -**Problem:** after 10 attempts it sets `manuallyStoppedRef = true`, which also disables the -NetInfo recovery path (`:455`) and the AppState path (1E). It should keep trying. - -**Fix:** -- Remove `manuallyStoppedRef.current = true` on exhaustion. Keep backoff **capped** at - `MAX_RECONNECT_MS` (30s); attempts continue indefinitely while a session is intended. -- After N attempts, downgrade to a low-frequency keep-trying state and post - "Connection lost — still retrying" **once** (guard with a ref) instead of giving up. -- Re-login on reconnect now goes through the **1B auth manager**, not inline. -- `manuallyStoppedRef` is set only by `stopStreaming()` (`:219`) — genuine user stop. - -### 1D. Zombie WebSocket detection *(A2; client side; backend done in Phase 0)* -**File:** `app/src/hooks/useAudioStreamer.ts:326-334` (+ `onmessage` `:348-370`) -**Problem:** the 25s heartbeat is fire-and-forget; a silently dead TCP keeps -`readyState === OPEN` so nothing reconnects until NetInfo happens to fire. - -**Fix (client):** add `lastPongRef`; stamp it on `{type:'pong'}` in `ws.onmessage`; in the -heartbeat interval, if `now - lastPongRef > 2×HEARTBEAT_MS`, `ws.close()` (→ `onclose` → -existing `attemptReconnect`). Reset `lastPongRef` on `onopen`. (Backend reply = Phase 0.) - -### 1E. App-lifecycle (foreground / relaunch) reconnect *(A4)* -**Files:** `app/src/hooks/useAudioStreamer.ts` (WS) + `app/src/hooks/useAutoReconnect.ts` (BLE) -**Problem:** no `AppState` listener anywhere. On iOS the JS runtime suspends in background; -on resume nothing re-establishes the socket. (True iOS *background streaming* needs -background modes and is **out of scope** — we handle *resume on foreground* + warm relaunch.) - -**Fix:** -- **WS:** `AppState` listener mirroring the NetInfo effect (`:452-465`): on `'active'`, if - `currentUrlRef` set, not manually stopped, and `readyState !== OPEN` → `attemptReconnect()`. -- **BLE:** `AppState` listener: on `'active'`, if a `lastKnownDeviceId` exists and isn't - connected, reset `triedAutoReconnectForCurrentId=false` so the launch-reconnect effect - (`:171-205`) re-fires. -- `onDeviceConnect`'s existing audio-pipeline auto-restart (`index.tsx:87-97`) means BLE - recovery transitively resumes streaming. - -### 1F. Connection Doctor + Settings home *(UX C1, C2, A4-UX; consumes 1B)* -**File:** `app/src/components/BackendStatus.tsx` (today: single `fetch(${baseUrl}/health)`) -**Problem:** any failure collapses to a bare "Network request failed" — no diagnosis, no -recovery path (this is the originating incident: backend reachable by browser/curl, app -just says "network failed"). - -**Fix:** -- Replace the single `/health` fetch with a **probe ladder** → classified, actionable - messages: *offline* → "You're offline"; *tailnet unreachable* → **"Can't reach your - tailnet — is Tailscale connected?"** [Open Tailscale] [Scan QR]; *backend down* → - **"Backend is down"** [Start backend] (Track 2); *ok* → connected. -- **QR scan reachable in every state** (un-gate it). Manual URL typing = last resort. -- Move backend URL + auth into a dedicated **Settings / Connection** screen; collapse the - main surface to a compact status chip when healthy. Set-once, not constantly present. -- (Optional) weak native VPN-interface hint (`utun*`/`tun*`) as a secondary signal — "a VPN - is up", not "Tailscale specifically" (no official API exists; documented limitation). - ---- - -## Track 2 — App network control (after Track 1 auth manager) - -### 2A. Richer QR bundle *(UX Phase 3; web + app)* -- `backends/advanced/webui/src/pages/System.tsx`: QR payload becomes JSON - `{ backendUrl, serviceManagerUrl, smToken }` (SM token from the System page's existing SM - config). Keep the copyable plain URL for fallback. -- `app/src/components/QRScanner.tsx` + `app/src/utils/urlConversion.ts`: parse **both** a - plain URL (legacy) and the JSON bundle; persist SM URL + token (token in SecureStore). - -### 2B. Service-manager client + Network Overview *(UX Phase 4, fixes C3)* -**New** `app/src/services/serviceManager.ts` with **two auto-selected transports**: -- **via backend proxy** `/api/admin/services` (existing JWT) when backend is up — no SM token; -- **direct to SM** (`/node`, `/cluster`, `/services`) with stored token when backend is down; -- SM URL auto-derived from backend host `:8775`, confirmed via the **unauthed** `/health`. - -**Network Overview screen:** status table of nodes/services (backend, ASR, speaker, -wakeword, …) from `/cluster` + `/services`, each with running/health. - -> Architectural note: the SM solves *"backend down"*, **not** *"Tailscale down"* (reaching -> the SM still needs the tailnet — that case is handled by 1F's guidance). - -### 2C. Remote start / restart *(UX Phase 5, fixes C4)* -In the overview, a down backend shows **Start** → direct `POST /services/backend/start` to -the SM (the only path that works when the backend is down) → poll `GET /operations/{id}` → -reflect progress. Same pattern restarts any service. - ---- - -## Track 3 — Device clients (independent of the app) - -### 3A. Wearable client — backend-WS reconnect + mid-session JWT refresh *(B1/B2)* -**Files:** `extras/local-wearable-client/backend_sender.py:212-296`, `main.py:245-256` -**Problem:** `stream_to_backend` opens the WS once; on WS error the wrapper just logs -(`main.py:255-256`) and the backend stays dead until the whole BLE session cycles. No -mid-session token refresh. - -**Fix:** -- Wrap the `websockets.connect` block (`backend_sender.py:234-296`) in a capped - exponential-backoff reconnect loop. On a WS drop while the audio generator is still - producing, re-fetch the JWT (`get_jwt_token`, `:217` — this **is** the mid-session refresh, - solving B2) and re-dial **without ending the BLE session**. -- Drain-vs-reconnect: the queue-backed generator (`main.py:246-251`) yields `None` only at - true session end → treat `None` as "stop", WS exceptions as "reconnect". Drop audio during - the gap (matches current outage behavior; document the choice). -- Reset backoff after a healthy connection duration (mirror app's `MIN_HEALTHY_DURATION`). - -### 3B. (optional) headless `run()` BLE backoff parity *(B3, low priority)* -**File:** `extras/local-wearable-client/main.py:584-621` — headless path rescans at a fixed -`scan_interval` with no backoff. The launchd default is the menu app (`menu_app.py:468-488`, -already has backoff), so add the same exponential backoff for parity only if cheap. - -### 3C. HAVPE relay — backend-WS reconnect *(C1)* -**File:** `extras/havpe-relay/relay_core.py:300-376` (`run_device_session`) -**Problem:** `websockets.connect` (`:301`) is opened once; any WS close tears down the -session (`:364-376`) and waits passively for the ESP32 to re-dial. -**Fix:** wrap WS-connect + task group (`:301-360`) in an inner capped-backoff retry loop, -keeping the device `reader`/`writer` alive across WS reconnects. Re-fetch the token -(`get_jwt_token`, `:282`) on each re-dial. Break only on device disconnect -(`IncompleteReadError`/reader EOF) or idle timeout — not on a WS blip. - -### 3D. HAVPE relay — ESPHome API reconnect within a session *(C2, lower priority)* -**Files:** `extras/havpe-relay/relay_core.py:304-321`, `device_controller.py:33-44` -**Problem:** one API connect timeout → audio-only for the entire session; no reconnect. -**Fix:** use `aioesphomeapi`'s `ReconnectLogic` (or a lightweight periodic re-`connect()`) -so button/dial/LED/speaker recover mid-session. - -> Note: device→relay audio TCP is server-side (relay can't initiate); the firmware re-dials -> and the idle-timeout (`relay_core.py:45`) reaps zombies. No change needed there. - ---- - -## Suggested implementation order - -1. **Phase 0** backend `pong` — tiny, standalone, unblocks 1D. -2. **1A** BLE device-forget — highest mechanical value, no auth dependency. -3. **1B** central auth manager — foundation; removes daily auth pain (A1–A3 UX). -4. **1C → 1D → 1E** mechanical WS reconnect on top of the auth manager. -5. **1F** Connection Doctor + Settings home — fixes the originating "useless message" incident. -6. **2A → 2B → 2C** network-control story, each building on the last. -- **Track 3 (3A, 3C, then 3D/3B)** in parallel throughout — no app dependency. - -Each item is independently shippable. - ---- - -## Verification - -**App mechanical (1A,1C,1D,1E):** dev client on a physical phone (BLE needs hardware): -- 1A: connect device, force-quit, walk out of range, relaunch in range → reconnects and is - **not** forgotten if the first attempt fails (`lastKnownDeviceId` persists). -- 1C: backend down for >10 reconnect attempts, then back → client resumes without a manual - restart (previously gave up permanently). -- 1D: start streaming, kill the backend's TCP path abruptly (drop wifi / `kill -STOP`) so no - close frame → within ~50s client detects missed pongs, closes, reconnects on return. -- 1E: background the app (iOS), foreground → WS re-establishes; toggle BLE off/on while - backgrounded then foreground → device reconnects. - -**App auth/UX (1B,1F):** log out → email retained, one-tap re-login; let a token expire (1h -or shortened TTL) and hit a non-streaming action → silent refresh, no forced relogin; point -the app at an unreachable tailnet → Connection Doctor says *tailnet unreachable* with -actions, not "Network request failed". - -**Network control (2A–2C):** scan the JSON QR → backend URL + SM URL + token persisted; take -the backend down → Overview shows it down with **Start** → SM starts it → status flips healthy. - -**Device clients (3A,3C):** `cd extras/local-wearable-client && uv run python main.py run` -against a local backend; mid-stream `./restart.sh` → client re-dials and resumes without -dropping BLE; run >1h (or shorten token TTL) to confirm mid-session re-auth. HAVPE: ESP32 -dialed in, restart backend mid-session → relay re-dials, ESP32 stays connected; kill/restore -the ESPHome API → button/LED recover (3D). - -**Regression:** `cd tests && make test-quick` after the Phase 0 `pong` change. diff --git a/backends/advanced/scripts/delete_all_conversations_api.py b/backends/advanced/scripts/delete_all_conversations_api.py deleted file mode 100755 index 25cd2b3b..00000000 --- a/backends/advanced/scripts/delete_all_conversations_api.py +++ /dev/null @@ -1,131 +0,0 @@ -#!/usr/bin/env python3 -""" -Script to delete all conversations using the API endpoints. -This uses the proper API authentication and endpoints. -""" - -import argparse -import asyncio -import os -import sys -from pathlib import Path - -import aiohttp -from dotenv import load_dotenv - -# Load environment variables -load_dotenv(Path(__file__).parent.parent / ".env") - - -async def get_auth_token(): - """Get admin authentication token.""" - admin_email = os.getenv("ADMIN_EMAIL") - admin_password = os.getenv("ADMIN_PASSWORD") - - if not admin_email or not admin_password: - print("Error: ADMIN_EMAIL and ADMIN_PASSWORD must be set in .env file") - sys.exit(1) - - base_url = "http://localhost:8000" - - async with aiohttp.ClientSession() as session: - # Login to get token - login_data = {"username": admin_email, "password": admin_password} - - async with session.post( - f"{base_url}/auth/jwt/login", data=login_data - ) as response: - if response.status != 200: - print(f"Failed to login: {response.status}") - text = await response.text() - print(f"Response: {text}") - sys.exit(1) - - result = await response.json() - return result["access_token"] - - -async def delete_all_conversations(skip_prompt=False): - """Delete all conversations using the API.""" - - base_url = "http://localhost:8000" - - # Get auth token - print("Getting admin authentication token...") - token = await get_auth_token() - - headers = {"Authorization": f"Bearer {token}"} - - async with aiohttp.ClientSession() as session: - # First, get all conversations - print("Fetching all conversations...") - async with session.get( - f"{base_url}/api/conversations", headers=headers - ) as response: - if response.status != 200: - print(f"Failed to fetch conversations: {response.status}") - text = await response.text() - print(f"Response: {text}") - return - - data = await response.json() - - # Extract conversations from nested structure - conversations_dict = data.get("conversations", {}) - conversations = [] - for client_id, client_conversations in conversations_dict.items(): - conversations.extend(client_conversations) - - print(f"Found {len(conversations)} conversations") - - if len(conversations) == 0: - print("No conversations to delete") - return - - # Confirm deletion unless --yes flag is used - if not skip_prompt: - response = input( - f"Are you sure you want to delete ALL {len(conversations)} conversations? (yes/no): " - ) - if response.lower() != "yes": - print("Deletion cancelled") - return - - # Delete each conversation - deleted_count = 0 - failed_count = 0 - - for conv in conversations: - audio_uuid = conv.get("audio_uuid") - if not audio_uuid: - print(f"Skipping conversation without audio_uuid: {conv.get('_id')}") - continue - - # Delete the conversation - async with session.delete( - f"{base_url}/api/conversations/{audio_uuid}", headers=headers - ) as response: - if response.status == 200: - deleted_count += 1 - print( - f"Deleted conversation {audio_uuid} ({deleted_count}/{len(conversations)})" - ) - else: - failed_count += 1 - text = await response.text() - print(f"Failed to delete {audio_uuid}: {response.status} - {text}") - - print(f"\nDeletion complete:") - print(f" Successfully deleted: {deleted_count}") - print(f" Failed: {failed_count}") - - -if __name__ == "__main__": - # Parse command line arguments - parser = argparse.ArgumentParser(description="Delete all conversations") - parser.add_argument( - "--yes", "-y", action="store_true", help="Skip confirmation prompt" - ) - args = parser.parse_args() - - asyncio.run(delete_all_conversations(skip_prompt=args.yes)) diff --git a/backends/advanced/scripts/laptop_client.py b/backends/advanced/scripts/laptop_client.py deleted file mode 100644 index a848469f..00000000 --- a/backends/advanced/scripts/laptop_client.py +++ /dev/null @@ -1,275 +0,0 @@ -import argparse -import asyncio -import json -import logging - -import aiohttp -import websockets -import websockets.exceptions -from easy_audio_interfaces.extras.local_audio import InputMicStream -from wyoming.audio import AudioChunk, AudioStart, AudioStop - -logger = logging.getLogger(__name__) -logging.basicConfig(level=logging.INFO) - -# Default WebSocket settings -DEFAULT_HOST = "localhost" -DEFAULT_PORT = 8000 -DEFAULT_ENDPOINT = "/ws?codec=pcm" - -# Audio format will be determined from the InputMicStream instance - - -def build_websocket_uri( - host: str, - port: int, - endpoint: str, - token: str | None = None, - device_name: str = "laptop", -) -> str: - """Build WebSocket URI with JWT token authentication.""" - base_uri = f"ws://{host}:{port}{endpoint}" - params = [] - if token: - params.append(f"token={token}") - if device_name: - params.append(f"device_name={device_name}") - - if params: - base_uri += "?" + "&".join(params) - return base_uri - - -async def authenticate_with_credentials( - host: str, port: int, username: str, password: str -) -> str: - """Authenticate with username/password and return JWT token.""" - auth_url = f"http://{host}:{port}/auth/jwt/login" - - # Prepare form data for authentication - form_data = aiohttp.FormData() - form_data.add_field("username", username) - form_data.add_field("password", password) - - try: - async with aiohttp.ClientSession() as session: - async with session.post(auth_url, data=form_data) as response: - if response.status == 200: - result = await response.json() - token = result.get("access_token") - if token: - logger.info(f"Successfully authenticated user '{username}'") - return token - else: - raise Exception("No access token received from server") - elif response.status == 400: - error_detail = await response.text() - raise Exception( - f"Authentication failed: Invalid credentials - {error_detail}" - ) - else: - error_detail = await response.text() - raise Exception( - f"Authentication failed with status {response.status}: {error_detail}" - ) - except aiohttp.ClientError as e: - raise Exception(f"Failed to connect to authentication server: {e}") - - -def validate_auth_args(args): - """Validate that exactly one authentication method is provided.""" - has_token = bool(args.token) - has_credentials = bool(args.username and args.password) - - if not has_token and not has_credentials: - raise ValueError( - "Authentication required: Please provide either --token OR both --username and --password" - ) - - if has_token and has_credentials: - raise ValueError( - "Conflicting authentication methods: Please provide either --token OR --username/--password, not both" - ) - - if args.username and not args.password: - raise ValueError( - "Username provided but password missing: Both --username and --password are required" - ) - - if args.password and not args.username: - raise ValueError( - "Password provided but username missing: Both --username and --password are required" - ) - - -async def send_wyoming_event(websocket, wyoming_event): - """Send a Wyoming protocol event over WebSocket. - - Based on how the backend processes Wyoming events, they expect: - 1. JSON header line ending with \n - 2. Optional binary payload if payload_length > 0 - - This replicates Wyoming's async_write_event behavior for WebSocket transport. - """ - # Get the event data from Wyoming event - event_data = wyoming_event.event() - - # Build event dict like Wyoming's async_write_event does - event_dict = event_data.to_dict() - event_dict["version"] = "1.0.0" # Wyoming adds version - - # Add payload_length if payload exists (critical for audio chunks!) - if event_data.payload: - event_dict["payload_length"] = len(event_data.payload) - - # Send JSON header - json_header = json.dumps(event_dict) + "\n" - await websocket.send(json_header) - logger.debug( - f"Sent Wyoming event: {event_data.type} (payload_length: {event_dict.get('payload_length', 0)})" - ) - - # Send binary payload if exists - if event_data.payload: - await websocket.send(event_data.payload) - logger.debug(f"Sent audio payload: {len(event_data.payload)} bytes") - - -async def main(): - # Parse command line arguments - parser = argparse.ArgumentParser( - description="Laptop audio client for OMI backend with dual authentication modes" - ) - parser.add_argument("--host", default=DEFAULT_HOST, help="WebSocket server host") - parser.add_argument( - "--port", type=int, default=DEFAULT_PORT, help="WebSocket server port" - ) - parser.add_argument( - "--endpoint", default=DEFAULT_ENDPOINT, help="WebSocket endpoint" - ) - - # Authentication options (mutually exclusive) - auth_group = parser.add_argument_group( - "authentication", "Choose one authentication method" - ) - auth_group.add_argument("--token", help="JWT authentication token") - auth_group.add_argument("--username", help="Username for login authentication") - auth_group.add_argument("--password", help="Password for login authentication") - - parser.add_argument( - "--device-name", default="laptop", help="Device name for client identification" - ) - args = parser.parse_args() - - # Validate authentication arguments - try: - validate_auth_args(args) - except ValueError as e: - logger.error(f"Authentication error: {e}") - parser.print_help() - return - - # Get or obtain authentication token - token = None - - if args.token: - # Use provided token directly - token = args.token - print(f"Using provided JWT token: {token[:8]}...") - - elif args.username and args.password: - # Authenticate with username/password to get token - print(f"Authenticating with username: {args.username}") - try: - token = await authenticate_with_credentials( - args.host, args.port, args.username, args.password - ) - print(f"Authentication successful! Received token: {token[:8]}...") - except Exception as e: - logger.error(f"Authentication failed: {e}") - return - - # Build WebSocket URI - ws_uri = build_websocket_uri( - args.host, args.port, args.endpoint, token, args.device_name - ) - print(f"Connecting to {ws_uri}") - print(f"Using device name: {args.device_name}") - - try: - async with websockets.connect(ws_uri) as websocket: - print("Connected to WebSocket") - - async def send_audio(): - """Capture audio from microphone and send Wyoming protocol audio events""" - try: - # Send audio-start event to begin session - async with InputMicStream(chunk_size=512) as stream: - audio_start = AudioStart( - rate=stream.sample_rate, - width=stream.sample_width, - channels=stream.channels, - ) - await send_wyoming_event(websocket, audio_start) - logger.info( - f"Sent audio-start event (rate={stream.sample_rate}, width={stream.sample_width}, channels={stream.channels})" - ) - while True: - try: - data = await stream.read() - if data and data.audio: - # Create Wyoming AudioChunk with format from stream - audio_chunk = AudioChunk( - audio=data.audio, - rate=stream.sample_rate, - width=stream.sample_width, - channels=stream.channels, - ) - await send_wyoming_event(websocket, audio_chunk) - logger.debug( - f"Sent audio chunk: {len(data.audio)} bytes" - ) - await asyncio.sleep( - 0.01 - ) # Small delay to prevent overwhelming - except websockets.exceptions.ConnectionClosed: - logger.info( - "WebSocket connection closed during audio sending" - ) - break - except Exception as e: - logger.error(f"Error sending audio: {e}") - break - - except Exception as e: - logger.error(f"Error in audio session: {e}") - finally: - # Send audio-stop event to end session - try: - audio_stop = AudioStop() - await send_wyoming_event(websocket, audio_stop) - logger.info("Sent audio-stop event") - except Exception as e: - logger.error(f"Error sending audio-stop: {e}") - - async def receive_messages(): - """Receive any messages from the WebSocket server""" - try: - async for message in websocket: - print(f"Received message: {message}") - except websockets.exceptions.ConnectionClosed: - logger.info("WebSocket connection closed during message receiving") - except Exception as e: - logger.error(f"Error receiving messages: {e}") - - # Run both audio sending and message receiving concurrently - await asyncio.gather(send_audio(), receive_messages()) - - except ConnectionRefusedError: - logger.error(f"Could not connect to {ws_uri}. Make sure the server is running.") - except Exception as e: - logger.error(f"Error connecting to WebSocket: {e}") - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/extras/ml-training/README.md b/extras/ml-training/README.md index 5d1acea2..badcd84f 100644 --- a/extras/ml-training/README.md +++ b/extras/ml-training/README.md @@ -1,6 +1,6 @@ -# ML Training Scripts +# ML Training Tools -Standalone CLI tools for exporting training data from Chronicle and fine-tuning models. These are **not** part of the backend runtime -- they're run manually on a workstation with GPU access. +Standalone CLI tools for exporting Chronicle annotations and training the event detection classifier. These are **not** part of the backend runtime -- they're run manually on a workstation. ## Contents @@ -14,33 +14,10 @@ Export accepted/rejected annotations from MongoDB and train an event detection c See `event-detection/README.md` for full usage. -### `whisper-adapter-finetuning/` - -Fine-tune a Whisper LoRA adapter for domain-specific ASR improvements (e.g., detecting non-speech events like sneezes, laughter). - -- `prepare_sneeze_data.py` - Prepare training data -- `train_sneeze.py` - LoRA adapter training script -- `evaluate_sneeze_model.py` - Evaluate trained adapter - -See `whisper-adapter-finetuning/README.md` for full usage. - -### `autoresearch-asr/` - -Autonomous LoRA fine-tuning loop for VibeVoice-ASR, adapted from [karpathy/autoresearch](https://github.com/karpathy/autoresearch). Give an AI agent the training setup and let it experiment overnight on Google Colab. - -- `prepare.py` - Fixed data loading, model caching, train/val/test split (DO NOT MODIFY) -- `evaluate.py` - Fixed evaluation harness: WER + SWER + boundary MAE (DO NOT MODIFY) -- `train.py` - The file the agent modifies: LoRA config, hyperparams, curriculum -- `program.md` - Agent instructions for the autonomous experiment loop -- `export_data.py` - Export training data from Chronicle API to VibeVoice format - -See `autoresearch-asr/program.md` for full usage. - ## Prerequisites ```bash pip install -r event-detection/requirements.txt -# For whisper adapter: see whisper-adapter-finetuning/README.md ``` ## Relationship to Backend diff --git a/extras/ml-training/whisper-adapter-finetuning/README.md b/extras/ml-training/whisper-adapter-finetuning/README.md deleted file mode 100644 index bf66b919..00000000 --- a/extras/ml-training/whisper-adapter-finetuning/README.md +++ /dev/null @@ -1,202 +0,0 @@ -# Whisper Sneeze Adapter Training - -This project fine-tunes OpenAI's Whisper model to transcribe sneezes in audio/video content using LoRA adapters. The model learns to recognize and transcribe sneezes as the token "SNEEZE" in transcriptions. - -## Prerequisites - -- Python 3.10+ -- CUDA-capable GPU (recommended for training) -- Access to Google Gemini API (for generating transcripts) - -## Installation - -1. Create a virtual environment: -```bash -python -m venv .venv -source .venv/bin/activate # On Windows: .venv\Scripts\activate -``` - -2. Install dependencies: -```bash -pip install torch torchaudio -pip install transformers datasets evaluate -pip install unsloth[colab-new] -pip install librosa soundfile jiwer -pip install tqdm -``` - -## Workflow - -### Step 1: Prepare Your Video - -1. Record or obtain a video file containing sneezes (e.g., `girls_sneezing.mp4` download with - ``` - yt-dlp -f "bestvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best" --merge-output-format mp4 -o "girls_sneezing.mp4" https://youtu.be/36b4248j5UE - ``` - -### Step 2: Generate Transcript with Gemini - -1. Upload your video to Google Gemini (or use Gemini API) -2. Request a transcript with sneezes marked using the format: `<sneeze>` -3. Generate a JSONL file named `sneeze_data.jsonl` with the following format: - -```jsonl -{"start": 0.0, "end": 5.0, "text": "Ugh, I really need to sneeze. Stuck? Yeah, it's right there."} -{"start": 5.0, "end": 11.0, "text": "Close one. <sneeze> Bless you. Thanks."} -{"start": 12.0, "end": 17.0, "text": "Ugh, I can feel it. I really need to sneeze so bad. Go on, let it out."} -``` - -**Format requirements:** -- Each line is a JSON object -- `start`: Start time in seconds (float) -- `end`: End time in seconds (float) -- `text`: Transcription text with sneezes marked as `<sneeze>` - -**Example Gemini prompt:** -``` -Please transcribe this video and create a JSONL file where each line contains: -- start: start time in seconds -- end: end time in seconds -- text: the transcription with sneezes marked as <sneeze> - -Format as JSONL (one JSON object per line). -``` - -### Step 3: Prepare Training Data - -Run the data preparation script to extract audio chunks and create train/test splits: - -```bash -python prepare_sneeze_data.py -``` - -This script will: -- Extract audio from your video file (`girls_sneezing.mp4`) -- Create audio chunks from the segments in `sneeze_data.jsonl` -- Save chunks to `sneeze_chunks/` directory -- Split data into `train.jsonl` (60%) and `test.jsonl` (40%) - -**Requirements:** -- `sneeze_data.jsonl` must exist in the project root -- Video file must be named `girls_sneezing.mp4` - -### Step 4: Train the Model - -Train the Whisper model with LoRA adapters: - -```bash -python train_sneeze.py -``` - -This will: -- Load the base Whisper Large v3 model -- Apply LoRA adapters (only trains 1-10% of parameters) -- Fine-tune on your sneeze data -- Save the adapter to `sneeze_lora_adapter_unsloth/` - -**Training configuration:** -- Model: `unsloth/whisper-large-v3` -- LoRA rank: 64 -- Batch size: 1 (with gradient accumulation: 4) -- Max steps: 200 -- Learning rate: 1e-4 - -**Note:** Training requires a GPU with sufficient VRAM. Adjust `load_in_4bit=True` in the script if you have limited memory. - -### Step 5: Evaluate the Model - -Evaluate the trained model on the test set: - -```bash -python evaluate_sneeze_model.py -``` - -This will: -- Load the base model and merge the LoRA adapter -- Run inference on test samples -- Calculate Word Error Rate (WER) -- Report sneeze detection recall and false positives - -## Results - -### Training Results - -Training was performed on a Tesla T4 GPU with the following configuration: -- **Model**: `unsloth/whisper-large-v3` -- **Trainable Parameters**: 31,457,280 of 1,574,947,840 (2.00%) -- **Training Time**: 12.04 minutes -- **Peak Memory Usage**: 8.896 GB (60.35% of max memory) -- **Training Samples**: 49 samples -- **Test Samples**: 4 samples - -**Training Loss Progression:** -| Step | Training Loss | Validation Loss | WER | -|------|---------------|-----------------|-----| -| 20 | 1.646100 | 1.869532 | 50.0% | -| 40 | 0.832500 | 1.004385 | 30.0% | -| 60 | 0.304600 | 0.354044 | 30.0% | -| 80 | 0.067700 | 0.051606 | 0.0% | -| 100 | 0.017600 | 0.162433 | 10.0% | -| 120 | 0.003400 | 0.006127 | 0.0% | -| 140 | 0.002000 | 0.004151 | 0.0% | -| 160 | 0.001400 | 0.003399 | 0.0% | -| 180 | 0.001300 | 0.003005 | 0.0% | -| 200 | 0.001000 | 0.002856 | 0.0% | - -**Final Metrics:** -- Final Training Loss: 0.001000 -- Final Validation Loss: 0.002856 -- Final Validation WER: 0.0% - -### Evaluation Results - -Evaluation was performed on 10 test samples (4 containing sneezes): - -**Overall Performance:** -- **Word Error Rate (WER)**: 0.3217 (32.17%) -- **Sneeze Recall**: 2/4 (50.0%) -- **False Positives**: 0 - -**Missed Sneezes:** -1. Reference: "Take your time, it'll come. SNEEZE Oh wow. Excuse me." - Prediction: "Take your time. It'll come. Oh, wow." - -2. Reference: "It's right there but... False alarm? No, it's stuck. SNEEZE Bless you." - Prediction: "It's right there, but... False alarm? No! It stopped..." - -**Analysis:** -- The model achieved perfect WER (0.0%) on the validation set during training, indicating good generalization on the training distribution. -- On the test set, the model achieved 50% sneeze recall, successfully detecting 2 out of 4 sneezes. -- No false positives were detected, showing the model is conservative in its sneeze predictions. -- The 32.17% WER on the test set suggests room for improvement, particularly in detecting sneezes in more varied contexts. - -## Project Structure - -``` -whisper-adapter-test/ -├── prepare_sneeze_data.py # Data preparation script -├── improved_sneeze_trainer.py # Training script -├── evaluate_sneeze_model.py # Evaluation script -├── sneeze_data.jsonl # Input transcript with sneezes -├── train.jsonl # Training manifest -├── test.jsonl # Test manifest -├── sneeze_chunks/ # Extracted audio chunks -└── sneeze_lora_adapter_unsloth/ # Trained adapter (created after training) -``` - -## Output Files - -- `train.jsonl`: Training dataset manifest -- `test.jsonl`: Test dataset manifest -- `sneeze_chunks/`: Directory with extracted audio chunks -- `sneeze_lora_adapter_unsloth/`: Trained LoRA adapter weights - -## Notes - -- The model replaces `<sneeze>` tags with `SNEEZE` during training -- LoRA adapters are memory-efficient and only update a small portion of model weights -- The evaluation script merges the adapter into the base model for inference - -## Conclusion - -Despite training on only 13 examples and evaluating on 10 test samples, the model achieved significant progress in sneeze detection. With just this small dataset, we were able to fine-tune the Whisper model to recognize and transcribe sneezes with 50% recall and zero false positives. This demonstrates the effectiveness of LoRA adapters for efficient fine-tuning on specialized tasks with limited data. diff --git a/extras/ml-training/whisper-adapter-finetuning/evaluate_sneeze_model.py b/extras/ml-training/whisper-adapter-finetuning/evaluate_sneeze_model.py deleted file mode 100644 index 4dcceaf2..00000000 --- a/extras/ml-training/whisper-adapter-finetuning/evaluate_sneeze_model.py +++ /dev/null @@ -1,123 +0,0 @@ -import json -import os - -import jiwer -import librosa -import torch -from peft import PeftModel -from tqdm import tqdm -from transformers import WhisperForConditionalGeneration, WhisperProcessor - -# --- CONFIGURATION (MUST MATCH YOUR TRAINING) --- -BASE_MODEL_ID = "openai/whisper-large-v3" -ADAPTER_PATH = "sneeze_lora_adapter_unsloth" # The folder Unsloth created -TEST_MANIFEST = "test.jsonl" - - -def main(): - # 1. Setup Device - device = "cuda" if torch.cuda.is_available() else "cpu" - print(f"Using device: {device}") - - # 2. Load Base Model (Large v3) - print(f"Loading base model: {BASE_MODEL_ID}") - processor = WhisperProcessor.from_pretrained(BASE_MODEL_ID) - model = WhisperForConditionalGeneration.from_pretrained( - BASE_MODEL_ID, torch_dtype=torch.float16 if device == "cuda" else torch.float32 - ) - - # 3. Load and MERGE Adapter - if os.path.exists(ADAPTER_PATH): - print(f"Loading LoRA adapter from: {ADAPTER_PATH}") - model = PeftModel.from_pretrained(model, ADAPTER_PATH) - print("Merging LoRA weights...") - model = model.merge_and_unload() - else: - print(f"❌ ERROR: Adapter {ADAPTER_PATH} not found!") - return - - model.to(device) - model.eval() - - # 4. Run Evaluation - evaluate_dataset(model, processor, device, TEST_MANIFEST) - - -def evaluate_dataset(model, processor, device, manifest_path): - if not os.path.exists(manifest_path): - print(f"Manifest {manifest_path} not found.") - return - - samples = [] - with open(manifest_path, "r") as f: - for line in f: - samples.append(json.loads(line)) - - print(f"Testing on {len(samples)} samples...") - - predictions = [] - references = [] - sneeze_stats = {"total": 0, "detected": 0, "fp": 0} - - for sample in tqdm(samples): - path = sample["audio"] - ref_text = sample["text"].replace("<sneeze>", "SNEEZE") - - try: - audio, _ = librosa.load(path, sr=16000) - except: - continue - - # Process audio - inputs = processor(audio, sampling_rate=16000, return_tensors="pt") - input_features = inputs.input_features.to(device) - - # Handle the dtype for half precision (if on GPU) - if device == "cuda": - input_features = input_features.half() - - # Generate - with torch.no_grad(): - generated_ids = model.generate( - input_features=input_features, # Use input_features, not inputs - language="en", - task="transcribe", - max_new_tokens=256, - ) - - pred = processor.batch_decode(generated_ids, skip_special_tokens=True)[ - 0 - ].strip() - - predictions.append(pred) - references.append(ref_text) - - # Stats - has_sneeze_ref = "SNEEZE" in ref_text - has_sneeze_pred = "SNEEZE" in pred - - if has_sneeze_ref: - sneeze_stats["total"] += 1 - if has_sneeze_pred: - sneeze_stats["detected"] += 1 - else: - print(f"\n❌ MISSED SNEEZE\nRef: {ref_text}\nPrd: {pred}") - elif has_sneeze_pred: - sneeze_stats["fp"] += 1 - print(f"\n⚠️ FALSE POSITIVE\nRef: {ref_text}\nPrd: {pred}") - - # Results - wer = jiwer.wer(references, predictions) - print("\n" + "=" * 40) - print(f"Word Error Rate: {wer:.4f}") - if sneeze_stats["total"] > 0: - recall = (sneeze_stats["detected"] / sneeze_stats["total"]) * 100 - print( - f"Sneeze Recall: {sneeze_stats['detected']}/{sneeze_stats['total']} ({recall:.1f}%)" - ) - print(f"False Positives: {sneeze_stats['fp']}") - print("=" * 40) - - -if __name__ == "__main__": - main() diff --git a/extras/ml-training/whisper-adapter-finetuning/prepare_sneeze_data.py b/extras/ml-training/whisper-adapter-finetuning/prepare_sneeze_data.py deleted file mode 100644 index 69dcbda9..00000000 --- a/extras/ml-training/whisper-adapter-finetuning/prepare_sneeze_data.py +++ /dev/null @@ -1,80 +0,0 @@ -import json -import os -import random - -import librosa -import numpy as np -import soundfile as sf - - -def prepare_data(): - jsonl_path = "sneeze_data.jsonl" - video_path = "girls_sneezing.mp4" - output_dir = "sneeze_chunks" - - if not os.path.exists(output_dir): - os.makedirs(output_dir) - - # Load full audio - print(f"Loading {video_path}...") - try: - y, sr = librosa.load(video_path, sr=16000) - except Exception as e: - print(f"Error loading video: {e}") - return - - segments = [] - with open(jsonl_path, "r") as f: - for line in f: - if line.strip(): - segments.append(json.loads(line)) - - print(f"Found {len(segments)} segments.") - - dataset_entries = [] - - for i, seg in enumerate(segments): - start_time = seg["start"] - end_time = seg["end"] - text = seg["text"] - - # Calculate sample indices - start_sample = int(start_time * sr) - end_sample = int(end_time * sr) - - # Extract audio - chunk = y[start_sample:end_sample] - - # Save to file - chunk_filename = f"chunk_{i:03d}.wav" - chunk_path = os.path.join(output_dir, chunk_filename) - sf.write(chunk_path, chunk, sr) - - dataset_entries.append({"audio": chunk_path, "text": text}) - print(f"Saved {chunk_filename}: {text[:30]}...") - - # Shuffle and Split - random.seed(42) - random.shuffle(dataset_entries) - - split_idx = int(len(dataset_entries) * 0.6) - train_data = dataset_entries[:split_idx] - test_data = dataset_entries[split_idx:] - - print(f"Training samples: {len(train_data)}") - print(f"Testing samples: {len(test_data)}") - - # Save split manifests - with open("train.jsonl", "w") as f: - for entry in train_data: - f.write(json.dumps(entry) + "\n") - - with open("test.jsonl", "w") as f: - for entry in test_data: - f.write(json.dumps(entry) + "\n") - - print("Data preparation complete.") - - -if __name__ == "__main__": - prepare_data() diff --git a/extras/ml-training/whisper-adapter-finetuning/sneeze_data.jsonl b/extras/ml-training/whisper-adapter-finetuning/sneeze_data.jsonl deleted file mode 100644 index 418ee305..00000000 --- a/extras/ml-training/whisper-adapter-finetuning/sneeze_data.jsonl +++ /dev/null @@ -1,23 +0,0 @@ -{"start": 0.0, "end": 5.0, "text": "Ugh, I really need to sneeze. Stuck? Yeah, it's right there."} -{"start": 5.0, "end": 11.0, "text": "Close one. <sneeze> Bless you. Thanks."} -{"start": 12.0, "end": 17.0, "text": "Ugh, I can feel it. I really need to sneeze so bad. Go on, let it out."} -{"start": 17.0, "end": 23.0, "text": "It's right there but... False alarm? No, it's stuck. <sneeze> Bless you."} -{"start": 24.0, "end": 29.0, "text": "Ugh, my nose. I... I really need to sneeze so bad. Do it then."} -{"start": 29.0, "end": 36.0, "text": "No, it's not coming. That's the worst. Stuck. <sneeze>"} -{"start": 36.0, "end": 42.0, "text": "Ugh, my nose, I want to sneeze so bad. You okay? Is it stuck?"} -{"start": 42.0, "end": 48.0, "text": "Nope, nothing. Teasing you. <sneeze>"} -{"start": 48.0, "end": 54.0, "text": "Ugh, my nose. I need to sneeze so bad. Go on, let it out."} -{"start": 54.0, "end": 60.0, "text": "Oh, it's stuck. It's teasing you. <sneeze> Bless you."} -{"start": 60.0, "end": 66.0, "text": "Ugh, I really... I really need to sneeze so bad. Go on, just let it out."} -{"start": 66.0, "end": 72.0, "text": "Ugh, it's stuck. Oh come on. <sneeze>"} -{"start": 72.0, "end": 78.0, "text": "Ugh, I really need to sneeze so bad. Go on, let it out. It's just stuck."} -{"start": 78.0, "end": 84.0, "text": "I can feel it right there. <sneeze> Oh finally."} -{"start": 84.0, "end": 90.0, "text": "Ugh, my nose is so itchy. I need to sneeze so badly. Do it. Let it out."} -{"start": 90.0, "end": 96.0, "text": "No. It's... it's stuck. Almost? <sneeze>"} -{"start": 96.0, "end": 102.0, "text": "Ugh, I want to sneeze so bad. It's right there, just not coming out."} -{"start": 102.0, "end": 108.0, "text": "No. Still stuck! <sneeze>"} -{"start": 108.0, "end": 114.0, "text": "Ugh, I really... I really need to sneeze so bad. Here it comes?"} -{"start": 114.0, "end": 120.0, "text": "Nope, it's uh, stuck. Why won't it come out? <sneeze>"} -{"start": 120.0, "end": 126.0, "text": "Ugh, I swear I need to sneeze so badly. Ugh, nope, still there."} -{"start": 126.0, "end": 131.0, "text": "Take your time, it'll come. <sneeze> Oh wow. Excuse me."} -{"start": 132.0, "end": 140.0, "text": "Don't forget to check out Patreon dot com slash AI sneeze for exclusive sneezing content and early access. Try it for free! <sneeze> <sneeze>"} diff --git a/extras/ml-training/whisper-adapter-finetuning/train_sneeze.py b/extras/ml-training/whisper-adapter-finetuning/train_sneeze.py deleted file mode 100644 index 1775582d..00000000 --- a/extras/ml-training/whisper-adapter-finetuning/train_sneeze.py +++ /dev/null @@ -1,295 +0,0 @@ -import json -import os -from dataclasses import dataclass -from typing import Any, Dict, List, Union - -import evaluate -import librosa -import numpy as np -import torch -import tqdm -from datasets import Dataset -from transformers import ( - Seq2SeqTrainer, - Seq2SeqTrainingArguments, - WhisperForConditionalGeneration, -) -from unsloth import FastModel, is_bf16_supported - -# Configuration -MODEL_ID = "unsloth/whisper-large-v3" -OUTPUT_DIR = "sneeze_lora_adapter_unsloth" -TRAIN_MANIFEST = "train.jsonl" - - -def prepare_dataset(): - """Load training data from jsonl""" - audio_paths = [] - texts = [] - - # Check if file exists - if not os.path.exists(TRAIN_MANIFEST): - raise FileNotFoundError(f"{TRAIN_MANIFEST} not found.") - - with open(TRAIN_MANIFEST, "r") as f: - for line in f: - try: - entry = json.loads(line) - # Normalize text - text = entry["text"].replace("<sneeze>", "SNEEZE") - - audio_paths.append(entry["audio"]) - texts.append(text) - - # Simple oversampling for target keyword - if "SNEEZE" in text: - for _ in range(5): - audio_paths.append(entry["audio"]) - texts.append(text) - except Exception as e: - print(f"Skipping bad line: {e}") - - print(f"Loaded {len(audio_paths)} training samples.") - - data = {"audio_path": audio_paths, "text": texts} - - return Dataset.from_dict(data) - - -def main(): - print(f"Loading model with Unsloth: {MODEL_ID}") - - # Load model using Unsloth's FastModel - model, tokenizer = FastModel.from_pretrained( - model_name=MODEL_ID, - dtype=None, # Auto detection - load_in_4bit=False, # Set to True for 4bit quantization (lower memory) - auto_model=WhisperForConditionalGeneration, - whisper_language="English", - whisper_task="transcribe", - # token = "hf_...", # Use if needed for gated models - ) - - # Apply LoRA adapters using Unsloth (only updates 1-10% of parameters) - model = FastModel.get_peft_model( - model, - r=64, # Suggested: 8, 16, 32, 64, 128 - target_modules=["q_proj", "v_proj"], - lora_alpha=64, - lora_dropout=0, # 0 is optimized - bias="none", # "none" is optimized - use_gradient_checkpointing="unsloth", # 30% less VRAM, fits 2x larger batch sizes - random_state=3407, - use_rslora=False, - loftq_config=None, - task_type=None, # MUST be None for Whisper - ) - - # Configure generation settings - model.generation_config.language = "<|en|>" - model.generation_config.task = "transcribe" - model.config.suppress_tokens = [] - model.generation_config.forced_decoder_ids = None - - # Load dataset - dataset = prepare_dataset() - - def formatting_prompts_func(example): - """Process audio and text for training""" - try: - # Load audio file - audio_array, sr = librosa.load(example["audio_path"], sr=16000) - except Exception as e: - print(f"Error loading {example['audio_path']}: {e}") - return None - - # Extract features using tokenizer's feature extractor - features = tokenizer.feature_extractor(audio_array, sampling_rate=16000) - - # Tokenize text - tokenized_text = tokenizer.tokenizer(example["text"]) - - return { - "input_features": features.input_features[0], - "labels": tokenized_text.input_ids, - } - - print("Processing dataset...") - train_data = [] - for example in tqdm.tqdm(dataset, desc="Processing audio"): - result = formatting_prompts_func(example) - if result is not None: - train_data.append(result) - - print(f"Successfully processed {len(train_data)} samples") - - # Split into train/test - split_idx = max(1, int(len(train_data) * 0.94)) - train_dataset = train_data[:split_idx] - test_dataset = train_data[split_idx:] - - print(f"Train samples: {len(train_dataset)}, Test samples: {len(test_dataset)}") - - # Setup WER metric for evaluation - metric = evaluate.load("wer") - - def compute_metrics(pred): - pred_logits = pred.predictions[0] - label_ids = pred.label_ids - - # Replace -100 with pad_token_id - label_ids[label_ids == -100] = tokenizer.pad_token_id - - pred_ids = np.argmax(pred_logits, axis=-1) - - pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True) - label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=True) - - wer = 100 * metric.compute(predictions=pred_str, references=label_str) - - return {"wer": wer} - - @dataclass - class DataCollatorSpeechSeq2SeqWithPadding: - processor: Any - - def __call__( - self, features: List[Dict[str, Union[List[int], torch.Tensor]]] - ) -> Dict[str, torch.Tensor]: - input_features = [ - {"input_features": feature["input_features"]} for feature in features - ] - batch = self.processor.feature_extractor.pad( - input_features, return_tensors="pt" - ) - - label_features = [{"input_ids": feature["labels"]} for feature in features] - labels_batch = self.processor.tokenizer.pad( - label_features, return_tensors="pt" - ) - - labels = labels_batch["input_ids"].masked_fill( - labels_batch.attention_mask.ne(1), -100 - ) - - if ( - (labels[:, 0] == self.processor.tokenizer.bos_token_id) - .all() - .cpu() - .item() - ): - labels = labels[:, 1:] - - batch["labels"] = labels - - return batch - - data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=tokenizer) - - # Show memory stats before training - if torch.cuda.is_available(): - gpu_stats = torch.cuda.get_device_properties(0) - start_gpu_memory = round( - torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3 - ) - max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3) - print(f"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.") - print(f"{start_gpu_memory} GB of memory reserved.") - - # Setup trainer with Seq2SeqTrainer - trainer = Seq2SeqTrainer( - model=model, - train_dataset=train_dataset, - data_collator=data_collator, - eval_dataset=test_dataset if len(test_dataset) > 0 else None, - tokenizer=tokenizer.feature_extractor, - compute_metrics=compute_metrics, - args=Seq2SeqTrainingArguments( - # predict_with_generate=True, - per_device_train_batch_size=1, - gradient_accumulation_steps=4, - warmup_steps=5, - # num_train_epochs=1, # Set for full training run - max_steps=200, - learning_rate=1e-4, - logging_steps=10, - optim="adamw_8bit", - fp16=not is_bf16_supported(), - bf16=is_bf16_supported(), - weight_decay=0.001, - remove_unused_columns=False, # Required for PEFT - lr_scheduler_type="linear", - label_names=["labels"], - eval_steps=20, - eval_strategy="steps" if len(test_dataset) > 0 else "no", - seed=3407, - output_dir=OUTPUT_DIR, - report_to="none", - ), - ) - - print("Starting training...") - trainer_stats = trainer.train() - - # Show final memory stats - if torch.cuda.is_available(): - used_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3) - used_memory_for_lora = round(used_memory - start_gpu_memory, 3) - used_percentage = round(used_memory / max_memory * 100, 3) - lora_percentage = round(used_memory_for_lora / max_memory * 100, 3) - print(f"{trainer_stats.metrics['train_runtime']} seconds used for training.") - print( - f"{round(trainer_stats.metrics['train_runtime']/60, 2)} minutes used for training." - ) - print(f"Peak reserved memory = {used_memory} GB.") - print(f"Peak reserved memory for training = {used_memory_for_lora} GB.") - print(f"Peak reserved memory % of max memory = {used_percentage} %.") - print( - f"Peak reserved memory for training % of max memory = {lora_percentage} %." - ) - - # Save the model - print(f"Saving adapter to {OUTPUT_DIR}") - model.save_pretrained(OUTPUT_DIR) - tokenizer.save_pretrained(OUTPUT_DIR) - - print("Training complete!") - - -def run_inference(audio_file: str, model_path: str = OUTPUT_DIR): - """Run inference with the trained model""" - from transformers import pipeline - - print(f"Loading model from {model_path}") - - # Load the fine-tuned model - model, tokenizer = FastModel.from_pretrained( - model_name=model_path, - dtype=None, - load_in_4bit=False, - auto_model=WhisperForConditionalGeneration, - ) - - # Set model to inference mode - FastModel.for_inference(model) - model.eval() - - # Create pipeline - whisper = pipeline( - "automatic-speech-recognition", - model=model, - tokenizer=tokenizer.tokenizer, - feature_extractor=tokenizer.feature_extractor, - processor=tokenizer, - return_language=True, - torch_dtype=torch.float16, - ) - - # Transcribe - result = whisper(audio_file) - print(f"Transcription: {result['text']}") - return result - - -if __name__ == "__main__": - main() diff --git a/extras/speaker-recognition/scripts/enroll_from_backup.py b/extras/speaker-recognition/scripts/enroll_from_backup.py deleted file mode 100755 index d641a653..00000000 --- a/extras/speaker-recognition/scripts/enroll_from_backup.py +++ /dev/null @@ -1,181 +0,0 @@ -#!/usr/bin/env python3 -""" -Script to enroll all speakers from the backup directory. - -This script reads the backup speaker data and re-enrolls all speakers -using the /enroll/batch endpoint. -""" - -import asyncio -import logging -import sys -from pathlib import Path - -import aiohttp - -# Add the src directory to Python path -sys.path.insert(0, str(Path(__file__).parent.parent / "src")) - -# Configure logging -logging.basicConfig( - level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s" -) -log = logging.getLogger(__name__) - -# Configuration -BACKUP_DIR = Path(__file__).parent.parent / "speaker_data_backup" -ENROLLMENT_AUDIO_DIR = BACKUP_DIR / "enrollment_audio" / "1" # User ID 1 -API_BASE_URL = "http://localhost:8085" -TIMEOUT = 300 # 5 minutes per enrollment - - -async def enroll_speaker_from_backup( - session: aiohttp.ClientSession, speaker_dir: Path -) -> bool: - """Enroll a single speaker from backup directory.""" - - try: - # Extract speaker info from directory name - speaker_id = speaker_dir.name - - # Extract speaker name from ID (remove user_1_ prefix and timestamp suffix) - speaker_name = speaker_id.replace("user_1_", "").rsplit("_", 1)[0] - - log.info(f"Enrolling speaker: {speaker_name} (ID: {speaker_id})") - - # Find all WAV files in the directory - audio_files = list(speaker_dir.glob("*.wav")) - - if not audio_files: - log.error(f"No WAV files found for {speaker_name}, skipping") - return False - - log.info(f" Found {len(audio_files)} audio file(s)") - for audio_file in audio_files: - log.info(f" - {audio_file.name}") - - # Prepare multipart form data - form_data = aiohttp.FormData() - form_data.add_field("speaker_id", speaker_id) - form_data.add_field("speaker_name", speaker_name) - - # Add all audio files - for audio_path in audio_files: - with open(audio_path, "rb") as f: - form_data.add_field( - "files", - f.read(), - filename=audio_path.name, - content_type="audio/wav", - ) - - # Send enrollment request - log.info(f"Sending enrollment request for {speaker_name}...") - async with session.post( - f"{API_BASE_URL}/enroll/batch", - data=form_data, - timeout=aiohttp.ClientTimeout(total=TIMEOUT), - ) as response: - if response.status == 200: - result = await response.json() - log.info(f"✅ Successfully enrolled {speaker_name}: {result}") - return True - else: - error_text = await response.text() - log.error( - f"❌ Failed to enroll {speaker_name} (HTTP {response.status}): {error_text}" - ) - return False - - except Exception as e: - log.error(f"❌ Error enrolling {speaker_dir.name}: {e}") - return False - - -async def check_service_health() -> bool: - """Check if the speaker service is healthy.""" - try: - async with aiohttp.ClientSession() as session: - async with session.get(f"{API_BASE_URL}/health") as response: - if response.status == 200: - health_info = await response.json() - log.info(f"Service health: {health_info}") - return True - else: - log.error(f"Service health check failed: HTTP {response.status}") - return False - except Exception as e: - log.error(f"Failed to connect to service: {e}") - return False - - -async def main(): - """Main enrollment function.""" - log.info("Starting speaker enrollment from backup directory") - log.info(f"Backup directory: {BACKUP_DIR}") - log.info(f"Enrollment audio directory: {ENROLLMENT_AUDIO_DIR}") - - # Check if backup directory exists - if not ENROLLMENT_AUDIO_DIR.exists(): - log.error(f"Backup directory not found: {ENROLLMENT_AUDIO_DIR}") - return 1 - - # Check service health - log.info("Checking speaker service health...") - if not await check_service_health(): - log.error("Speaker service is not healthy, aborting") - return 1 - - # Find all speaker directories - speaker_dirs = [] - for item in ENROLLMENT_AUDIO_DIR.iterdir(): - if item.is_dir() and item.name.startswith("user_1_"): - speaker_dirs.append(item) - - if not speaker_dirs: - log.error("No speaker directories found in backup") - return 1 - - log.info(f"Found {len(speaker_dirs)} speakers to enroll") - - # Enroll each speaker - success_count = 0 - failed_count = 0 - - async with aiohttp.ClientSession() as session: - for speaker_dir in sorted(speaker_dirs): - success = await enroll_speaker_from_backup(session, speaker_dir) - if success: - success_count += 1 - else: - failed_count += 1 - - # Small delay between enrollments - await asyncio.sleep(1) - - # Summary - log.info(f"\n{'='*50}") - log.info(f"Enrollment Summary:") - log.info(f" Successfully enrolled: {success_count}") - log.info(f" Failed: {failed_count}") - log.info(f" Total: {len(speaker_dirs)}") - log.info(f"{'='*50}") - - if failed_count > 0: - log.warning("Some enrollments failed. Check logs above for details.") - return 1 - else: - log.info("All speakers enrolled successfully! 🎉") - return 0 - - -if __name__ == "__main__": - try: - exit_code = asyncio.run(main()) - sys.exit(exit_code) - except KeyboardInterrupt: - log.info("Enrollment interrupted by user") - sys.exit(1) - except Exception as e: - log.error(f"Unexpected error: {e}") - sys.exit(1) diff --git a/extras/speaker-recognition/scripts/simple_enroll.py b/extras/speaker-recognition/scripts/simple_enroll.py deleted file mode 100644 index 6ec5dbcc..00000000 --- a/extras/speaker-recognition/scripts/simple_enroll.py +++ /dev/null @@ -1,65 +0,0 @@ -#!/usr/bin/env python3 -"""Simple enrollment script without global variable issues.""" - -import sys - -import requests - -SERVICE_URL = "http://localhost:8085" - - -def list_speakers(): - """List enrolled speakers.""" - try: - response = requests.get(f"{SERVICE_URL}/speakers") - if response.status_code == 200: - data = response.json() - speakers = data.get("speakers", {}) - if not speakers: - print("No speakers enrolled") - else: - print(f"Enrolled speakers ({len(speakers)}):") - for speaker_id, info in speakers.items(): - print(f" {speaker_id}: {info.get('name', 'Unknown')}") - return True - else: - print(f"Error: {response.status_code}") - return False - except Exception as e: - print(f"Error: {e}") - return False - - -def enroll_speaker(file_path, speaker_id, speaker_name): - """Enroll a speaker from audio file.""" - try: - with open(file_path, "rb") as f: - files = {"file": f} - data = {"speaker_id": speaker_id, "speaker_name": speaker_name} - response = requests.post( - f"{SERVICE_URL}/enroll/upload", files=files, data=data - ) - - if response.status_code == 200: - result = response.json() - action = "Updated" if result.get("updated") else "Enrolled" - print(f"✅ {action} speaker: {speaker_name} (ID: {speaker_id})") - return True - else: - print(f"❌ Error: {response.status_code} - {response.text}") - return False - except Exception as e: - print(f"❌ Error: {e}") - return False - - -if __name__ == "__main__": - if len(sys.argv) == 2 and sys.argv[1] == "--list": - list_speakers() - elif len(sys.argv) == 4: - file_path, speaker_id, speaker_name = sys.argv[1], sys.argv[2], sys.argv[3] - enroll_speaker(file_path, speaker_id, speaker_name) - else: - print("Usage:") - print(" python simple_enroll.py --list") - print(" python simple_enroll.py <audio_file> <speaker_id> <speaker_name>") diff --git a/extras/speaker-recognition/scripts/test_diarize.py b/extras/speaker-recognition/scripts/test_diarize.py deleted file mode 100644 index dffbfafc..00000000 --- a/extras/speaker-recognition/scripts/test_diarize.py +++ /dev/null @@ -1,65 +0,0 @@ -#!/usr/bin/env python3 -"""Test the diarize-and-identify endpoint.""" - -import sys - -import requests - -SERVICE_URL = "http://localhost:8085" - - -def test_diarize_and_identify(audio_file): - """Test diarization and identification.""" - try: - with open(audio_file, "rb") as f: - files = {"file": f} - data = { - "min_duration": 0.5, - "similarity_threshold": 0.65, - "identify_only_enrolled": False, - } - - print(f"Testing diarization and identification on: {audio_file}") - response = requests.post( - f"{SERVICE_URL}/diarize-and-identify", files=files, data=data - ) - - if response.status_code == 200: - result = response.json() - print("\n✅ Success!") - - # Show summary - summary = result.get("summary", {}) - print(f"\nSummary:") - print(f" Duration: {summary.get('total_duration', 0):.2f}s") - print(f" Segments: {summary.get('num_segments', 0)}") - print(f" Diarized Speakers: {summary.get('num_diarized_speakers', 0)}") - print(f" Identified: {summary.get('identified_speakers', [])}") - print(f" Unknown: {summary.get('unknown_speakers', [])}") - - # Show segments - segments = result.get("segments", []) - print(f"\nSegments:") - for i, seg in enumerate(segments): - identified = seg.get("identified_as", "Unknown") - confidence = seg.get("confidence", 0) - print( - f" {i+1}. {seg['start']:.2f}s-{seg['end']:.2f}s: {seg['speaker']} → {identified} ({confidence:.3f})" - ) - - return True - else: - print(f"❌ Error: {response.status_code} - {response.text}") - return False - - except Exception as e: - print(f"❌ Error: {e}") - return False - - -if __name__ == "__main__": - if len(sys.argv) != 2: - print("Usage: python test_diarize.py <audio_file>") - sys.exit(1) - - test_diarize_and_identify(sys.argv[1]) diff --git a/extras/speaker-recognition/scripts/test_workflow.py b/extras/speaker-recognition/scripts/test_workflow.py deleted file mode 100755 index 8aa1c6c2..00000000 --- a/extras/speaker-recognition/scripts/test_workflow.py +++ /dev/null @@ -1,220 +0,0 @@ -#!/usr/bin/env python3 -""" -Test script to demonstrate the complete speaker enrollment and identification workflow. - -This script shows how to: -1. Enroll speakers with sample audio -2. Run diarization and identification on a conversation -3. Display results - -Usage: - python test_workflow.py -""" - -import os -import sys -from pathlib import Path - -import requests - -SPEAKER_SERVICE_URL = os.getenv("SPEAKER_SERVICE_URL", "http://localhost:8085") - - -def test_enrollment(): - """Test speaker enrollment with sample files.""" - print("\n=== Testing Speaker Enrollment ===") - - # Check if we have sample audio files - sample_dir = Path(__file__).parent / "training/youtube-transcript/data/audio" - if not sample_dir.exists(): - print(f"❌ Sample audio directory not found: {sample_dir}") - print( - "Please ensure you have audio files in the training/youtube-transcript/data/audio/ directory" - ) - return False - - audio_files = list(sample_dir.glob("*.wav")) - if not audio_files: - print(f"❌ No WAV files found in {sample_dir}") - return False - - print(f"Found {len(audio_files)} audio files") - - # Enroll first speaker using first audio file - if len(audio_files) >= 1: - file1 = audio_files[0] - print(f"\nEnrolling Speaker 1 from: {file1.name}") - - with open(file1, "rb") as f: - files = {"file": (file1.name, f, "audio/wav")} - data = {"speaker_id": "test_speaker_1", "speaker_name": "Test Speaker 1"} - - response = requests.post( - f"{SPEAKER_SERVICE_URL}/enroll/upload", files=files, data=data - ) - if response.status_code == 200: - print("✅ Successfully enrolled Test Speaker 1") - else: - print(f"❌ Failed to enroll: {response.status_code} - {response.text}") - return False - - # Enroll second speaker if we have another file - if len(audio_files) >= 2: - file2 = audio_files[1] - print(f"\nEnrolling Speaker 2 from: {file2.name}") - - with open(file2, "rb") as f: - files = {"file": (file2.name, f, "audio/wav")} - data = {"speaker_id": "test_speaker_2", "speaker_name": "Test Speaker 2"} - - response = requests.post( - f"{SPEAKER_SERVICE_URL}/enroll/upload", files=files, data=data - ) - if response.status_code == 200: - print("✅ Successfully enrolled Test Speaker 2") - else: - print(f"❌ Failed to enroll: {response.status_code} - {response.text}") - - return True - - -def test_diarize_and_identify(): - """Test the diarize-and-identify endpoint.""" - print("\n=== Testing Diarize and Identify ===") - - # Use one of the sample files for testing - sample_dir = Path(__file__).parent / "training/youtube-transcript/data/audio" - audio_files = list(sample_dir.glob("*.wav")) - - if not audio_files: - print("❌ No audio files available for testing") - return False - - # Use the first file for diarization test - test_file = audio_files[0] - print(f"Testing with: {test_file.name}") - - with open(test_file, "rb") as f: - files = {"file": (test_file.name, f, "audio/wav")} - data = { - "min_duration": 0.5, - "similarity_threshold": 0.65, - "identify_only_enrolled": False, - } - - print("Sending request to /diarize-and-identify...") - response = requests.post( - f"{SPEAKER_SERVICE_URL}/diarize-and-identify", files=files, data=data - ) - - if response.status_code == 200: - result = response.json() - print("\n✅ Diarization and identification successful!") - - # Display summary - summary = result.get("summary", {}) - print(f"\nSummary:") - print(f" Total Duration: {summary.get('total_duration', 0):.2f}s") - print(f" Segments: {summary.get('num_segments', 0)}") - print(f" Diarized Speakers: {summary.get('num_diarized_speakers', 0)}") - print(f" Identified Speakers: {summary.get('identified_speakers', [])}") - print(f" Unknown Speakers: {summary.get('unknown_speakers', [])}") - print(f" Similarity Threshold: {summary.get('similarity_threshold', 0)}") - - # Display first few segments - segments = result.get("segments", []) - if segments: - print(f"\nFirst {min(5, len(segments))} segments:") - for i, seg in enumerate(segments[:5]): - print(f"\n Segment {i+1}:") - print(f" Time: {seg['start']:.2f}s - {seg['end']:.2f}s") - print(f" Speaker: {seg['speaker']}") - print(f" Identified As: {seg.get('identified_as', 'Unknown')}") - print(f" Confidence: {seg.get('confidence', 0):.3f}") - print(f" Status: {seg.get('status', 'unknown')}") - - return True - else: - print(f"❌ Request failed: {response.status_code} - {response.text}") - return False - - -def list_speakers(): - """List all enrolled speakers.""" - print("\n=== Enrolled Speakers ===") - - response = requests.get(f"{SPEAKER_SERVICE_URL}/speakers") - if response.status_code == 200: - data = response.json() - speakers = data.get("speakers", {}) - - if not speakers: - print("No speakers enrolled") - else: - for speaker_id, info in speakers.items(): - print(f"\nID: {speaker_id}") - print(f" Name: {info.get('name', 'Unknown')}") - print(f" FAISS Index: {info.get('faiss_index', 'N/A')}") - - return True - else: - print(f"❌ Failed to list speakers: {response.status_code}") - return False - - -def cleanup(): - """Clean up test speakers.""" - print("\n=== Cleaning Up Test Data ===") - - # Delete test speakers - for speaker_id in ["test_speaker_1", "test_speaker_2"]: - response = requests.delete(f"{SPEAKER_SERVICE_URL}/speakers/{speaker_id}") - if response.status_code == 200: - print(f"✅ Deleted {speaker_id}") - elif response.status_code == 404: - print(f"ℹ️ {speaker_id} not found (already deleted)") - - -def main(): - print("Speaker Recognition Workflow Test") - print("=" * 50) - - # Check service health - try: - response = requests.get(f"{SPEAKER_SERVICE_URL}/health", timeout=5) - if response.status_code != 200: - print(f"❌ Speaker service not responding at {SPEAKER_SERVICE_URL}") - return 1 - health = response.json() - print(f"✅ Service running (Device: {health.get('device', 'unknown')})") - except Exception as e: - print(f"❌ Cannot connect to speaker service: {e}") - print("Make sure the service is running: docker compose up speaker-recognition") - return 1 - - # Run tests - try: - # Test enrollment - if not test_enrollment(): - print("\n❌ Enrollment test failed") - return 1 - - # List speakers - list_speakers() - - # Test diarize and identify - if not test_diarize_and_identify(): - print("\n❌ Diarize and identify test failed") - return 1 - - print("\n✅ All tests passed!") - - finally: - # Always cleanup - cleanup() - - return 0 - - -if __name__ == "__main__": - sys.exit(main())