llmwalk 0.1.2__tar.gz

llmwalk-0.1.2/.gitignore

@@ -0,0 +1,176 @@
+# Created by https://www.toptal.com/developers/gitignore/api/python
+# Edit at https://www.toptal.com/developers/gitignore?templates=python
+
+### Python ###
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+### Python Patch ###
+# Poetry local configuration file - https://python-poetry.org/docs/configuration/#local-configuration
+poetry.toml
+
+# ruff
+.ruff_cache/
+
+# LSP config files
+pyrightconfig.json
+
+# End of https://www.toptal.com/developers/gitignore/api/python

llmwalk-0.1.2/AGENTS.md

@@ -0,0 +1,12 @@
+## Project notes for Codex/agents
+
+This repo is a single-file `uv` inline script.
+
+- Entry point: `llmtree.py`
+- Dependencies are declared in the `# /// script` block at the top of `llmtree.py`
+- Prefer running via `uv` so dependencies resolve correctly:
+  - `uv run llmtree.py -- --help`
+  - `uv run llmtree.py -- -p "Your prompt here"`
+
+Avoid using `pip install` directly for this repo unless you have a specific reason; the intended workflow is `uv run` with the inline dependency block.
+

llmwalk-0.1.2/PKG-INFO

@@ -0,0 +1,42 @@
+Metadata-Version: 2.4
+Name: llmwalk
+Version: 0.1.2
+Summary: Explore the answer-space for any prompt and any MLX-supported model.
+Requires-Python: >=3.10
+Requires-Dist: mlx-lm==0.28.4
+Requires-Dist: rich==14.2.0
+Requires-Dist: sortedcontainers==2.4.0
+Description-Content-Type: text/markdown
+
+# llmwalk
+
+Explore the answer-space for any prompt and any MLX-supported model. See
+<https://huggingface.co/mlx-community/models> for supported models.
+
+![Usage example gif](example1.gif)
+
+Instead of sampling from the possible tokens each step, llmwalk branches out
+and completes all of the branches the sampler would consider based on
+`--top-k`, `--top-p` and `--temperature`, ranking the results by probability
+as it goes.
+
+The tree is walked prioritising the most likely branches, until it finds `-n`
+branches and then it stops. It doesn't enumerate all possibilities, just enough
+to know for sure it has found the `-n` most likely branches.
+
+## Usage
+
+- `uvx llmwalk -p "In what year was Barack Obama born?"`
+- `uvx llmwalk -p "Write a haiku about compilers" -n 5`
+- `uvx llmwalk -p "Give me one word: " --top-k 200 --temperature 0.7`
+
+## Options
+
+- `-p, --prompt TEXT`: Prompt to score (wrapped with the model’s chat template).
+- `-m, --model MODEL`: MLX-LM model identifier or path (default: `mlx-community/Llama-3.2-1B-Instruct-4bit`), supported models can be found at <https://huggingface.co/mlx-community/models>
+- `-n N`: Number of answers to show. The search stops once it has `N` finished answers and no unfinished branch can beat the worst of those `N`.
+- `--min-probability FLOAT`: Any branch whose cumulative probability falls below this is marked finished (`low_probability`) and not expanded further.
+- `--top-k INT`: At each step, expand at most `k` next tokens (highest probability).
+- `--top-p FLOAT`: Nucleus cutoff applied *within the top-k tokens* at each step (keep adding tokens until cumulative probability ≥ `p`).
+- `--temperature FLOAT`: Softmax temperature applied when computing per-step probabilities (`1.0` is the model distribution; must be `> 0`).
+- `--stats-interval SECONDS`: How often to refresh the live view (`<= 0` disables periodic refresh; still renders at start/end).

llmwalk-0.1.2/README.md

