save all this

Author: aduarterengifo

.direnv/flake-profile

@@ -0,0 +1 @@
+flake-profile-2-link

.direnv/flake-profile-2-link

@@ -0,0 +1 @@
+/nix/store/6iabdxd5vr4sgrqsclls76f5qd9yqrym-nix-shell-env

CLAUDE.md

@@ -0,0 +1,51 @@
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+
+## Project Overview
+
+This is a Rust/Python hybrid project that implements a Byte Pair Encoding (BPE) tokenizer. The core tokenization algorithms are implemented in Rust for performance, with Python bindings provided via PyO3. The project focuses on efficient text tokenization with support for special tokens and multiprocessing.
+
+## Key Components Architecture
+
+- **Rust Core (`src/lib.rs`)**: Contains the main BPE merge algorithm (`rusty_merge`) and helper functions. Uses PyO3 to expose functions to Python. The merge function implements an optimized version of the BPE algorithm with efficient pair counting and token updating.
+
+- **Python Interface (`test.py`)**: Main training script that orchestrates the BPE training process. Handles file chunking, multiprocessing for pre-tokenization, and calls into Rust for the merge operations. Uses regex for tokenization patterns and implements heap-based optimization.
+
+- **Build System**: Uses maturin for building the Rust-Python extension, configured for stable ABI (abi3-py38) for compatibility across Python versions.
+
+## Development Commands
+
+### Building the Extension
+```bash
+# Build the Rust extension and install in development mode
+maturin develop
+```
+
+### Running Tests/Training
+```bash
+# Run the main BPE training script
+python test.py
+```
+
+### Dependencies Management
+```bash
+# Install Python dependencies
+uv sync
+
+# Add new dependencies
+uv add <package>
+```
+
+### Development Environment
+```bash
+# Enter Nix development shell (if using Nix)
+nix develop
+```
+
+## Key Implementation Details
+
+- The Rust implementation maintains two key data structures: `pair_to_count` (frequency of byte pairs) and `pair_to_toks` (which tokens contain each pair)
+- The Python side handles file I/O, chunking for parallel processing, and pre-tokenization using regex patterns
+- Special tokens like `<|endoftext|>` are handled separately in the tokenization pipeline
+- The merge algorithm efficiently updates pair counts when tokens are merged, avoiding full recomputation

Cargo.lock

@@ -9,4 +9,11 @@ name = "rusty_tokey"
 crate-type = ["cdylib"]
 
 [dependencies]
-pyo3 = "0.25.0"
+once_cell = "1.21.3"
+rand = "0.9.0"
+regex = "1.11.1"
+
+[dependencies.pyo3]
+version = "0.25.0"
+# "abi3-py38" tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.8
+features = ["abi3-py38"]

Cargo.toml

@@ -0,0 +1,33 @@
+# rusty-tokey
+
+```
+                 _         _        _                
+ _ __ _   _ ___ | |_ _   _| |_ ___ | | _____ _   _ 
+| '__| | | / __|| __| | | | __/ _ \| |/ / _ \ | | |
+| |  | |_| \__ \| |_| |_| | || (_) |   <  __/ |_| |
+|_|   \__,_|___/ \__|\__, |\__\___/|_|\_\___|\__, |
+                     |___/                   |___/ 
+```
+
+Byte Pair Encoding tokenizer. Rust core with Python bindings.
+
+## Build
+
+```bash
+maturin develop
+```
+
+## Run
+
+```bash
+python test.py
+```
+
+## What it does
+
+- Pre-tokenizes text using regex patterns
+- Counts byte pair frequencies 
+- Merges most frequent pairs iteratively
+- Outputs vocabulary and merge operations
+
+Core algorithm in Rust for speed. Python handles I/O and orchestration.

README.md

@@ -9,12 +9,14 @@
   outputs = { self, nixpkgs, flake-utils, ... }:
     flake-utils.lib.eachDefaultSystem (system:
       let
-        pkgs = import nixpkgs { inherit system; };
+        pkgs = import nixpkgs { inherit system; config.allowUnfree = true; };
       in {
+        pkgs.config.allowUnfree = true;
         devShells.default = pkgs.mkShell {
           buildInputs = [
             pkgs.python312
             pkgs.uv
+            pkgs.claude-code
           ];
           shellHook = ''
             echo "Python: $(python --version)"

flake.lock

@@ -11,5 +11,10 @@ classifiers = [
     "Programming Language :: Python :: Implementation :: PyPy",
 ]
 dynamic = ["version"]
+dependencies = [
+    "maturin>=1.8.7",
+    "regex>=2024.11.6",
+]
 [tool.maturin]
 features = ["pyo3/extension-module"]
+python-source = "python"