jvamvas's picture
Escape HTML symbols before rendering
60aa5ef
Raw
History Blame Contribute Delete
9.15 kB
import html
from typing import Dict, List, Tuple
from pathlib import Path
from transformers import AutoTokenizer
import tiktoken
# UZH color palette
UZH_COLORS = [
"#BACBFF", # UZH Blue V1
"#DBF4F9", # UZH Cyan V1
"#ECF6D6", # UZH Apple V1
"#FFF4DA", # UZH Gold V1
"#FFDBCC", # UZH Orange V1
"#FBC6D4", # UZH Berry V1
"#C2C2C2", # UZH Grey V1
"#FAFAFA", # UZH Light Grey V1
"#7596FF", # UZH Blue V2
"#B7E9F4", # UZH Cyan V2
"#DBEDAD", # UZH Apple V2
"#FFE9B5", # UZH Gold V2
"#FEB799", # UZH Orange V2
"#F78CAA", # UZH Berry V2
"#A3A3A3", # UZH Grey V2
"#EFEFEF", # UZH Light Grey V2
"#9CB4FF", # UZH Blue V3
"#9FDCEA", # UZH Cyan V3
"#CFE58A", # UZH Apple V3
"#FFDC8A", # UZH Gold V3
"#FD9A73", # UZH Orange V3
"#F26A93", # UZH Berry V3
"#8A8A8A", # UZH Grey V3
"#E1E1E1", # UZH Light Grey V3
"#7D9BFF", # UZH Blue V4
"#75C9DB", # UZH Cyan V4
"#BEDA6C", # UZH Apple V4
"#FFD06A", # UZH Gold V4
"#F87A52", # UZH Orange V4
"#EA4D7F", # UZH Berry V4
"#757575", # UZH Grey V4
"#D2D2D2", # UZH Light Grey V4
"#5F82F4", # UZH Blue V5
"#4FB3C7", # UZH Cyan V5
"#A9CB52", # UZH Apple V5
"#FBC24E", # UZH Gold V5
"#EE643C", # UZH Orange V5
"#DC3B6E", # UZH Berry V5
"#5F5F5F", # UZH Grey V5
"#C2C2C2", # UZH Light Grey V5
"#456CE0", # UZH Blue V6
"#369DB3", # UZH Cyan V6
"#94BA3E", # UZH Apple V6
"#E9B138", # UZH Gold V6
"#DE512D", # UZH Orange V6
"#CB2C5F", # UZH Berry V6
"#4A4A4A", # UZH Grey V6
"#B3B3B3", # UZH Light Grey V6
]
def load_hf_tokenizer(name: str) -> Tuple[str, object]:
"""
Load a single HuggingFace tokenizer.
Args:
name: The name of the tokenizer to load
Returns:
Tuple of (tokenizer_name, tokenizer_object)
"""
try:
tokenizer = AutoTokenizer.from_pretrained(
name,
use_fast=True,
model_max_length=1000000,
clean_up_tokenization_spaces=True,
legacy=False
)
except Exception as e:
tokenizer = AutoTokenizer.from_pretrained(
name,
model_max_length=1000000,
clean_up_tokenization_spaces=True,
legacy=False
)
return name, tokenizer
def load_openai_tokenizer(name: str) -> Tuple[str, object]:
"""
Load a single OpenAI tokenizer.
Args:
name: The name of the tokenizer to load
Returns:
Tuple of (tokenizer_name, tokenizer_object)
"""
return name, tiktoken.encoding_for_model(name)
def load_gsw_tokenizer() -> Tuple[str, object]:
"""
Load a Swiss German (GSW) tokenizer from local vocabulary files in gsw_tokenizer directory.
Returns:
Tuple of (tokenizer_name, tokenizer_object)
"""
tokenizer = AutoTokenizer.from_pretrained("jvamvas/swissbert-gsw-vocab")
return "swissbert-gsw", tokenizer
def load_tokenizers() -> Dict[str, object]:
"""
Load all tokenizers.
Returns:
Dictionary mapping tokenizer names to tokenizer objects
"""
tokenizers = {}
# Load OpenAI tokenizers first
openai_names = ["gpt-4o"]
for name in openai_names:
tokenizer_name, tokenizer = load_openai_tokenizer(name)
tokenizers[tokenizer_name] = tokenizer
# Load HuggingFace tokenizers in specified order
hf_names = [
"meta-llama/Llama-4-Scout-17B-16E-Instruct",
"deepseek-ai/DeepSeek-V3-0324",
"ZurichNLP/swissbert",
"google/gemma-3-27b-it",
"mistralai/Mistral-Nemo-Instruct-2407",
"CohereLabs/aya-expanse-8b",
"Qwen/Qwen3-4B-Base",
]
for name in hf_names:
tokenizer_name, tokenizer = load_hf_tokenizer(name)
tokenizers[tokenizer_name] = tokenizer
return tokenizers
# Mapping of model names to display names
MODEL_DISPLAY_NAMES = {
"meta-llama/Llama-4-Scout-17B-16E-Instruct": "Llama 4",
"deepseek-ai/DeepSeek-V3-0324": "DeepSeek V3",
"ZurichNLP/swissbert": "SwissBERT 🇨🇭",
"mistralai/Mistral-Nemo-Instruct-2407": "Mistral NeMo",
"google/gemma-3-27b-it": "Gemma 3",
"gpt-4o": "ChatGPT (gpt-4o)",
"CohereLabs/aya-expanse-8b": "Aya Expanse",
"Qwen/Qwen3-4B-Base": "Qwen3 4B Base",
}
def tokenize(s: str, tokenizer) -> List[str]:
"""
Tokenize a string using any tokenizer from load_hf_tokenizers() or load_openai_tokenizers().
For SwissBERT tokenizer, compares both SwissBERT and SwissBERT-GSW tokenizations and returns the shorter one.
Args:
s: The string to tokenize
tokenizer: A tokenizer from load_hf_tokenizers() or load_openai_tokenizers()
Returns:
A list of tokens, with special tokens removed and any tail token markers (## or @@) removed
"""
# Special handling for SwissBERT tokenizer
if hasattr(tokenizer, "name_or_path") and "swissbert" in tokenizer.name_or_path.lower():
# Get SwissBERT-GSW tokenizer
_, gsw_tokenizer = load_gsw_tokenizer()
# Get tokenizations from both tokenizers§
swissbert_tokens = _tokenize_with_tokenizer(s, tokenizer)
gsw_tokens = _tokenize_with_tokenizer(s, gsw_tokenizer)
# Return the shorter tokenization
shorter_tokens = swissbert_tokens if len(swissbert_tokens) <= len(gsw_tokens) else gsw_tokens
if len(shorter_tokens) > 0 and shorter_tokens[0].startswith(" "):
shorter_tokens[0] = shorter_tokens[0][1:]
return shorter_tokens
return _tokenize_with_tokenizer(s, tokenizer)
def _tokenize_with_tokenizer(s: str, tokenizer) -> List[str]:
"""
Internal helper function to tokenize a string with a given tokenizer.
Args:
s: The string to tokenize
tokenizer: A tokenizer object
Returns:
A list of tokens, with special tokens removed and any tail token markers (## or @@) removed
"""
if hasattr(tokenizer, "tokenize"):
encoded = tokenizer.encode(s, add_special_tokens=False)
if hasattr(tokenizer, "name_or_path") and any(name in tokenizer.name_or_path.lower() for name in ["llama", "deepseek", "mistral", "aya", "qwen"]):
tokens = [tokenizer.decode([token_id], skip_special_tokens=False) for token_id in encoded]
else:
tokens = tokenizer.convert_ids_to_tokens(encoded)
filtered_tokens = []
for t in tokens:
if t.startswith("<") or t.startswith("["):
continue
elif "Ġ" in t:
filtered_tokens.append(t.replace("Ġ", " "))
elif "Ċ" in t:
filtered_tokens.append(t.replace("Ċ", " "))
elif t.startswith("▁"):
filtered_tokens.append(" " + t[1:])
else:
filtered_tokens.append(t)
return [t.rstrip("##").rstrip("@@") for t in filtered_tokens]
elif hasattr(tokenizer, "encode"):
token_ids = tokenizer.encode(s)
return [tokenizer.decode([token_id]) for token_id in token_ids]
else:
raise ValueError("Unsupported tokenizer type")
def get_uzh_color(index):
"""Get a color from the UZH color palette based on index."""
return UZH_COLORS[index % len(UZH_COLORS)]
def visualize_tokens(text: str, tokenizers: Dict[str, object]):
"""
Tokenize text with each tokenizer and visualize the tokens with colors.
Colors are consistent across tokenizers for the same token sequences.
Colors are deterministic based on token content.
Args:
text: The input text to tokenize
tokenizers: Dictionary of tokenizers
Returns:
Dictionary mapping tokenizer names to HTML visualizations
"""
results = {}
# First pass: collect all unique tokens across all tokenizers
all_tokens = set()
for tokenizer in tokenizers.values():
tokens = tokenize(text, tokenizer)
all_tokens.update(tokens)
# Generate colors for all unique tokens using hash-based approach
token_colors = {}
for token in all_tokens:
# Use hash of token to get a deterministic index
token_hash = hash(token)
# Ensure positive index and wrap around to color list length
index = abs(token_hash) % len(UZH_COLORS)
token_colors[token] = get_uzh_color(index)
# Second pass: create visualizations using the consistent colors
for name, tokenizer in tokenizers.items():
tokens = tokenize(text, tokenizer)
# Create a colored visualization
html_output = ""
# Build the HTML with colored spans for each token
for token in tokens:
color = token_colors[token]
escaped_token = html.escape(token)
html_output += f'<span style="background-color: {color}; padding: 2px; margin: 1px; border-radius: 3px;">{escaped_token}</span>'
results[name] = html_output
return results