Full Code of willkurt/token-explorer for AI

Repository: willkurt/token-explorer
Branch: main
Commit: b63a7db3aed7
Files: 27
Total size: 46.2 KB

Directory structure:
gitextract_3pk0mwhm/

├── .gitignore
├── LICENSE
├── README.md
├── config.toml
├── demo_prompt.txt
├── main.py
├── prompts/
│   ├── about.txt
│   └── date_demo.txt
├── pyproject.toml
├── simple_prompt.txt
├── src/
│   ├── .cache/
│   │   ├── state_token_map_27f9e2691b02faa4fd2065ea2c552e76.pkl
│   │   ├── state_token_map_5380414885d63dd961e6535f1b045780.pkl
│   │   ├── state_token_map_684ecf3a4ef13bc80396e045a63f3efa.pkl
│   │   ├── state_token_map_b679a835c11dec6145cf6fe92b0909bf.pkl
│   │   └── state_token_map_f20b8ac4e8144c064419a5fe50482eb5.pkl
│   ├── __init__.py
│   ├── explorer.py
│   ├── simpleguide.py
│   └── utils.py
├── struct/
│   ├── eu_date.txt
│   ├── iso_date.txt
│   ├── month_name.txt
│   ├── name.txt
│   └── us_date.txt
└── tests/
    ├── __init__.py
    ├── test_explorer.py
    └── test_simpleguide.py

================================================
FILE CONTENTS
================================================

================================================
FILE: .gitignore
================================================
# Python-generated files
__pycache__/
*.py[oc]
build/
dist/
wheels/
*.egg-info

# Virtual environments
.venv


================================================
FILE: LICENSE
================================================
MIT License

Copyright (c) 2025 Will Kurt

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.


================================================
FILE: README.md
================================================
# Token Explorer

Token Explorer allows you to interactively explore the token generation process of an LLM, using a "video game" style interface. You can use either arrow keys or vim-style navigation (h/j/k/l) along with WASD keys.

Token explore allows you to:
- Choose a starting prompt, or provide your own text file.
- Step through generation one token at a time using either:
  * Arrow keys to navigate, pop and append tokens
  * Vim-style keys: h/l to pop/append, j/k to move up/down
- View the token probabilities and entropies.
- Add a copy of your current prompt to the list of prompts.
- Cycle through the prompts by pressing `w` and `s`.
- Add and remove prompts from the list with `a` and `d`.
- Automatically uses the best available device (CUDA > MPS > CPU).


## Running the app

