internal rework 1

docs/alignement_principles.md

@@ -0,0 +1,71 @@
+### Alignment between STT Tokens and Diarization Segments 
+
+- Example 1: The punctuation from STT and the speaker change from Diariation come in the prediction `t`
+- Example 2: The punctuation from STT comes from prediction `t`, but the speaker change from Diariation come in the prediction `t-1`
+- Example 3: The punctuation from STT comes from prediction `t-1`, but the speaker change from Diariation come in the prediction `t`
+
+> `#` Is the split between the `t-1` prediction and t prediction.  
+
+
+## Example 1:
+```text
+punctuations_segments : __#_______.__________________!____
+diarization_segments:
+SPK1                    __#____________
+SPK2                      #            ___________________
+-->
+ALIGNED SPK1            __#_______.
+ALIGNED SPK2              #        __________________!____
+
+t-1 output:
+SPK1:                   __#
+SPK2: NO
+DIARIZATION BUFFER: NO
+
+t output:
+SPK1:                       __#__.
+SPK2:                             __________________!____
+DIARIZATION BUFFER: No
+```
+
+## Example 2:
+```text
+punctuations_segments : _____#__.___________
+diarization_segments:
+SPK1                    ___  #
+SPK2                       __#______________
+-->
+ALIGNED SPK1            _____#__.
+ALIGNED SPK2                 #   ___________
+
+t-1 output:
+SPK1:                   ___  #
+SPK2:
+DIARIZATION BUFFER:        __#
+
+t output:
+SPK1:                      __#__.
+SPK2:                            ___________
+DIARIZATION BUFFER: No
+```
+
+## Example 3:
+```text
+punctuations_segments : ___.__#__________
+diarization_segments:
+SPK1                    ______#__
+SPK2                          #  ________
+-->
+ALIGNED SPK1            ___.  #
+ALIGNED SPK2                __#__________
+
+t-1 output:
+SPK1:                   ___.  #
+SPK2:
+DIARIZATION BUFFER:         __#
+
+t output:
+SPK1:                         #
+SPK2:                       __#___________
+DIARIZATION BUFFER: NO
+```

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "whisperlivekit"
-version = "0.2.14"
+version = "0.2.14.post4"
 description = "Real-time speech-to-text with speaker diarization using Whisper"
 readme = "README.md"
 authors = [

whisperlivekit/TokensAlignment.py

@@ -0,0 +1,168 @@
+class TokensAlignment:
+
+    def __init__(self, state_light, silence=None, args=None):
+        self.state_light = state_light
+        self.silence = silence
+        self.args = args
+
+        self._tokens_index = 0
+        self._diarization_index = 0
+        self._translation_index = 0
+
+    def update(self):
+        pass
+
+
+    def compute_punctuations_segments(self):
+        punctuations_breaks = []
+        new_tokens = self.state.tokens[self.state.last_validated_token:]
+        for i in range(len(new_tokens)):
+            token = new_tokens[i]
+            if token.is_punctuation():
+                punctuations_breaks.append({
+                    'token_index': i,
+                    'token': token,
+                    'start': token.start,
+                    'end': token.end,
+                })
+        punctuations_segments = []
+        for i, break_info in enumerate(punctuations_breaks):
+            start = punctuations_breaks[i - 1]['end'] if i > 0 else 0.0
+            end = break_info['end']
+            punctuations_segments.append({
+                'start': start,
+                'end': end,
+                'token_index': break_info['token_index'],
+                'token': break_info['token']
+            })
+        return punctuations_segments
+
+    def concatenate_diar_segments(self):
+        diarization_segments = self.state.diarization_segments
+
+if __name__ == "__main__":
+    from whisperlivekit.timed_objects import State, ASRToken, SpeakerSegment, Transcript, Silence
+    
+    # Reconstruct the state from the backup data
+    tokens = [
+        ASRToken(start=1.38, end=1.48, text=' The'),
+        ASRToken(start=1.42, end=1.52, text=' description'),
+        ASRToken(start=1.82, end=1.92, text=' technology'),
+        ASRToken(start=2.54, end=2.64, text=' has'),
+        ASRToken(start=2.7, end=2.8, text=' improved'),
+        ASRToken(start=3.24, end=3.34, text=' so'),
+        ASRToken(start=3.66, end=3.76, text=' much'),
+        ASRToken(start=4.02, end=4.12, text=' in'),
+        ASRToken(start=4.08, end=4.18, text=' the'),
+        ASRToken(start=4.26, end=4.36, text=' past'),
+        ASRToken(start=4.48, end=4.58, text=' few'),
+        ASRToken(start=4.76, end=4.86, text=' years'),
+        ASRToken(start=5.76, end=5.86, text='.'),
+        ASRToken(start=5.72, end=5.82, text=' Have'),
+        ASRToken(start=5.92, end=6.02, text=' you'),
+        ASRToken(start=6.08, end=6.18, text=' noticed'),
+        ASRToken(start=6.52, end=6.62, text=' how'),
+        ASRToken(start=6.8, end=6.9, text=' accurate'),
+        ASRToken(start=7.46, end=7.56, text=' real'),
+        ASRToken(start=7.72, end=7.82, text='-time'),
+        ASRToken(start=8.06, end=8.16, text=' speech'),
+        ASRToken(start=8.48, end=8.58, text=' to'),
+        ASRToken(start=8.68, end=8.78, text=' text'),
+        ASRToken(start=9.0, end=9.1, text=' is'),
+        ASRToken(start=9.24, end=9.34, text=' now'),
+        ASRToken(start=9.82, end=9.92, text='?'),
+        ASRToken(start=9.86, end=9.96, text=' Absolutely'),
+        ASRToken(start=11.26, end=11.36, text='.'),
+        ASRToken(start=11.36, end=11.46, text=' I'),
+        ASRToken(start=11.58, end=11.68, text=' use'),
+        ASRToken(start=11.78, end=11.88, text=' it'),
+        ASRToken(start=11.94, end=12.04, text=' all'),
+        ASRToken(start=12.08, end=12.18, text=' the'),
+        ASRToken(start=12.32, end=12.42, text=' time'),
+        ASRToken(start=12.58, end=12.68, text=' for'),
+        ASRToken(start=12.78, end=12.88, text=' taking'),
+        ASRToken(start=13.14, end=13.24, text=' notes'),
+        ASRToken(start=13.4, end=13.5, text=' during'),
+        ASRToken(start=13.78, end=13.88, text=' meetings'),
+        ASRToken(start=14.6, end=14.7, text='.'),
+        ASRToken(start=14.82, end=14.92, text=' It'),
+        ASRToken(start=14.92, end=15.02, text="'s"),
+        ASRToken(start=15.04, end=15.14, text=' amazing'),
+        ASRToken(start=15.5, end=15.6, text=' how'),
+        ASRToken(start=15.66, end=15.76, text=' it'),
+        ASRToken(start=15.8, end=15.9, text=' can'),
+        ASRToken(start=15.96, end=16.06, text=' recognize'),
+        ASRToken(start=16.58, end=16.68, text=' different'),
+        ASRToken(start=16.94, end=17.04, text=' speakers'),
+        ASRToken(start=17.82, end=17.92, text=' and'),
+        ASRToken(start=18.0, end=18.1, text=' even'),
+        ASRToken(start=18.42, end=18.52, text=' add'),
+        ASRToken(start=18.74, end=18.84, text=' punct'),
+        ASRToken(start=19.02, end=19.12, text='uation'),
+        ASRToken(start=19.68, end=19.78, text='.'),
+        ASRToken(start=20.04, end=20.14, text=' Yeah'),
+        ASRToken(start=20.5, end=20.6, text=','),
+        ASRToken(start=20.6, end=20.7, text=' but'),
+        ASRToken(start=20.76, end=20.86, text=' sometimes'),
+        ASRToken(start=21.42, end=21.52, text=' noise'),
+        ASRToken(start=21.82, end=21.92, text=' can'),
+        ASRToken(start=22.08, end=22.18, text=' still'),
+        ASRToken(start=22.38, end=22.48, text=' cause'),
+        ASRToken(start=22.72, end=22.82, text=' mistakes'),
+        ASRToken(start=23.74, end=23.84, text='.'),
+        ASRToken(start=23.96, end=24.06, text=' Does'),
+        ASRToken(start=24.16, end=24.26, text=' this'),
+        ASRToken(start=24.4, end=24.5, text=' system'),
+        ASRToken(start=24.76, end=24.86, text=' handle'),
+        ASRToken(start=25.12, end=25.22, text=' that'),
+        ASRToken(start=25.38, end=25.48, text=' well'),
+        ASRToken(start=25.68, end=25.78, text='?'),
+        ASRToken(start=26.4, end=26.5, text=' It'),
+        ASRToken(start=26.5, end=26.6, text=' does'),
+        ASRToken(start=26.7, end=26.8, text=' a'),
+        ASRToken(start=27.08, end=27.18, text=' pretty'),
+        ASRToken(start=27.12, end=27.22, text=' good'),
+        ASRToken(start=27.34, end=27.44, text=' job'),
+        ASRToken(start=27.64, end=27.74, text=' filtering'),
+        ASRToken(start=28.1, end=28.2, text=' noise'),
+        ASRToken(start=28.64, end=28.74, text=','),
+        ASRToken(start=28.78, end=28.88, text=' especially'),
+        ASRToken(start=29.3, end=29.4, text=' with'),
+        ASRToken(start=29.51, end=29.61, text=' models'),
+        ASRToken(start=29.99, end=30.09, text=' that'),
+        ASRToken(start=30.21, end=30.31, text=' use'),
+        ASRToken(start=30.51, end=30.61, text=' voice'),
+        ASRToken(start=30.83, end=30.93, text=' activity'),
+    ]
+    
+    diarization_segments = [
+        SpeakerSegment(start=1.3255040645599365, end=4.3255040645599365, speaker=0),
+        SpeakerSegment(start=4.806154012680054, end=9.806154012680054, speaker=0),
+        SpeakerSegment(start=9.806154012680054, end=10.806154012680054, speaker=1),
+        SpeakerSegment(start=11.168735027313232, end=14.168735027313232, speaker=1),
+        SpeakerSegment(start=14.41029405593872, end=17.41029405593872, speaker=1),
+        SpeakerSegment(start=17.52983808517456, end=19.52983808517456, speaker=1),
+        SpeakerSegment(start=19.64953374862671, end=20.066200415293377, speaker=1),
+        SpeakerSegment(start=20.066200415293377, end=22.64953374862671, speaker=2),
+        SpeakerSegment(start=23.012792587280273, end=25.012792587280273, speaker=2),
+        SpeakerSegment(start=25.495875597000122, end=26.41254226366679, speaker=2),
+        SpeakerSegment(start=26.41254226366679, end=30.495875597000122, speaker=0),
+    ]
+    
+    state = State(
+        tokens=tokens,
+        last_validated_token=72,
+        last_speaker=-1,
+        last_punctuation_index=71,
+        translation_validated_segments=[],
+        buffer_translation=Transcript(start=0, end=0, speaker=-1),
+        buffer_transcription=Transcript(start=None, end=None, speaker=-1),
+        diarization_segments=diarization_segments,
+        end_buffer=31.21587559700018,
+        end_attributed_speaker=30.495875597000122,
+        remaining_time_transcription=0.4,
+        remaining_time_diarization=0.7,
+        beg_loop=1763627603.968919
+    )
+    
+    alignment = TokensAlignment(state)

whisperlivekit/audio_processor.py

@@ -4,12 +4,12 @@ from time import time, sleep
 import math
 import logging
 import traceback
-from whisperlivekit.timed_objects import ASRToken, Silence, Line, FrontData, State, Transcript, ChangeSpeaker
+from whisperlivekit.timed_objects import ASRToken, Silence, Line, FrontData, State, StateLight, Transcript, ChangeSpeaker
 from whisperlivekit.core import TranscriptionEngine, online_factory, online_diarization_factory, online_translation_factory
 from whisperlivekit.silero_vad_iterator import FixedVADIterator
 from whisperlivekit.results_formater import format_output
 from whisperlivekit.ffmpeg_manager import FFmpegManager, FFmpegState
-
+from whisperlivekit.TokensAlignment import TokensAlignment
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
@@ -31,13 +31,29 @@ def cut_at(cumulative_pcm, cut_sec):
 
 async def get_all_from_queue(queue):
     items = []
-    try:
-        while True:
-            item = queue.get_nowait()
-            items.append(item)
-    except asyncio.QueueEmpty:
-        pass
-    return items
+
+    first_item = await queue.get()
+    queue.task_done()
+    if first_item is SENTINEL:
+        return first_item
+    if isinstance(first_item, Silence):
+        return first_item
+    items.append(first_item)
+    
+    while True:
+        if not queue._queue:
+            break
+        next_item = queue._queue[0]
+        if next_item is SENTINEL:
+            break
+        if isinstance(next_item, Silence):
+            break
+        items.append(await queue.get())
+        queue.task_done()
+    if isinstance(items[0], np.ndarray):
+        return np.concatenate(items)
+    else: #translation
+        return items
 
 class AudioProcessor:
     """
@@ -65,24 +81,21 @@ class AudioProcessor:
 
         # State management
         self.is_stopping = False
-        self.silence = False
+        self.silence = True
         self.silence_duration = 0.0
+        self.start_silence = None
+        self.last_silence_dispatch_time = None
         self.state = State()
+        self.state_light = StateLight()
         self.lock = asyncio.Lock()
         self.sep = " "  # Default separator
         self.last_response_content = FrontData()
         self.last_detected_speaker = None
         self.speaker_languages = {}
-        self.diarization_before_transcription = False
 
-        self.segments = []
-        
+        self.tokens_alignment = TokensAlignment(self.state_light, self.args, self.sep)
+        self.beg_loop = None
 
-        if self.diarization_before_transcription:
-            self.cumulative_pcm = []
-            self.last_start = 0.0
-            self.last_end = 0.0
-        
         # Models and processing
         self.asr = models.asr
         self.vac_model = models.vac_model
@@ -109,7 +122,8 @@ class AudioProcessor:
         self.diarization_queue = asyncio.Queue() if self.args.diarization else None
         self.translation_queue = asyncio.Queue() if self.args.target_language else None
         self.pcm_buffer = bytearray()
-
+        self.total_pcm_samples = 0
+        self.end_buffer = 0.0
         self.transcription_task = None
         self.diarization_task = None
         self.translation_task = None
@@ -128,18 +142,57 @@ class AudioProcessor:
         if models.translation_model:
             self.translation = online_translation_factory(self.args, models.translation_model)
 
+    async def _push_silence_event(self, silence_buffer: Silence):
+        if self.transcription_queue:
+            await self.transcription_queue.put(silence_buffer)
+        if self.args.diarization and self.diarization_queue:
+            await self.diarization_queue.put(silence_buffer)
+        if self.translation_queue:
+            await self.translation_queue.put(silence_buffer)
+
+    async def _begin_silence(self):
+        if self.silence:
+            return
+        self.silence = True
+        now = time()
+        self.start_silence = now
+        self.last_silence_dispatch_time = now
+        await self._push_silence_event(Silence(is_starting=True))
+
+    async def _end_silence(self):
+        if not self.silence:
+            return
+        now = time()
+        duration = now - (self.last_silence_dispatch_time if self.last_silence_dispatch_time else self.beg_loop)
+        await self._push_silence_event(Silence(duration=duration, has_ended=True))
+        self.last_silence_dispatch_time = now
+        self.silence = False
+        self.start_silence = None
+        self.last_silence_dispatch_time = None
+
+    async def _enqueue_active_audio(self, pcm_chunk: np.ndarray):
+        if pcm_chunk is None or pcm_chunk.size == 0:
+            return
+        if self.transcription_queue:
+            await self.transcription_queue.put(pcm_chunk.copy())
+        if self.args.diarization and self.diarization_queue:
+            await self.diarization_queue.put(pcm_chunk.copy())
+        self.silence_duration = 0.0
+
+    def _slice_before_silence(self, pcm_array, chunk_sample_start, silence_sample):
+        if silence_sample is None:
+            return None
+        relative_index = int(silence_sample - chunk_sample_start)
+        if relative_index <= 0:
+            return None
+        split_index = min(relative_index, len(pcm_array))
+        if split_index <= 0:
+            return None
+        return pcm_array[:split_index]
+
     def convert_pcm_to_float(self, pcm_buffer):
         """Convert PCM buffer in s16le format to normalized NumPy array."""
         return np.frombuffer(pcm_buffer, dtype=np.int16).astype(np.float32) / 32768.0
-
-    async def add_dummy_token(self):
-        """Placeholder token when no transcription is available."""
-        async with self.lock:
-            current_time = time() - self.state.beg_loop
-            self.state.tokens.append(ASRToken(
-                start=current_time, end=current_time + 1,
-                text=".", speaker=-1, is_dummy=True
-            ))
             
     async def get_current_state(self):
         """Get current state."""
@@ -147,12 +200,12 @@ class AudioProcessor:
             current_time = time()
             
             remaining_transcription = 0
-            if self.state.end_buffer > 0:
-                remaining_transcription = max(0, round(current_time - self.state.beg_loop - self.state.end_buffer, 1))
+            if self.end_buffer > 0:
+                remaining_transcription = max(0, round(current_time - self.beg_loop - self.end_buffer, 1))
                 
             remaining_diarization = 0
             if self.state.tokens:
-                latest_end = max(self.state.end_buffer, self.state.tokens[-1].end if self.state.tokens else 0)
+                latest_end = max(self.end_buffer, self.state.tokens[-1].end if self.state.tokens else 0)
                 remaining_diarization = max(0, round(latest_end - self.state.end_attributed_speaker, 1))
                 
             self.state.remaining_time_transcription = remaining_transcription
@@ -203,7 +256,7 @@ class AudioProcessor:
                 await asyncio.sleep(0.2)
 
         logger.info("FFmpeg stdout processing finished. Signaling downstream processors if needed.")
-        if not self.diarization_before_transcription and self.transcription_queue:
+        if self.transcription_queue:
             await self.transcription_queue.put(SENTINEL)
         if self.diarization:
             await self.diarization_queue.put(SENTINEL)
@@ -216,37 +269,47 @@ class AudioProcessor:
         
         while True:
             try:
-                item = await self.transcription_queue.get()
+                # item = await self.transcription_queue.get()
+                item = await get_all_from_queue(self.transcription_queue)
                 if item is SENTINEL:
                     logger.debug("Transcription processor received sentinel. Finishing.")
-                    self.transcription_queue.task_done()
                     break
 
                 asr_internal_buffer_duration_s = len(getattr(self.transcription, 'audio_buffer', [])) / self.transcription.SAMPLING_RATE
-                transcription_lag_s = max(0.0, time() - self.state.beg_loop - self.state.end_buffer)
+                transcription_lag_s = max(0.0, time() - self.beg_loop - self.end_buffer)
                 asr_processing_logs = f"internal_buffer={asr_internal_buffer_duration_s:.2f}s | lag={transcription_lag_s:.2f}s |"
-                if type(item) is Silence:
-                    asr_processing_logs += f" + Silence of = {item.duration:.2f}s"
+                stream_time_end_of_current_pcm = cumulative_pcm_duration_stream_time
+                new_tokens = []
+                current_audio_processed_upto = self.end_buffer
+
+                if isinstance(item, Silence):
+                    if item.is_starting:
+                        new_tokens, current_audio_processed_upto = await asyncio.to_thread(
+                            self.transcription.start_silence
+                        )
+                        asr_processing_logs += f" + Silence starting"
+                    if item.has_ended:
+                        asr_processing_logs += f" + Silence of = {item.duration:.2f}s"
+                        cumulative_pcm_duration_stream_time += item.duration
+                        current_audio_processed_upto = cumulative_pcm_duration_stream_time
+                        self.transcription.end_silence(item.duration, self.state.tokens[-1].end if self.state.tokens else 0)
                     if self.state.tokens:
                         asr_processing_logs += f" | last_end = {self.state.tokens[-1].end} |"
                     logger.info(asr_processing_logs)
-                    cumulative_pcm_duration_stream_time += item.duration
-                    self.transcription.insert_silence(item.duration, self.state.tokens[-1].end if self.state.tokens else 0)
-                    continue
+                    new_tokens = new_tokens or []
+                    current_audio_processed_upto = max(current_audio_processed_upto, stream_time_end_of_current_pcm)
                 elif isinstance(item, ChangeSpeaker):
                     self.transcription.new_speaker(item)
+                    continue
                 elif isinstance(item, np.ndarray):
                     pcm_array = item
-                
-                logger.info(asr_processing_logs)
-                
-                duration_this_chunk = len(pcm_array) / self.sample_rate
-                cumulative_pcm_duration_stream_time += duration_this_chunk
-                stream_time_end_of_current_pcm = cumulative_pcm_duration_stream_time
+                    logger.info(asr_processing_logs)
+                    cumulative_pcm_duration_stream_time += len(pcm_array) / self.sample_rate
+                    stream_time_end_of_current_pcm = cumulative_pcm_duration_stream_time
+                    self.transcription.insert_audio_chunk(pcm_array, stream_time_end_of_current_pcm)
+                    new_tokens, current_audio_processed_upto = await asyncio.to_thread(self.transcription.process_iter)
+                    new_tokens = new_tokens or []
 
-                self.transcription.insert_audio_chunk(pcm_array, stream_time_end_of_current_pcm)
-                new_tokens, current_audio_processed_upto = await asyncio.to_thread(self.transcription.process_iter)
-                
                 _buffer_transcript = self.transcription.get_buffer()
                 buffer_text = _buffer_transcript.text
 
@@ -255,7 +318,7 @@ class AudioProcessor:
                     if buffer_text.startswith(validated_text):
                         _buffer_transcript.text = buffer_text[len(validated_text):].lstrip()
 
-                candidate_end_times = [self.state.end_buffer]
+                candidate_end_times = [self.end_buffer]
 
                 if new_tokens:
                     candidate_end_times.append(new_tokens[-1].end)
@@ -268,14 +331,14 @@ class AudioProcessor:
                 async with self.lock:
                     self.state.tokens.extend(new_tokens)
                     self.state.buffer_transcription = _buffer_transcript
-                    self.state.end_buffer = max(candidate_end_times)
-                
+                    self.end_buffer = max(candidate_end_times)
+                    self.state_light.new_tokens = new_tokens
+                    self.state_light.new_tokens += 1
+                    self.state_light.new_tokens_buffer = _buffer_transcript
+
                 if self.translation_queue:
                     for token in new_tokens:
-                        await self.translation_queue.put(token)
-                        
-                self.transcription_queue.task_done()
-                
+                        await self.translation_queue.put(token)                
             except Exception as e:
                 logger.warning(f"Exception in transcription_processor: {e}")
                 logger.warning(f"Traceback: {traceback.format_exc()}")
@@ -292,68 +355,25 @@ class AudioProcessor:
         logger.info("Transcription processor task finished.")
 
 
-    async def diarization_processor(self, diarization_obj):
-        """Process audio chunks for speaker diarization."""
-        if self.diarization_before_transcription:
-            self.current_speaker = 0
-            await self.transcription_queue.put(ChangeSpeaker(speaker=self.current_speaker, start=0.0))
+    async def diarization_processor(self):
         while True:
             try:
-                item = await self.diarization_queue.get()
+                item = await get_all_from_queue(self.diarization_queue)
                 if item is SENTINEL:
-                    logger.debug("Diarization processor received sentinel. Finishing.")
-                    self.diarization_queue.task_done()
                     break
                 elif type(item) is Silence:
-                    diarization_obj.insert_silence(item.duration)
+                    if item.has_ended:
+                        self.diarization.insert_silence(item.duration)
                     continue
-                elif isinstance(item, np.ndarray):
-                    pcm_array = item
-                else:
-                    raise Exception('item should be pcm_array') 
-                
-                
-                
-                # Process diarization
-                await diarization_obj.diarize(pcm_array)
-                if self.diarization_before_transcription:
-                    segments = diarization_obj.get_segments()
-                    self.cumulative_pcm.append(pcm_array)
-                    if segments:
-                        last_segment = segments[-1]                    
-                        if last_segment.speaker != self.current_speaker:
-                            cut_sec = last_segment.start - self.last_end
-                            to_transcript, self.cumulative_pcm = cut_at(self.cumulative_pcm, cut_sec)
-                            await self.transcription_queue.put(to_transcript)
-                            
-                            self.current_speaker = last_segment.speaker
-                            await self.transcription_queue.put(ChangeSpeaker(speaker=self.current_speaker, start=last_segment.start))
-                            
-                            cut_sec = last_segment.end - last_segment.start
-                            to_transcript, self.cumulative_pcm = cut_at(self.cumulative_pcm, cut_sec)
-                            await self.transcription_queue.put(to_transcript)                            
-                            self.last_start = last_segment.start
-                            self.last_end = last_segment.end
-                        else:
-                            cut_sec = last_segment.end - self.last_end
-                            to_transcript, self.cumulative_pcm = cut_at(self.cumulative_pcm, cut_sec)
-                            await self.transcription_queue.put(to_transcript)
-                            self.last_end = last_segment.end
-                elif not self.diarization_before_transcription:           
-                    async with self.lock:
-                        self.state.tokens = diarization_obj.assign_speakers_to_tokens(
-                            self.state.tokens,
-                            use_punctuation_split=self.args.punctuation_split
-                        )
-                if len(self.state.tokens) > 0:
-                    self.state.end_attributed_speaker = max(self.state.tokens[-1].end, self.state.end_attributed_speaker)
-                self.diarization_queue.task_done()
+
+                self.diarization.insert_audio_chunk(item)
+                diarization_segments = await self.diarization.diarize()
+                self.state_light.new_diarization = diarization_segments
+                self.state_light.new_diarization_index += 1
                 
             except Exception as e:
                 logger.warning(f"Exception in diarization_processor: {e}")
                 logger.warning(f"Traceback: {traceback.format_exc()}")
-                if 'pcm_array' in locals() and pcm_array is not SENTINEL:
-                    self.diarization_queue.task_done()
         logger.info("Diarization processor task finished.")
 
     async def translation_processor(self):
@@ -362,51 +382,24 @@ class AudioProcessor:
         # in the future we want to have different languages for each speaker etc, so it will be more complex.
         while True:
             try:
-                item = await self.translation_queue.get() #block until at least 1 token
-                if item is SENTINEL:
+                tokens_to_process = await get_all_from_queue(self.translation_queue)
+                if tokens_to_process is SENTINEL:
                     logger.debug("Translation processor received sentinel. Finishing.")
                     self.translation_queue.task_done()
                     break
-                elif type(item) is Silence:
-                    self.translation.insert_silence(item.duration)
-                    continue
-                
-                # get all the available tokens for translation. The more words, the more precise
-                tokens_to_process = [item]
-                additional_tokens = await get_all_from_queue(self.translation_queue)
-                
-                sentinel_found = False
-                for additional_token in additional_tokens:
-                    if additional_token is SENTINEL:
-                        sentinel_found = True
-                        break
-                    elif type(additional_token) is Silence:
-                        self.translation.insert_silence(additional_token.duration)
-                        continue
-                    else:
-                        tokens_to_process.append(additional_token)                
+                elif type(tokens_to_process) is Silence:
+                    if tokens_to_process.has_ended:
+                        self.translation.insert_silence(tokens_to_process.duration)
+                    continue               
                 if tokens_to_process:
                     self.translation.insert_tokens(tokens_to_process)
                     translation_validated_segments, buffer_translation = await asyncio.to_thread(self.translation.process)
                     async with self.lock:
                         self.state.translation_validated_segments = translation_validated_segments
                         self.state.buffer_translation = buffer_translation
-                self.translation_queue.task_done()
-                for _ in additional_tokens:
-                    self.translation_queue.task_done()
-                
-                if sentinel_found:
-                    logger.debug("Translation processor received sentinel in batch. Finishing.")
-                    break
-                
             except Exception as e:
                 logger.warning(f"Exception in translation_processor: {e}")
                 logger.warning(f"Traceback: {traceback.format_exc()}")
-                if 'token' in locals() and item is not SENTINEL:
-                    self.translation_queue.task_done()
-                if 'additional_tokens' in locals():
-                    for _ in additional_tokens:
-                        self.translation_queue.task_done()
         logger.info("Translation processor task finished.")
 
     async def results_formatter(self):
@@ -420,7 +413,7 @@ class AudioProcessor:
                     continue
 
                 state = await self.get_current_state()
-                
+                self.tokens_alignment.compute_punctuations_segments()
                 lines, undiarized_text = format_output(
                     state,
                     self.silence,
@@ -507,7 +500,7 @@ class AudioProcessor:
             processing_tasks_for_watchdog.append(self.transcription_task)
             
         if self.diarization:
-            self.diarization_task = asyncio.create_task(self.diarization_processor(self.diarization))
+            self.diarization_task = asyncio.create_task(self.diarization_processor())
             self.all_tasks_for_cleanup.append(self.diarization_task)
             processing_tasks_for_watchdog.append(self.diarization_task)
         
@@ -586,8 +579,8 @@ class AudioProcessor:
     async def process_audio(self, message):
         """Process incoming audio data."""
 
-        if not self.state.beg_loop:
-            self.state.beg_loop = time()
+        if not self.beg_loop:
+            self.beg_loop = time()
 
         if not message:
             logger.info("Empty audio message received, initiating stop sequence.")
@@ -639,40 +632,30 @@ class AudioProcessor:
         pcm_array = self.convert_pcm_to_float(self.pcm_buffer[:aligned_chunk_size])
         self.pcm_buffer = self.pcm_buffer[aligned_chunk_size:]
 
-        res = None
-        end_of_audio = False
-        silence_buffer = None
+        num_samples = len(pcm_array)
+        chunk_sample_start = self.total_pcm_samples
+        chunk_sample_end = chunk_sample_start + num_samples
 
+        res = None
         if self.args.vac:
             res = self.vac(pcm_array)
 
         if res is not None:
-            if res.get("end", 0) > res.get("start", 0):
-                end_of_audio = True
-            elif self.silence: #end of silence
-                self.silence = False
-                silence_buffer = Silence(duration=time() - self.start_silence)
-
-        if silence_buffer:
-            if not self.diarization_before_transcription and self.transcription_queue:
-                await self.transcription_queue.put(silence_buffer)
-            if self.args.diarization and self.diarization_queue:
-                await self.diarization_queue.put(silence_buffer)
-            if self.translation_queue:
-                await self.translation_queue.put(silence_buffer)
+            silence_detected = res.get("end", 0) > res.get("start", 0)
+            if silence_detected and not self.silence:
+                pre_silence_chunk = self._slice_before_silence(
+                    pcm_array, chunk_sample_start, res.get("end")
+                )
+                if pre_silence_chunk is not None and pre_silence_chunk.size > 0:
+                    await self._enqueue_active_audio(pre_silence_chunk)
+                await self._begin_silence()
+            elif self.silence:
+                await self._end_silence()
 
         if not self.silence:
-            if not self.diarization_before_transcription and self.transcription_queue:
-                await self.transcription_queue.put(pcm_array.copy())
+            await self._enqueue_active_audio(pcm_array)
 
-            if self.args.diarization and self.diarization_queue:
-                await self.diarization_queue.put(pcm_array.copy())
-
-            self.silence_duration = 0.0
-
-            if end_of_audio:
-                self.silence = True
-                self.start_silence = time()
+        self.total_pcm_samples = chunk_sample_end
 
         if not self.args.transcription and not self.args.diarization:
             await asyncio.sleep(0.1)

whisperlivekit/core.py

@@ -52,8 +52,8 @@ class TranscriptionEngine:
 
         transcription_common_params = {
             "warmup_file": None,
-            "min_chunk_size": 0.5,
-            "model_size": "tiny",
+            "min_chunk_size": 0.1,
+            "model_size": "base",
             "model_cache_dir": None,
             "model_dir": None,
             "model_path": None,

whisperlivekit/diarization/diart_backend.py

@@ -26,7 +26,7 @@ class DiarizationObserver(Observer):
     """Observer that logs all data emitted by the diarization pipeline and stores speaker segments."""
     
     def __init__(self):
-        self.speaker_segments = []
+        self.diarization_segments = []
         self.processed_time = 0
         self.segment_lock = threading.Lock()
         self.global_time_offset = 0.0
@@ -48,7 +48,7 @@ class DiarizationObserver(Observer):
                 for speaker, label in annotation._labels.items():
                     for start, end in zip(label.segments_boundaries_[:-1], label.segments_boundaries_[1:]):
                         print(f"  {speaker}: {start:.2f}s-{end:.2f}s")
-                        self.speaker_segments.append(SpeakerSegment(
+                        self.diarization_segments.append(SpeakerSegment(
                             speaker=speaker,
                             start=start + self.global_time_offset,
                             end=end + self.global_time_offset
@@ -59,14 +59,14 @@ class DiarizationObserver(Observer):
     def get_segments(self) -> List[SpeakerSegment]:
         """Get a copy of the current speaker segments."""
         with self.segment_lock:
-            return self.speaker_segments.copy()
+            return self.diarization_segments.copy()
     
     def clear_old_segments(self, older_than: float = 30.0):
         """Clear segments older than the specified time."""
         with self.segment_lock:
             current_time = self.processed_time
-            self.speaker_segments = [
-                segment for segment in self.speaker_segments 
+            self.diarization_segments = [
+                segment for segment in self.diarization_segments 
                 if current_time - segment.end < older_than
             ]
     
@@ -178,7 +178,6 @@ class DiartDiarization:
         
         self.pipeline = SpeakerDiarization(config=config)        
         self.observer = DiarizationObserver()
-        self.lag_diart = None
         
         if use_microphone:
             self.source = MicrophoneAudioSource(block_duration=block_duration)
@@ -217,32 +216,6 @@ class DiartDiarization:
         if self.custom_source:
             self.custom_source.close()
 
-    def assign_speakers_to_tokens(self, tokens: list, use_punctuation_split: bool = False) -> float:
-        """
-        Assign speakers to tokens based on timing overlap with speaker segments.
-        Uses the segments collected by the observer.
-        
-        If use_punctuation_split is True, uses punctuation marks to refine speaker boundaries.
-        """
-        segments = self.observer.get_segments()
-        
-        # Debug logging
-        logger.debug(f"assign_speakers_to_tokens called with {len(tokens)} tokens")
-        logger.debug(f"Available segments: {len(segments)}")
-        for i, seg in enumerate(segments[:5]):  # Show first 5 segments
-            logger.debug(f"  Segment {i}: {seg.speaker} [{seg.start:.2f}-{seg.end:.2f}]")
-        
-        if not self.lag_diart and segments and tokens:
-            self.lag_diart = segments[0].start - tokens[0].start
-        
-        if not use_punctuation_split:
-            for token in tokens:
-                for segment in segments:
-                    if not (segment.end <= token.start + self.lag_diart or segment.start >= token.end + self.lag_diart):
-                        token.speaker = extract_number(segment.speaker) + 1
-        else:
-            tokens = add_speaker_to_tokens(segments, tokens)
-        return tokens
         
 def concatenate_speakers(segments):
     segments_concatenated = [{"speaker": 1, "begin": 0.0, "end": 0.0}]

whisperlivekit/diarization/sortformer_backend.py

@@ -94,11 +94,11 @@ class SortformerDiarizationOnline:
             model_name: Pre-trained model name (default: "nvidia/diar_streaming_sortformer_4spk-v2")
         """
         self.sample_rate = sample_rate
-        self.speaker_segments = []
+        self.diarization_segments = []
+        self.diar_segments = []
         self.buffer_audio = np.array([], dtype=np.float32)
         self.segment_lock = threading.Lock()
         self.global_time_offset = 0.0
-        self.processed_time = 0.0
         self.debug = False
                 
         self.diar_model = shared_model.diar_model
@@ -155,12 +155,10 @@ class SortformerDiarizationOnline:
         )
         self.streaming_state.fifo_lengths = torch.zeros((batch_size,), dtype=torch.long, device=device)
         self.streaming_state.mean_sil_emb = torch.zeros((batch_size, self.diar_model.sortformer_modules.fc_d_model), device=device)
