fix word alignment

app/services/alignment.py

@@ -8,7 +8,6 @@ from typing import List, Tuple, Optional
 from dataclasses import dataclass
 
 from app.core.config import get_settings
-
 from app.services.transcription import WordTimestamp
 from app.services.diarization import SpeakerSegment
 from app.schemas.models import TranscriptSegment
@@ -34,120 +33,133 @@ class AlignmentService:
     Uses word-center-based algorithm for accurate speaker-to-text mapping.
     """
     
-    # Pause threshold for splitting segments (seconds)
     PAUSE_THRESHOLD = 1.0
+    CENTER_TOL = 0.12
+    OVERLAP_TH = 0.5
+    DIA_MERGE_GAP = 0.25
     
     @staticmethod
     def get_word_center(word: WordTimestamp) -> float:
         """Calculate the center time of a word."""
         return (word.start + word.end) / 2
     
+    
+    @staticmethod
+    def overlap_ratio(w_start, w_end, s_start, s_end):
+        overlap = max(0.0, min(w_end, s_end) - max(w_start, s_start))
+        dur = max(1e-6, w_end - w_start)
+        return overlap / dur
+    
+    
+    # Diarization merge
+    @classmethod
+    def merge_dia_segments(cls, segments: List[SpeakerSegment]) -> List[SpeakerSegment]:
+        if not segments:
+            return []
+
+        segments = sorted(segments, key=lambda s: s.start)
+        merged = [segments[0]]
+
+        for s in segments[1:]:
+            p = merged[-1]
+            if s.speaker == p.speaker and (s.start - p.end) <= cls.DIA_MERGE_GAP:
+                p.end = s.end
+            else:
+                merged.append(s)
+
+        return merged
+    
+    
     @classmethod
-    def find_speaker_at_time(
+    def find_speaker_center(
         cls,
         time: float,
-        speaker_segments: List[SpeakerSegment]
+        speaker_segments: List[SpeakerSegment],
     ) -> Optional[str]:
-        """
-        Find which speaker is speaking at a given time.
-        
-        Args:
-            time: Time point in seconds
-            speaker_segments: List of speaker segments from diarization
-            
-        Returns:
-            Speaker label or None if no speaker found
-        """
+
         for seg in speaker_segments:
-            if seg.start <= time <= seg.end:
+            if seg.start - cls.CENTER_TOL <= time <= seg.end + cls.CENTER_TOL:
                 return seg.speaker
         return None
     
+
     @classmethod
-    def find_closest_speaker(
+    def find_speaker_center(
         cls,
         time: float,
-        speaker_segments: List[SpeakerSegment]
-    ) -> str:
-        """
-        Find the closest speaker to a given time (for gaps/silence).
-        
-        Args:
-            time: Time point in seconds
-            speaker_segments: List of speaker segments
-            
-        Returns:
-            Closest speaker label or "Unknown"
-        """
+        speaker_segments: List[SpeakerSegment],
+    ) -> Optional[str]:
+
+        for seg in speaker_segments:
+            if seg.start - cls.CENTER_TOL <= time <= seg.end + cls.CENTER_TOL:
+                return seg.speaker
+        return None
+
+    @staticmethod
+    def find_closest_speaker(time: float, speaker_segments: List[SpeakerSegment]) -> str:
         if not speaker_segments:
             return "Unknown"
-        
-        min_distance = float('inf')
-        closest_speaker = "Unknown"
-        
+
+        min_dist = float("inf")
+        closest = "Unknown"
+
         for seg in speaker_segments:
-            # Distance to segment start or end
-            dist_to_start = abs(time - seg.start)
-            dist_to_end = abs(time - seg.end)
-            min_seg_dist = min(dist_to_start, dist_to_end)
-            
-            if min_seg_dist < min_distance:
-                min_distance = min_seg_dist
-                closest_speaker = seg.speaker
-        
-        return closest_speaker
+            d = min(abs(time - seg.start), abs(time - seg.end))
+            if d < min_dist:
+                min_dist = d
+                closest = seg.speaker
+
+        return closest
+    
     
     @classmethod
     def assign_speakers_to_words(
         cls,
         words: List[WordTimestamp],
-        speaker_segments: List[SpeakerSegment]
+        speaker_segments: List[SpeakerSegment],
     ) -> List[WordWithSpeaker]:
-        """
-        Step 3c: Assign speakers to each word based on word center time.
-        
-        Args:
-            words: List of words with timestamps from transcription
-            speaker_segments: List of speaker segments from diarization
-            
-        Returns:
-            List of words with speaker assignments
-        """
+
+        words = [w for w in words if w.word and w.word.strip()]
+
         if not speaker_segments:
-            # No diarization available, assign all to "Speaker 1"
-            logger.warning("No speaker segments available, using single speaker")
+            logger.warning("No diarization, fallback single speaker")
             return [
-                WordWithSpeaker(
-                    word=w.word,
-                    start=w.start,
-                    end=w.end,
-                    speaker="Speaker 1"
-                )
+                WordWithSpeaker(w.word, w.start, w.end, "Speaker 1")
                 for w in words
             ]
-        
-        words_with_speakers = []
-        
+
+        speaker_segments = cls.merge_dia_segments(speaker_segments)
+
+        results = []
+
         for word in words:
-            # Calculate word center time
-            center_time = cls.get_word_center(word)
-            
-            # Find speaker at this time
-            speaker = cls.find_speaker_at_time(center_time, speaker_segments)
-            
-            # If no direct match, find closest speaker
+            center = cls.get_word_center(word)
+
+            # 1. CENTER
+            speaker = cls.find_speaker_center(center, speaker_segments)
+
             if speaker is None:
-                speaker = cls.find_closest_speaker(center_time, speaker_segments)
-            
-            words_with_speakers.append(WordWithSpeaker(
-                word=word.word,
-                start=word.start,
-                end=word.end,
-                speaker=speaker
-            ))
-        
-        logger.debug(f"Assigned speakers to {len(words_with_speakers)} words")
-        return words_with_speakers
+                # 2. OVERLAP
+                best_ratio = 0
+                best_spk = None
+
+                for seg in speaker_segments:
+                    r = cls.overlap_ratio(word.start, word.end, seg.start, seg.end)
+                    if r > best_ratio:
+                        best_ratio = r
+                        best_spk = seg.speaker
+
+                if best_ratio >= cls.OVERLAP_TH:
+                    speaker = best_spk
+                else:
+                    # 3. CLOSEST
+                    speaker = cls.find_closest_speaker(center, speaker_segments)
+
+            results.append(
+                WordWithSpeaker(word.word, word.start, word.end, speaker)
+            )
+
+        return results
     
     @classmethod
     def reconstruct_segments(

app/services/processor.py

@@ -13,6 +13,9 @@ import torch
 
 from app.core.config import get_settings
 from app.services.transcription import TranscriptionService
+from app.services.alignment import AlignmentService
+from app.services.transcription import WordTimestamp
+
 from app.services.diarization import DiarizationService, SpeakerSegment, DiarizationResult
 
 logger = logging.getLogger(__name__)
@@ -42,92 +45,6 @@ class ProcessingResult:
     txt_content: str = ""
     csv_content: str = ""
 
-def assign_speaker_to_word(word_start, diarization_segments):
-    if not diarization_segments:
-        return "SPEAKER_0"
-
-    for seg in diarization_segments:
-        if seg.start - 0.05 <= word_start <= seg.end + 0.05:
-            return seg.speaker
-
-    return min(
-        diarization_segments,
-        key=lambda s: abs((s.start + s.end) / 2 - word_start)
-    ).speaker
-
-
-
-def _safe_extract_word(w: dict) -> str:
-    """
-    Robust extractor for finetuned Whisper outputs.
-    Handles:
-      - str
-      - dict {text, confidence, ...}
-      - None
-    """
-    raw = w.get("word", "")
-
-    if isinstance(raw, dict):
-        return str(raw.get("text", "")).strip()
-
-    if isinstance(raw, str):
-        return raw.strip()
-
-    return str(raw).strip()
-
-def group_words_into_segments(
-    words: List[dict],
-    diarization_segments: List[SpeakerSegment],
-    speaker_map: Dict[str, str],
-    roles: Dict[str, str],
-    max_word_gap_s: float = 0.6,
-) -> List[TranscriptSegment]:
-
-    segments: List[TranscriptSegment] = []
-    current: Optional[TranscriptSegment] = None
-
-    for w in words or []:
-        text = _safe_extract_word(w)
-        if not text:
-            continue
-
-        w_start = float(w["start"])
-        w_end = float(w["end"])
-
-        speaker_raw = assign_speaker_to_word(w_start, diarization_segments)
-        speaker = speaker_map.get(speaker_raw, speaker_raw)
-        role = roles.get(speaker, "KH")
-
-        if current is None:
-            current = TranscriptSegment(
-                start=w_start,
-                end=w_end,
-                speaker=speaker,
-                role=role,
-                text=text,
-            )
-            continue
-
-        gap = w_start - current.end
-
-        if speaker == current.speaker and gap <= max_word_gap_s:
-            current.text += " " + text
-            current.end = max(current.end, w_end)
-        else:
-            segments.append(current)
-            current = TranscriptSegment(
-                start=w_start,
-                end=w_end,
-                speaker=speaker,
-                role=role,
-                text=text,
-            )
-
-    if current:
-        segments.append(current)
-
-    return segments
-
 
 def normalize_asr_result(result):
     """
