app.py

import os
import gradio as gr
import torch
import itertools # For color cycling
import tiktoken # For GPT-4 tokenizer
from transformers import AutoTokenizer, AutoModel # For Llama3 tokenizer

# Bytelatent imports (assuming they are in the python path)
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

# --- 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."""
    # (Keep the function from the previous version - no changes needed)
    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
    color_cycler = itertools.cycle(colors)
    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))
        current_token_index += length
    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 highlighted_data


def create_tiktoken_highlight_data(prompt, colors):
    """Generates data for gr.HighlightedText based on tiktoken (gpt-4) tokens."""
    # (Keep the function from the previous version - no changes needed)
    try:
        enc = tiktoken.get_encoding("cl100k_base")
        tiktoken_ids = enc.encode(prompt)
        highlighted_data = []
        color_cycler = itertools.cycle(colors)
        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))
        print(f"Tiktoken processing complete. Found {len(tiktoken_ids)} tokens.")
        return highlighted_data
    except ImportError:
         print("Error: tiktoken library not found. Please install it: pip install tiktoken")
         return [("tiktoken library not installed.", "Error")]
    except Exception as tiktoken_err:
        print(f"Error during tiktoken processing: {tiktoken_err}")
        return [(f"Error processing with tiktoken: {str(tiktoken_err)}", "Error")]


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
        # This might download the tokenizer files on the first run
        # May require `huggingface-cli login` if model is private or gated
        print(f"Loading Llama 3 tokenizer: {model_name}")
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        print("Llama 3 tokenizer loaded.")

        # Encode the prompt
        llama_token_ids = tokenizer.encode(prompt)

        highlighted_data = []
        color_cycler = itertools.cycle(colors)

        for i, token_id in enumerate(llama_token_ids):
            try:
                # Decode individual token. Llama/SentencePiece tokenizers usually handle this well.
                token_text = tokenizer.decode([token_id])
                # Special case: Handle potential leading space added by sentencepiece during decode
                # if token_text.startswith(' '): # Check if this improves visualization
                #      token_text = token_text[1:] # Remove leading space visual artifact? Test this.
            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))

        print(f"Llama 3 processing complete. Found {len(llama_token_ids)} tokens.")
        return highlighted_data

    except ImportError:
         print("Error: transformers or sentencepiece library not found. Please install them: pip install transformers sentencepiece")
         return [("transformers/sentencepiece library not installed.", "Error")]
    except OSError as e:
        # Handle errors like model not found, network issues, authentication needed
        print(f"Error loading Llama 3 tokenizer '{model_name}': {e}")
        if "authentication" in str(e).lower():
             return [(f"Authentication required for Llama 3 tokenizer '{model_name}'. Use `huggingface-cli login`.", "Error")]
        else:
            return [(f"Could not load Llama 3 tokenizer '{model_name}'. Check model name and network. Error: {e}", "Error")]
    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")]