-        self.streaming_state.n_sil_frames = torch.zeros((batch_size,), dtype=torch.long, device=device)
-        
-        # Initialize total predictions tensor
+        self.streaming_state.n_sil_frames = torch.zeros((batch_size,), dtype=torch.long, device=device)        
         self.total_preds = torch.zeros((batch_size, 0, self.diar_model.sortformer_modules.n_spk), device=device)
 
-    def insert_silence(self, silence_duration: float):
+    def insert_silence(self, silence_duration: Optional[float]):
         """
         Insert silence period by adjusting the global time offset.
         
@@ -171,248 +169,111 @@ class SortformerDiarizationOnline:
             self.global_time_offset += silence_duration
         logger.debug(f"Inserted silence of {silence_duration:.2f}s, new offset: {self.global_time_offset:.2f}s")
 
-    async def diarize(self, pcm_array: np.ndarray):
+    def insert_audio_chunk(self, pcm_array: np.ndarray):
+        if self.debug:
+            self.audio_buffer.append(pcm_array.copy())
+        self.buffer_audio = np.concatenate([self.buffer_audio, pcm_array.copy()])
+  
+
+    async def diarize(self):
         """
         Process audio data for diarization in streaming fashion.
         
         Args:
             pcm_array: Audio data as numpy array
         """
-        try:
-            if self.debug:
-                self.audio_buffer.append(pcm_array.copy())
 
-            threshold = int(self.chunk_duration_seconds * self.sample_rate)
+        threshold = int(self.chunk_duration_seconds * self.sample_rate)
+        
+        if not len(self.buffer_audio) >= threshold:
+            return []
+        
+        audio = self.buffer_audio[:threshold]
+        self.buffer_audio = self.buffer_audio[threshold:]
+        
+        device = self.diar_model.device
+        audio_signal_chunk = torch.tensor(audio, device=device).unsqueeze(0)
+        audio_signal_length_chunk = torch.tensor([audio_signal_chunk.shape[1]], device=device)
+        
+        processed_signal_chunk, processed_signal_length_chunk = self.audio2mel.get_features(
+            audio_signal_chunk, audio_signal_length_chunk
+        )
+        processed_signal_chunk = processed_signal_chunk.to(device)
+        processed_signal_length_chunk = processed_signal_length_chunk.to(device)
+        
+        if self._previous_chunk_features is not None:
+            to_add = self._previous_chunk_features[:, :, -99:].to(device)
+            total_features = torch.concat([to_add, processed_signal_chunk], dim=2).to(device)
+        else:
+            total_features = processed_signal_chunk.to(device)
+        
+        self._previous_chunk_features = processed_signal_chunk.to(device)
+        
+        chunk_feat_seq_t = torch.transpose(total_features, 1, 2).to(device)
+        
+        with torch.inference_mode():
+            left_offset = 8 if self._chunk_index > 0 else 0
+            right_offset = 8
             
-            self.buffer_audio = np.concatenate([self.buffer_audio, pcm_array.copy()])
-            if not len(self.buffer_audio) >= threshold:
-                return
-            
-            audio = self.buffer_audio[:threshold]
-            self.buffer_audio = self.buffer_audio[threshold:]
-            
-            device = self.diar_model.device
-            audio_signal_chunk = torch.tensor(audio, device=device).unsqueeze(0)
-            audio_signal_length_chunk = torch.tensor([audio_signal_chunk.shape[1]], device=device)
-            
-            processed_signal_chunk, processed_signal_length_chunk = self.audio2mel.get_features(
-                audio_signal_chunk, audio_signal_length_chunk
-            )
-            processed_signal_chunk = processed_signal_chunk.to(device)
-            processed_signal_length_chunk = processed_signal_length_chunk.to(device)
-            
-            if self._previous_chunk_features is not None:
-                to_add = self._previous_chunk_features[:, :, -99:].to(device)
-                total_features = torch.concat([to_add, processed_signal_chunk], dim=2).to(device)
-            else:
-                total_features = processed_signal_chunk.to(device)
-            
-            self._previous_chunk_features = processed_signal_chunk.to(device)
-            
-            chunk_feat_seq_t = torch.transpose(total_features, 1, 2).to(device)
-            
-            with torch.inference_mode():
-                left_offset = 8 if self._chunk_index > 0 else 0
-                right_offset = 8
-                
-                self.streaming_state, self.total_preds = self.diar_model.forward_streaming_step(
-                    processed_signal=chunk_feat_seq_t,
-                    processed_signal_length=torch.tensor([chunk_feat_seq_t.shape[1]]).to(device),
-                    streaming_state=self.streaming_state,
-                    total_preds=self.total_preds,
-                    left_offset=left_offset,
-                    right_offset=right_offset,
-                )
-                
-            # Convert predictions to speaker segments
-            self._process_predictions()
-            
-            self._chunk_index += 1
-            
-        except Exception as e:
-            logger.error(f"Error in diarize: {e}")
-            raise
-            
-        # TODO: Handle case when stream ends with partial buffer (accumulated_duration > 0 but < chunk_duration_seconds)
+            self.streaming_state, self.total_preds = self.diar_model.forward_streaming_step(
+                processed_signal=chunk_feat_seq_t,
+                processed_signal_length=torch.tensor([chunk_feat_seq_t.shape[1]]).to(device),
+                streaming_state=self.streaming_state,
+                total_preds=self.total_preds,
+                left_offset=left_offset,
+                right_offset=right_offset,
+            )                
+        new_segments = self._process_predictions()
+        
+        self._chunk_index += 1
+        return new_segments
 
     def _process_predictions(self):
         """Process model predictions and convert to speaker segments."""
-        try:
-            preds_np = self.total_preds[0].cpu().numpy()
-            active_speakers = np.argmax(preds_np, axis=1)
-            
-            if self._len_prediction is None:
-                self._len_prediction = len(active_speakers)
-            
-            # Get predictions for current chunk
-            frame_duration = self.chunk_duration_seconds / self._len_prediction
-            current_chunk_preds = active_speakers[-self._len_prediction:]
-            
-            with self.segment_lock:
-                # Process predictions into segments
-                base_time = self._chunk_index * self.chunk_duration_seconds + self.global_time_offset
-                
-                for idx, spk in enumerate(current_chunk_preds):
-                    start_time = base_time + idx * frame_duration
-                    end_time = base_time + (idx + 1) * frame_duration
-                    
-                    # Check if this continues the last segment or starts a new one
-                    if (self.speaker_segments and 
-                        self.speaker_segments[-1].speaker == spk and 
-                        abs(self.speaker_segments[-1].end - start_time) < frame_duration * 0.5):
-                        # Continue existing segment
-                        self.speaker_segments[-1].end = end_time
-                    else:
-                        
-                        # Create new segment
-                        self.speaker_segments.append(SpeakerSegment(
-                            speaker=spk,
-                            start=start_time,
-                            end=end_time
-                        ))
-                
-                # Update processed time
-                self.processed_time = max(self.processed_time, base_time + self.chunk_duration_seconds)
-                
-                logger.debug(f"Processed chunk {self._chunk_index}, total segments: {len(self.speaker_segments)}")
-                
-        except Exception as e:
-            logger.error(f"Error processing predictions: {e}")
-
-    def assign_speakers_to_tokens(self, tokens: list, use_punctuation_split: bool = False) -> list:
-        """
-        Assign speakers to tokens based on timing overlap with speaker segments.
+        preds_np = self.total_preds[0].cpu().numpy()
+        active_speakers = np.argmax(preds_np, axis=1)
         