Token Explore uses `uv` for project management. Please see the [uv docs](https://docs.astral.sh/uv/getting-started/installation/) for more information.

Once you have `uv` installed, you can install the dependencies and run the app with:

```bash 
uv run main.py
```

In the model has a default prompt, but you can provide any text file as an argument to the app.

```bash
uv run main.py --input <path_to_file>
```

If you're using regex structs, you can precompile them with:

```bash
uv run main.py --precompile
```

## Usage

When you start the app you will see your prompt as well as a table of the top 30 tokens and their probabilities.

![Starting prompt](./imgs/starting_prompt.png)

The idea of Token Explorer is to make it very easy to explore the space of possible prompts and token generations from an LLM. To use Token Explorer it's best to treat your keyboard like you're playing a video game: put your left hand on WASD and your right hand on the arrow keys.

### Basic Usage

Use the up and down arrow keys to navigate the table. Use 'k'/'j' keys to select the current token so LLM can start generate the next one.

![Navigating the table](./imgs/navigating_table.png)

Here you can see I've highlighted the token "very" which has a probability of 0.03. Pressing the right arrow key or 'l' will append this token to the end of your prompt. Then it will display the next set of tokens.

![Appending a token](./imgs/appending_token.png)

If you want to go back and reselect the token, you can use the left arrow key or 'h' to pop the token back off the end of the prompt.

To **quit** the app, you can press `ctrl+q`.

You can also save your current prompt by pressing `x`. This will save the prompt to the `prompts` folder.

### Adding prompts

One of the goals of Token Explorer is to make it easy to play around with alternate methods of prompting. To faciliate this, Token Explorer allows you to duplicate your current prompt and add it to the list of prompts by pressing 'd'. In this image below we've added a copy of our current prompt to the list of prompts and are now at propmt 2 of 2:

![Adding a prompt](./imgs/add_a_prompt.png)

You can cycle through the prompts by pressing 'w' and 's', making it easy to try out different possible paths for your prompt, all while acting like you are the models sampler!

If you want to experiment with dramatically different prompts, you should write these out in a text file and pass them as an argument to the app.

## Visualization Layers

Token Explorer has a few different visualization layers that you can toggle on and off.

### Token Probabilities

It can be very helpful to see the probabilities of each token when generated, in part so we can see where our model might be going wrong. You can press `e` to toggle the probability view.

![Probabilities](./imgs/probabilities.png)

In the image above we've used the entire output of an LLM as our prompt. This allows us to understand better what the model was reasoning about when it generated the output. Notice for example that the model was basically certain the answer was 72.

### Token Entropies

Another way to understand the model's reasoning is to look at the entropy of the token distribution. Entropy represents how uncertain it is about the next token chosen. The highest (normalized) entropy is 1 (which means all tokens look like reasonable choices). The lowest is 0 (which means the model is certain about the next token).

You can simply press `e` again to enable the entropy view.

![Entropy](./imgs/entropies.png)

Pressing `e` again will return you back to the default view.

## Example Workflow

Let's try to understand our GSM8K prompt a bit better. The plaintext prompt is:

```
Question: Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?
Reasoning: Natalia sold 48/2 = <<48/2=24>>24 clips in May. Natalia sold 48+24 = <<48+24=72>>72 clips altogether in April and May.
Answer: 72
```

First let's understand the model's answer. We'll start by loading the prompt into Token Explorer, and then backup using the `left` arrow key until we get to the answer token.

![Backing up to the answer](./imgs/backing_up_to_answer.png)

Here we can see that the model was basically certain about the answer, which makes sense given that the prompt is a simple arithmetic problem. As we can see, the model assigns a probability of essentially 1 to the answer starting with '7'. Recall that we could also see this visually be looking at the 'probabilities' layer.

It looks like our model is doing great, but let's go back to the entropy layer to see if we can find places to explore. Notice that the token 'Natalia' has higher entropy than the other tokens, which means the model is more uncertain about which token to choose next.

![Entropy](./imgs/natalia.png)

I'm curious what's happening there. I want to back up, but at the same time, don't want to lose my place in the prompt. I can use the `d` copy my prompt as a new prompt.

![Swapping prompts](./imgs/workflow1.png)

Now I can rewind until the token 'Natalia' and see if we can understand what's happening there, while still preserving my place in the prompt.

When we look at the probabilities for the next token, we can see that the model is confused between a range of choices.

![Probabilities](./imgs/workflow2.png)

I'm curious if this part was important for our getting the correct answer. To explore we'll:

- Create a copy of this point in the prompt with 'd'
- use 'right' arrow to fill until the end.

Here's the probability view for the new prompt:

![Exploring](./imgs/workflow3.png)

You can see both that we *did* still get the correct answer, and that the path I chose was fairly low probability for a bit. So we've learned something interesting! Even if we perturb the prompt to a low-probability path in the middle of it's reasoning, we still get the correct answer!

## Experimental Support For Structured Outputs

Token Explorer now supports structured outputs!

Currently Token Explorer comes with 5 precompiled regex structs:
- eu_date
- iso_date
- us_date
- month_name
- name

You'll see your current struct at the bottom of the prompt like so:

![Struct](./imgs/current_struct.png)

By pressing `Shift+r` you can cycle through the structs.

![Toggle struct](./imgs/toggle_struct.png)

When you press `r` the structure will be activated. This will change it's color to green and you will only see tokens that adhere to that structure.

![Struct active](./imgs/struct_active.png)

The structure will eventually reach a final state where the structure is satisfied (though there may be multiple ways to satisfy the structure). In this case the named structure will be highlighted in red.

![Struct final](./imgs/struct_final.png)

Finally if you complete the defined structure, the structure will automatically be toggled off, allowing you to choose from among the unstructured tokens.

![Struct complete](./imgs/struct_complete.png)

### Adding structure

You can add your own structure by putting a .txt file in the `struct` folder. The file should contain a Python regex pattern.

**Warning:** This is an experimental feature and the regex pattern can currently take a loooong time to compile (if they're sophisticated). Once compiled they take up very little memory, and are fast to run. To precompile your structs, you can use the `--precompile` flag.

The precompile structs are stored in `src/.cache/` as `.pkl` files. 5 of them are precompiled in the repo for easy demoing.


## Configuration

The configuration is done in the `config.toml` file. The only thing you might want to change is the `model` section, which defaults to `Qwen/Qwen2.5-0.5B`. However Token Explorer is *far* from optimized for performance, so it's best to use a smaller model for now.


================================================
FILE: config.toml
================================================
# Model Configuration
[model]
name = "Qwen/Qwen2.5-0.5B"        # Model identifier

# Prompt Settings
[prompt]
example_prompt = "Once upon a time, there was a"
max_prompts = 9

# Display Settings
[display]
tokens_to_show = 30          # Number of tokens to display in preview


================================================
FILE: demo_prompt.txt
================================================
This is a story aobut

================================================
FILE: main.py
================================================
from ast import literal_eval
from itertools import cycle
from src.explorer import Explorer
from src.utils import entropy_to_color, probability_to_color
from textual.app import App, ComposeResult, Binding
from textual.containers import VerticalScroll
from textual.reactive import reactive
from textual.widgets import Footer, Header, Static, DataTable
from textwrap import dedent
import sys
import os
import argparse
import tomli 
from datetime import datetime


def load_config():
    try:
        with open("config.toml", "rb") as f:
            return tomli.load(f)
    except FileNotFoundError:
        print("Config file not found, using default values")
        return {
            "model": "Qwen/Qwen2.5-0.5B",
            "example_prompt": "Once upon a time, there was a",
            "tokens_to_show": 30,
            "max_prompts": 9
        }

config = load_config()
MODEL_NAME = config["model"]["name"]
EXAMPLE_PROMPT = config["prompt"]["example_prompt"]
TOKENS_TO_SHOW = config["display"]["tokens_to_show"]
MAX_PROMPTS = config["prompt"]["max_prompts"]

class TokenExplorer(App):
    """Main application class."""

    display_modes = cycle(["prompt", "prob", "entropy"])
    display_mode = reactive(next(display_modes))

    BINDINGS = [("e", "change_display_mode", "Mode"),
                ("left,h", "pop_token", "Back"),
                ("right,l", "append_token", "Add"),
                ("d", "add_prompt", "New"),
                ("a", "remove_prompt", "Del"),
                ("w", "increment_prompt", "Next"),
                ("s", "decrement_prompt", "Prev"),
                ("x", "save_prompt", "Save"),
                ("j", "select_next", "Down"),
                ("k", "select_prev", "Up"),
                ("r", "toggle_struct", "Toggle struct"),
                ("R", "next_struct", "Next struct")

                ]
    
    
    def __init__(self, prompt=EXAMPLE_PROMPT, precompile=False):
        super().__init__()
        # Add support for multiple prompts.
        self.prompts = [prompt]
        self.prompt_index = 0
        self.explorer = Explorer(MODEL_NAME)
        self.explorer.set_prompt(prompt)
        self.rows = self._top_tokens_to_rows(
            self.explorer.get_top_n_tokens(n=TOKENS_TO_SHOW)
            )
        self.selected_row = 0  # Track currently selected token row
        self.regex_structs = self._get_regex_structs()
        # this is the position of the stuct in the prompt
        self.struct_index = None
        # this is the position of the struct in the regex_structs list
        self.current_struct_index = 0
        if precompile:
            self.precompile_regex_structs()

    def precompile_regex_structs(self):
        print("Precompiling regex structs, this may take a while...")
        for name, regex in self.regex_structs:
            print(name)
            self.explorer.set_guide(regex)
            self.explorer.clear_guide()
        self.explorer.clear_guide()

    def _get_regex_structs(self):
        try:
            struct_files = []
            # Get all files in struct directory
            for file in os.listdir("struct"):
                if file.endswith(".txt"):
                    file_path = os.path.join("struct", file)
                    try:
                        with open(file_path, "r") as f:
                            # Get first line and strip whitespace
                            regex = f.readline().strip()
                            # Remove file extension and add tuple
                            name = os.path.splitext(file)[0]
                            struct_files.append((name, str(literal_eval(regex))))
                    except:
                        # Skip files that can't be read
                        continue
            return struct_files
        except FileNotFoundError:
            return []

    def _top_tokens_to_rows(self, tokens):
        return [("token_id", "token", "prob")] + [
            (token["token_id"], token["token"], token["probability"])
            for token in tokens
        ]
        
    def compose(self) -> ComposeResult:
        yield Header()
        with VerticalScroll():
            yield Static(id="results")
            yield DataTable(id="table")
        yield Footer()

    def _refresh_table(self):
        table = self.query_one(DataTable)
        self.rows = self._top_tokens_to_rows(
            self.explorer.get_top_n_tokens(n=TOKENS_TO_SHOW)
            )
        table.clear()
        table.add_rows(self.rows[1:])
        # Reset cursor to top
        self.selected_row = 0
        table.move_cursor(row=self.selected_row)
        self.query_one("#results", Static).update(self._render_prompt())
    

    def _render_structure_section(self):
        struct_section = ""
        if self.explorer.guide_is_finished():
            struct_section = f"[on red]{self.regex_structs[self.current_struct_index][0]}[/on]"
        elif self.struct_index is not None:
            struct_section = f"[on green]{self.regex_structs[self.current_struct_index][0]}[/on]"

        else:
            struct_section = f"[on grey]{self.regex_structs[self.current_struct_index][0]}[/on]"
        return struct_section
    
    def _render_prompt(self):
        if self.display_mode == "entropy":
            entropy_legend = "".join([
                f"[on {entropy_to_color(i/10)}] {i/10:.2f} [/on]"
                for i in range(11)
                ])
            prompt_legend = f"[bold]Token entropy:[/bold]{entropy_legend}"
            token_entropies = self.explorer.get_prompt_token_normalized_entropies()
            token_strings = self.explorer.get_prompt_tokens_strings()
            prompt_text = "".join(f"[on {entropy_to_color(entropy)}]{token}[/on]" for token, entropy in zip(token_strings, token_entropies))
        elif self.display_mode == "prob":
            prob_legend = "".join([
                f"[on {probability_to_color(i/10)}] {i/10:.2f} [/on]"
                for i in range(11)
                ])
            prompt_legend = f"[bold]Token prob:[/bold]{prob_legend}"
            token_probs = self.explorer.get_prompt_token_probabilities()
            token_strings = self.explorer.get_prompt_tokens_strings()
            prompt_text = "".join(f"[on {probability_to_color(prob)}]{token}[/on]" for token, prob in zip(token_strings, token_probs))
        else:
            prompt_text = self.explorer.get_prompt()
            prompt_legend = ""
        return dedent(f"""
{prompt_text}





{prompt_legend}
[bold]Prompt[/bold] {self.prompt_index+1}/{len(self.prompts)} tokens: {len(self.explorer.prompt_tokens)}
[bold]Struct[/bold] {self._render_structure_section()}
""")
    
    def on_mount(self) -> None:
        self.query_one("#results", Static).update(self._render_prompt())
        table = self.query_one(DataTable)
        table.add_columns(*self.rows[0])
        table.add_rows(self.rows[1:])
        table.cursor_type = "row"

    def action_next_struct(self):
        self.current_struct_index = (self.current_struct_index + 1) % len(self.regex_structs)
        self.query_one("#results", Static).update(self._render_prompt())

    def action_toggle_struct(self):
        if self.struct_index is None:
            # this is the theoretical index of the first
            # structure token when structured gen is activated
            # even though that token *doesn't* exist yet.
            # this track to help with backtracking.
            self.struct_index = len(self.explorer.get_prompt_tokens())
            self.explorer.set_guide(self.regex_structs[self.current_struct_index][1])
        else:
            self.struct_index = None
            self.explorer.clear_guide()
        self.query_one("#results", Static).update(self._render_prompt())
        self._refresh_table()
        
    def action_add_prompt(self):
        if len(self.prompts) < MAX_PROMPTS:
            self.prompts.append(self.explorer.get_prompt())
            self.prompt_index = (self.prompt_index + 1) % len(self.prompts)
            self.explorer.set_prompt(self.prompts[self.prompt_index])
            self.query_one("#results", Static).update(self._render_prompt())
            self._refresh_table()

    def action_remove_prompt(self):
        if len(self.prompts) > 1:
            self.prompts.pop(self.prompt_index)
            self.prompt_index = (self.prompt_index - 1) % len(self.prompts)
            self.explorer.set_prompt(self.prompts[self.prompt_index])
            self.query_one("#results", Static).update(self._render_prompt())
            self._refresh_table()
    
    def action_increment_prompt(self):
        self.prompt_index = (self.prompt_index + 1) % len(self.prompts)
        self.explorer.set_prompt(self.prompts[self.prompt_index])
        self.query_one("#results", Static).update(self._render_prompt())
        self._refresh_table()

    def action_decrement_prompt(self):
        self.prompt_index = (self.prompt_index - 1) % len(self.prompts)
        self.explorer.set_prompt(self.prompts[self.prompt_index])
        self.query_one("#results", Static).update(self._render_prompt())
        self._refresh_table()

    def action_change_display_mode(self):
        self.display_mode = next(self.display_modes)
        self.query_one("#results", Static).update(self._render_prompt())

    def action_save_prompt(self):
        with open(f"prompts/prompt_{self.prompt_index}_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.txt", "w") as f:
            f.write(self.explorer.get_prompt())

    def action_select_next(self):
        """Move selection down one row"""
        if self.selected_row < len(self.rows) - 2:  # -2 for header row
            self.selected_row += 1
            table = self.query_one(DataTable)
            table.move_cursor(row=self.selected_row)
            
    def action_select_prev(self):
        """Move selection up one row"""
        if self.selected_row > 0:
            self.selected_row -= 1
            table = self.query_one(DataTable)
            table.move_cursor(row=self.selected_row)

    def action_append_token(self):
        """Append currently selected token"""
        # TODO: here we need to distinguish between a dead and finished guide

        table = self.query_one(DataTable)
        if table.cursor_row is not None:
            if len(self.rows) > (table.cursor_row+1):
                self.explorer.append_token(self.rows[table.cursor_row+1][0])
                if self.explorer.guide_is_dead():
                    self.explorer.clear_guide()
                    self.struct_index = None
                self.prompts[self.prompt_index] = self.explorer.get_prompt()
                self._refresh_table()  # This will reset cursor position

    def action_pop_token(self):
        if len(self.explorer.get_prompt_tokens()) > 1:
            self.explorer.pop_token()
            if self.explorer.guide is not None:
                self.explorer.clear_guide()
                #  need to add logic for backtracking the guide
                self.explorer.set_guide(self.regex_structs[self.current_struct_index][1]
                                        ,ff_from=self.struct_index)
            self.prompts[self.prompt_index] = self.explorer.get_prompt()
            self.query_one("#results", Static).update(self._render_prompt())
            self._refresh_table()



if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Token Explorer Application')
    parser.add_argument('--input', '-i', type=str, help='Path to input text file')
    parser.add_argument('--precompile', '-p', action='store_true', help='Precompile regex structs')
    args = parser.parse_args()

    prompt = EXAMPLE_PROMPT
    if args.input:
        try:
            with open(args.input, 'r') as f:
                prompt = f.read()
        except FileNotFoundError:
            print(f"Error: Could not find input file '{args.input}'")
            sys.exit(1)
        except Exception as e:
            print(f"Error reading file: {e}")
            sys.exit(1)
    app = TokenExplorer(prompt, args.precompile)

    app.run()


================================================
FILE: prompts/about.txt
================================================
This is a directory where all of your quick save prompt end up, just in case you find a real gem!


================================================
FILE: prompts/date_demo.txt
================================================
Remember, the date of this this year's (2025) Fourth of July party is: 


================================================
FILE: pyproject.toml
================================================
[project]
name = "token-explorer"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.12"
dependencies = [
    "greenery>=4.2.2",
    "httpx>=0.28.1",
    "jupyter>=1.1.1",
    "pytest>=8.3.5",
    "textual>=2.1.2",
    "textual-dev>=1.7.0",
    "tomli>=2.2.1",
    "torch>=2.6.0",
    "transformers>=4.49.0",
]


================================================
FILE: simple_prompt.txt
================================================
Once upon a time, there was a

================================================
FILE: src/__init__.py
================================================


================================================
FILE: src/explorer.py
================================================
"""
This code is used to process an LLM one token at at time.

The Explorer class manages the prompt internally and handles all interactions with the LLM.
"""
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import numpy as np
from src.simpleguide import SimpleGuide
class Explorer:
    def __init__(self, model_name="Qwen/Qwen2.5-0.5B"):
        """
        Initialize the Explorer with a model name.
        
        Args:
            model_name: Name of the model to load (default "Qwen/Qwen2.5-0.5B")
        """
        self.model_name = model_name
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForCausalLM.from_pretrained(model_name)
        
        # Auto select device (CUDA > MPS > CPU)
        if torch.cuda.is_available():
            self.device = torch.device("cuda")
        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            self.device = torch.device("mps")
        else:
            self.device = torch.device("cpu")
        self.model = self.model.to(self.device)
        self.guide = None
        
        # Initialize with empty promp
        self.prompt_text = ""
        self.prompt_tokens = []
    

    def clear_guide(self):
        self.guide = None

    def set_guide(self, regex_struct,ff_from=None):
        self.clear_guide()
        self.guide = SimpleGuide(regex_struct, self.tokenizer)
        if ff_from is not None:
            for token in self.prompt_tokens[ff_from:]:
                self.guide.advance(token)


    def set_prompt(self, prompt_text):
        """
        Set the current prompt text and update the encoded tokens.
        
        Args:
            prompt_text: The prompt text to set
        """
        self.prompt_text = prompt_text
        self.prompt_tokens = self.tokenizer.encode(prompt_text)
        return self
    

    def get_prompt_token_probabilities(self):
        """
        Calculate the probability of each token in the sequence given its preceding context,
        using a single forward pass.
        
        Args:
            self: The Explorer object
        Returns:
            list: A list of probabilities for each token in the sequence
        """
        # Convert token IDs to tensor and create input
        input_ids = torch.tensor([self.prompt_tokens]).to(self.device)
        
        # Get the model's output in a single forward pass
        with torch.no_grad():
            outputs = self.model(input_ids)
            logits = outputs.logits[0]  # Shape: [sequence_length, vocab_size]
        
        # Calculate probabilities for each position
        token_probabilities = []
        
        # First token has no context, so we'll use None or some default
        token_probabilities.append(0.5)
        
        # For each position after the first
        for pos in range(len(self.prompt_tokens) - 1):
            # The logits at position 'pos' predict the token at position 'pos+1'
            position_logits = logits[pos]
            position_probs = torch.softmax(position_logits, dim=-1)
            
            # Get probability of the actual next token
            next_token_id = self.prompt_tokens[pos + 1]
            next_token_prob = position_probs[next_token_id].item()
            
            token_probabilities.append(next_token_prob)
        return token_probabilities
    
    def get_prompt_token_normalized_entropies(self):
        # Convert token IDs to tensor and create input
        input_ids = torch.tensor([self.prompt_tokens]).to(self.device)
        
        # Get the model's output in a single forward pass
        with torch.no_grad():
            outputs = self.model(input_ids)
            logits = outputs.logits[0]  # Shape: [sequence_length, vocab_size]
        
        # Calculate normalized entropy for each position
        normalized_entropies = []
        
        # First token has no context, so we'll use None or some default
        normalized_entropies.append(0.5)
        
        # For each position after the first
        for pos in range(len(self.prompt_tokens) - 1):
            # The logits at position 'pos' predict the token at position 'pos+1'
            position_logits = logits[pos]
            position_probs = torch.softmax(position_logits, dim=-1)
            
            # Calculate entropy: -sum(p * log(p))
            # We filter out zeros to avoid log(0) issues
            probs_np = position_probs.cpu().numpy()
            non_zero_probs = probs_np[probs_np > 0]
            entropy = -np.sum(non_zero_probs * np.log2(non_zero_probs))
            
            # Normalize by maximum possible entropy (log2 of vocabulary size)
            max_entropy = np.log2(len(position_probs))
            normalized_entropy = entropy / max_entropy
            
            normalized_entropies.append(normalized_entropy)
        
        return normalized_entropies


    def get_prompt(self):
        """
        Get the current prompt text.
        
        Returns:
            The current prompt text
        """
        return self.prompt_text
    
    def get_prompt_tokens(self):
        """
        Get the current encoded prompt tokens.
        
        Returns:
            List of token ids representing the current prompt
        """
        return self.prompt_tokens
    
    def get_prompt_tokens_strings(self):
        """
        Get the current prompt tokens as a string.
        """
        return [self.tokenizer.decode(token) for token in self.prompt_tokens]
    
    def pop_token(self):
        """
        NOTE: Need to handle the guide in this case.
        Remove and return the last token from the prompt tokens.
        If the prompt is empty, return None.
        
        Returns:
            The removed token id, or None if prompt was empty
        """
        if not self.prompt_tokens:
            return None
            
        # Pop last token and update prompt text
        last_token = self.prompt_tokens.pop()
        self.prompt_text = self.tokenizer.decode(self.prompt_tokens)
        
        return last_token
    
    def append_token(self, token_id):
        """
        Append a token to the current prompt tokens and update prompt text.
        
        Args:
            token_id: The token id to append
        """
        # Add token to prompt tokens
        self.prompt_tokens.append(token_id)
        
        # Update prompt text to match new tokens
        self.prompt_text = self.tokenizer.decode(self.prompt_tokens)
        if self.guide is not None:
            self.guide.advance(token_id)

        
        return self
    
    def guide_is_finished(self):
        if self.guide is not None:
            return self.guide.is_finished()
        return False
    
    def guide_is_dead(self):
        if self.guide is not None:
            return self.guide.is_dead()
        return False
    
    def get_top_n_tokens(self, n=5, search=""):
        """
        Get the top n most likely next tokens given the current prompt.
        Optionally filter tokens by a search string.
        
        Args:
            n: Number of top tokens to return (default 5)
            search: Optional string to filter tokens (default "")
            
        Returns:
            List of dicts containing token info and probabilities, sorted by probability
        """
        # Get model output for the encoded prompt
        with torch.no_grad():
            outputs = self.model(torch.tensor([self.prompt_tokens]).to(self.device))
            
        # Get logits for the next token
        next_token_logits = outputs.logits[0, -1, :]
        
        # Get probabilities using softmax
        next_token_probs = torch.nn.functional.softmax(next_token_logits, dim=0)

        if self.guide is not None:
            allowed_tokens = self.guide.get_tokens()
            allowed_tokens_mask = torch.zeros(len(next_token_probs), device=next_token_logits.device)
            allowed_tokens_mask[allowed_tokens] = 1.0
            next_token_probs =  next_token_probs * allowed_tokens_mask
            # renormalize the probabilities
            next_token_probs = next_token_probs / next_token_probs.sum()
        if search:
            # Filter tokens that contain the search string
            matching_tokens = []
            for idx, prob in enumerate(next_token_probs):
                token = self.tokenizer.decode(idx)
                if search.lower() in token.lower():
                    matching_tokens.append({
                        "token_id": idx,
                        "token": token,
                        "probability": prob.item()
                    })
            
            # Sort by probability and take top n
            matching_tokens.sort(key=lambda x: x["probability"], reverse=True)
            if self.guide is not None:
                # make sure that the token id is in the allowed tokens
                matching_tokens = [token for token in matching_tokens if token["token_id"] in allowed_tokens]
            return matching_tokens[:n]
        else:
            # Original behavior for no search string
            top_probs, top_indices = torch.topk(next_token_probs, n)
            
            results = []
            for prob, idx in zip(top_probs, top_indices):
                token = self.tokenizer.decode(idx)
                results.append({
                    "token": token,
                    "token_id": idx.item(),
                    "probability": prob.item()
                })
            if self.guide is not None:
                # make sure that the token id is in the allowed tokens
                results = [token for token in results if token["token_id"] in allowed_tokens]
            return results

"""
Attempting to replicate the basic api of outlines-core, but
we're going to try to reduce the memory footprint and make it more efficient.

"""

# Example usage
if __name__ == "__main__":
    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")  
    # test the RegexGuide
    guide = RegexGuide(r'a{1,5}', tokenizer)
    print("Tokens:", guide.get_tokens())
    guide.advance('a')
    print("Tokens:", guide.get_tokens())
    guide.advance('a')
    print("Tokens:", guide.get_tokens())
    explorer = Explorer()
    explorer.set_prompt("Once upon a time, there was a")
    
    print("Prompt:", explorer.get_prompt())
    print("Encoded prompt:", explorer.get_prompt_tokens())
    print("-----")
    print("Top tokens:", explorer.get_top_n_tokens())
    print("-----")
    print("Filtered tokens:", explorer.get_top_n_tokens(search="man"))
    print("-----")
    print("Appending token:", explorer.get_top_n_tokens(search="man")[0])
    explorer.append_token(explorer.get_top_n_tokens(search="man")[0]["token_id"])
    print("-----")
    print("Prompt:", explorer.get_prompt())
    print("Encoded prompt:", explorer.get_prompt_tokens())
    print("-----")
    print("Popping token:", explorer.pop_token())
    print("-----")
    print("Prompt:", explorer.get_prompt()) 
    print("Token probabilities:", explorer.get_prompt_token_probabilities())
    print("-----")
    print("Token entropies:", explorer.get_prompt_token_normalized_entropies())
    explorer.set_guide(r'a{1,5}')
    print("-----")
    print("Top tokens:", explorer.get_top_n_tokens())
    print("-----")
    print("Guide is finished:", explorer.guide.is_finished())



================================================
FILE: src/simpleguide.py
================================================
from transformers import AutoTokenizer
import re
import time
from greenery import parse
import pickle
import hashlib
import os



class TrieNode:
    def __init__(self):
        self.children = {}
        self.token_id = None
        
class TokenTrie:
    def __init__(self):
        self.root = TrieNode()
        
    def insert(self, token_meta):
        node = self.root
        word = token_meta['str']
        token_id = token_meta['id']
        for char in word:
            if char not in node.children:
                node.children[char] = TrieNode()
            node = node.children[char]
        node.token_id = token_id

    def collect_valid_tokens(self, state, fsm):
        node_state_stack = [(self.root, state)]
        valid_tokens = []
        while node_state_stack:
            node, s = node_state_stack.pop()
            for c, next_node in node.children.items():
                for cc, next_state in fsm.map[s].items():
                    if cc.accepts(c) and fsm.islive(next_state):
                        if next_node.token_id:
                            valid_tokens.append(next_node.token_id)
                        node_state_stack.append((next_node, next_state))
        return valid_tokens

class SimpleGuide:
    """
    A minimal guide for structured generation, based on the greenery library for regex parsing.
    """
    def __init__(self, regex_struct, tokenizer, no_cache=False, verbose=False):
        self.regex_struct = re.compile(regex_struct)
        start_time = time.time()
        self.pattern = parse(regex_struct)
        end_time = time.time()
        if verbose:
            print(f"Regex parsing time: {(end_time - start_time) * 1000:.2f}ms")
        start_time = time.time()
        self.fsm = self.pattern.to_fsm()
        end_time = time.time()
        if verbose:
            print(f"FSM construction time: {(end_time - start_time) * 1000:.2f}ms")
        self.tokenizer = tokenizer
        # get the string representation of the tokenizer vocabulary
        self.vocab = tokenizer.get_vocab()
        self.vocab_list = [{
            'id': value,
            'str': tokenizer.decode(value)
            } for value in self.vocab.values()]
        self.string_so_far = ""
        self.tokens_so_far = []
        self.finished = False
        self.build_token_str_trie()
        self.build_state_token_map(no_cache, verbose)
    
    def build_token_str_trie(self):
        self.token_str_trie = TokenTrie()
        for item in self.vocab_list:
            self.token_str_trie.insert(item)

    def build_state_token_map(self, no_cache=False, verbose=False):
        # Create a unique hash for this regex and tokenizer combination
        cache_key = hashlib.md5(
            f"{self.regex_struct}_{self.tokenizer.name_or_path}".encode()
        ).hexdigest()
        
        cache_dir = os.path.join(os.path.dirname(__file__), ".cache")
        cache_file = os.path.join(cache_dir, f"state_token_map_{cache_key}.pkl")
        
        # Try to load from cache first
        if os.path.exists(cache_file) and not no_cache:
            try:
                with open(cache_file, 'rb') as f:
                    self.state_token_map = pickle.load(f)
                return
            except Exception as e:
                print(f"Cache load failed: {e}")
        
        # If cache doesn't exist or fails, build the map
        start_time = time.time()
        self.state_token_map = {}
        for state in self.fsm.states:
            self.state_token_map[state] = self.token_str_trie.collect_valid_tokens(state, self.fsm)
        end_time = time.time()
        if verbose:
            print(f"State token map construction time: {(end_time - start_time) * 1000:.2f}ms")
        # Save to cache
        os.makedirs(cache_dir, exist_ok=True)
        if not no_cache:
            try:
                with open(cache_file, 'wb') as f:
                    pickle.dump(self.state_token_map, f)
            except Exception as e:
                print(f"Cache save failed: {e}")

    def get_current_state(self, candidate, state=None):
        if state is None:
            state = self.fsm.initial
        for c in candidate:
            valid_next = [state for cc, state in self.fsm.map[state].items() if cc.accepts(c) and self.fsm.islive(state)]
            if len(valid_next) == 0:
                return None
            state = valid_next[0]
        return state
    
    def is_potential_prefix(self, candidate):
        state = self.get_current_state(candidate)
        return state is not None
    
    def get_tokens(self):
        """
        Appends the token to the string_so_far (temporarily) and returns the ids of the tokens that match the current regex
        given the string so far.
        """
        # Here's where we can distinguish between in a finished state and when dead.
        if self.finished:
            return [self.tokenizer.eos_token_id]
        matching_tokens = self.state_token_map[self.get_current_state(self.string_so_far)]
        return matching_tokens 

    def advance(self, token_id):
        if token_id == self.tokenizer.eos_token_id:
            self.finished = True
            return self
        self.string_so_far += self.tokenizer.decode(token_id)
        self.tokens_so_far.append(token_id)
        return self

    def is_finished(self):
        # Might want to also check if it's dead.
        finished=self.get_current_state(self.string_so_far) in self.fsm.finals
        dead=not self.fsm.islive(self.get_current_state(self.string_so_far))
        return finished or dead
    
    def is_dead(self):
        current_state = self.get_current_state(self.string_so_far)
        live_states = [val for val in self.fsm.map[current_state].values() if self.fsm.islive(val)]
        return len(live_states) == 0
        return not self.fsm.islive(current_state)

def test_guide_loading():
    import sys
    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")
    print("Vocab size:", len(tokenizer.get_vocab()))
    start_time = time.time()
    #guide = SimpleGuide(r'(0?[1-9]|[12]\d|3[01])/(0?[1-9]|1[0-2])/\d{4}', tokenizer)
    #regex = r'\w{5} \w{5} \w{5}\n'
    regex = r'\s?(January|February|March|April|May|June|July|August|September|October|November|December)\s+(0?[1-9]|[12]\d|3[01]),\s+\d{4}'
    guide = SimpleGuide(regex, tokenizer, no_cache=True, verbose=True)
    end_time = time.time()
    loading_time = (end_time - start_time) * 1000  # Convert to milliseconds
    print(f"Guide loading time: {loading_time:.2f}ms")

    start_time = time.time()
    for _ in range(10000):
        guide.get_current_state("1")
    end_time = time.time()
    print(f"Current state time: {(end_time - start_time) * 1000 / 10000:.2f}ms")
    start_time = time.time()
    tokens = guide.get_tokens()
    end_time = time.time()
    #print([tokenizer.decode(token) for token in tokens])
    print(f"Tokens time: {(end_time - start_time) * 1000:.2f}ms")
    # get the size of the pickled guide
    print(f"SimpleGuide size: {sys.getsizeof(pickle.dumps(guide)) / 1024 / 1024:.2f}MB")

if __name__ == "__main__":
    test_guide_loading()


================================================
FILE: src/utils.py
================================================
def probability_to_color(probability, alpha=1.0):
    """
    Maps a probability value (0.0-1.0) to a color on a blue-red scale.
    Blue represents high probability (1.0)
    Red represents low probability (0.0)
    
    Args:
        probability (float): Probability value between 0.0 and 1.0
        alpha (float, optional): Alpha/opacity value between 0.0 and 1.0. Defaults to 1.0.
    
    Returns:
        str: RGBA color string (format: 'rgba(r, g, b, a)')
    """
    # Ensure probability is in valid range
    probability = max(0, min(1, probability))
    
    # Red component (high when probability is low)
    red = int(255 * (1 - probability))
    
    # Blue component (high when probability is high)
    blue = int(255 * probability)
    
    # Green component (kept at 0 for a cleaner red-blue gradient)
    green = 0
    
    # Return rgba string
    return f"rgba({red}, {green}, {blue}, {alpha})"

def entropy_to_color(entropy, alpha=1.0):
    """
    Maps a normalized entropy value (0.0-1.0) to a grayscale color.
    White (255,255,255) represents highest entropy (1.0)
    Black (0,0,0) represents lowest entropy (0.0)
    
    Args:
        entropy (float): Normalized entropy value between 0.0 and 1.0
        alpha (float, optional): Alpha/opacity value between 0.0 and 1.0. Defaults to 1.0.
    
    Returns:
        str: RGBA color string (format: 'rgba(r, g, b, a)')
    """
    # Ensure entropy is in valid range
    entropy = max(0, min(1, entropy))
    
    # For grayscale, all RGB components have the same value
    # Higher entropy = lighter color (closer to white)
    value = int(255 * entropy)
    
    # Return rgba string
    return f"rgba({value}, {value}, {value}, {alpha})"

================================================
FILE: struct/eu_date.txt
================================================
r'\s?(0?[1-9]|[12]\d|3[01])/(0?[1-9]|1[0-2])/\d{4}'

================================================
FILE: struct/iso_date.txt
================================================
r'\s?\d{4}-(?:0[1-9]|1[0-2])-(?:0[1-9]|[12]\d|3[01])'


================================================
FILE: struct/month_name.txt
================================================
r'\s?(January|February|March|April|May|June|July|August|September|October|November|December)\s+(0?[1-9]|[12]\d|3[01]),\s+\d{4}'

================================================
FILE: struct/name.txt
================================================
r'\s?[A-Z][a-z]{3,10}\s[A-Z][a-z]{4,16}'

================================================
FILE: struct/us_date.txt
================================================
r'\s?(0?[1-9]|1[0-2])/(0?[1-9]|[12]\d|3[01])/\d{4}'

================================================
FILE: tests/__init__.py
================================================


================================================
FILE: tests/test_explorer.py
================================================
from src.explorer import Explorer

def test_get_prompt_token_probabilities():
    explorer = Explorer()
    explorer.set_prompt("Hello, world")
    probabilities = explorer.get_prompt_token_probabilities()
    assert probabilities[0] == 0.5 # first token has no context
    assert len(probabilities) == len(explorer.prompt_tokens)

def test_get_top_n_tokens():
    explorer = Explorer()
    explorer.set_prompt("Hello, world")
    tokens = explorer.get_top_n_tokens(n=5)
    assert len(tokens) == 5
    assert tokens[0]["token"] == "!"

def test_guide():
    explorer = Explorer()
    explorer.set_prompt("Hello, world")
    explorer.set_guide("ba+")
    tokens = explorer.get_top_n_tokens(n=5)
    assert len(tokens) == 2 # actually only 2 valid tokens
    assert tokens[0]["token"][0] == "b"

def test_guide_append_token():
    explorer = Explorer()
    tokenizer = explorer.tokenizer
    explorer.set_prompt("Hello, world")
    explorer.set_guide("abc")
    tokens = explorer.get_top_n_tokens(n=5)
    assert tokens[0]["token"][0] == "a"
    explorer.append_token(tokenizer.encode("a")[0])
    tokens = explorer.get_top_n_tokens(n=5)
    assert tokens[0]["token"][0] == "b"

================================================
FILE: tests/test_simpleguide.py
================================================
from src.simpleguide import SimpleGuide
from transformers import AutoTokenizer

def test_get_tokens():
    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")
    # test a basic regex
    guide = SimpleGuide("a+", tokenizer)
    decoded_tokens = [tokenizer.decode(token) for token in guide.get_tokens()]
    for token in decoded_tokens:
        for i in range(len(token)):
            assert token[i] == "a"

def test_advance():
    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")
    guide = SimpleGuide("abbbc", tokenizer)
    guide.advance(tokenizer.encode("a")[0])
    decoded_tokens = [tokenizer.decode(token) for token in guide.get_tokens()]
    for token in decoded_tokens:
        assert token[0] == "b"

def test_is_finished_single_finish():
    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")
    guide = SimpleGuide("abc", tokenizer)
    guide.advance(tokenizer.encode("a")[0])
    assert not guide.is_finished()
    guide.advance(tokenizer.encode("b")[0])
    assert not guide.is_finished()
    guide.advance(tokenizer.encode("c")[0])
    assert guide.is_finished()

def test_is_finished_multiple_finish():
    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")
    guide = SimpleGuide("abc{1,2}", tokenizer)
    guide.advance(tokenizer.encode("a")[0])
    assert not guide.is_finished()
    guide.advance(tokenizer.encode("b")[0])
    assert not guide.is_finished()
    guide.advance(tokenizer.encode("c")[0])
    assert guide.is_finished()
    guide.advance(tokenizer.encode("c")[0])
    assert guide.is_finished()

def test_is_dead():
    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")
    guide = SimpleGuide("abc", tokenizer)
    guide.advance(tokenizer.encode("a")[0])
    assert not guide.is_dead()
    guide.advance(tokenizer.encode("b")[0])
    assert not guide.is_dead()
    guide.advance(tokenizer.encode("c")[0])
    assert guide.is_dead()

def test_is_dead_multiple_finish():
    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")
    guide = SimpleGuide("abc{1,2}", tokenizer)
    guide.advance(tokenizer.encode("a")[0])
    assert not guide.is_dead()
    guide.advance(tokenizer.encode("b")[0])
    assert not guide.is_dead()
    guide.advance(tokenizer.encode("c")[0])
    assert not guide.is_dead() and guide.is_finished()
    guide.advance(tokenizer.encode("c")[0])
    assert guide.is_dead()

def test_spaces():
    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")
    guide = SimpleGuide(" ab", tokenizer, no_cache=True)
    assert any([cc.accepts(' ') for cc in guide.fsm.map[0].keys()])
    assert tokenizer.encode(" ")[0] in guide.state_token_map[0]