import os import gradio as gr import torch import itertools # For color cycling import tiktoken # For GPT-4 tokenizer from transformers import AutoTokenizer # For Llama3 tokenizer - AutoModel usually not needed just for tokenizer # Bytelatent imports (assuming they are in the python path) try: from bytelatent.data.file_util import get_fs from bytelatent.generate_patcher import patcher_nocache from bytelatent.tokenizers.blt_tokenizer import BltTokenizer from bytelatent.plotting.entropy_figure_via_matplot_lib import plot_entropies from bytelatent.args import TrainArgs from download_blt_weights import main as ensure_present BLT_AVAILABLE = True except ImportError as e: print(f"Warning: Bytelatent libraries not found. Bytelatent functionality will be disabled. Error: {e}") BLT_AVAILABLE = False # Define dummy classes/functions if BLT is not available to avoid NameErrors later class BltTokenizer: pass class TrainArgs: pass def patcher_nocache(*args, **kwargs): return None def plot_entropies(*args, **kwargs): return None def ensure_present(*args, **kwargs): pass # --- Global Setup --- # Define colors for patches/tokens VIZ_COLORS = [ "#a6cee3", "#1f78b4", "#b2df8a", "#33a02c", "#fb9a99", "#e31a1c", "#fdbf6f", "#ff7f00", "#cab2d6", "#6a3d9a", "#ffff99", "#b15928" ] # Add more if you expect many segments LLAMA3_MODEL_NAME = "meta-llama/Meta-Llama-3-8B" # Or choose another variant like Instruct # --- Helper Functions --- def create_bytelatent_highlight_data(tokenizer, patch_lengths_tensor, tokens_tensor, colors): """Generates data for gr.HighlightedText based on bytelatent patches.""" if not BLT_AVAILABLE: return [("Bytelatent library not available.", "Error")] if patch_lengths_tensor is None or tokens_tensor is None or patch_lengths_tensor.numel() == 0: return None patch_lengths = patch_lengths_tensor.tolist() all_tokens = tokens_tensor.tolist() highlighted_data = [] current_token_index = 0 patch_count = 0 # Initialize patch count # color_cycler = itertools.cycle(colors) # Moved inside loop if needed per-patch for i, length in enumerate(patch_lengths): if length <= 0: continue patch_token_ids = all_tokens[current_token_index : current_token_index + length] if not patch_token_ids: continue try: patch_text = tokenizer.decode(patch_token_ids) except Exception as decode_err: print(f"Warning: Bytelatent patch decoding failed: {decode_err}") patch_text = f"[Decode Error: {len(patch_token_ids)} tokens]" patch_label = f"BL Patch {i+1}" highlighted_data.append((patch_text, patch_label)) patch_count += 1 # Increment count for each valid patch added current_token_index += length # Handle remainder separately, don't count it as a 'patch' if current_token_index != len(all_tokens): print(f"Warning: Bytelatent token mismatch. Consumed {current_token_index}, total {len(all_tokens)}") remaining_tokens = all_tokens[current_token_index:] if remaining_tokens: try: remaining_text = tokenizer.decode(remaining_tokens) except Exception: remaining_text = f"[Decode Error: {len(remaining_tokens)} remaining tokens]" highlighted_data.append((remaining_text, "BL Remainder")) # Return both highlighted data and the calculated patch count return highlighted_data, patch_count def create_tiktoken_highlight_data(prompt, colors): """Generates data for gr.HighlightedText based on tiktoken (gpt-4) tokens.""" try: enc = tiktoken.get_encoding("cl100k_base") tiktoken_ids = enc.encode(prompt) highlighted_data = [] # color_cycler = itertools.cycle(colors) # Moved inside loop if needed per-token for i, token_id in enumerate(tiktoken_ids): try: token_text = enc.decode([token_id]) except UnicodeDecodeError: try: token_bytes = enc.decode_single_token_bytes(token_id) token_text = f"[Bytes: {token_bytes.hex()}]" except Exception: token_text = "[Decode Error]" except Exception as e: print(f"Unexpected tiktoken decode error: {e}") token_text = "[Decode Error]" token_label = f"GPT4 Tk {i+1}" highlighted_data.append((token_text, token_label)) token_count = len(tiktoken_ids) print(f"Tiktoken processing complete. Found {token_count} tokens.") return highlighted_data, token_count except ImportError: print("Error: tiktoken library not found. Please install it: pip install tiktoken") return [("tiktoken library not installed.", "Error")], 0 except Exception as tiktoken_err: print(f"Error during tiktoken processing: {tiktoken_err}") return [(f"Error processing with tiktoken: {str(tiktoken_err)}", "Error")], 0 def create_llama3_highlight_data(prompt, colors, model_name=LLAMA3_MODEL_NAME): """Generates data for gr.HighlightedText based on Llama 3 tokenizer.""" try: # Load Llama 3 tokenizer from Hugging Face Hub print(f"Loading Llama 3 tokenizer: {model_name}") # Use trust_remote_code=True if required by the specific model revision tokenizer = AutoTokenizer.from_pretrained(model_name) #, trust_remote_code=True) print("Llama 3 tokenizer loaded.") # Encode the prompt llama_token_ids = tokenizer.encode(prompt) highlighted_data = [] # color_cycler = itertools.cycle(colors) # Moved inside loop if needed per-token for i, token_id in enumerate(llama_token_ids): try: # Decode individual token. token_text = tokenizer.decode([token_id]) except Exception as e: print(f"Unexpected Llama 3 decode error for token {token_id}: {e}") token_text = "[Decode Error]" token_label = f"Llama3 Tk {i+1}" # Clearer label prefix highlighted_data.append((token_text, token_label)) token_count = len(llama_token_ids) print(f"Llama 3 processing complete. Found {token_count} tokens.") return highlighted_data, token_count except ImportError: print("Error: transformers or sentencepiece library not found. Please install them: pip install transformers sentencepiece") return [("transformers/sentencepiece library not installed.", "Error")], 0 except OSError as e: # Handle errors like model not found, network issues, authentication needed print(f"Error loading Llama 3 tokenizer '{model_name}': {e}") error_msg = f"Could not load Llama 3 tokenizer '{model_name}'. Check model name and network." if "authentication" in str(e).lower(): error_msg = f"Authentication required for Llama 3 tokenizer '{model_name}'. Use `huggingface-cli login`." return [(f"{error_msg} Error: {e}", "Error")], 0 except Exception as llama_err: print(f"Error during Llama 3 processing: {llama_err}") import traceback traceback.print_exc() # Print full traceback for debugging return [(f"Error processing with Llama 3: {str(llama_err)}", "Error")], 0 # --- Main Processing Function --- def process_text(prompt: str, model_name: str = "blt-1b"): """ Processes the input prompt using ByteLatent, Tiktoken, and Llama 3, returning visualizations, counts, and status. Args: prompt: The input text string from the Gradio interface. model_name: The name of the bytelatent model to use. Returns: A tuple containing: - Matplotlib Figure for the entropy plot (or None). - List of tuples for bytelatent gr.HighlightedText (or None). - Integer count of bytelatent patches. - List of tuples for tiktoken gr.HighlightedText (or None). - Integer count of tiktoken tokens. - List of tuples for Llama 3 gr.HighlightedText (or None). - Integer count of Llama 3 tokens. - Status/Error message string. """ fig = None bl_highlighted_data = None tk_highlighted_data = None llama_highlighted_data = None bl_count = 0 tk_count = 0 llama_count = 0 status_message = "Starting processing..." # --- 1. Tiktoken Processing (Independent) --- status_message += "\nProcessing with Tiktoken (gpt-4)..." tk_highlighted_data, tk_count_calc = create_tiktoken_highlight_data(prompt, VIZ_COLORS) if tk_highlighted_data and tk_highlighted_data[0][1] == "Error": status_message += f"\nTiktoken Error: {tk_highlighted_data[0][0]}" tk_count = 0 # Ensure count is 0 on error else: tk_count = tk_count_calc # Assign calculated count status_message += f"\nTiktoken processing successful ({tk_count} tokens)." # --- 2. Llama 3 Processing (Independent) --- status_message += "\nProcessing with Llama 3 tokenizer..." llama_highlighted_data, llama_count_calc = create_llama3_highlight_data(prompt, VIZ_COLORS) if llama_highlighted_data and llama_highlighted_data[0][1] == "Error": status_message += f"\nLlama 3 Error: {llama_highlighted_data[0][0]}" llama_count = 0 # Ensure count is 0 on error else: llama_count = llama_count_calc # Assign calculated count status_message += f"\nLlama 3 processing successful ({llama_count} tokens)." # --- 3. Bytelatent Processing --- if BLT_AVAILABLE: try: status_message += f"\nLoading Bytelatent entropy model for '{model_name}'..." # (Bytelatent loading code remains the same) consolidated_path = os.path.join("hf-weights", model_name) train_args_path = os.path.join(consolidated_path, "params.json") if not os.path.exists(train_args_path): raise FileNotFoundError(f"BLT training args not found at {train_args_path}.") fs = get_fs(train_args_path); train_args = TrainArgs.model_validate_json(fs.read_text(train_args_path)) bl_tokenizer = train_args.data.tokenizer_args.build(); assert isinstance(bl_tokenizer, BltTokenizer) patcher_args = train_args.data.patcher_args.model_copy(deep=True); patcher_args.realtime_patching = True device = "cuda" if torch.cuda.is_available() else "cpu"; print(f"Using BLT device: {device}") patcher_args.patching_device = device; patcher_args.device = device entropy_model_dir = os.path.join(consolidated_path, "entropy_model") if not os.path.exists(entropy_model_dir): raise FileNotFoundError(f"Entropy model directory not found at {entropy_model_dir}.") patcher_args.entropy_model_checkpoint_dir = entropy_model_dir; bl_patcher = patcher_args.build() status_message += "\nBytelatent entropy model loaded." # --- Processing --- status_message += "\nRunning Bytelatent entropy model patching..." print(f"Processing prompt with entropy model: '{prompt}'") prompt_bytes = prompt.encode('utf-8') max_bytes = 512 # Define max bytes if len(prompt_bytes) > max_bytes: print(f"Warning: Prompt exceeds {max_bytes} bytes ({len(prompt_bytes)}). Truncating for entropy model.") # Find the byte position that corresponds to the last full character within the limit # This avoids splitting a multi-byte character try: last_char_pos = prompt_bytes[:max_bytes].rfind(b' ') # Simple whitespace split point find, might not be perfect if last_char_pos == -1: # If no space, truncate hard (less ideal) prompt_bl = prompt_bytes[:max_bytes].decode('utf-8', errors='ignore') else: prompt_bl = prompt_bytes[:last_char_pos].decode('utf-8', errors='ignore') except Exception: # Fallback to simple truncation on decode errors prompt_bl = prompt_bytes[:max_bytes].decode('utf-8', errors='ignore') status_message += f"\nWarning: Prompt truncated to approx {len(prompt_bl.encode('utf-8'))} bytes for Bytelatent entropy model." else: prompt_bl = prompt results = patcher_nocache([prompt_bl], tokenizer=bl_tokenizer, patcher=bl_patcher) if not results: print("Bytelatent entropy processing returned no results.") status_message += "\nBytelatent entropy model warning: Processing completed, but no results were generated." bl_highlighted_data = [("No patches generated.", "Info")] bl_count = 0 else: batch_patch_lengths, batch_scores, batch_tokens = results patch_lengths, scores, tokens = batch_patch_lengths[0], batch_scores[0], batch_tokens[0] # --- Visualization Data Generation --- try: decoded_output_for_plot = bl_tokenizer.decode(tokens.tolist()) except Exception as decode_err: print(f"Warning: Error decoding full sequence for plot: {decode_err}") decoded_output_for_plot = prompt_bl # Use truncated prompt for plot if decode fails fig = plot_entropies(patch_lengths, scores, decoded_output_for_plot, threshold=bl_patcher.threshold) bl_highlighted_data, bl_count_calc = create_bytelatent_highlight_data(bl_tokenizer, patch_lengths, tokens, VIZ_COLORS) bl_count = bl_count_calc # Assign calculated count status_message += f"\nBytelatent entropy model processing and visualization successful ({bl_count} patches)." print("Bytelatent Entropy model processing and decoding complete.") except FileNotFoundError as e: print(f"Bytelatent Error: {e}") status_message += f"\nBytelatent FileNotFoundError: {str(e)}" bl_highlighted_data = [(f"Bytelatent Error: {e}", "Error")] bl_count = 0 except Exception as e: print(f"An unexpected Bytelatent error occurred: {e}") import traceback traceback.print_exc() status_message += f"\nBytelatent Unexpected Error: {str(e)}" bl_highlighted_data = [(f"Bytelatent Error: {e}", "Error")] bl_count = 0 else: status_message += "\nBytelatent processing skipped (library not found)." bl_highlighted_data = [("Bytelatent library not available.", "Error")] bl_count = 0 fig = None # Ensure fig is None if BLT is skipped # Return all generated data and the final status message return fig, bl_highlighted_data, bl_count, tk_highlighted_data, tk_count, llama_highlighted_data, llama_count, status_message # --- Gradio Interface --- # Create color maps for HighlightedText dynamically MAX_EXPECTED_SEGMENTS = 2000 # Increased max segments further just in case common_error_map = {"Error": "#FF0000", "Info": "#808080"} # Red for errors, Gray for info bytelatent_color_map = {f"BL Patch {i+1}": color for i, color in zip(range(MAX_EXPECTED_SEGMENTS), itertools.cycle(VIZ_COLORS))} bytelatent_color_map["BL Remainder"] = "#AAAAAA"; bytelatent_color_map.update(common_error_map) tiktoken_color_map = {f"GPT4 Tk {i+1}": color for i, color in zip(range(MAX_EXPECTED_SEGMENTS), itertools.cycle(VIZ_COLORS))} tiktoken_color_map.update(common_error_map) llama3_color_map = {f"Llama3 Tk {i+1}": color for i, color in zip(range(MAX_EXPECTED_SEGMENTS), itertools.cycle(VIZ_COLORS))} llama3_color_map.update(common_error_map) with gr.Blocks(theme=gr.themes.Origin()) as iface: gr.Markdown("# BLT's Entropy-based Patcher vs. Tokenizer Visualisation") gr.Markdown( "Enter text to visualize its segmentation according to different methods:\n" "1. **Byte Latent Transformer (BLT):** Entropy-based patching plot and patched text (_for this space ONLY_ - limited to ~512 bytes).\n" "2. **Tiktoken (GPT-4):** Text segmented by `cl100k_base` tokens.\n" f"3. **Llama 3:** Text segmented by the `{LLAMA3_MODEL_NAME}` tokenizer." ) with gr.Row(): with gr.Column(scale=1): # Input Column prompt_input = gr.Textbox( label="Input Prompt", value="Daenerys Targaryen is in Game of Thrones, a fantasy epic by George R.R. Martin.", placeholder="Enter text here...", max_length=512, # Allow even longer input, Bytelatent will truncate lines=5, info="For this space ONLY, processing is limited to ~512 bytes." ) submit_button = gr.Button("Generate Visualizations", variant="primary") status_output = gr.Textbox(label="Processing Status", interactive=False, lines=7) # Increased lines slightly with gr.Column(scale=2): # Output Column # --- Bytelatent Output Area --- with gr.Row(equal_height=False): # Use Row to place title and count together gr.Markdown("### BLT Entropy Patcher Output (`blt_main_entropy_100m_512w`)") bl_count_output = gr.Number(label="Patch Count", value=0, interactive=False, scale=1, step=1) # Added Number output highlighted_output_bl = gr.HighlightedText( label="BLT's Entropy-based Patches", color_map=bytelatent_color_map, show_legend=False, # Legend can get very long # show_label=False, # Hide the HighlightedText label as we have the markdown title show_inline_category=False, # container=False, # Reduces vertical space slightly ) plot_output = gr.Plot(label="Entropy vs. Token Index", show_label=True) # --- Tiktoken Output Area --- with gr.Row(equal_height=False): gr.Markdown("### Tiktoken Output (`cl100k_base`)") tk_count_output = gr.Number(label="Token Count", value=0, interactive=False, scale=1, step=1) # Added Number output highlighted_output_tk = gr.HighlightedText( label="Tiktoken Segmented Text", color_map=tiktoken_color_map, show_legend=False, show_inline_category=False, # show_label=False, # container=False, ) # --- Llama 3 Output Area --- with gr.Row(equal_height=False): gr.Markdown(f"### Llama 3 Output (`{LLAMA3_MODEL_NAME}`)") llama_count_output = gr.Number(label="Token Count", value=0, interactive=False, scale=1, step=1) # Added Number output highlighted_output_llama = gr.HighlightedText( label="Llama 3 Segmented Text", color_map=llama3_color_map, show_legend=False, show_inline_category=False, # show_label=False, # container=False, ) # Define the action for the button click submit_button.click( fn=process_text, inputs=prompt_input, # Ensure order matches the 8 return values of process_text outputs=[ plot_output, # fig highlighted_output_bl, # bl_highlighted_data bl_count_output, # bl_count <-- New highlighted_output_tk, # tk_highlighted_data tk_count_output, # tk_count <-- New highlighted_output_llama,# llama_highlighted_data llama_count_output, # llama_count <-- New status_output # status_message ] ) # --- Launch the Gradio App --- if __name__ == "__main__": print("Checking required libraries...") try: import tiktoken print("- tiktoken found.") except ImportError: print("WARNING: 'tiktoken' not found. GPT-4 visualization will fail. Install with: pip install tiktoken") try: import transformers import sentencepiece print("- transformers found.") print("- sentencepiece found.") except ImportError: print("WARNING: 'transformers' or 'sentencepiece' not found. Llama 3 visualization will fail. Install with: pip install transformers sentencepiece") if BLT_AVAILABLE: print("- Bytelatent libraries found.") # Ensure bytelatent model is present only if library is available try: print("Ensuring Bytelatent model 'blt-1b' weights are present...") ensure_present(["blt-1b"]) print("Bytelatent model check complete.") except Exception as blt_dl_err: print(f"WARNING: Failed to ensure Bytelatent model presence: {blt_dl_err}") else: print("INFO: Bytelatent libraries not found, skipping related functionality.") print(f"Attempting to use Llama 3 Tokenizer: {LLAMA3_MODEL_NAME}. Ensure you have access (e.g., via `huggingface-cli login` if needed).") print("Launching Gradio interface...") iface.launch()