-        Args:
-            tokens: List of tokens with timing information
-            use_punctuation_split: Whether to use punctuation for boundary refinement
-            
-        Returns:
-            List of tokens with speaker assignments
-            Last speaker_segment
-        """
+        if self._len_prediction is None:
+            self._len_prediction = len(active_speakers) #12
+        
+        frame_duration = self.chunk_duration_seconds / self._len_prediction
+        current_chunk_preds = active_speakers[-self._len_prediction:]
+        
+        new_segments = []
+
         with self.segment_lock:
-            segments = self.speaker_segments.copy()
-        
-        if not segments or not tokens:
-            logger.debug("No segments or tokens available for speaker assignment")
-            return tokens
-        
-        logger.debug(f"Assigning speakers to {len(tokens)} tokens using {len(segments)} segments")
-        use_punctuation_split = False
-        if not use_punctuation_split:
-            # Simple overlap-based assignment
-            for token in tokens:
-                token.speaker = -1  # Default to no speaker
-                for segment in segments:
-                    # Check for timing overlap
-                    if not (segment.end <= token.start or segment.start >= token.end):
-                        token.speaker = segment.speaker + 1  # Convert to 1-based indexing
-                        break
-        else:
-            # Use punctuation-aware assignment (similar to diart_backend)
-            tokens = self._add_speaker_to_tokens_with_punctuation(segments, tokens)
-        
-        return tokens
-
-    def _add_speaker_to_tokens_with_punctuation(self, segments: List[SpeakerSegment], tokens: list) -> list:
-        """
-        Assign speakers to tokens with punctuation-aware boundary adjustment.
-        
-        Args:
-            segments: List of speaker segments
-            tokens: List of tokens to assign speakers to
-            
-        Returns:
-            List of tokens with speaker assignments
-        """
-        punctuation_marks = {'.', '!', '?'}
-        punctuation_tokens = [token for token in tokens if token.text.strip() in punctuation_marks]
-        
-        # Convert segments to concatenated format
-        segments_concatenated = self._concatenate_speakers(segments)
-        
-        # Adjust segment boundaries based on punctuation
-        for ind, segment in enumerate(segments_concatenated):
-            for i, punctuation_token in enumerate(punctuation_tokens):
-                if punctuation_token.start > segment['end']:
-                    after_length = punctuation_token.start - segment['end']
-                    before_length = segment['end'] - punctuation_tokens[i - 1].end if i > 0 else float('inf')
-                    
-                    if before_length > after_length:
-                        segment['end'] = punctuation_token.start
-                        if i < len(punctuation_tokens) - 1 and ind + 1 < len(segments_concatenated):
-                            segments_concatenated[ind + 1]['begin'] = punctuation_token.start
-                    else:
-                        segment['end'] = punctuation_tokens[i - 1].end if i > 0 else segment['end']
-                        if i < len(punctuation_tokens) - 1 and ind - 1 >= 0:
-                            segments_concatenated[ind - 1]['begin'] = punctuation_tokens[i - 1].end
-                    break
-        
-        # Ensure non-overlapping tokens
-        last_end = 0.0
-        for token in tokens:
-            start = max(last_end + 0.01, token.start)
-            token.start = start
-            token.end = max(start, token.end)
-            last_end = token.end
-        
-        # Assign speakers based on adjusted segments
-        ind_last_speaker = 0
-        for segment in segments_concatenated:
-            for i, token in enumerate(tokens[ind_last_speaker:]):
-                if token.end <= segment['end']:
-                    token.speaker = segment['speaker']
-                    ind_last_speaker = i + 1
-                elif token.start > segment['end']:
-                    break
-        
-        return tokens
-
-    def _concatenate_speakers(self, segments: List[SpeakerSegment]) -> List[dict]:
-        """
-        Concatenate consecutive segments from the same speaker.
-        
-        Args:
-            segments: List of speaker segments
-            
-        Returns:
-            List of concatenated speaker segments
-        """
-        if not segments:
-            return []
-            
-        segments_concatenated = [{"speaker": segments[0].speaker + 1, "begin": segments[0].start, "end": segments[0].end}]
-        
-        for segment in segments[1:]:
-            speaker = segment.speaker + 1
-            if segments_concatenated[-1]['speaker'] != speaker:
-                segments_concatenated.append({"speaker": speaker, "begin": segment.start, "end": segment.end})
-            else:
-                segments_concatenated[-1]['end'] = segment.end
+            base_time = self._chunk_index * self.chunk_duration_seconds + self.global_time_offset
+            current_spk = current_chunk_preds[0]
+            start_time = round(base_time, 2)
+            for idx, spk in enumerate(current_chunk_preds):
+                current_time = round(base_time + idx * frame_duration, 2)
+                if spk != current_spk:
+                    new_segments.append(SpeakerSegment(
+                        speaker=current_spk,
+                        start=start_time,
+                        end=current_time
+                    ))
+                    start_time = current_time
+                    current_spk = spk
+            new_segments.append(
+                SpeakerSegment(
+                        speaker=current_spk,
+                        start=start_time,
+                        end=current_time
+            )
+            )
+        return new_segments
                 
-        return segments_concatenated
-
     def get_segments(self) -> List[SpeakerSegment]:
         """Get a copy of the current speaker segments."""
         with self.segment_lock:
-            return self.speaker_segments.copy()
-
-    def clear_old_segments(self, older_than: float = 30.0):
-        """Clear segments older than the specified time."""
-        with self.segment_lock:
-            current_time = self.processed_time
-            self.speaker_segments = [
-                segment for segment in self.speaker_segments 
-                if current_time - segment.end < older_than
-            ]
-            logger.debug(f"Cleared old segments, remaining: {len(self.speaker_segments)}")
+            return self.diarization_segments.copy()
 
     def close(self):
         """Close the diarization system and clean up resources."""
         logger.info("Closing SortformerDiarization")
         with self.segment_lock:
-            self.speaker_segments.clear()
+            self.diarization_segments.clear()
         
         if self.debug:
             concatenated_audio = np.concatenate(self.audio_buffer)
@@ -438,7 +299,7 @@ if __name__ == '__main__':
     
     async def main():
         """TEST ONLY."""
-        an4_audio = 'audio_test.mp3'
+        an4_audio = 'diarization_audio.wav'
         signal, sr = librosa.load(an4_audio, sr=16000)
         signal = signal[:16000*30]
 
@@ -450,13 +311,15 @@ if __name__ == '__main__':
         print("Speaker 0: 0:25 - 0:30")
         print("=" * 50)
         
-        diarization = SortformerDiarization(sample_rate=16000)        
+        diarization_backend = SortformerDiarization()
+        diarization = SortformerDiarizationOnline(shared_model = diarization_backend)     
         chunk_size = 1600
         
         for i in range(0, len(signal), chunk_size):
             chunk = signal[i:i+chunk_size]
-            await diarization.diarize(chunk)
+            new_segments = await diarization.diarize(chunk)
             print(f"Processed chunk {i // chunk_size + 1}")
+            print(new_segments)
         
         segments = diarization.get_segments()
         print("\nDiarization results:")

whisperlivekit/diarization/sortformer_backend_offline.py

@@ -1,205 +0,0 @@
-import numpy as np
-import torch
-import logging
-
-from nemo.collections.asr.models import SortformerEncLabelModel
-from nemo.collections.asr.modules import AudioToMelSpectrogramPreprocessor
-import librosa
-
-logger = logging.getLogger(__name__)
-
-def load_model():
-
-    diar_model = SortformerEncLabelModel.from_pretrained("nvidia/diar_streaming_sortformer_4spk-v2")
-    diar_model.eval()
-
-    if torch.cuda.is_available():
-        diar_model.to(torch.device("cuda"))
-
-    #we target 1 second lag for the moment. chunk_len could be reduced.
-    diar_model.sortformer_modules.chunk_len = 10
-    diar_model.sortformer_modules.subsampling_factor = 10 #8 would be better ideally
-
-    diar_model.sortformer_modules.chunk_right_context = 0 #no.
-    diar_model.sortformer_modules.chunk_left_context = 10 #big so it compensiate the problem with no padding later.
-
-    diar_model.sortformer_modules.spkcache_len = 188
-    diar_model.sortformer_modules.fifo_len = 188
-    diar_model.sortformer_modules.spkcache_update_period = 144
-    diar_model.sortformer_modules.log = False
-    diar_model.sortformer_modules._check_streaming_parameters()
-
-
-    audio2mel = AudioToMelSpectrogramPreprocessor(
-            window_size= 0.025, 
-            normalize="NA",
-            n_fft=512,
-            features=128,
-            pad_to=0) #pad_to 16 works better than 0. On test audio, we detect a third speaker for 1 second with pad_to=0. To solve that : increase left context to 10.
-
-    return diar_model, audio2mel
-
-diar_model, audio2mel = load_model()
-
-class StreamingSortformerState:
-    """
-    This class creates a class instance that will be used to store the state of the
-    streaming Sortformer model.
-
-    Attributes:
-        spkcache (torch.Tensor): Speaker cache to store embeddings from start
-        spkcache_lengths (torch.Tensor): Lengths of the speaker cache
-        spkcache_preds (torch.Tensor): The speaker predictions for the speaker cache parts
-        fifo (torch.Tensor): FIFO queue to save the embedding from the latest chunks
-        fifo_lengths (torch.Tensor): Lengths of the FIFO queue
-        fifo_preds (torch.Tensor): The speaker predictions for the FIFO queue parts
-        spk_perm (torch.Tensor): Speaker permutation information for the speaker cache
-        mean_sil_emb (torch.Tensor): Mean silence embedding
-        n_sil_frames (torch.Tensor): Number of silence frames
-    """
-
-    spkcache = None  # Speaker cache to store embeddings from start
-    spkcache_lengths = None  #
-    spkcache_preds = None  # speaker cache predictions
-    fifo = None  # to save the embedding from the latest chunks
-    fifo_lengths = None
-    fifo_preds = None
-    spk_perm = None
-    mean_sil_emb = None
-    n_sil_frames = None
-
-
-def init_streaming_state(self, batch_size: int = 1, async_streaming: bool = False, device: torch.device = None):
-    """
-    Initializes StreamingSortformerState with empty tensors or zero-valued tensors.
-
-    Args:
-        batch_size (int): Batch size for tensors in streaming state
-        async_streaming (bool): True for asynchronous update, False for synchronous update
-        device (torch.device): Device for tensors in streaming state
-
-    Returns:
-        streaming_state (SortformerStreamingState): initialized streaming state
-    """
-    streaming_state = StreamingSortformerState()
-    if async_streaming:
-        streaming_state.spkcache = torch.zeros((batch_size, self.spkcache_len, self.fc_d_model), device=device)
-        streaming_state.spkcache_preds = torch.zeros((batch_size, self.spkcache_len, self.n_spk), device=device)
-        streaming_state.spkcache_lengths = torch.zeros((batch_size,), dtype=torch.long, device=device)
-        streaming_state.fifo = torch.zeros((batch_size, self.fifo_len, self.fc_d_model), device=device)
-        streaming_state.fifo_lengths = torch.zeros((batch_size,), dtype=torch.long, device=device)
-    else:
-        streaming_state.spkcache = torch.zeros((batch_size, 0, self.fc_d_model), device=device)
-        streaming_state.fifo = torch.zeros((batch_size, 0, self.fc_d_model), device=device)
-    streaming_state.mean_sil_emb = torch.zeros((batch_size, self.fc_d_model), device=device)
-    streaming_state.n_sil_frames = torch.zeros((batch_size,), dtype=torch.long, device=device)
-    return streaming_state
-
-
-def process_diarization(chunks):
-    """ 
-    what it does:
-    1. Preprocessing: Applies dithering and pre-emphasis (high-pass filter) if enabled
-    2. STFT: Computes the Short-Time Fourier Transform using:
-        - the window of window_size=0.025 --> size of a window : 400 samples
-        - the hop parameter : n_window_stride = 0.01 -> every 160 samples, a new window
-    3. Magnitude Calculation: Converts complex STFT output to magnitude spectrogram
-    4. Mel Conversion: Applies Mel filterbanks (128 filters in this case) to get Mel spectrogram
-    5. Logarithm: Takes the log of the Mel spectrogram (if `log=True`)
-    6. Normalization: Skips normalization since `normalize="NA"`
-    7. Padding: Pads the time dimension to a multiple of `pad_to` (default 16)    
-    """
-    previous_chunk = None
-    l_chunk_feat_seq_t = []
-    for chunk in chunks:
-        audio_signal_chunk = torch.tensor(chunk).unsqueeze(0).to(diar_model.device)
-        audio_signal_length_chunk = torch.tensor([audio_signal_chunk.shape[1]]).to(diar_model.device)
-        processed_signal_chunk, processed_signal_length_chunk = audio2mel.get_features(audio_signal_chunk, audio_signal_length_chunk)
-        if previous_chunk is not None:
-            to_add = previous_chunk[:, :, -99:]
-            total = torch.concat([to_add, processed_signal_chunk], dim=2)
-        else:
-            total = processed_signal_chunk
-        previous_chunk = processed_signal_chunk
-        l_chunk_feat_seq_t.append(torch.transpose(total, 1, 2))
-
-    batch_size = 1
-    streaming_state = init_streaming_state(diar_model.sortformer_modules,
-        batch_size = batch_size,
-        async_streaming = True,
-        device = diar_model.device
-    )
-    total_preds = torch.zeros((batch_size, 0, diar_model.sortformer_modules.n_spk), device=diar_model.device)
-
-    chunk_duration_seconds = diar_model.sortformer_modules.chunk_len * diar_model.sortformer_modules.subsampling_factor * diar_model.preprocessor._cfg.window_stride
-
-    l_speakers = [
-        {'start_time': 0,
-        'end_time': 0,
-        'speaker': 0
-        }
-    ]
-    len_prediction = None
-    left_offset = 0
-    right_offset = 8
-    for i, chunk_feat_seq_t in enumerate(l_chunk_feat_seq_t):
-        with torch.inference_mode():
-                streaming_state, total_preds = diar_model.forward_streaming_step(
-                    processed_signal=chunk_feat_seq_t,
-                    processed_signal_length=torch.tensor([chunk_feat_seq_t.shape[1]]),
-                    streaming_state=streaming_state,
-                    total_preds=total_preds,
-                    left_offset=left_offset,
-                    right_offset=right_offset,
-                )
-                left_offset = 8
-                preds_np = total_preds[0].cpu().numpy()
-                active_speakers = np.argmax(preds_np, axis=1)
-                if len_prediction is None:
-                    len_prediction = len(active_speakers) # we want to get the len of 1 prediction
-                frame_duration = chunk_duration_seconds / len_prediction
-                active_speakers = active_speakers[-len_prediction:]
-                for idx, spk in enumerate(active_speakers):
-                    if spk != l_speakers[-1]['speaker']:
-                        l_speakers.append(
-                            {'start_time': (i * chunk_duration_seconds + idx * frame_duration),
-                            'end_time': (i * chunk_duration_seconds + (idx + 1) * frame_duration),
-                            'speaker': spk
-                        })                    
-                    else:
-                        l_speakers[-1]['end_time'] = i * chunk_duration_seconds + (idx + 1) * frame_duration
-                    
-        
-        """
-        Should print
-        [{'start_time': 0, 'end_time': 8.72, 'speaker': 0}, 
-        {'start_time': 8.72, 'end_time': 18.88, 'speaker': 1},
-        {'start_time': 18.88, 'end_time': 24.96, 'speaker': 2},
-        {'start_time': 24.96, 'end_time': 31.68, 'speaker': 0}]
-        """
-    for speaker in l_speakers:
-        print(f"Speaker {speaker['speaker']}: {speaker['start_time']:.2f}s - {speaker['end_time']:.2f}s")    
-    
-
-if __name__ == '__main__':
-    
-    an4_audio = 'audio_test.mp3'
-    signal, sr = librosa.load(an4_audio, sr=16000)
-    signal = signal[:16000*30]
-    # signal = signal[:-(len(signal)%16000)]
-
-    print("\n" + "=" * 50)
-    print("Expected ground truth:")
-    print("Speaker 0: 0:00 - 0:09")
-    print("Speaker 1: 0:09 - 0:19") 
-    print("Speaker 2: 0:19 - 0:25")
-    print("Speaker 0: 0:25 - 0:30")
-    print("=" * 50)
-
-    chunk_size = 16000  # 1 second
-    chunks = []
-    for i in range(0, len(signal), chunk_size):
-        chunk = signal[i:i+chunk_size]
-        chunks.append(chunk)
-    
-    process_diarization(chunks)

whisperlivekit/local_agreement/online_asr.py

@@ -151,21 +151,32 @@ class OnlineASRProcessor:
         """Append an audio chunk (a numpy array) to the current audio buffer."""
         self.audio_buffer = np.append(self.audio_buffer, audio)
 
-    def insert_silence(self, silence_duration, offset):
-        """
-        If silences are > 5s, we do a complete context clear. Otherwise, we just insert a small silence and shift the last_attend_frame
-        """
-        # if self.transcript_buffer.buffer:
-        #     self.committed.extend(self.transcript_buffer.buffer)
-        #     self.transcript_buffer.buffer = []
-            
-        if True: #silence_duration < 3: #we want the last audio to be treated to not have a gap. could also be handled in the future in ends_with_silence.
-            gap_silence = np.zeros(int(16000 * silence_duration), dtype=np.int16)
-            self.insert_audio_chunk(gap_silence)
+    def start_silence(self):
+        if self.audio_buffer.size == 0:
+            return [], self.get_audio_buffer_end_time()
+        return self.process_iter()
+
+    def end_silence(self, silence_duration: Optional[float], offset: float):
+        if not silence_duration or silence_duration <= 0:
+            return
+
+        long_silence = silence_duration >= 5
+        if not long_silence:
+            gap_samples = int(self.SAMPLING_RATE * silence_duration)
+            if gap_samples > 0:
+                gap_silence = np.zeros(gap_samples, dtype=np.float32)
+                self.insert_audio_chunk(gap_silence)
         else:
             self.init(offset=silence_duration + offset)
+
         self.global_time_offset += silence_duration
 
+    def insert_silence(self, silence_duration, offset):
+        """
+        Backwards compatibility shim for legacy callers that still use insert_silence.
+        """
+        self.end_silence(silence_duration, offset)
+
     def prompt(self) -> Tuple[str, str]:
         """
         Returns a tuple: (prompt, context), where:

whisperlivekit/parse_args.py

@@ -81,14 +81,14 @@ def parse_args():
     parser.add_argument(
         "--min-chunk-size",
         type=float,
-        default=0.5,
+        default=0.1,
         help="Minimum audio chunk size in seconds. It waits up to this time to do processing. If the processing takes shorter time, it waits, otherwise it processes the whole segment that was received by this time.",
     )
     
     parser.add_argument(
         "--model",
         type=str,
-        default="small",
+        default="base",
         dest='model_size',
         help="Name size of the Whisper model to use (default: tiny). Suggested values: tiny.en,tiny,base.en,base,small.en,small,medium.en,medium,large-v1,large-v2,large-v3,large,large-v3-turbo. The model is automatically downloaded from the model hub if not present in model cache dir.",
     )

whisperlivekit/remove_silences.py

@@ -36,7 +36,6 @@ def blank_to_silence(tokens):
                             start=token.start,
                             end=token.end,
                             speaker=-2,
-                            probability=0.95
                         )
                 else:
                     if silence_token: #there was silence but no more
@@ -70,7 +69,6 @@ def no_token_to_silence(tokens):
                     start=last_end,
                     end=token.start,
                     speaker=-2,
-                    probability=0.95
                     )
                 new_tokens.append(silence_token)
         
