Add tokens for each category

src/shared.py

@@ -7,18 +7,19 @@ from typing import Optional
 from dataclasses import dataclass, field
 from enum import Enum
 
+
+START_SEGMENT_TEMPLATE = 'START_{}_TOKEN'
+END_SEGMENT_TEMPLATE = 'END_{}_TOKEN'
+
 class CustomTokens(Enum):
     EXTRACT_SEGMENTS_PREFIX = 'EXTRACT_SEGMENTS: '
 
+    # Preprocessing tokens
     URL = 'URL_TOKEN'
     HYPHENATED_URL = 'HYPHENATED_URL_TOKEN'
     NUMBER_PERCENTAGE = 'NUMBER_PERCENTAGE_TOKEN'
     NUMBER = 'NUMBER_TOKEN'
 
-    START_SEGMENT = 'START_SEGMENT_TOKEN'
-    END_SEGMENT = 'END_SEGMENT_TOKEN'
-    NO_SEGMENT = 'NO_SEGMENT_FOUND'
-
     SHORT_HYPHENATED = 'SHORT_HYPHENATED_TOKEN'
     LONG_WORD = 'LONG_WORD_TOKEN'
 
@@ -28,6 +29,20 @@ class CustomTokens(Enum):
     LAUGHTER = '[Laughter]'
 
     PROFANITY = 'PROFANITY_TOKEN'
+    
+    # Segment tokens
+    NO_SEGMENT = 'NO_SEGMENT_TOKEN'
+
+    START_SPONSOR = START_SEGMENT_TEMPLATE.format('SPONSOR')
+    END_SPONSOR = END_SEGMENT_TEMPLATE.format('SPONSOR')
+
+    START_SELFPROMO = START_SEGMENT_TEMPLATE.format('SELFPROMO')
+    END_SELFPROMO = END_SEGMENT_TEMPLATE.format('SELFPROMO')
+
+    START_INTERACTION = START_SEGMENT_TEMPLATE.format('INTERACTION')
+    END_INTERACTION = END_SEGMENT_TEMPLATE.format('INTERACTION')
+
+    BETWEEN_SEGMENTS = 'BETWEEN_SEGMENTS_TOKEN'
 
     @classmethod
     def custom_tokens(cls):