lyrics-transcriber 0.41.0__py3-none-any.whl → 0.43.0__py3-none-any.whl

lyrics_transcriber/correction/corrector.py

@@ -1,18 +1,30 @@
-from typing import List, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union, Dict, Any
 import logging
 from pathlib import Path
+from copy import deepcopy
 
+from lyrics_transcriber.correction.handlers.levenshtein import LevenshteinHandler
+from lyrics_transcriber.correction.handlers.llm import LLMHandler
 from lyrics_transcriber.correction.handlers.no_space_punct_match import NoSpacePunctuationMatchHandler
 from lyrics_transcriber.correction.handlers.relaxed_word_count_match import RelaxedWordCountMatchHandler
+from lyrics_transcriber.correction.handlers.repeat import RepeatCorrectionHandler
+from lyrics_transcriber.correction.handlers.sound_alike import SoundAlikeHandler
 from lyrics_transcriber.correction.handlers.syllables_match import SyllablesMatchHandler
-from lyrics_transcriber.types import GapSequence, LyricsData, TranscriptionResult, CorrectionResult, LyricsSegment, WordCorrection, Word
+from lyrics_transcriber.correction.handlers.word_count_match import WordCountMatchHandler
+from lyrics_transcriber.types import (
+    CorrectionStep,
+    GapSequence,
+    LyricsData,
+    TranscriptionResult,
+    CorrectionResult,
+    LyricsSegment,
+    WordCorrection,
+    Word,
+)
 from lyrics_transcriber.correction.anchor_sequence import AnchorSequenceFinder
 from lyrics_transcriber.correction.handlers.base import GapCorrectionHandler
-from lyrics_transcriber.correction.handlers.word_count_match import WordCountMatchHandler
 from lyrics_transcriber.correction.handlers.extend_anchor import ExtendAnchorHandler
-from lyrics_transcriber.correction.handlers.sound_alike import SoundAlikeHandler
-from lyrics_transcriber.correction.handlers.levenshtein import LevenshteinHandler
-from lyrics_transcriber.correction.handlers.repeat import RepeatCorrectionHandler
+from lyrics_transcriber.utils.word_utils import WordUtils
 
 
 class LyricsCorrector:
@@ -24,25 +36,54 @@ class LyricsCorrector:
         self,
         cache_dir: Union[str, Path],
         handlers: Optional[List[GapCorrectionHandler]] = None,
+        enabled_handlers: Optional[List[str]] = None,
         anchor_finder: Optional[AnchorSequenceFinder] = None,
         logger: Optional[logging.Logger] = None,
     ):
         self.logger = logger or logging.getLogger(__name__)
         self._anchor_finder = anchor_finder