@@ -280,9 +197,8 @@ class Processor:
         temperature: float = 0.0,
         best_of: int = 5,
         log_prob_threshold = -1.2,
-        compression_ratio_threshold = 2.4,
-
-        initial_prompt: Optional[str] = None,
+        compression_ratio_threshold = 2.4
+        
     ) -> ProcessingResult:
 
         import asyncio
@@ -374,8 +290,13 @@ class Processor:
 
         logger.info(f"roles(mapped) = {roles}")
 
-
-            
+        # context window
+        windows = build_context_windows(
+            duration,
+            win_s=settings.CONTEXT_WINDOW_S,
+            overlap_s=settings.CONTEXT_OVERLAP_S,
+        )
+  
         # 7: Transcribe
         vad_options = None
         if vad_filter:
@@ -439,22 +360,47 @@ class Processor:
 
             if not words and not text:
                 continue
+            
+            words = []
+            for w in words:
+                try:
+                    words.append(
+                        WordTimestamp(
+                            word=str(w.get("word", "")).strip(),
+                            start=float(w.get("start", 0)) + w_start,
+                            end=float(w.get("end", 0)) + w_start,
+                        )
+                    )
+                except:
+                    pass
+
+            if not words:
+                continue
 
-            # ===== MAP WORD → SPEAKER =====
-            word_segments = group_words_into_segments(
-                words=words,
-                diarization_segments=refined_segments,
-                speaker_map=speaker_map,
-                roles=roles,
-                max_word_gap_s=0.6
+            # ===== ALIGNMENT =====
+            aligned_segments = AlignmentService.align_precision(
+                words,
+                refined_segments
             )
 
-            for seg in word_segments:
-                processed_segments.append(seg)
+            # ===== MAP WORD → ROLE =====
+            for seg in aligned_segments:
+                label = speaker_map.get(seg.speaker, seg.speaker)
+                role = roles.get(label, "KH")
+
+                processed_segments.append(
+                    TranscriptSegment(
+                        start=seg.start,
+                        end=seg.end,
+                        speaker=label,
+                        role=role,
+                        text=seg.text
+                    )
+                )
 
             if text:
                 prev_prompt += " " + text
-            
+                
 
         if not processed_segments:
             processed_segments = [
@@ -468,15 +414,9 @@ class Processor:
             ]
         
         processed_segments = cls._merge_adjacent_segments(
-            processed_segments,
-            max_gap_s=0.7
+            processed_segments
         )
 
-        processed_segments = [
-            s for s in processed_segments
-            if cls._is_meaningful_segment(s)
-        ]
-
         processed_segments = cls._filter_segments_with_context(processed_segments)
         processing_time = time.time() - t0
         

app/services/transcription.py

@@ -83,9 +83,10 @@ class TranscriptionService:
     
     @classmethod
     def is_loaded(cls, model_name: str = None) -> bool:
+        """Check if a model is loaded."""
         if model_name is None:
             model_name = settings.default_whisper_model
-        """Check if a model is loaded."""
+
         cache_key = f"{model_name}_{settings.resolved_compute_type}"
         return cache_key in cls._models