added openai and cohere models support along with token visuvalizations

app.py

@@ -4,11 +4,24 @@ import ast
 from collections import Counter
 import re
 import plotly.graph_objs as go
+import html
+import random
+import tiktoken
+import anthropic
 
 model_path = "models/"
 
 # Available models
-MODELS = ["Meta-Llama-3.1-8B", "gemma-2b"]
+MODELS = ["Meta-Llama-3.1-8B", "gemma-2b", "gpt-3.5-turbo","gpt-4","gpt-4o" "Claude-3-Sonnet"]
+openai_models = ["gpt-3.5-turbo","gpt-4","gpt-4o"]
+# Color palette visible on both light and dark themes
+COLOR_PALETTE = [
+    "#e6194B", "#3cb44b", "#ffe119", "#4363d8",
+    "#f58231", "#911eb4", "#42d4f4", "#f032e6",
+    "#bfef45", "#fabed4", "#469990", "#dcbeff",
+    "#9A6324", "#fffac8", "#800000", "#aaffc3",
+    "#808000", "#ffd8b1", "#000075", "#a9a9a9"
+]
 
 def create_vertical_histogram(data, title):
     labels, values = zip(*data) if data else ([], [])
@@ -25,32 +38,80 @@ def create_vertical_histogram(data, title):
     )
     return fig
 
-def process_text(text:str,model_name):
-    tokenizer = AutoTokenizer.from_pretrained(model_path + model_name)
-    token_ids = tokenizer.encode(text, add_special_tokens=True)
-    tokens = tokenizer.convert_ids_to_tokens(token_ids)
-    return text,tokens,token_ids,
+def process_text(text: str, model_name: str, api_key: str = None):
+    if model_name in ["Meta-Llama-3.1-8B", "gemma-2b"]:
+        tokenizer = AutoTokenizer.from_pretrained(model_path + model_name)
+        token_ids = tokenizer.encode(text, add_special_tokens=True)
+        tokens = tokenizer.convert_ids_to_tokens(token_ids)
+    elif model_name in openai_models:
+        encoding = tiktoken.encoding_for_model(model_name=model_name)
+        token_ids = encoding.encode(text)
+        tokens = [encoding.decode([id]) for id in token_ids]
+    elif model_name == "Claude-3-Sonnet":
+        if not api_key:
+            raise ValueError("API key is required for Claude models")
+        client = anthropic.Anthropic(api_key=api_key)
+        tokenizer = client.get_tokenizer()
+        token_ids = tokenizer.encode(text).ids
+        tokens = [tokenizer.decode([id]) for id in token_ids]
+    else:
+        raise ValueError(f"Unsupported model: {model_name}")
+    
+    return text, tokens, token_ids
 
-def process_ids(ids:str,model_name):
-    tokenizer = AutoTokenizer.from_pretrained(model_path + model_name)
+def process_ids(ids: str, model_name: str, api_key: str = None):
     token_ids = ast.literal_eval(ids)
-    text = tokenizer.decode(token_ids)
-    tokens = tokenizer.convert_ids_to_tokens(token_ids)
-    return text,tokens,token_ids
+    if model_name in ["Meta-Llama-3.1-8B", "gemma-2b"]:
+        tokenizer = AutoTokenizer.from_pretrained(model_path + model_name)
+        text = tokenizer.decode(token_ids)
+        tokens = tokenizer.convert_ids_to_tokens(token_ids)
+    elif model_name == openai_models:
+        encoding = tiktoken.encoding_for_model(model_name=model_name)
+        text = encoding.decode(token_ids)
+        tokens = [encoding.decode([id]) for id in token_ids]
+    elif model_name == "Claude-3-Sonnet":
+        client = anthropic.Anthropic(api_key=api_key)
+        tokenizer = client.get_tokenizer()
+        text = tokenizer.decode(token_ids)
+        tokens = [tokenizer.decode([id]) for id in token_ids]
+    else:
+        raise ValueError(f"Unsupported model: {model_name}")
+    
+    return text, tokens, token_ids
 
-def process_input(input_type, input_value, model_name):
-    
+def get_token_color(token, token_colors):
+    if token.startswith('<') and token.endswith('>'):
+        return "#42d4f4"  # Cyan for special tokens
+    elif token == '▁' or token == ' ':
+        return "#3cb44b"  # Green for space tokens
+    elif not token.isalnum():
+        return "#f032e6"  # Magenta for special characters
+    else:
+        if token not in token_colors:
+            token_colors[token] = random.choice(COLOR_PALETTE)
+        return token_colors[token]
+
+def create_html_tokens(tokens):
+    html_output = '<div style="font-family: monospace; border: 1px solid #ccc; padding: 10px; border-radius: 5px; background-color: #f9f9f9; white-space: pre-wrap; word-break: break-all;">'
+    token_colors = {}
+    for token in tokens:
+        color = get_token_color(token, token_colors)
+        escaped_token = html.escape(token)
+        html_output += f'<span style="background-color: {color}; color: black; padding: 2px 4px; margin: 1px; border-radius: 3px; display: inline-block;">{escaped_token}</span>'
+    html_output += '</div>'
+    return html_output
+
+def process_input(input_type, input_value, model_name, api_key):
     if input_type == "Text":