# --- 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 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).
            - List of tuples for tiktoken gr.HighlightedText (or None).
            - List of tuples for Llama 3 gr.HighlightedText (or None).
            - Status/Error message string.
    """
    fig = None
    bl_highlighted_data = None
    tk_highlighted_data = None
    llama_highlighted_data = None
    status_message = "Starting processing..."

    # --- 1. Tiktoken Processing (Independent) ---
    status_message += "\nProcessing with Tiktoken (gpt-4)..."
    tk_highlighted_data = 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]}"
    else:
         status_message += "\nTiktoken processing successful."

    # --- 2. Llama 3 Processing (Independent) ---
    status_message += "\nProcessing with Llama 3 tokenizer..."
    llama_highlighted_data = 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]}"
    else:
         status_message += "\nLlama 3 processing successful."

    # --- 3. Bytelatent Processing ---
    try:
        status_message += f"\nLoading entropy model for '{model_name}'..."
        # (Bytelatent loading code remains the same as previous version)
        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"Bytelatent 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 Bytelatent 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"Bytelatent 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 model loaded."

        # --- Processing ---
        status_message += "\nRunning Bytelatent patching..."
        print(f"Processing prompt with Bytelatent: '{prompt}'")
        # Limit prompt length for bytelatent if necessary
        prompt_bytes = prompt.encode('utf-8')
        if len(prompt_bytes) > 512:
             print(f"Warning: Prompt exceeds 512 bytes ({len(prompt_bytes)}). Truncating for Bytelatent.")
             prompt_bl = prompt_bytes[:512].decode('utf-8', errors='ignore')
             status_message += "\nWarning: Prompt truncated to 512 bytes for Bytelatent."
        else:
             prompt_bl = prompt

        results = patcher_nocache([prompt_bl], tokenizer=bl_tokenizer, patcher=bl_patcher)

        if not results:
            print("Bytelatent processing returned no results.")
            status_message += "\nBytelatent Warning: Processing completed, but no results were generated."
        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 = create_bytelatent_highlight_data(bl_tokenizer, patch_lengths, tokens, VIZ_COLORS)
            status_message += "\nBytelatent processing and visualization successful."
            print("Bytelatent processing and decoding complete.")

    except FileNotFoundError as e:
        print(f"Bytelatent Error: {e}")
        status_message += f"\nBytelatent FileNotFoundError: {str(e)}"
    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)}"

    # Return all generated data and the final status message
    return fig, bl_highlighted_data, tk_highlighted_data, llama_highlighted_data, status_message


# --- Gradio Interface ---

# Create color maps for HighlightedText dynamically
MAX_EXPECTED_SEGMENTS = 1000 # Increase max expected segments further
common_error_map = {"Error": "#FF0000"} # Red for errors

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"] = "#808080"; 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.Soft()) as iface:
    gr.Markdown("# BLT's Entropy Patcher Visualisation") # Updated Title
    gr.Markdown(
        "Enter text to visualize its segmentation according to different tokenizers:\n"
        "1.  **BLT:** Entropy plot and text segmented by dynamic patches (Input limited to 512 bytes).\n"
        "2.  **Tiktoken (GPT-4):** Text segmented by `cl100k_base` tokens.\n"
        "3.  **Llama 3:** Text segmented by the `meta-llama/Meta-Llama-3-8B` 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=2048, # Allow even longer input, Bytelatent will truncate
                lines=5,
                info="Processing is limited to the first 512 bytes of the input."
            )
            submit_button = gr.Button("Generate Visualizations", variant="primary")
            status_output = gr.Textbox(label="Processing Status", interactive=False, lines=5)

         with gr.Column(scale=2): # Output Column
             gr.Markdown("### BLT's Entropy Patcher Output (`100m`)")
             highlighted_output_bl = gr.HighlightedText(
                 label="Bytelatent Patched Text",
                 color_map=bytelatent_color_map,
                 show_legend=False, # Legend can get very long, disable for compactness
                 show_inline_category=False,
             )
             plot_output = gr.Plot(label="Bytelatent Entropy vs. Token Index")

             gr.Markdown("### Tiktoken Output (`cl100k_base` for GPT-4)")
             highlighted_output_tk = gr.HighlightedText(
                 label="Tiktoken Segmented Text",
                 color_map=tiktoken_color_map,
                 show_legend=False,
                 show_inline_category=False,
             )

             gr.Markdown(f"### Llama 3 Output (`{LLAMA3_MODEL_NAME}`)")
             highlighted_output_llama = gr.HighlightedText(
                 label="Llama 3 Segmented Text",
                 color_map=llama3_color_map,
                 show_legend=False,
                 show_inline_category=False,
             )

    # Define the action for the button click
    submit_button.click(
        fn=process_text,
        inputs=prompt_input,
        # Ensure order matches the 5 return values of process_text
        outputs=[
            plot_output,
            highlighted_output_bl,
            highlighted_output_tk,
            highlighted_output_llama,
            status_output
            ]
    )

# --- Launch the Gradio App ---
if __name__ == "__main__":
    print("Please ensure 'tiktoken', 'transformers', and 'sentencepiece' are installed (`pip install tiktoken transformers sentencepiece`)")
    print(f"Attempting to use Llama 3 Tokenizer: {LLAMA3_MODEL_NAME}. Ensure you have access (e.g., via `huggingface-cli login` if needed).")
    ensure_present(["blt-1b"]) # Ensure bytelatent model is present
    iface.launch()