@@ -90,7 +88,6 @@ def ends_with_silence(tokens, beg_loop, vac_detected_silence):
                     start=tokens[-1].end,
                     end=current_time,
                     speaker=-2,
-                    probability=0.95
                 )
             )
     return tokens

whisperlivekit/result_diarization.md

@@ -0,0 +1,60 @@
+########### WHAT IS PRODUCED: ########### 
+
+SPEAKER 1 0:00:04 - 0:00:06
+Transcription technology has improved so much in the past
+
+SPEAKER 1 0:00:07 - 0:00:12
+years. Have you noticed how accurate real-time speech detects is now?
+
+SPEAKER 2 0:00:12 - 0:00:12
+Absolutely
+
+SPEAKER 1 0:00:13 - 0:00:13
+.
+
+SPEAKER 2 0:00:14 - 0:00:14
+I
+
+SPEAKER 1 0:00:14 - 0:00:17
+use it all the time for taking notes during meetings.
+
+SPEAKER 2 0:00:17 - 0:00:17
+It
+
+SPEAKER 1 0:00:17 - 0:00:22
+'s amazing how it can recognize different speakers, and even add punctuation.
+
+SPEAKER 2 0:00:22 - 0:00:22
+Yeah
+
+SPEAKER 1 0:00:23 - 0:00:26
+, but sometimes noise can still cause mistakes.
+
+SPEAKER 3 0:00:26 - 0:00:27
+Does
+
+SPEAKER 1 0:00:27 - 0:00:28
+this system handle that
+
+SPEAKER 1 0:00:29 - 0:00:29
+?
+
+SPEAKER 3 0:00:29 - 0:00:29
+It
+
+SPEAKER 1 0:00:29 - 0:00:33
+does a pretty good job filtering noise, especially with models that use voice activity
+
+########### WHAT SHOULD BE PRODUCED: ########### 
+
+SPEAKER 1 0:00:04 - 0:00:12
+Transcription technology has improved so much in the past years. Have you noticed how accurate real-time speech detects is now?
+
+SPEAKER 2 0:00:12 - 0:00:22
+Absolutely. I use it all the time for taking notes during meetings. It's amazing how it can recognize different speakers, and even add punctuation.
+
+SPEAKER 3 0:00:22 - 0:00:28
+Yeah, but sometimes noise can still cause mistakes. Does this system handle that well?
+
+SPEAKER 1 0:00:29 - 0:00:29
+It does a pretty good job filtering noise, especially with models that use voice activity