@@ -0,0 +1,32 @@
+# llmwalk
+
+Explore the answer-space for any prompt and any MLX-supported model. See
+<https://huggingface.co/mlx-community/models> for supported models.
+
+![Usage example gif](example1.gif)
+
+Instead of sampling from the possible tokens each step, llmwalk branches out
+and completes all of the branches the sampler would consider based on
+`--top-k`, `--top-p` and `--temperature`, ranking the results by probability
+as it goes.
+
+The tree is walked prioritising the most likely branches, until it finds `-n`
+branches and then it stops. It doesn't enumerate all possibilities, just enough
+to know for sure it has found the `-n` most likely branches.
+
+## Usage
+
+- `uvx llmwalk -p "In what year was Barack Obama born?"`
+- `uvx llmwalk -p "Write a haiku about compilers" -n 5`
+- `uvx llmwalk -p "Give me one word: " --top-k 200 --temperature 0.7`
+
+## Options
+
+- `-p, --prompt TEXT`: Prompt to score (wrapped with the model’s chat template).
+- `-m, --model MODEL`: MLX-LM model identifier or path (default: `mlx-community/Llama-3.2-1B-Instruct-4bit`), supported models can be found at <https://huggingface.co/mlx-community/models>
+- `-n N`: Number of answers to show. The search stops once it has `N` finished answers and no unfinished branch can beat the worst of those `N`.
+- `--min-probability FLOAT`: Any branch whose cumulative probability falls below this is marked finished (`low_probability`) and not expanded further.
+- `--top-k INT`: At each step, expand at most `k` next tokens (highest probability).
+- `--top-p FLOAT`: Nucleus cutoff applied *within the top-k tokens* at each step (keep adding tokens until cumulative probability ≥ `p`).
+- `--temperature FLOAT`: Softmax temperature applied when computing per-step probabilities (`1.0` is the model distribution; must be `> 0`).
+- `--stats-interval SECONDS`: How often to refresh the live view (`<= 0` disables periodic refresh; still renders at start/end).

llmwalk-0.1.2/llmwalk/__init__.py

@@ -0,0 +1,3 @@
+__all__ = ["__version__"]
+
+__version__ = "0.1.0"

llmwalk-0.1.2/llmwalk/__main__.py

@@ -0,0 +1,6 @@
+from __future__ import annotations
+
+from .cli import main
+
+if __name__ == "__main__":
+    main()

llmwalk-0.1.2/llmwalk/cli.py

