Merge 3b372875c1 into bc1fca39f3

.gitignore

@@ -1,6 +1,7 @@
 .venv/
 __pycache__/
 *.pyc
+*.so
 rustbpe/target/
 dev-ignore/
 report.md

nanochat/gpt.py

@@ -8,7 +8,7 @@ Notable features:
 - norm after token embedding
 - no learnable params in rmsnorm
 - no bias in linear layers
-- Multi-Query Attention (MQA) support for more efficient inference
+- Group-Query Attention (GQA) support for more efficient inference
 """
 
 import math
@@ -29,7 +29,7 @@ class GPTConfig:
     vocab_size: int = 50304
     n_layer: int = 12
     n_head: int = 6 # number of query heads
-    n_kv_head: int = 6 # number of key/value heads (MQA)
+    n_kv_head: int = 6 # number of key/value heads (GQA)
     n_embd: int = 768
 
 

nanochat/tokenizer.py

@@ -149,7 +149,7 @@ class HuggingFaceTokenizer:
 # -----------------------------------------------------------------------------
 # Tokenizer based on rustbpe + tiktoken combo
 import pickle
-import rustbpe
+from nanochat import rustbpe
 import tiktoken
 
 class RustBPETokenizer:

pyproject.toml

@@ -23,7 +23,7 @@ requires = ["maturin>=1.7,<2.0"]
 build-backend = "maturin"
 
 [tool.maturin]
-module-name = "rustbpe"
+module-name = "nanochat.rustbpe"
 bindings = "pyo3"
 python-source = "."
 manifest-path = "rustbpe/Cargo.toml"
@@ -69,6 +69,12 @@ gpu = [
 ]  
 
 [tool.uv]
+cache-keys = [
+    { file = "pyproject.toml" },
+    { file = "rustbpe/src/**/*.rs" },
+    { file = "rustbpe/Cargo.toml" },
+    { file = "rustbpe/Cargo.lock" }
+]
 conflicts = [
     [  
         { extra = "cpu" },