whisperlivekit/results_formater.py

@@ -1,7 +1,9 @@
 
 import logging
+import re
 from whisperlivekit.remove_silences import handle_silences
-from whisperlivekit.timed_objects import Line, format_time
+from whisperlivekit.timed_objects import Line, format_time, SpeakerSegment
+from typing import List
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
@@ -9,22 +11,16 @@ logger.setLevel(logging.DEBUG)
 CHECK_AROUND = 4
 DEBUG = False
 
-
-def is_punctuation(token):
-    if token.is_punctuation():
-        return True
-    return False
-
 def next_punctuation_change(i, tokens):
     for ind in range(i+1, min(len(tokens), i+CHECK_AROUND+1)):
-        if is_punctuation(tokens[ind]):
+        if tokens[ind].is_punctuation():
             return ind        
     return None
 
 def next_speaker_change(i, tokens, speaker):
     for ind in range(i-1, max(0, i-CHECK_AROUND)-1, -1):
         token = tokens[ind]
-        if is_punctuation(token):
+        if token.is_punctuation():
             break
         if token.speaker != speaker:
             return ind, token.speaker
@@ -50,7 +46,106 @@ def append_token_to_last_line(lines, sep, token):
             lines[-1].end = token.end
         if not lines[-1].detected_language and token.detected_language:
             lines[-1].detected_language = token.detected_language
-            
+
+def extract_number(s) -> int:
+    """Extract number from speaker string (for diart compatibility)."""
+    if isinstance(s, int):
+        return s
+    m = re.search(r'\d+', str(s))
+    return int(m.group()) if m else 0
+
+def concatenate_speakers(segments: List[SpeakerSegment]) -> List[dict]:
+    """Concatenate consecutive segments from the same speaker."""
+    if not segments:
+        return []
+    
+    # Get speaker number from first segment
+    first_speaker = extract_number(segments[0].speaker)
+    segments_concatenated = [{"speaker": first_speaker + 1, "begin": segments[0].start, "end": segments[0].end}]
+    
+    for segment in segments[1:]:
+        speaker = extract_number(segment.speaker) + 1
+        if segments_concatenated[-1]['speaker'] != speaker:
+            segments_concatenated.append({"speaker": speaker, "begin": segment.start, "end": segment.end})
+        else:
+            segments_concatenated[-1]['end'] = segment.end
+    
+    return segments_concatenated
+
+def add_speaker_to_tokens_with_punctuation(segments: List[SpeakerSegment], tokens: list) -> list:
+    """Assign speakers to tokens with punctuation-aware boundary adjustment."""
+    punctuation_marks = {'.', '!', '?'}
+    punctuation_tokens = [token for token in tokens if token.text.strip() in punctuation_marks]
+    segments_concatenated = concatenate_speakers(segments)
+    
+    for ind, segment in enumerate(segments_concatenated):
+        for i, punctuation_token in enumerate(punctuation_tokens):
+            if punctuation_token.start > segment['end']:
+                after_length = punctuation_token.start - segment['end']
+                before_length = segment['end'] - punctuation_tokens[i - 1].end if i > 0 else float('inf')
+                if before_length > after_length:
+                    segment['end'] = punctuation_token.start
+                    if i < len(punctuation_tokens) - 1 and ind + 1 < len(segments_concatenated):
+                        segments_concatenated[ind + 1]['begin'] = punctuation_token.start
+                else:
+                    segment['end'] = punctuation_tokens[i - 1].end if i > 0 else segment['end']
+                    if i < len(punctuation_tokens) - 1 and ind - 1 >= 0:
+                        segments_concatenated[ind - 1]['begin'] = punctuation_tokens[i - 1].end
+                break
+    
+    # Ensure non-overlapping tokens
+    last_end = 0.0
+    for token in tokens:
+        start = max(last_end + 0.01, token.start)
+        token.start = start
+        token.end = max(start, token.end)
+        last_end = token.end
+    
+    # Assign speakers based on adjusted segments
+    ind_last_speaker = 0
+    for segment in segments_concatenated:
+        for i, token in enumerate(tokens[ind_last_speaker:]):
+            if token.end <= segment['end']:
+                token.speaker = segment['speaker']
+                ind_last_speaker = i + 1
+            elif token.start > segment['end']:
+                break
+    
+    return tokens
+
+def assign_speakers_to_tokens(tokens: list, segments: List[SpeakerSegment], use_punctuation_split: bool = False) -> list:
+    """
+    Assign speakers to tokens based on timing overlap with speaker segments.
+    
+    Args:
+        tokens: List of tokens with timing information
+        segments: List of speaker segments
+        use_punctuation_split: Whether to use punctuation for boundary refinement
+        
+    Returns:
+        List of tokens with speaker assignments
+    """
+    if not segments or not tokens:
+        logger.debug("No segments or tokens available for speaker assignment")
+        return tokens
+    
+    logger.debug(f"Assigning speakers to {len(tokens)} tokens using {len(segments)} segments")
+    
+    if not use_punctuation_split:
+        # Simple overlap-based assignment
+        for token in tokens:
+            token.speaker = -1  # Default to no speaker
+            for segment in segments:
+                # Check for timing overlap
+                if not (segment.end <= token.start or segment.start >= token.end):
+                    speaker_num = extract_number(segment.speaker)
+                    token.speaker = speaker_num + 1  # Convert to 1-based indexing
+                    break
+    else:
+        # Use punctuation-aware assignment
+        tokens = add_speaker_to_tokens_with_punctuation(segments, tokens)
+    
+    return tokens
 
 def format_output(state, silence, args, sep):
     diarization = args.diarization
@@ -58,63 +153,72 @@ def format_output(state, silence, args, sep):
     tokens = state.tokens
     translation_validated_segments = state.translation_validated_segments # Here we will attribute the speakers only based on the timestamps of the segments
     last_validated_token = state.last_validated_token
-    
-    previous_speaker = 1
+     
+    last_speaker = abs(state.last_speaker)
     undiarized_text = []
     tokens = handle_silences(tokens, state.beg_loop, silence)
+    
+    # Assign speakers to tokens based on segments stored in state
+    if False and diarization and state.diarization_segments:
+        use_punctuation_split = args.punctuation_split if hasattr(args, 'punctuation_split') else False
+        tokens = assign_speakers_to_tokens(tokens, state.diarization_segments, use_punctuation_split=use_punctuation_split)
     for i in range(last_validated_token, len(tokens)):
         token = tokens[i]
         speaker = int(token.speaker)
         token.corrected_speaker = speaker
-        if not diarization:
+        if True or not diarization:
             if speaker == -1: #Speaker -1 means no attributed by diarization. In the frontend, it should appear under 'Speaker 1'
                 token.corrected_speaker = 1
                 token.validated_speaker = True
         else:
-                if is_punctuation(token):
-                    state.last_punctuation_index = i  
-                
-                if state.last_punctuation_index == i-1:
-                    if token.speaker != previous_speaker:
+            if token.speaker == -1:
+                undiarized_text.append(token.text)
+            elif token.is_punctuation():
+                state.last_punctuation_index = i
+                token.corrected_speaker = last_speaker
+                token.validated_speaker = True
+            elif state.last_punctuation_index == i-1:
+                if token.speaker != last_speaker:
+                    token.corrected_speaker = token.speaker
+                    token.validated_speaker = True
+                    # perfect, diarization perfectly aligned
+                else:
+                    speaker_change_pos, new_speaker = next_speaker_change(i, tokens, speaker)
+                    if speaker_change_pos:
+                        # Corrects delay:
+                        # That was the idea. <Okay> haha |SPLIT SPEAKER| that's a good one 
+                        # should become:
+                        # That was the idea. |SPLIT SPEAKER| <Okay> haha that's a good one 
+                        token.corrected_speaker = new_speaker
                         token.validated_speaker = True
-                        # perfect, diarization perfectly aligned
-                        last_punctuation = None
-                    else:
-                        speaker_change_pos, new_speaker = next_speaker_change(i, tokens, speaker)
-                        if speaker_change_pos:
-                            # Corrects delay:
-                            # That was the idea. <Okay> haha |SPLIT SPEAKER| that's a good one 
-                            # should become:
-                            # That was the idea. |SPLIT SPEAKER| <Okay> haha that's a good one 
-                            token.corrected_speaker = new_speaker
-                            token.validated_speaker = True
-                elif speaker != previous_speaker:
-                    if not (speaker == -2 or previous_speaker == -2):
-                        if next_punctuation_change(i, tokens):
-                            # Corrects advance:
-                            # Are you |SPLIT SPEAKER| <okay>? yeah, sure. Absolutely 
-                            # should become:
-                            # Are you <okay>? |SPLIT SPEAKER| yeah, sure. Absolutely 
-                            token.corrected_speaker = previous_speaker
-                            token.validated_speaker = True
-                        else: #Problematic, except if the language has no punctuation. We append to previous line, except if disable_punctuation_split is set to True.
-                            if not disable_punctuation_split:
-                                token.corrected_speaker = previous_speaker
-                                token.validated_speaker = False
+            elif speaker != last_speaker:
+                if not (speaker == -2 or last_speaker == -2):
+                    if next_punctuation_change(i, tokens):
+                        # Corrects advance:
+                        # Are you |SPLIT SPEAKER| <okay>? yeah, sure. Absolutely 
+                        # should become:
+                        # Are you <okay>? |SPLIT SPEAKER| yeah, sure. Absolutely 
+                        token.corrected_speaker = last_speaker
+                        token.validated_speaker = True
+                    else: #Problematic, except if the language has no punctuation. We append to previous line, except if disable_punctuation_split is set to True.
+                        if not disable_punctuation_split:
+                            token.corrected_speaker = last_speaker
+                            token.validated_speaker = False
         if token.validated_speaker:
             state.last_validated_token = i
-        previous_speaker = token.corrected_speaker  
+        state.last_speaker = token.corrected_speaker  
 
-    previous_speaker = 1
+    last_speaker = 1
     
     lines = []
     for token in tokens:
-        if int(token.corrected_speaker) != int(previous_speaker):
-            lines.append(new_line(token))
-        else:
-            append_token_to_last_line(lines, sep, token)
+        if token.corrected_speaker != -1:
+            if int(token.corrected_speaker) != int(last_speaker):
+                lines.append(new_line(token))
+            else:
+                append_token_to_last_line(lines, sep, token)
 
-        previous_speaker = token.corrected_speaker        
+        last_speaker = token.corrected_speaker        
 
     if lines:
         unassigned_translated_segments = []

whisperlivekit/simul_whisper/backend.py

