Update segmenter.py

segmenter.py

@@ -7,7 +7,8 @@ logger = logging.getLogger(__name__)
 
 class TextSegmenter:
     def __init__(self):
-        self.speakers = ["speaker1", "speaker2", "speaker3", "speaker4"]
+        # Changed speakers to Nari DIA's expected tags
+        self.speakers = ["S1", "S2"] 
         self.current_speaker_index = 0
     
     def segment_and_assign_speakers(
@@ -45,11 +46,11 @@ class TextSegmenter:
     
     def _segment_by_dialogue(self, text: str) -> List[Tuple[str, str]]:
         """Segment by detecting dialogue patterns."""
-        # Look for dialogue markers like quotes, dashes, etc.
         lines = text.split('\n')
         segments = []
         current_segment = []
-        current_speaker = self.speakers[0]
+        # Start with the first speaker in the list
+        current_speaker = self.speakers[0] 
         
         for line in lines:
             line = line.strip()
@@ -79,55 +80,49 @@ class TextSegmenter:
     
     def _segment_auto(self, text: str) -> List[Tuple[str, str]]:
         """Automatic segmentation using multiple heuristics."""
-        # Try to detect natural breaks
         segments = []
         
-        # Split by double newlines first (paragraphs)
         paragraphs = [p.strip() for p in text.split('\n\n') if p.strip()]
         
         if len(paragraphs) > 1:
-            # Use paragraph-based segmentation
             return self._segment_by_paragraphs(text)
         
-        # Fall back to sentence-based segmentation for long text
         sentences = self._split_into_sentences(text)
         if len(sentences) > 10:
             return self._segment_by_sentence_groups(sentences)
         
-        # For short text, just alternate every few sentences
         return self._segment_simple(text)
     
     def _split_into_sentences(self, text: str) -> List[str]:
         """Split text into sentences."""
         # Simple sentence splitting
-        sentences = re.split(r'[.!?]+', text)
+        # Use a more robust regex to avoid splitting on abbreviations (e.g., "Mr.")
+        # This is a common simple improvement, though full NLP libraries are best for complex cases.
+        sentences = re.split(r'(?<=[.!?])\s+', text) # Split after . ! ? followed by space
         return [s.strip() for s in sentences if s.strip()]
     
     def _segment_by_sentence_groups(self, sentences: List[str]) -> List[Tuple[str, str]]:
         """Group sentences and assign to different speakers."""
         segments = []
-        group_size = max(2, len(sentences) // 8)  # Aim for reasonable segment sizes
+        group_size = max(2, len(sentences) // 8) 
         
         for i in range(0, len(sentences), group_size):
             group = sentences[i:i + group_size]
             speaker = self.speakers[i // group_size % len(self.speakers)]
-            text_segment = '. '.join(group) + '.'
+            text_segment = ' '.join(group) # No need to add '.' if already present from sentence splitting
             segments.append((speaker, text_segment))
         
         return segments
     
     def _segment_simple(self, text: str) -> List[Tuple[str, str]]:
         """Simple segmentation for short texts."""
-        # Just split roughly in half or thirds
         words = text.split()
         total_words = len(words)
         
         if total_words < 50:
-            # Too short to split meaningfully
-            return [("speaker1", text)]
+            return [(self.speakers[0], text)] # Assign to S1
         
-        # Split into 2-3 segments
-        num_segments = min(3, max(2, total_words // 100))
+        num_segments = min(len(self.speakers), max(2, total_words // 100)) # Limit segments by available speakers
         segment_size = total_words // num_segments
         
         segments = []