@@ -0,0 +1,412 @@
+from __future__ import annotations
+
+import argparse
+import heapq
+import sys
+import time
+from dataclasses import dataclass
+from datetime import datetime
+from functools import lru_cache
+
+import mlx.core as mx
+from mlx.nn import Module
+from mlx_lm import load
+from mlx_lm.models.cache import KVCache
+from mlx_lm.tokenizer_utils import TokenizerWrapper
+from rich.console import Console, Group
+from rich.live import Live
+from rich.style import Style
+from rich.table import Table
+from rich.text import Text
+from sortedcontainers import SortedList
+
+args: argparse.Namespace
+
+_BAND_COLORS = [
+    "#7f7f7f",  # 0-10%: grey
+    "#ff3b30",  # 10-20%: red
+    "#ff6a00",  # 20-30%: orange
+    "#ff8c00",  # 30-40%: dark orange
+    "#ffb000",  # 40-50%: amber
+    "#ffd000",  # 50-60%: yellow
+    "#d7e500",  # 60-70%: yellow-green
+    "#a8e600",  # 70-80%: greenish
+    "#4cd964",  # 80-90%: green
+    "#00c853",  # 90-100%: bright green
+]
+_BAND_STYLES = [Style(color=c) for c in _BAND_COLORS]
+
+
+@dataclass
+class OutputToken:
+    token: int
+    prob: float
+
+
+@dataclass(eq=False)
+class Branch:
+    parent: Branch | None
+    token: OutputToken | None
+    probability: float = 1.0
+    finish_reason: str | None = None
+    cache: list[KVCache] | None = None
+
+    def answer_tokens(self) -> list[OutputToken]:
+        toks: list[OutputToken] = []
+        cur: Branch | None = self
+        while cur is not None and cur.token is not None:
+            toks.append(cur.token)
+            cur = cur.parent
+        toks.reverse()
+        return toks
+
+
+def _clone_kv_cache(c: KVCache) -> KVCache:
+    cloned = KVCache()
+    cloned.offset = c.offset
+    cloned.keys = mx.array(c.keys) if c.keys is not None else None
+    cloned.values = mx.array(c.values) if c.values is not None else None
+    return cloned
+
+
+def _clone_prompt_cache(cache: list[KVCache]) -> list[KVCache]:
+    return [_clone_kv_cache(c) for c in cache]
+
+
+def _top_tokens_from_logprobs(logprobs: mx.array) -> list[OutputToken]:
+    vocab = int(logprobs.shape[0])
+    k = min(args.top_k, vocab)
+    part = mx.argpartition(logprobs, vocab - k)
+    top_idx = part[vocab - k :]
+    top_lp = mx.take(logprobs, top_idx)
+    order = mx.argsort(top_lp)[::-1]
+    sorted_indices = mx.take(top_idx, order)
+
+    if args.temperature == 1.0:
+        probs = mx.exp(mx.take(logprobs, sorted_indices))
+    else:
+        lse = mx.logsumexp(logprobs / args.temperature, axis=-1)
+        probs = mx.exp(mx.take(logprobs, sorted_indices) / args.temperature - lse)
+
+    mx.eval(sorted_indices, probs)
+    token_ids: list[int] = sorted_indices.astype(mx.int64).tolist()
+    token_probs: list[float] = mx.reshape(probs, (-1,)).tolist()
+
+    output_tokens: list[OutputToken] = []
+    cum_prob = 0.0
+    for token_id, prob in zip(token_ids, token_probs):  # type: ignore[call-arg]
+        if output_tokens and cum_prob >= args.top_p:
+            break
+        output_tokens.append(OutputToken(token=token_id, prob=float(prob)))
+        cum_prob += float(prob)
+
+    return output_tokens
+
+
+class PromptTreeSearch:
+    model: Module
+    tokenizer: TokenizerWrapper
+    prompt: list[int]
+    _frontier: list[tuple[float, int, Branch]]
+    _finished_eos: SortedList[Branch]
+    _heap_counter: int = 0
+    _stopped: bool = False
+
+    tokens: int = 0
+    pruned: int = 0
+
+    _low_watermark: float | None = None
+    _start: datetime | None = None
+    _end: datetime | None = None
+
+    def __init__(self, model: Module, tokenizer: TokenizerWrapper, prompt: list[int]) -> None:
+        self.model = model
+        self.tokenizer = tokenizer
+        self.prompt = prompt
+        self._frontier = []
+        self._finished_eos = SortedList(key=lambda b: -b.probability)
+
+        root = Branch(parent=None, token=None)
+        self.branches = SortedList(key=lambda b: -b.probability)
+        self.branches.add(root)
+        self._push_frontier(root)
+
+    @lru_cache(maxsize=65536)
+    def decode_token(self, token_id: int) -> str:
+        return self.tokenizer.decode([token_id], skip_special_tokens=True)  # type: ignore[call-arg]
+
+    def _run_model(self, cache: list[KVCache], input_ids: list[int]) -> mx.array:
+        inputs = mx.array([input_ids], mx.int32)
+        logits = self.model(inputs, cache=cache)[:, -1, :]
+        logits = logits.astype(mx.float32)
+        logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
+        return mx.reshape(logprobs, (-1,))
+
+    @property
+    def active(self) -> int:
+        return len(self._frontier)
+
+    def top_branches(self, n: int) -> list[Branch]:
+        return list(self.branches[:n])
+
+    def _push_frontier(self, branch: Branch) -> None:
+        self._heap_counter += 1
+        heapq.heappush(self._frontier, (-branch.probability, self._heap_counter, branch))
+
+    def _update_low_watermark(self) -> None:
+        if len(self._finished_eos) < args.n:
+            self._low_watermark = None
+            return
+        self._low_watermark = self._finished_eos[args.n - 1].probability
+
+    def stop(self) -> None:
+        self._stopped = True
+
+    def should_stop(self) -> bool:
+        if self._stopped:
+            return True
+        if not self._frontier:
+            return True
+        if self._low_watermark is None:
+            return False
+        best_prob = -self._frontier[0][0]
+        return best_prob < self._low_watermark
+
+    def begin(self) -> None:
+        self._start = datetime.now()
+
+    def step(self) -> None:
+        if self.should_stop():
+            return
+
+        _, _, branch = heapq.heappop(self._frontier)
+
+        if self._low_watermark is not None and branch.probability < self._low_watermark:
+            self.pruned += 1
+            branch.finish_reason = "pruned"
+            branch.cache = None
+            return
+
+        if branch.token is None:  # root branch
+            cache_after = self.model.make_cache() if hasattr(self.model, "make_cache") else []  # type: ignore[assignment]
+            logprobs = self._run_model(cache_after, self.prompt)
+        else:
+            if branch.cache is None:
+                raise RuntimeError("Branch cache missing while expanding leaf")
+            cache_after = _clone_prompt_cache(branch.cache)
+            logprobs = self._run_model(cache_after, [branch.token.token])
+
+        self.branches.remove(branch)
+
+        new_branches: list[Branch] = []
+        frontier_add: list[Branch] = []
+        eos_add: list[Branch] = []
+        for tok in _top_tokens_from_logprobs(logprobs):
+            new_prob = branch.probability * tok.prob
+
+            if new_prob < args.min_probability:
+                self.pruned += 1
+                new_branch = Branch(
+                    parent=branch,
+                    token=tok,
+                    probability=new_prob,
+                    finish_reason="low_probability",
+                )
+                new_branches.append(new_branch)
+                continue
+
+            if tok.token in self.tokenizer.eos_token_ids:
+                new_branch = Branch(
+                    parent=branch,
+                    token=tok,
+                    probability=new_prob,
+                    finish_reason="eos_token",
+                )
+                new_branches.append(new_branch)
+                eos_add.append(new_branch)
+                continue
+
+            new_branch = Branch(
+                parent=branch,
+                token=tok,
+                probability=new_prob,
+                cache=cache_after,
+            )
+            new_branches.append(new_branch)
+            frontier_add.append(new_branch)
+
+        for b in new_branches:
+            self.branches.add(b)
+        for b in frontier_add:
+            self._push_frontier(b)
+        for b in eos_add:
+            self._finished_eos.add(b)
+        if eos_add:
+            self._update_low_watermark()
+        self.tokens += 1
+
+        branch.cache = None
+
+
+def style_for_token_probability(prob: float) -> Style:
+    if prob != prob:  # NaN
+        prob = 0.0
+    elif prob < 0.0:
+        prob = 0.0
+    elif prob > 1.0:
+        prob = 1.0
+
+    band = min(int(prob * 10), 9)  # 0..9
+    return _BAND_STYLES[band]
+
+
+def render_probability_legend() -> Text:
+    legend = Text("Legend: ", style="bold", no_wrap=True, overflow="ellipsis")
+    for i in range(9, -1, -1):
+        style = style_for_token_probability((i + 0.5) / 10)
+        if i == 9:
+            label = "90%+"
+        elif i == 0:
+            label = "0–10%"
+        else:
+            label = f"{i * 10}%+"
+
+        legend.append("■", style=style)
+        legend.append(f" {label}")
+        if i != 0:
+            legend.append("  ")
+
+    return legend
+
+
+_PROBABILITY_LEGEND = render_probability_legend()
+
+
+def render_branches(walker: PromptTreeSearch) -> Table:
+    table = Table(expand=True)
+    table.add_column("Fin", justify="center", no_wrap=True, width=3)
+    table.add_column("Prob.", justify="right", no_wrap=True, width=8)
+    table.add_column("Answer", ratio=1)
+
+    branches = walker.top_branches(args.n)
+    for i in range(args.n):
+        if i >= len(branches):
+            table.add_row("", "", "")
+            continue
+
+        branch = branches[i]
+        answer_text = Text()
+        for tok in branch.answer_tokens():
+            piece = walker.decode_token(tok.token)
+            if not piece:
+                continue
+            answer_text.append(piece, style=style_for_token_probability(tok.prob))
+        probability_text = f"{branch.probability * 100:6.2f}%"
+        status: Text
+        if branch.finish_reason == "eos_token":
+            status = Text("✓", style="green")
+        elif branch.finish_reason == "low_probability":
+            status = Text("✓", style="yellow")
+        elif branch.finish_reason == "pruned":
+            status = Text("✓", style="dim")
+        else:
+            status = Text("")
+
+        table.add_row(status, probability_text, answer_text)
+
+    return table
+
+
+def render_stats_bar(walker: PromptTreeSearch) -> Table:
+    elapsed = (datetime.now() - walker._start).total_seconds() if walker._start else 0.0
+    tps = walker.tokens / elapsed if elapsed > 0 else 0.0
+    left = f"active {walker.active}  pruned {walker.pruned} tps {tps:0.1f}"
+    grid = Table.grid(expand=True)
+    grid.add_column(ratio=1)
+    grid.add_column(justify="right", no_wrap=True)
+    grid.add_row(
+        Text(left, overflow="ellipsis", no_wrap=True),
+        Text(
+            f"top_k={args.top_k}  top_p={args.top_p}  temp={args.temperature}",
+            no_wrap=True,
+        ),
+    )
+    return grid
+
+
+def render_view(walker: PromptTreeSearch) -> Group:
+    return Group(
+        _PROBABILITY_LEGEND,
+        render_branches(walker),
+        render_stats_bar(walker),
+    )
+
+
+def run() -> None:
+    load_resp = load(args.model)
+    model = load_resp[0]
+    tokenizer = load_resp[1]
+
+    prompt = tokenizer.apply_chat_template(  # type: ignore[call-arg]
+        [{"role": "user", "content": args.prompt}],
+        add_generation_prompt=True,
+    )
+
+    console = Console()
+
+    walker = PromptTreeSearch(model, tokenizer, prompt)
+    walker.begin()
+
+    try:
+        with Live(console=console, transient=False) as live:
+            interval = max(0.1, args.stats_interval)
+            next_render = time.monotonic()
+            live.update(render_view(walker))
+            while not walker.should_stop():
+                walker.step()
+                if args.stats_interval > 0 and time.monotonic() >= next_render:
+                    live.update(render_view(walker))
+                    next_render = time.monotonic() + interval
+            live.update(render_view(walker))
+    except KeyboardInterrupt:
+        walker.stop()
+
+
+def parse_args(argv: list[str] | None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-p", "--prompt", default="What is 2+2?", help="Prompt to score")
+    parser.add_argument("-m", "--model", default="mlx-community/Llama-3.2-1B-Instruct-4bit")
+    parser.add_argument("-n", default=10, type=int, help="Number of answers to show")
+    parser.add_argument("--min-probability", type=float, default=0.0001)
+    parser.add_argument("--top-k", dest="top_k", default=50, type=int)
+    parser.add_argument(
+        "--top-p",
+        dest="top_p",
+        default=1.0,
+        type=float,
+        help="Nucleus sampling threshold (0 < p <= 1)",
+    )
+    parser.add_argument("--temperature", type=float, default=1.0, help="Softmax temperature (> 0)")
+    parser.add_argument(
+        "--stats-interval",
+        type=float,
+        default=0.1,
+        help="Seconds between live stats bar updates (<=0 disables)",
+    )
+
+    raw = list(sys.argv[1:] if argv is None else argv)
+    filtered = [a for a in raw if a != "--"]
+    parsed = parser.parse_args(filtered)
+
+    if parsed.temperature <= 0:
+        parser.error("--temperature must be > 0")
+    if not (0 < parsed.top_p <= 1):
+        parser.error("--top-p must be in the range (0, 1]")
+
+    return parsed
+
+
+def main(argv: list[str] | None = None) -> None:
+    global args
+    args = parse_args(argv)
+    run()

llmwalk-0.1.2/pyproject.toml

@@ -0,0 +1,26 @@
+[build-system]
+requires = ["hatchling>=1.22.0"]
+build-backend = "hatchling.build"
+
+[project]
+name = "llmwalk"
+version = "0.1.2"
+description = "Explore the answer-space for any prompt and any MLX-supported model."
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+  "mlx-lm==0.28.4",
+  "rich==14.2.0",
+  "sortedcontainers==2.4.0",
+]
+
+[project.scripts]
+llmwalk = "llmwalk.cli:main"
+
+[tool.hatch.build]
+exclude = [
+  "example1.gif",
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["llmwalk"]