@@ -18,7 +18,7 @@ from whisperlivekit.backend_support import (
 
 import torch
 from whisperlivekit.simul_whisper.config import AlignAttConfig
-from whisperlivekit.simul_whisper.simul_whisper import PaddedAlignAttWhisper
+from whisperlivekit.simul_whisper.simul_whisper import AlignAtt
 
 logger = logging.getLogger(__name__)
 
@@ -34,6 +34,8 @@ if HAS_FASTER_WHISPER:
 else:
     WhisperModel = None
 
+MIN_DURATION_REAL_SILENCE = 5
+
 class SimulStreamingOnlineProcessor:
     SAMPLING_RATE = 16000
 
@@ -56,23 +58,29 @@ class SimulStreamingOnlineProcessor:
 
     def load_new_backend(self):
         model = self.asr.get_new_model_instance()
-        self.model = PaddedAlignAttWhisper(
+        self.model = AlignAtt(
             cfg=self.asr.cfg,
             loaded_model=model,
             mlx_encoder=self.asr.mlx_encoder,
             fw_encoder=self.asr.fw_encoder,
             )
 
-    def insert_silence(self, silence_duration, offset):
+    def start_silence(self):
+        tokens, processed_upto = self.process_iter(is_last=True)
+        return tokens, processed_upto
+
+    def end_silence(self, silence_duration, offset):
         """
-        If silences are > 5s, we do a complete context clear. Otherwise, we just insert a small silence and shift the last_attend_frame
+        If silences are > MIN_DURATION_REAL_SILENCE, we do a complete context clear. Otherwise, we just insert a small silence and shift the last_attend_frame
         """
-        if silence_duration < 5:
-            gap_silence = torch.zeros(int(16000*silence_duration))
-            self.model.insert_audio(gap_silence)
-            # self.global_time_offset += silence_duration
-        else:
-            self.process_iter(is_last=True) #we want to totally process what remains in the buffer.
+        self.end += silence_duration
+        long_silence = silence_duration >= MIN_DURATION_REAL_SILENCE
+        if not long_silence:
+            gap_len = int(16000 * silence_duration)
+            if gap_len > 0:
+                gap_silence = torch.zeros(gap_len)
+                self.model.insert_audio(gap_silence)
+        if long_silence:
             self.model.refresh_segment(complete=True)
             self.model.global_time_offset = silence_duration + offset
 
@@ -300,7 +308,7 @@ class SimulStreamingASR():
         if warmup_audio is not None:
             warmup_audio = torch.from_numpy(warmup_audio).float()
             if self.fast_encoder:                
-                temp_model = PaddedAlignAttWhisper(
+                temp_model = AlignAtt(
                     cfg=self.cfg,
                     loaded_model=whisper_model,
                     mlx_encoder=self.mlx_encoder,

whisperlivekit/simul_whisper/generation_progress.py

@@ -1,43 +0,0 @@
-class Tokens:
-    def __init__(self, tokens):
-        self.tokens = tokens
-
-#    def clone(self):
-#        return Tokens(self.tokens.clone())
-
-    def __str__(self):
-        return str(self.tokens.tolist())
-
-    def __repr__(self):
-        return self.__str__()
-
-class BeamTokens(Tokens):
-    def __init__(self, tokens, beam_size):
-        self.tokens = tokens
-        self.beam_size = beam_size
-
-    def clone(self):
-        return BeamTokens(self.tokens.clone())
-
-    def __str__(self):
-        return f"BeamTokens({self.tokens.tolist()}, beam_size={self.beam_size})"
-
-    def __repr__(self):
-        return self.__str__()
-
-    def as_text(self, tokenizer):
-        return tokenizer.decode(self.tokens)
-
-class Logits(Tokens):
-    def __init__(self, logits):
-        super().__init__(logits)
-
-#    def clone(self):
-#        return Logits(self.tokens.clone(), self.beam_size)
-
-    def __str__(self):
-#        return "abc"
-        return f"Logits({self.tokens.shape})"
-
-    def __repr__(self):
-        return self.__str__()

whisperlivekit/simul_whisper/simul_whisper.py

@@ -1,17 +1,16 @@
-# This code was originally in simul_whisper/transcriber/simul_whisper.py . It is adapted a lot for SimulStreaming.
-
 import os
 import logging
 
 import torch
 import torch.nn.functional as F
+import numpy as np
 
-from whisperlivekit.whisper import load_model, DecodingOptions, tokenizer
+from whisperlivekit.whisper import DecodingOptions, tokenizer
 from .config import AlignAttConfig
 from whisperlivekit.timed_objects import ASRToken
 from whisperlivekit.whisper.audio import log_mel_spectrogram, TOKENS_PER_SECOND, pad_or_trim, N_SAMPLES, N_FRAMES
 from whisperlivekit.whisper.timing import median_filter
-from whisperlivekit.whisper.decoding import GreedyDecoder, BeamSearchDecoder, SuppressTokens, detect_language
+from whisperlivekit.whisper.decoding import GreedyDecoder, BeamSearchDecoder, SuppressTokens
 from .beam import BeamPyTorchInference
 from .eow_detection import fire_at_boundary, load_cif
 import os
@@ -22,28 +21,37 @@ from whisperlivekit.backend_support import (
     faster_backend_available,
 )
 
-import numpy as np
 from ..timed_objects import PUNCTUATION_MARKS
-from .generation_progress import *
 
 DEC_PAD = 50257
 logger = logging.getLogger(__name__)
 
-
-HAS_MLX_WHISPER = False
-HAS_FASTER_WHISPER = False
-
 if mlx_backend_available():
     from mlx_whisper.audio import log_mel_spectrogram as mlx_log_mel_spectrogram
     from mlx_whisper.transcribe import pad_or_trim as mlx_pad_or_trim
-    HAS_MLX_WHISPER = True
 
 if faster_backend_available():
     from faster_whisper.audio import pad_or_trim as fw_pad_or_trim
     from faster_whisper.feature_extractor import FeatureExtractor
-    HAS_FASTER_WHISPER = True
 
-class PaddedAlignAttWhisper:
+USE_MLCORE = False
+
+
+def load_coreml_encoder():
+    try:
+        from coremltools.models import MLModel
+    except ImportError:
+        logger.warning("coremltools is not installed")
+        return None
+    COREML_ENCODER_PATH = os.environ.get("MLCORE_ENCODER_PATH", "whisperlivekit/whisper/whisper_encoder.mlpackage")
+    _coreml_encoder = MLModel(COREML_ENCODER_PATH)
+    spec = _coreml_encoder.get_spec()
+    _coreml_input_name = spec.description.input[0].name if spec.description.input else "mel"
+    _coreml_output_name = spec.description.output[0].name if spec.description.output else None
+    return _coreml_encoder, _coreml_input_name, _coreml_output_name
+
+
+class AlignAtt:
     def __init__(
             self, 
             cfg: AlignAttConfig,
@@ -55,9 +63,13 @@ class PaddedAlignAttWhisper:
         
         self.model = loaded_model
         self.mlx_encoder = mlx_encoder
-        self.fw_encoder = fw_encoder
+        self.fw_encoder = fw_encoder            
         if fw_encoder:
             self.fw_feature_extractor = FeatureExtractor(feature_size=self.model.dims.n_mels)
+        self.coreml_encoder_tuple = None
+        if USE_MLCORE:
+            self.coreml_encoder_tuple = load_coreml_encoder()
+        self.use_mlcore = self.coreml_encoder_tuple is not None
         
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
         
@@ -393,15 +405,28 @@ class PaddedAlignAttWhisper:
         else:
             input_segments = self.segments[0]
 
-        # if self.cfg.language == "auto" and self.reset_tokenizer_to_auto_next_call:
-        #     logger.debug("Resetting tokenizer to auto for new sentence.")
-        #     self.create_tokenizer(None)
-        #     self.detected_language = None
-        #     self.init_tokens()
-        #     self.reset_tokenizer_to_auto_next_call = False
-
-        # NEW : we can use a different encoder, before using standart whisper for cross attention with the hooks on the decoder
         beg_encode = time()
+        if self.use_mlcore:
+            coreml_encoder, coreml_input_name, coreml_output_name = self.coreml_encoder_tuple
+            mel_padded = log_mel_spectrogram(
+                input_segments,
+                n_mels=self.model.dims.n_mels,
+                padding=N_SAMPLES,
+                device="cpu",
+            ).unsqueeze(0)
+            mel = pad_or_trim(mel_padded, N_FRAMES)
+            content_mel_len = int((mel_padded.shape[2] - mel.shape[2]) / 2)
+            mel_np = np.ascontiguousarray(mel.numpy())
+            ml_inputs = {coreml_input_name or "mel": mel_np}
+            coreml_outputs = coreml_encoder.predict(ml_inputs)
+            if coreml_output_name and coreml_output_name in coreml_outputs:
+                encoder_feature_np = coreml_outputs[coreml_output_name]
+            else:
+                encoder_feature_np = next(iter(coreml_outputs.values()))
+            encoder_feature = torch.as_tensor(
+                np.array(encoder_feature_np),
+                device=self.device,
+            )
         if self.mlx_encoder:
             mlx_mel_padded = mlx_log_mel_spectrogram(audio=input_segments.detach(), n_mels=self.model.dims.n_mels, padding=N_SAMPLES)
             mlx_mel = mlx_pad_or_trim(mlx_mel_padded, N_FRAMES, axis=-2)
@@ -616,10 +641,9 @@ class PaddedAlignAttWhisper:
             timestamp_idx += len(word_tokens)
 
             timestamp_entry = ASRToken(
-                    start=current_timestamp,
-                    end=current_timestamp + 0.1,
+                    start=round(current_timestamp, 2),
+                    end=round(current_timestamp + 0.1, 2),
                     text= word,
-                    probability=0.95,
                     speaker=self.speaker,
                     detected_language=self.detected_language
                 ).with_offset(

whisperlivekit/timed_objects.py

@@ -8,15 +8,15 @@ def format_time(seconds: float) -> str:
     """Format seconds as HH:MM:SS."""
     return str(timedelta(seconds=int(seconds)))
 
-
 @dataclass
-class TimedText:
+class Timed:
     start: Optional[float] = 0
     end: Optional[float] = 0
+
+@dataclass
+class TimedText(Timed):
     text: Optional[str] = ''
     speaker: Optional[int] = -1
-    probability: Optional[float] = None
-    is_dummy: Optional[bool] = False
     detected_language: Optional[str] = None
     
     def is_punctuation(self):
@@ -51,7 +51,7 @@ class ASRToken(TimedText):
     
     def with_offset(self, offset: float) -> "ASRToken":
         """Return a new token with the time offset added."""
-        return ASRToken(self.start + offset, self.end + offset, self.text, self.speaker, self.probability, detected_language=self.detected_language)
+        return ASRToken(self.start + offset, self.end + offset, self.text, self.speaker, detected_language=self.detected_language)
 
 @dataclass
 class Sentence(TimedText):
@@ -72,21 +72,21 @@ class Transcript(TimedText):
     ) -> "Transcript":
         sep = sep if sep is not None else ' '
         text = sep.join(token.text for token in tokens)
-        probability = sum(token.probability for token in tokens if token.probability) / len(tokens) if tokens else None
         if tokens:
             start = offset + tokens[0].start
             end = offset + tokens[-1].end
         else:
             start = None
             end = None
-        return cls(start, end, text, probability=probability)
+        return cls(start, end, text)
 
 
 @dataclass
-class SpeakerSegment(TimedText):
+class SpeakerSegment(Timed):
     """Represents a segment of audio attributed to a specific speaker.
     No text nor probability is associated with this segment.
     """
+    speaker: Optional[int] = -1
     pass
 
 @dataclass
@@ -123,7 +123,9 @@ class Translation(TimedText):
 
 @dataclass
 class Silence():
-    duration: float
+    duration: Optional[float] = None
+    is_starting: bool = False
+    has_ended: bool = False
     
     
 @dataclass
@@ -178,12 +180,24 @@ class ChangeSpeaker:
 class State():
     tokens: list = field(default_factory=list)
     last_validated_token: int = 0
+    last_speaker: int = 1
     last_punctuation_index: Optional[int] = None
     translation_validated_segments: list = field(default_factory=list)
     buffer_translation: str = field(default_factory=Transcript)
     buffer_transcription: str = field(default_factory=Transcript)
+    diarization_segments: list = field(default_factory=list)
     end_buffer: float = 0.0
     end_attributed_speaker: float = 0.0
     remaining_time_transcription: float = 0.0
     remaining_time_diarization: float = 0.0
-    beg_loop: Optional[int] = None
+
+
+@dataclass  
+class StateLight():
+    new_tokens: list = field(default_factory=list)
+    new_translation: list = field(default_factory=list)
+    new_diarization: list = field(default_factory=list)
+    new_tokens_buffer: list = field(default_factory=list) #only when local agreement
+    new_tokens_index = 0
+    new_translation_index = 0
+    new_diarization_index = 0

whisperlivekit/web/live_transcription.js

@@ -391,12 +391,11 @@ function renderLinesWithBuffer(
 
       if (idx === lines.length - 1) {
         if (!isFinalizing && item.speaker !== -2) {
-          if (remaining_time_transcription > 0) {
             speakerLabel += `<span class="label_transcription"><span class="spinner"></span>Transcription lag <span id='timeInfo'><span class="lag-transcription-value">${fmt1(
               remaining_time_transcription
             )}</span>s</span></span>`;
-          }
-          if (buffer_diarization && remaining_time_diarization > 0) {
+
+          if (buffer_diarization && remaining_time_diarization) {
             speakerLabel += `<span class="label_diarization"><span class="spinner"></span>Diarization lag<span id='timeInfo'><span class="lag-diarization-value">${fmt1(
               remaining_time_diarization
             )}</span>s</span></span>`;

whisperlivekit/whisper/__init__.py

@@ -4,17 +4,18 @@ import json
 import os
 import urllib
 import warnings
-from typing import List, Optional, Union, Dict
+from typing import Dict, List, Optional, Union
 
 import torch
 from tqdm import tqdm
 from pathlib import Path
+from torch import Tensor
 
-from .audio import load_audio, log_mel_spectrogram, pad_or_trim
-from .decoding import DecodingOptions, DecodingResult, decode, detect_language
-from .model import ModelDimensions, Whisper
-from .transcribe import transcribe
-from .version import __version__
+from whisperlivekit.whisper.audio import load_audio, log_mel_spectrogram, pad_or_trim
+from whisperlivekit.whisper.decoding import DecodingOptions, DecodingResult, decode, detect_language
+from whisperlivekit.whisper.model import ModelDimensions, Whisper
+from whisperlivekit.whisper.transcribe import transcribe
+from whisperlivekit.whisper.version import __version__
 
 _MODELS = {
     "tiny.en": "https://openaipublic.azureedge.net/main/whisper/models/d3dd57d32accea0b295c96e26691aa14d8822fac7d9d27d5dc00b4ca2826dd03/tiny.en.pt",
@@ -233,13 +234,97 @@ def _convert_hf_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, tor
     return converted if converted else state_dict
 
 
+def _load_lora_state(lora_path: str):
+    safe_path = os.path.join(lora_path, "adapter_model.safetensors")
+    bin_path = os.path.join(lora_path, "adapter_model.bin")
+    if os.path.isfile(safe_path):
+        try:
+            from safetensors.torch import load_file
+        except ImportError as exc:
+            raise ImportError(
+                "Loading LoRA adapters stored as .safetensors requires the `safetensors` package."
+            ) from exc
+        return load_file(safe_path)
+    if os.path.isfile(bin_path):
+        return torch.load(bin_path, map_location="cpu")
+    raise FileNotFoundError(
+        f"No adapter weights found under {lora_path}. Expected adapter_model.safetensors or adapter_model.bin."
+    )
+
+
+def _collapse_hf_module_name(module: str):
+    if module.startswith("base_model."):
+        module = module[len("base_model.") :]
+    if module.startswith("model.model."):
+        module = module[len("model.") :]
+    if not module.startswith("model."):
+        module = f"model.{module}"
+    return module
+
+
+def _apply_lora_adapter(state_dict: Dict[str, Tensor], lora_path: Optional[str]):
+    if not lora_path:
+        return
+
+    config_path = os.path.join(lora_path, "adapter_config.json")
+    if not os.path.isfile(config_path):
+        raise FileNotFoundError(f"Missing adapter_config.json inside {lora_path}")
+    with open(config_path, "r", encoding="utf-8") as handle:
+        config = json.load(handle)
+    if config.get("peft_type") != "LORA":
+        raise ValueError("Only LoRA adapters are supported.")
+
+    r = config.get("r")
+    alpha = config.get("lora_alpha") or config.get("alpha")
+    if not r or not alpha:
+        raise ValueError("LoRA config must include `r` and `lora_alpha`.")
+    scaling = alpha / r
+
+    adapter_state = _load_lora_state(lora_path)
+    lora_layers: Dict[str, Dict[str, Tensor]] = {}
+    for key, tensor in adapter_state.items():
+        if key.endswith("lora_A.weight"):
+            module = key[: -len(".lora_A.weight")]
+            lora_layers.setdefault(module, {})["A"] = tensor
+        elif key.endswith("lora_B.weight"):
+            module = key[: -len(".lora_B.weight")]
+            lora_layers.setdefault(module, {})["B"] = tensor
+
+    if not lora_layers:
+        raise ValueError(f"No LoRA tensors found in {lora_path}")
+
+    for module, parts in lora_layers.items():
+        if "A" not in parts or "B" not in parts:
+            raise ValueError(f"Incomplete LoRA tensors for module '{module}'")
+
+        hf_module = _collapse_hf_module_name(module)
+        hf_weight_key = f"{hf_module}.weight"
+
+        delta = parts["B"] @ parts["A"]
+        delta = delta * scaling
+
+        converted = _convert_hf_state_dict({hf_weight_key: delta})
+        if not converted:
+            raise KeyError(f"Failed to map LoRA module '{module}' into Whisper state dict.")
+        target_name, delta_tensor = next(iter(converted.items()))
+        if target_name not in state_dict:
+            raise KeyError(
+                f"LoRA module '{module}' mapped to '{target_name}', but the base model has no such parameter."
+            )
+
+        state_dict[target_name] = state_dict[target_name] + delta_tensor.to(
+            dtype=state_dict[target_name].dtype, device=state_dict[target_name].device
+        )
+
+
 def load_model(
     name: str,
     device: Optional[Union[str, torch.device]] = None,
     download_root: str = None,
     in_memory: bool = False,
-    decoder_only=False,
-    custom_alignment_heads=None
+    decoder_only: bool = False,
+    custom_alignment_heads: Optional[str] = None,
+    lora_path: Optional[str] = None,
 ) -> Whisper:
     """
     Load a Whisper ASR model
@@ -255,6 +340,8 @@ def load_model(
         path to download the model files; by default, it uses "~/.cache/whisper"
     in_memory: bool
         whether to preload the model weights into host memory
+    lora_path: str
+        optional directory containing PEFT LoRA adapter weights (adapter_config + adapter_model)
 
     Returns
     -------
@@ -302,6 +389,7 @@ def load_model(
     else:
         state_dict = checkpoint
     state_dict = _convert_hf_state_dict(state_dict)
+    _apply_lora_adapter(state_dict, lora_path)
 
     if dims_cfg is not None:
         dims = ModelDimensions(**dims_cfg)
@@ -329,3 +417,47 @@ def load_model(
         model.set_alignment_heads(alignment_heads)
 
     return model.to(device)
+
+
+def convert_encoder_to_coreml(
+    model_name = "base",
+    output_path= "whisper_encoder.mlpackage",
+    dummy_frames = 3000, #Number of time frames to use for the dummy mel input during tracing
+    precision = "float16",
+):
+   
+    import coremltools as ct
+    model = load_model(model_name, device="cpu", decoder_only=False)
+    encoder = model.encoder.eval().cpu()
+
+    dummy_input = torch.randn(
+        1,
+        model.dims.n_mels,
+        dummy_frames,
+        dtype=next(encoder.parameters()).dtype,
+    )
+
+    with torch.no_grad():
+        traced_encoder = torch.jit.trace(encoder, dummy_input)
+
+    precision_map = {
+        "float16": ct.precision.FLOAT16,
+        "fp16": ct.precision.FLOAT16,
+        "float32": ct.precision.FLOAT32,
+        "fp32": ct.precision.FLOAT32,
+    }
+    coreml_precision = precision_map[precision.lower()]
+
+    mlmodel = ct.convert(
+        traced_encoder,
+        inputs=[ct.TensorType(name="mel", shape=dummy_input.shape)],
+        convert_to= "mlprogram",
+        compute_precision=coreml_precision,
+    )
+
+    output_path = Path(output_path)
+    mlmodel.save(str(output_path))
+    return output_path
+
+# if __name__ == "__main__":
+#     convert_encoder_to_coreml(model_name="tiny", output_path="whisper_encoder.mlpackage", dummy_frames=3000, precision="float16", convert_to="mlprogram")