admin-valentin/WhisperLiveKit
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

diarization now works at word - not chunk - level!

Browse source
This commit is contained in:
Quentin Fuxa 2025-02-24 00:35:42 +01:00
parent d89622b9c2
commit 2ced4fef20
4 changed files with 36 additions and 44 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
src/diarization/diarization_online.py
View file

@ -81,11 +81,10 @@ class DiartDiarization:
def close(self): def close(self):
self.source.close() self.source.close()
def assign_speakers_to_chunks(self, chunks: list) -> list: def assign_speakers_to_tokens(self, end_attributed_speaker, tokens: list) -> list:
end_attributed_speaker = 0 for token in tokens:
for chunk in chunks:
for segment in self.segment_speakers: for segment in self.segment_speakers:
if not (segment["end"] <= chunk["beg"] or segment["beg"] >= chunk["end"]): if not (segment["end"] <= token.start or segment["beg"] >= token.end):
chunk["speaker"] = extract_number(segment["speaker"]) + 1 token.speaker = extract_number(segment["speaker"]) + 1
end_attributed_speaker = chunk["end"] end_attributed_speaker = max(token.end, end_attributed_speaker)
return end_attributed_speaker return end_attributed_speaker

2
src/whisper_streaming/online_asr.py
View file

@ -202,7 +202,7 @@ class OnlineASRProcessor:
logger.debug( logger.debug(
f"Length of audio buffer now: {len(self.audio_buffer)/self.SAMPLING_RATE:.2f} seconds" f"Length of audio buffer now: {len(self.audio_buffer)/self.SAMPLING_RATE:.2f} seconds"
) )
return self.concatenate_tokens(committed_tokens) return committed_tokens
def chunk_completed_sentence(self): def chunk_completed_sentence(self):
""" """

3
src/whisper_streaming/timed_objects.py
View file

@ -5,7 +5,8 @@ from typing import Optional
class TimedText: class TimedText:
start: Optional[float] start: Optional[float]
end: Optional[float] end: Optional[float]
text: str text: Optional[str] = ''
speaker: Optional[int] = -1
@dataclass @dataclass
class ASRToken(TimedText): class ASRToken(TimedText):

64
whisper_fastapi_online_server.py
View file

@ -11,6 +11,7 @@ from fastapi.responses import HTMLResponse
from fastapi.middleware.cors import CORSMiddleware from fastapi.middleware.cors import CORSMiddleware
from src.whisper_streaming.whisper_online import backend_factory, online_factory, add_shared_args from src.whisper_streaming.whisper_online import backend_factory, online_factory, add_shared_args
from src.whisper_streaming.timed_objects import ASRToken
import math import math
import logging import logging
@ -47,7 +48,7 @@ parser.add_argument(
parser.add_argument( parser.add_argument(
"--diarization", "--diarization",
type=bool, type=bool,
default=False, default=True,
help="Whether to enable speaker diarization.", help="Whether to enable speaker diarization.",
) )
@ -157,7 +158,9 @@ async def websocket_endpoint(websocket: WebSocket):
full_transcription = "" full_transcription = ""
beg = time() beg = time()
beg_loop = time() beg_loop = time()
chunk_history = [] # Will store dicts: {beg, end, text, speaker} tokens = []
end_attributed_speaker = 0
sep = online.asr.sep
while True: while True:
try: try:
@ -177,7 +180,6 @@ async def websocket_endpoint(websocket: WebSocket):
logger.warning("FFmpeg read timeout. Restarting...") logger.warning("FFmpeg read timeout. Restarting...")
await restart_ffmpeg() await restart_ffmpeg()
full_transcription = "" full_transcription = ""
chunk_history = []
beg = time() beg = time()
continue # Skip processing and read from new process continue # Skip processing and read from new process
@ -202,63 +204,53 @@ async def websocket_endpoint(websocket: WebSocket):
if args.transcription: if args.transcription:
logger.info(f"{len(online.audio_buffer) / online.SAMPLING_RATE} seconds of audio will be processed by the model.") logger.info(f"{len(online.audio_buffer) / online.SAMPLING_RATE} seconds of audio will be processed by the model.")
online.insert_audio_chunk(pcm_array) online.insert_audio_chunk(pcm_array)
transcription = online.process_iter() new_tokens = online.process_iter()
if transcription.start: tokens.extend(new_tokens)
chunk_history.append({ full_transcription += sep.join([t.text for t in new_tokens])
"beg": transcription.start,
"end": transcription.end,
"text": transcription.text,
"speaker": -1
})
full_transcription += transcription.text if transcription else ""
buffer = online.get_buffer() buffer = online.get_buffer()
if buffer in full_transcription: # With VAC, the buffer is not updated until the next chunk is processed if buffer in full_transcription: # With VAC, the buffer is not updated until the next chunk is processed
buffer = "" buffer = ""
else: else:
chunk_history.append({ tokens.append(
"beg": time() - beg_loop, ASRToken(
"end": time() - beg_loop + 1, start = time() - beg_loop,
"text": '', end = time() - beg_loop + 0.5))
"speaker": -1 sleep(0.5)
})
sleep(1)
buffer = '' buffer = ''
if args.diarization: if args.diarization:
await diarization.diarize(pcm_array) await diarization.diarize(pcm_array)
end_attributed_speaker = diarization.assign_speakers_to_chunks(chunk_history) end_attributed_speaker = diarization.assign_speakers_to_tokens(end_attributed_speaker, tokens)
previous_speaker = -10
current_speaker = -10
lines = [] lines = []
last_end_diarized = 0 last_end_diarized = 0
previous_speaker = -1 for token in tokens:
for ind, ch in enumerate(chunk_history): speaker = token.speaker
speaker = ch.get("speaker")
if args.diarization: if args.diarization:
if speaker == -1 or speaker == 0: if speaker == -1 or speaker == 0:
if ch['end'] < end_attributed_speaker: if token.end < end_attributed_speaker:
speaker = previous_speaker speaker = previous_speaker
else: else:
speaker = 0 speaker = 0
else: else:
last_end_diarized = max(ch['end'], last_end_diarized) last_end_diarized = max(token.end, last_end_diarized)
if speaker != current_speaker: if speaker != previous_speaker:
lines.append( lines.append(
{ {
"speaker": speaker, "speaker": speaker,
"text": ch['text'], "text": token.text,
"beg": format_time(ch['beg']), "beg": format_time(token.start),
"end": format_time(ch['end']), "end": format_time(token.end),
"diff": round(ch['end'] - last_end_diarized, 2) "diff": round(token.end - last_end_diarized, 2)
} }
) )
current_speaker = speaker previous_speaker = speaker
else: else:
lines[-1]["text"] += ch['text'] lines[-1]["text"] += sep + token.text
lines[-1]["end"] = format_time(ch['end']) lines[-1]["end"] = format_time(token.end)
lines[-1]["diff"] = round(ch['end'] - last_end_diarized, 2) lines[-1]["diff"] = round(token.end - last_end_diarized, 2)
response = {"lines": lines, "buffer": buffer} response = {"lines": lines, "buffer": buffer}
await websocket.send_json(response) await websocket.send_json(response)