Add ability to color code.

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import json
+import random
 from read_files import Tokenizer  # Make sure to include this file
 
 def load_tokenizer(path):
@@ -9,19 +10,74 @@ def load_tokenizer(path):
     merges = {tuple(map(int, k.split(','))): v for k, v in serialized_merges.items()}
     return Tokenizer(merges)
 
+def generate_color():
+    """Generate a random pastel color"""
+    hue = random.random()
+    saturation = 0.3 + random.random() * 0.2
+    value = 0.9 + random.random() * 0.1
+    
+    # Convert HSV to RGB
+    import colorsys
+    rgb = colorsys.hsv_to_rgb(hue, saturation, value)
+    return f"rgb({int(rgb[0]*255)}, {int(rgb[1]*255)}, {int(rgb[2]*255)})"
+
 # Load tokenizer
 tokenizer = load_tokenizer('tokenizer.json')
 
 def encode_text(text):
     """Encode text to tokens"""
+    # Get the encoded tokens
     encoded = tokenizer.encode(text)
     decoded = tokenizer.decode(encoded)
-    return str(encoded), len(encoded), decoded, text == decoded
+    
+    # Create color-coded HTML
+    colors = {}
+    html_parts = []
+    current_pos = 0
+    
+    # Track each token's bytes and their position in the original text
+    token_bytes = []
+    for token in encoded:
+        if token < 256:
+            token_bytes.append(bytes([token]))
+        else:
+            # Recursively expand merged tokens
+            def expand_token(t):
+                if t < 256:
+                    return bytes([t])
+                pair = next((k for k, v in tokenizer.merges.items() if v == t), None)
+                if pair:
+                    return expand_token(pair[0]) + expand_token(pair[1])
+                return b''
+            
+            token_bytes.append(expand_token(token))
+    
+    # Convert bytes to text segments and color-code them
+    current_text = ''
+    for i, token_byte in enumerate(token_bytes):
+        try:
+            token_text = token_byte.decode('utf-8')
+            if token_text:
+                if encoded[i] not in colors:
+                    colors[encoded[i]] = generate_color()
+                color = colors[encoded[i]]
+                html_parts.append(f'<span style="background-color: {color};">{token_text}</span>')
+        except UnicodeDecodeError:
+            continue
+    
+    colored_text = ''.join(html_parts)
+    
+    return (
+        str(encoded),
+        len(encoded),
+        decoded,
+        text == decoded,
+        colored_text
+    )
 
 def decode_tokens(token_string):
     """Decode token sequence back to text"""
     try:
-        # Convert string representation of tokens to list of integers
         tokens = [int(t.strip()) for t in token_string.replace('[', '').replace(']', '').split(',')]
         decoded = tokenizer.decode(tokens)
         return decoded, len(tokens)
@@ -43,9 +99,14 @@ with gr.Blocks(title="Marathi BPE Tokenizer") as iface:
         with gr.Row():
             token_ids = gr.Textbox(label="Token IDs")
             token_count = gr.Number(label="Token Count")
+        
+        with gr.Row():
             decoded_text = gr.Textbox(label="Decoded Text")
             roundtrip_success = gr.Checkbox(label="Successful Round-trip")
         
+        with gr.Row():
+            colored_tokens = gr.HTML(label="Tokenized Text (Color Coded)")
+        
         # Add example inputs for encoding
         gr.Examples(
             examples=[
@@ -73,7 +134,7 @@ with gr.Blocks(title="Marathi BPE Tokenizer") as iface:
         # Add example inputs for decoding
         gr.Examples(
             examples=[
-                ["[256, 257, 258, 259]"],  # Add some actual token sequences here
+                ["[256, 257, 258, 259]"],
                 ["[260, 261, 262, 263]"],
             ],
             inputs=input_tokens
@@ -83,7 +144,7 @@ with gr.Blocks(title="Marathi BPE Tokenizer") as iface:
     encode_btn.click(
         fn=encode_text,
         inputs=input_text,
-        outputs=[token_ids, token_count, decoded_text, roundtrip_success]
+        outputs=[token_ids, token_count, decoded_text, roundtrip_success, colored_tokens]
     )
     
     decode_btn.click(
@@ -94,4 +155,4 @@ with gr.Blocks(title="Marathi BPE Tokenizer") as iface:
 
 # Launch the app
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()