-        self._cache_dir = cache_dir
-
-        # Default handlers in order of preference
-        self.handlers = handlers or [
-            # WordCountMatchHandler(logger=self.logger),
-            # RelaxedWordCountMatchHandler(logger=self.logger),
-            # NoSpacePunctuationMatchHandler(logger=self.logger),
-            # SyllablesMatchHandler(logger=self.logger),
-            ExtendAnchorHandler(logger=self.logger),
-            # RepeatCorrectionHandler(logger=self.logger),
-            # SoundAlikeHandler(logger=self.logger),
-            # LevenshteinHandler(logger=self.logger),
+        self._cache_dir = Path(cache_dir)
+
+        # Define default enabled handlers - excluding LLM, Repeat, SoundAlike, and Levenshtein
+        DEFAULT_ENABLED_HANDLERS = [
+            "ExtendAnchorHandler",
+            "WordCountMatchHandler",
+            "SyllablesMatchHandler",
+            "RelaxedWordCountMatchHandler",
+            "NoSpacePunctuationMatchHandler",
         ]
 
+        # Create all handlers but respect enabled_handlers if provided
+        all_handlers = [
+            ("ExtendAnchorHandler", ExtendAnchorHandler(logger=self.logger)),
+            ("WordCountMatchHandler", WordCountMatchHandler(logger=self.logger)),
+            ("SyllablesMatchHandler", SyllablesMatchHandler(logger=self.logger)),
+            ("RelaxedWordCountMatchHandler", RelaxedWordCountMatchHandler(logger=self.logger)),
+            ("NoSpacePunctuationMatchHandler", NoSpacePunctuationMatchHandler(logger=self.logger)),
+            ("LLMHandler", LLMHandler(logger=self.logger, cache_dir=self._cache_dir)),
+            ("RepeatCorrectionHandler", RepeatCorrectionHandler(logger=self.logger)),
+            ("SoundAlikeHandler", SoundAlikeHandler(logger=self.logger)),
+            ("LevenshteinHandler", LevenshteinHandler(logger=self.logger)),
+        ]
+
+        # Store all handler information
+        self.all_handlers = [
+            {
+                "id": handler_id,
+                "name": handler_id,
+                "description": handler.__class__.__doc__ or "",
+                "enabled": handler_id in (enabled_handlers if enabled_handlers is not None else DEFAULT_ENABLED_HANDLERS),
+            }
+            for handler_id, handler in all_handlers
+        ]
+
+        if handlers:
+            self.handlers = handlers
+        else:
+            # Use provided enabled_handlers if available, otherwise use defaults
+            handler_filter = enabled_handlers if enabled_handlers is not None else DEFAULT_ENABLED_HANDLERS
+            self.handlers = [h[1] for h in all_handlers if h[0] in handler_filter]
+
     @property
     def anchor_finder(self) -> AnchorSequenceFinder:
         """Lazy load the anchor finder instance, initializing it if not already set."""
@@ -50,39 +91,52 @@ class LyricsCorrector:
             self._anchor_finder = AnchorSequenceFinder(cache_dir=self._cache_dir, logger=self.logger)
         return self._anchor_finder
 
-    def run(self, transcription_results: List[TranscriptionResult], lyrics_results: List[LyricsData]) -> CorrectionResult:
+    def run(
+        self,
+        transcription_results: List[TranscriptionResult],
+        lyrics_results: Dict[str, LyricsData],
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> CorrectionResult:
         """Execute the correction process."""
         if not transcription_results:
             self.logger.error("No transcription results available")
             raise ValueError("No primary transcription data available")
 
+        # Store reference lyrics for use in word map
+        self.reference_lyrics = lyrics_results
+
         # Get primary transcription
         primary_transcription = sorted(transcription_results, key=lambda x: x.priority)[0].result
         transcribed_text = " ".join(" ".join(w.text for w in segment.words) for segment in primary_transcription.segments)
-        reference_texts = {lyrics.source: lyrics.lyrics for lyrics in lyrics_results}
 
         # Find anchor sequences and gaps
         self.logger.debug("Finding anchor sequences and gaps")
-        anchor_sequences = self.anchor_finder.find_anchors(transcribed_text, reference_texts)
-        gap_sequences = self.anchor_finder.find_gaps(transcribed_text, anchor_sequences, reference_texts)
+        anchor_sequences = self.anchor_finder.find_anchors(transcribed_text, lyrics_results, primary_transcription)
+        gap_sequences = self.anchor_finder.find_gaps(transcribed_text, anchor_sequences, lyrics_results, primary_transcription)
+
+        # Store anchor sequences for use in correction handlers
+        self._anchor_sequences = anchor_sequences
 
-        # Process corrections
-        corrections, corrected_segments = self._process_corrections(primary_transcription.segments, gap_sequences)
+        # Process corrections with metadata
+        corrections, corrected_segments, correction_steps, word_id_map, segment_id_map = self._process_corrections(
+            primary_transcription.segments, gap_sequences, metadata=metadata
+        )
 
         # Calculate correction ratio
         total_words = sum(len(segment.words) for segment in corrected_segments)
         corrections_made = len(corrections)
         correction_ratio = 1 - (corrections_made / total_words if total_words > 0 else 0)
 
+        # Get the currently enabled handler IDs using full class names
+        enabled_handlers = [handler.__class__.__name__ for handler in self.handlers]
+
         return CorrectionResult(
             original_segments=primary_transcription.segments,
             corrected_segments=corrected_segments,
-            corrected_text="\n".join(segment.text for segment in corrected_segments) + "\n",
             corrections=corrections,
             corrections_made=corrections_made,
             confidence=correction_ratio,
-            transcribed_text=transcribed_text,
-            reference_texts=reference_texts,
+            reference_lyrics=lyrics_results,
             anchor_sequences=anchor_sequences,
             resized_segments=[],
             gap_sequences=gap_sequences,
@@ -91,7 +145,12 @@ class LyricsCorrector:
                 "gap_sequences_count": len(gap_sequences),
                 "total_words": total_words,
                 "correction_ratio": correction_ratio,
+                "available_handlers": self.all_handlers,
+                "enabled_handlers": enabled_handlers,
             },
+            correction_steps=correction_steps,
+            word_id_map=word_id_map,
+            segment_id_map=segment_id_map,
         )
 
     def _preserve_formatting(self, original: str, new_word: str) -> str:
@@ -102,8 +161,8 @@ class LyricsCorrector:
         return leading_space + new_word.strip() + trailing_space
 
     def _process_corrections(
-        self, segments: List[LyricsSegment], gap_sequences: List[GapSequence]
-    ) -> Tuple[List[WordCorrection], List[LyricsSegment]]:
+        self, segments: List[LyricsSegment], gap_sequences: List[GapSequence], metadata: Optional[Dict[str, Any]] = None
+    ) -> Tuple[List[WordCorrection], List[LyricsSegment], List[CorrectionStep], Dict[str, str], Dict[str, str]]:
         """Process corrections using handlers.
 
         The correction flow works as follows:
@@ -121,84 +180,109 @@ class LyricsCorrector:
         b) Applying those corrections to the original text (segment-centric)
         """
         self.logger.info(f"Starting correction process with {len(gap_sequences)} gaps")