-        text,tokens,token_ids = process_text(text=input_value,model_name=model_name)
+        text, tokens, token_ids = process_text(text=input_value, model_name=model_name, api_key=api_key)
     elif input_type == "Token IDs":
-        text,tokens,token_ids = process_ids(ids=input_value,model_name=model_name)
+        text, tokens, token_ids = process_ids(ids=input_value, model_name=model_name, api_key=api_key)
 
     character_count = len(text)
     word_count = len(text.split())
     
-    
-    space_count = sum(1 for token in tokens if token == '▁')
-    special_char_count = sum(1 for token in tokens if not token.isalnum() and token != '▁')
+    space_count = sum(1 for token in tokens if token in ['▁', ' '])
+    special_char_count = sum(1 for token in tokens if not token.isalnum() and token not in ['▁', ' '])
     
     words = re.findall(r'\b\w+\b', text.lower())
     special_chars = re.findall(r'[^\w\s]', text)
@@ -71,7 +132,9 @@ def process_input(input_type, input_value, model_name):
     analysis += f"Special character tokens: {special_char_count}\n"
     analysis += f"Other tokens: {len(tokens) - space_count - special_char_count}"
     
-    return analysis, text,tokens, str(token_ids), words_hist, special_chars_hist, numbers_hist
+    html_tokens = create_html_tokens(tokens)
+    
+    return analysis, text, html_tokens, str(token_ids), words_hist, special_chars_hist, numbers_hist
 
 def text_example():
     return "Hello, world! This is an example text input for tokenization."
@@ -85,8 +148,9 @@ with gr.Blocks() as iface:
     
     with gr.Row():
         input_type = gr.Radio(["Text", "Token IDs"], label="Input Type", value="Text")
-        model_name = gr.Dropdown(choices=MODELS, label="Select Model",value=MODELS[0])
+        model_name = gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[0])
     
+    api_key = gr.Textbox(label="API Key Claude models)", type="password")
     input_text = gr.Textbox(lines=5, label="Input")
     
     with gr.Row():
@@ -96,8 +160,8 @@ with gr.Blocks() as iface:
     submit_button = gr.Button("Process")
     
     analysis_output = gr.Textbox(label="Analysis", lines=6)
-    text_output = gr.Textbox(label="Text",lines=6)
-    tokens_output = gr.Textbox(label="Tokens", lines=3)
+    text_output = gr.Textbox(label="Text", lines=6)
+    tokens_output = gr.HTML(label="Tokens")
     token_ids_output = gr.Textbox(label="Token IDs", lines=2)
     
     with gr.Row():
@@ -117,8 +181,8 @@ with gr.Blocks() as iface:
     
     submit_button.click(
         process_input,
-        inputs=[input_type, input_text, model_name],
-        outputs=[analysis_output,text_output ,tokens_output, token_ids_output, words_plot, special_chars_plot, numbers_plot]
+        inputs=[input_type, input_text, model_name, api_key],
+        outputs=[analysis_output, text_output, tokens_output, token_ids_output, words_plot, special_chars_plot, numbers_plot]
     )
 
 if __name__ == "__main__":

requirements.txt

@@ -1,11 +1,13 @@
 aiofiles==23.2.1
 annotated-types==0.7.0
+anthropic==0.34.1
 anyio==4.4.0
 certifi==2024.7.4
 charset-normalizer==3.3.2
 click==8.1.7
 contourpy==1.2.1
 cycler==0.12.1
+distro==1.9.0
 fastapi==0.112.2
 ffmpy==0.4.0
 filelock==3.15.4
@@ -20,6 +22,7 @@ huggingface-hub==0.24.6
 idna==3.8
 importlib_resources==6.4.4
 Jinja2==3.1.4
+jiter==0.5.0
 kiwisolver==1.4.5
 markdown-it-py==3.0.0
 MarkupSafe==2.1.5
@@ -51,6 +54,7 @@ six==1.16.0
 sniffio==1.3.1
 starlette==0.38.2
 tenacity==9.0.0
+tiktoken==0.7.0
 tokenizers==0.19.1
 tomlkit==0.12.0
 tqdm==4.66.5