diff --git a/components/SettingsMenu.tsx b/components/SettingsMenu.tsx index fc00f33..bc0d843 100644 --- a/components/SettingsMenu.tsx +++ b/components/SettingsMenu.tsx @@ -28,6 +28,33 @@ const MENU_LINKS = [ { label: "Follow on X", href: X_PROFILE_URL, Icon: XIcon }, ] as const; +const TRANSCRIPTION_SETTINGS_STORAGE_KEY = "rescript.transcription.settings"; +const DEFAULT_TRANSCRIPTION_SETTINGS = { + maxSpeakers: 2, + onsetThreshold: 0.7, +}; + +function readTranscriptionSettings() { + if (typeof window === "undefined" || !window.localStorage) { + return DEFAULT_TRANSCRIPTION_SETTINGS; + } + try { + const raw = window.localStorage.getItem(TRANSCRIPTION_SETTINGS_STORAGE_KEY); + if (!raw) return DEFAULT_TRANSCRIPTION_SETTINGS; + const parsed = JSON.parse(raw) as Partial; + return { + maxSpeakers: Number.isFinite(parsed.maxSpeakers) + ? Math.max(1, Math.round(Number(parsed.maxSpeakers))) + : DEFAULT_TRANSCRIPTION_SETTINGS.maxSpeakers, + onsetThreshold: Number.isFinite(parsed.onsetThreshold) + ? Math.max(0, Math.min(1, Number(parsed.onsetThreshold))) + : DEFAULT_TRANSCRIPTION_SETTINGS.onsetThreshold, + }; + } catch { + return DEFAULT_TRANSCRIPTION_SETTINGS; + } +} + /** * Top-bar settings popover. Houses appearance, transcript source, and social * links for now — structure is section-based so more prefs can land here later. @@ -37,6 +64,30 @@ export default function SettingsMenu() { const panelId = useId(); const { appearance, setAppearance } = useAppearance(); const { enabled: telemetry, setEnabled: setTelemetry } = useTelemetryPref(); + const [transcriptionSettings, setTranscriptionSettings] = useState(() => + readTranscriptionSettings() + ); + + const commitTranscriptionSettings = ( + patch: Partial + ) => { + const next = { ...transcriptionSettings, ...patch }; + const normalized = { + maxSpeakers: Math.max(1, Math.round(Number(next.maxSpeakers))), + onsetThreshold: Math.max(0, Math.min(1, Number(next.onsetThreshold))), + }; + setTranscriptionSettings(normalized); + try { + if (typeof window !== "undefined") { + window.localStorage.setItem( + TRANSCRIPTION_SETTINGS_STORAGE_KEY, + JSON.stringify(normalized) + ); + } + } catch { + // localStorage may be unavailable in private browsing or storage-full modes. + } + }; return ( +
+

+ Transcription +

+
+ + +
+
+
{MENU_LINKS.map(({ label, href, Icon }) => ( ; + return { + maxSpeakers: Number.isFinite(parsed.maxSpeakers) + ? Math.max(1, Math.round(Number(parsed.maxSpeakers))) + : DEFAULT_TRANSCRIPTION_SETTINGS.maxSpeakers, + onsetThreshold: Number.isFinite(parsed.onsetThreshold) + ? Math.max(0, Math.min(1, Number(parsed.onsetThreshold))) + : DEFAULT_TRANSCRIPTION_SETTINGS.onsetThreshold, + }; + } catch { + return DEFAULT_TRANSCRIPTION_SETTINGS; + } +} + /** Stop an in-flight ASR job (e.g. after importing a transcript). */ export function cancelTranscription() { activeWorker?.terminate(); @@ -90,13 +117,22 @@ export function useTranscriber() { ); }; + const settings = readTranscriptionSettings(); + // Transfer, not copy: the worker takes ownership of the PCM and `audio` is // detached here. Nothing on the main thread reads it afterwards — the // waveform draws from the envelope the store built in setAudio — and on a // long recording the copy this replaces was hundreds of megabytes held for // the length of the run. workerRef.current.postMessage( - { audio, duration, model, language: transcriptLanguage }, + { + audio, + duration, + model, + language: transcriptLanguage, + maxSpeakers: settings.maxSpeakers, + onsetThreshold: settings.onsetThreshold, + }, [audio.buffer] ); }, []); diff --git a/lib/types.ts b/lib/types.ts index cca974f..0fbc567 100644 --- a/lib/types.ts +++ b/lib/types.ts @@ -104,4 +104,8 @@ export interface WorkerRequest { * and is passed through for Parakeet alignment (Parakeet ASR still auto-detects). */ language: import("./languages").TranscriptLanguage; + /** Optional speaker-cap setting for diarization post-processing. */ + maxSpeakers?: number; + /** Optional diarization onset confidence threshold. */ + onsetThreshold?: number; } diff --git a/workers/transcription.worker.ts b/workers/transcription.worker.ts index 8da612f..fdee348 100644 --- a/workers/transcription.worker.ts +++ b/workers/transcription.worker.ts @@ -132,6 +132,10 @@ if (/apple/i.test(navigator.vendor)) { const DIARIZATION_MODEL = "onnx-community/pyannote-segmentation-3.0"; const VAD_MODEL = "onnx-community/silero-vad"; +/** Limit diarized speaker IDs to the requested UI-friendly headroom. */ +const MAX_ALLOWED_SPEAKERS = 2; +/** Raise the diarization onset confidence bar above the default soft-slice margin. */ +const ONSET_THRESHOLD = 0.7; /** Viterbi needs a frames x tokens lattice, so alignment runs in bounded batches. */ const ALIGN_BATCH_MAX_S = 20; /** Context either side of a batch, so edge words are not clipped. */ @@ -984,14 +988,18 @@ async function refineWordTimestamps( * class indices are only meaningful within a pass, which is what the overlap and * `stitchDiarizationWindows` are for. */ -async function diarize(audio: Float32Array): Promise { +async function diarize( + audio: Float32Array, + onsetThreshold = ONSET_THRESHOLD +): Promise { const { processor, model } = await getDiarizer(); // post_process_speaker_diarization is specific to the PyAnnote processor // and is not part of the generic Processor typings. const pyannote = processor as unknown as { post_process_speaker_diarization: ( logits: unknown, - numSamples: number + numSamples: number, + opts?: { onsetThreshold?: number; onset_threshold?: number } ) => DiarizationSegment[][]; }; @@ -1007,7 +1015,11 @@ async function diarize(audio: Float32Array): Promise { windows.push({ offsetS: startSample / VAD_SAMPLE_RATE, durationS: slice.length / VAD_SAMPLE_RATE, - segments: pyannote.post_process_speaker_diarization(logits, slice.length)[0] ?? [], + segments: + pyannote.post_process_speaker_diarization(logits, slice.length, { + onsetThreshold, + onset_threshold: onsetThreshold, + })[0] ?? [], }); postLive({ type: "progress", @@ -1019,7 +1031,11 @@ async function diarize(audio: Float32Array): Promise { } /** Assign a speaker to each word from the diarization segments. */ -function assignSpeakers(words: Word[], segments: DiarizationSegment[]) { +function assignSpeakers( + words: Word[], + segments: DiarizationSegment[], + maxSpeakers = MAX_ALLOWED_SPEAKERS +) { // Segment id 0 is "no speaker" (silence/noise); ignore it. const speech = segments.filter((s) => s.id !== 0); if (speech.length === 0) { @@ -1033,6 +1049,7 @@ function assignSpeakers(words: Word[], segments: DiarizationSegment[]) { const byStart = [...speech].sort((a, b) => a.start - b.start); let cursor = 0; + const normalizedMax = Math.max(1, Math.round(maxSpeakers)); const idMap = new Map(); // pyannote id -> sequential index for (const w of words) { const mid = (w.start + w.end) / 2; @@ -1057,7 +1074,10 @@ function assignSpeakers(words: Word[], segments: DiarizationSegment[]) { } } const raw = seg ? seg.id : -1; - if (raw >= 0 && !idMap.has(raw)) idMap.set(raw, idMap.size); + if (raw >= 0 && !idMap.has(raw)) { + const speakerId = (parseInt(String(raw), 10) % normalizedMax) + 1; + idMap.set(raw, speakerId); + } w.speaker = raw >= 0 ? (idMap.get(raw) as number) : 0; } } @@ -1107,12 +1127,14 @@ function wordsFromParakeet( async function finishWithDiarization( words: Word[], - audio: Float32Array + audio: Float32Array, + onsetThreshold = ONSET_THRESHOLD, + maxSpeakers = MAX_ALLOWED_SPEAKERS ): Promise { try { post({ type: "progress", message: "Identifying speakers…", value: 0 }); - const segments = await diarize(audio); - assignSpeakers(words, segments); + const segments = await diarize(audio, onsetThreshold); + assignSpeakers(words, segments, maxSpeakers); } catch (err) { console.warn("Speaker diarization failed; using a single speaker.", err); } @@ -1122,7 +1144,9 @@ async function finishWithDiarization( async function runParakeet( audio: Float32Array, duration: number, - transcriptLanguage: TranscriptLanguage + transcriptLanguage: TranscriptLanguage, + maxSpeakers = MAX_ALLOWED_SPEAKERS, + onsetThreshold = ONSET_THRESHOLD ): Promise { // Overlap diarizer (+ language-matched aligner) with Parakeet load. getDiarizer().catch(() => {}); @@ -1201,14 +1225,16 @@ async function runParakeet( duration, transcriptLanguage ); - return finishWithDiarization(words, audio); + return finishWithDiarization(words, audio, onsetThreshold, maxSpeakers); } async function runWhisper( audio: Float32Array, duration: number, choice: WhisperModel, - transcriptLanguage: TranscriptLanguage + transcriptLanguage: TranscriptLanguage, + maxSpeakers = MAX_ALLOWED_SPEAKERS, + onsetThreshold = ONSET_THRESHOLD ): Promise { // Overlap Whisper + Silero downloads; diarizer and language-matched aligner // warm in the background so both are cached by the time the transcript lands. @@ -1402,20 +1428,39 @@ async function runWhisper( transcriptLanguage ); - return finishWithDiarization(words, audio); + return finishWithDiarization(words, audio, onsetThreshold, maxSpeakers); } self.onmessage = async (event: MessageEvent) => { - const { audio, duration, model, language } = event.data; + const { audio, duration, model, language, maxSpeakers, onsetThreshold } = event.data; try { const choice: ModelId = model ?? "base"; const transcriptLanguage: TranscriptLanguage = language ?? "en"; + const normalizedMaxSpeakers = Number.isFinite(maxSpeakers) + ? Math.max(1, Math.round(Number(maxSpeakers))) + : MAX_ALLOWED_SPEAKERS; + const normalizedThreshold = Number.isFinite(onsetThreshold) + ? Math.max(0, Math.min(1, Number(onsetThreshold))) + : ONSET_THRESHOLD; let words: Word[]; if (isParakeetModel(choice)) { - words = await runParakeet(audio, duration, transcriptLanguage); + words = await runParakeet( + audio, + duration, + transcriptLanguage, + normalizedMaxSpeakers, + normalizedThreshold + ); } else if (isWhisperModel(choice)) { - words = await runWhisper(audio, duration, choice, transcriptLanguage); + words = await runWhisper( + audio, + duration, + choice, + transcriptLanguage, + normalizedMaxSpeakers, + normalizedThreshold + ); } else { throw new Error(`Unknown speech model: ${String(choice)}`); }