-
-        # First pass: Process all gaps
-        all_corrections = self._process_gaps(gap_sequences)
-
-        # Second pass: Apply corrections to segments
-        corrected_segments = self._apply_corrections_to_segments(segments, all_corrections)
-
-        self.logger.info(f"Correction process complete. Made {len(all_corrections)} corrections")
-        return all_corrections, corrected_segments
-
-    def _process_gaps(self, gap_sequences: List[GapSequence]) -> List[WordCorrection]:
-        """Process each gap using available handlers until all words are corrected or no handlers remain."""
+        correction_steps = []
         all_corrections = []
-        # return all_corrections
-
-        for gap in gap_sequences:
-            self.logger.debug(f"Processing gap: {gap.text}")
-            high_confidence_positions = set()  # Track positions that have high confidence corrections
-            corrected_positions = set()  # Track all corrected positions regardless of confidence
-
-            # Try each handler until gap is fully corrected
+        word_id_map = {}
+        segment_id_map = {}
+
+        # Create word map for handlers - include both transcribed and reference words
+        word_map = {w.id: w for s in segments for w in s.words}  # Transcribed words
+
+        # Add reference words from all sources
+        for source, lyrics_data in self.reference_lyrics.items():
+            for segment in lyrics_data.segments:
+                for word in segment.words:
+                    if word.id not in word_map:  # Don't overwrite transcribed words
+                        word_map[word.id] = word
+
+        # Base handler data that all handlers need
+        base_handler_data = {
+            "word_map": word_map,
+            "anchor_sequences": self._anchor_sequences,
+            "audio_file_hash": metadata.get("audio_file_hash") if metadata else None,
+        }
+
+        for i, gap in enumerate(gap_sequences, 1):
+            self.logger.info(f"Processing gap {i}/{len(gap_sequences)} at position {gap.transcription_position}")
+
+            # Get the actual words for logging
+            gap_words = [word_map[word_id] for word_id in gap.transcribed_word_ids]
+            self.logger.debug(f"Gap text: '{' '.join(w.text for w in gap_words)}'")
+
+            # Try each handler in order
             for handler in self.handlers:
-                # Skip if all words have high confidence corrections
-                uncorrected_positions = set(range(gap.transcription_position, gap.transcription_position + gap.length))
-                uncorrected_positions -= corrected_positions  # Skip any corrected positions
+                handler_name = handler.__class__.__name__
+                can_handle, handler_data = handler.can_handle(gap, base_handler_data)
 
-                if not uncorrected_positions:
-                    self.logger.debug("All words have been corrected, skipping remaining handlers")
-                    break
+                if can_handle:
+                    # Merge base handler data with specific handler data
+                    handler_data = {**base_handler_data, **(handler_data or {})}
 
-                self.logger.debug(f"Trying handler {handler.__class__.__name__}")
+                    corrections = handler.handle(gap, handler_data)
+                    if corrections:
+                        self.logger.info(f"Handler {handler_name} made {len(corrections)} corrections")
+                        # Track affected IDs
+                        affected_word_ids = [w.id for w in self._get_affected_words(gap, segments)]
+                        affected_segment_ids = [s.id for s in self._get_affected_segments(gap, segments)]
 
-                # Pass previous corrections to RepeatCorrectionHandler
-                if isinstance(handler, RepeatCorrectionHandler):
-                    handler.set_previous_corrections(all_corrections)
+                        # Apply corrections and get updated segments
+                        updated_segments = self._apply_corrections_to_segments(self._get_affected_segments(gap, segments), corrections)
 
-                can_handle, handler_data = handler.can_handle(gap)
-                if can_handle:
-                    self.logger.debug(f"{handler.__class__.__name__} can handle gap")
-                    # Only pass handler_data if it's not empty
-                    corrections = handler.handle(gap, handler_data if handler_data else None)
-                    if corrections:
-                        # Add corrections to gap and track corrected positions
+                        # Update ID maps
+                        for correction in corrections:
+                            if correction.word_id and correction.corrected_word_id:
+                                word_id_map[correction.word_id] = correction.corrected_word_id
+
+                        # Map segment IDs
+                        for old_seg, new_seg in zip(self._get_affected_segments(gap, segments), updated_segments):
+                            segment_id_map[old_seg.id] = new_seg.id
+
+                        # Create correction step
+                        step = CorrectionStep(
+                            handler_name=handler_name,
+                            affected_word_ids=affected_word_ids,
+                            affected_segment_ids=affected_segment_ids,
+                            corrections=corrections,
+                            segments_before=self._get_affected_segments(gap, segments),
+                            segments_after=updated_segments,
+                            created_word_ids=[w.id for w in self._get_new_words(updated_segments, affected_word_ids)],
+                            deleted_word_ids=[id for id in affected_word_ids if not self._word_exists(id, updated_segments)],
+                        )
+                        correction_steps.append(step)
+                        all_corrections.extend(corrections)
+
+                        # Log correction details
                         for correction in corrections:
-                            # Skip if this position was already corrected
-                            if correction.original_position in corrected_positions:
-                                continue
-
-                            gap.add_correction(correction)
-                            corrected_positions.add(correction.original_position)
-                            # Track positions with high confidence corrections (>= 0.9)
-                            if correction.confidence >= 0.9:
-                                high_confidence_positions.add(correction.original_position)
-
-                        # Filter out corrections for already corrected positions
-                        new_corrections = [c for c in corrections if c.original_position in corrected_positions]
-                        if new_corrections:
-                            self.logger.debug(
-                                f"{handler.__class__.__name__} made {len(new_corrections)} corrections: "
-                                f"{[f'{c.original_word}->{c.corrected_word}' for c in new_corrections]}"
+                            self.logger.info(
+                                f"Made correction: '{correction.original_word}' -> '{correction.corrected_word}' "
+                                f"(confidence: {correction.confidence:.2f}, reason: {correction.reason})"
                             )
-                            all_corrections.extend(new_corrections)
+                        break  # Stop trying other handlers once we've made corrections
+                    else:
+                        self.logger.debug(f"Handler {handler_name} found no corrections needed")
+                else:
+                    self.logger.debug(f"Handler {handler_name} cannot handle gap")
 
-                        # Log remaining uncorrected words
-                        if not gap.is_fully_corrected:
-                            uncorrected = [word for pos, word in gap.uncorrected_words if pos not in corrected_positions]
-                            if uncorrected:
-                                self.logger.debug(f"Uncorrected words remaining: {', '.join(uncorrected)}")
+        # Create final result with correction history
+        corrected_segments = self._apply_all_corrections(segments, all_corrections)
+        self.logger.info(f"Correction process completed with {len(all_corrections)} total corrections")
+        return all_corrections, corrected_segments, correction_steps, word_id_map, segment_id_map
 
-            if not gap.corrections:
-                self.logger.warning("No handler could handle the gap")
+    def _get_new_words(self, segments: List[LyricsSegment], original_word_ids: List[str]) -> List[Word]:
+        """Find words that were created during correction."""
+        return [w for s in segments for w in s.words if w.id not in original_word_ids]
 
-        return all_corrections
+    def _word_exists(self, word_id: str, segments: List[LyricsSegment]) -> bool:
+        """Check if a word ID still exists in the segments."""
+        return any(w.id == word_id for s in segments for w in s.words)
 
     def _apply_corrections_to_segments(self, segments: List[LyricsSegment], corrections: List[WordCorrection]) -> List[LyricsSegment]:
         """Apply corrections to create new segments."""
-        correction_map = {}
+        # Create word ID map for quick lookup
+        word_map = {w.id: w for s in segments for w in s.words}
+
         # Group corrections by original_position to handle splits
+        correction_map = {}
         for c in corrections:
             if c.original_position not in correction_map:
                 correction_map[c.original_position] = []
@@ -207,7 +291,7 @@ class LyricsCorrector:
         corrected_segments = []
         current_word_idx = 0
 
-        for segment_idx, segment in enumerate(segments):
+        for segment in segments:
             corrected_words = []
             for word in segment.words:
                 if current_word_idx in correction_map:
@@ -226,28 +310,30 @@ class LyricsCorrector:
 
                             # Update corrected_position as we create new words
                             correction.corrected_position = len(corrected_words)
-                            corrected_words.append(
-                                Word(
-                                    text=self._preserve_formatting(correction.original_word, correction.corrected_word),
-                                    start_time=start_time,
-                                    end_time=end_time,
-                                    confidence=correction.confidence,
-                                )
+                            new_word = Word(
+                                id=correction.corrected_word_id or WordUtils.generate_id(),
+                                text=self._preserve_formatting(correction.original_word, correction.corrected_word),
+                                start_time=start_time,
+                                end_time=end_time,
+                                confidence=correction.confidence,
+                                created_during_correction=True,
                             )
+                            corrected_words.append(new_word)
                     else:
                         # Handle single word replacement
                         correction = word_corrections[0]
                         if not correction.is_deletion:
                             # Update corrected_position
                             correction.corrected_position = len(corrected_words)
-                            corrected_words.append(
-                                Word(
-                                    text=self._preserve_formatting(correction.original_word, correction.corrected_word),
-                                    start_time=word.start_time,
-                                    end_time=word.end_time,
-                                    confidence=correction.confidence,
-                                )
+                            new_word = Word(
+                                id=correction.corrected_word_id or WordUtils.generate_id(),
+                                text=self._preserve_formatting(correction.original_word, correction.corrected_word),
+                                start_time=word.start_time,
+                                end_time=word.end_time,
+                                confidence=correction.confidence,
+                                created_during_correction=True,
                             )
+                            corrected_words.append(new_word)
                 else:
                     corrected_words.append(word)
                 current_word_idx += 1
@@ -255,6 +341,7 @@ class LyricsCorrector:
             if corrected_words:
                 corrected_segments.append(
                     LyricsSegment(
+                        id=segment.id,  # Preserve original segment ID
                         text=" ".join(w.text for w in corrected_words),
                         words=corrected_words,
                         start_time=segment.start_time,
@@ -263,3 +350,33 @@ class LyricsCorrector:
                 )
 
         return corrected_segments
+
+    def _get_affected_segments(self, gap: GapSequence, segments: List[LyricsSegment]) -> List[LyricsSegment]:
+        """Get segments that contain words from the gap sequence."""
+        affected_segments = []
+        gap_word_ids = set(gap.transcribed_word_ids)
+
+        for segment in segments:
+            # Check if any words in this segment are part of the gap
+            if any(w.id in gap_word_ids for w in segment.words):
+                affected_segments.append(segment)
+            elif affected_segments:  # We've passed the gap
+                break
+
+        return affected_segments
+
+    def _get_affected_words(self, gap: GapSequence, segments: List[LyricsSegment]) -> List[Word]:
+        """Get words that are part of the gap sequence."""
+        # Create a map of word IDs to Word objects for quick lookup
+        word_map = {w.id: w for s in segments for w in s.words}
+
+        # Get the actual Word objects using the IDs
+        return [word_map[word_id] for word_id in gap.transcribed_word_ids]
+
+    def _apply_all_corrections(self, segments: List[LyricsSegment], corrections: List[WordCorrection]) -> List[LyricsSegment]:
+        """Apply all corrections to create final corrected segments."""
+        # Make a deep copy to avoid modifying original segments
+        working_segments = deepcopy(segments)
+
+        # Apply corrections in order
+        return self._apply_corrections_to_segments(working_segments, corrections)

lyrics_transcriber/correction/handlers/base.py

@@ -12,23 +12,41 @@ class GapCorrectionHandler(ABC):
         self.logger = logger or logging.getLogger(__name__)
 
     @abstractmethod
-    def can_handle(self, gap: GapSequence) -> Tuple[bool, Dict[str, Any]]:
-        """Determine if this handler can process the given gap.
-
+    def can_handle(self, gap: GapSequence, data: Optional[Dict[str, Any]] = None) -> Tuple[bool, Dict[str, Any]]:
+        """Check if this handler can process the given gap.
+        
+        Args:
+            gap: The gap sequence to check
+            data: Optional dictionary containing additional data like word_map
+            
         Returns:
-            Tuple containing:
-            - bool: Whether this handler can process the gap
-            - dict: Data computed during can_handle that will be needed by handle().
-                   Empty dict if no data needs to be passed.
+            Tuple of (can_handle, handler_data)
         """
         pass
 
     @abstractmethod
     def handle(self, gap: GapSequence, data: Optional[Dict[str, Any]] = None) -> List[WordCorrection]:
-        """Process a gap and return any corrections.
-
+        """Process the gap and return any corrections.
+        
         Args:
             gap: The gap sequence to process
-            data: Optional data dictionary returned by can_handle()
+            data: Optional dictionary containing additional data like word_map
+            
+        Returns:
+            List of corrections to apply
         """
         pass
+
+    def _validate_data(self, data: Optional[Dict[str, Any]]) -> bool:
+        """Validate that required data is present.
+        
+        Args:
+            data: The data dictionary to validate
+            
+        Returns:
+            True if data is valid, False otherwise
+        """
+        if not data or "word_map" not in data:
+            self.logger.error("No word_map provided in data")
+            return False
+        return True

lyrics_transcriber/correction/handlers/extend_anchor.py

@@ -1,7 +1,7 @@
 from typing import List, Optional, Tuple, Dict, Any
 import logging
 
-from lyrics_transcriber.types import GapSequence, WordCorrection
+from lyrics_transcriber.types import GapSequence, WordCorrection, Word
 from lyrics_transcriber.correction.handlers.base import GapCorrectionHandler
 from lyrics_transcriber.correction.handlers.word_operations import WordOperations
 
@@ -40,26 +40,31 @@ class ExtendAnchorHandler(GapCorrectionHandler):
     """
 
     def __init__(self, logger: Optional[logging.Logger] = None):
-        super().__init__(logger)
         self.logger = logger or logging.getLogger(__name__)
 
-    def can_handle(self, gap: GapSequence) -> Tuple[bool, Dict[str, Any]]:
-        # Must have reference words
-        if not gap.reference_words:
-            self.logger.debug("No reference words available.")
+    def can_handle(self, gap: GapSequence, data: Optional[Dict[str, Any]] = None) -> Tuple[bool, Dict[str, Any]]:
+        """Check if this gap can be handled by extending anchor sequences."""
+        # Check if we have anchor sequences
+        if not data or "anchor_sequences" not in data:
+            self.logger.debug("No anchor sequences available")
             return False, {}
 
-        # Gap must have words
-        if not gap.words:
-            self.logger.debug("No words in the gap to process.")
+        # Must have reference word IDs
+        if not gap.reference_word_ids:
+            self.logger.debug("No reference word IDs available.")
             return False, {}
 
-        # At least one word must match between gap and any reference source
+        # Gap must have word IDs
+        if not gap.transcribed_word_ids:
+            self.logger.debug("No word IDs in the gap to process.")
+            return False, {}
+
+        # At least one word ID must match between gap and any reference source
         # in the same position
         has_match = any(
-            i < len(ref_words) and gap.words[i].lower() == ref_words[i].lower()
-            for ref_words in gap.reference_words.values()
-            for i in range(min(len(gap.words), len(ref_words)))
+            i < len(ref_word_ids) and gap.transcribed_word_ids[i] == ref_word_ids[i]
+            for ref_word_ids in gap.reference_word_ids.values()
+            for i in range(min(len(gap.transcribed_word_ids), len(ref_word_ids)))
         )
 
         self.logger.debug(f"Can handle gap: {has_match}")
@@ -68,16 +73,32 @@ class ExtendAnchorHandler(GapCorrectionHandler):
     def handle(self, gap: GapSequence, data: Optional[Dict[str, Any]] = None) -> List[WordCorrection]:
         corrections = []
 
+        # Get word lookup map from data
+        word_map = data.get("word_map", {})
+        if not word_map:
+            self.logger.error("No word_map provided in data")
+            return []
+
         # Process each word in the gap that has a corresponding reference position
-        for i, word in enumerate(gap.words):
+        for i, word_id in enumerate(gap.transcribed_word_ids):
+            # Get the actual word object
+            if word_id not in word_map:
+                self.logger.error(f"Word ID {word_id} not found in word_map")
+                continue
+            word = word_map[word_id]
+
             # Find reference sources that have a matching word at this position
             matching_sources = [
-                source for source, ref_words in gap.reference_words.items() if i < len(ref_words) and word.lower() == ref_words[i].lower()
+                source for source, ref_word_ids in gap.reference_word_ids.items() if i < len(ref_word_ids) and word_id == ref_word_ids[i]
             ]
 
+            if not matching_sources:
+                self.logger.debug(f"Skipping word '{word.text}' at position {i} - no matching references")
+                continue
+
             if matching_sources:
                 # Word matches reference(s) at this position - validate it
-                confidence = len(matching_sources) / len(gap.reference_words)
+                confidence = len(matching_sources) / len(gap.reference_word_ids)
                 sources = ", ".join(matching_sources)
 
                 # Get base reference positions
@@ -88,24 +109,26 @@ class ExtendAnchorHandler(GapCorrectionHandler):
                 for source in matching_sources:
                     if source in base_reference_positions:
                         # Find this word's position in the reference text
-                        ref_words = gap.reference_words[source]
-                        for ref_idx, ref_word in enumerate(ref_words):
-                            if ref_word.lower() == word.lower():
+                        ref_word_ids = gap.reference_word_ids[source]
+                        for ref_idx, ref_word_id in enumerate(ref_word_ids):
+                            if ref_word_id == word_id:
                                 reference_positions[source] = base_reference_positions[source] + ref_idx
                                 break
 
                 corrections.append(
                     WordOperations.create_word_replacement_correction(
-                        original_word=word,
-                        corrected_word=word,  # Same word, just validating
+                        original_word=word.text,
+                        corrected_word=word.text,
                         original_position=gap.transcription_position + i,
                         source=sources,
                         confidence=confidence,
-                        reason="ExtendAnchorHandler: Matched reference source(s)",
+                        reason="Matched reference source(s)",
                         reference_positions=reference_positions,
+                        handler="ExtendAnchorHandler",
+                        original_word_id=word_id,
+                        corrected_word_id=word_id,
                     )
                 )
-                self.logger.debug(f"Validated word '{word}' with confidence {confidence} from sources: {sources}")
-            # No else clause - non-matching words are left unchanged
+                self.logger.debug(f"Validated word '{word.text}' with confidence {confidence} from sources: {sources}")
 
         return corrections