Merge 3b372875c1 into 4a87a0d19f

Fix comment: rotary embeddings final dimension size

.gitignore

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

nanochat/gpt.py

@@ -244,7 +244,7 @@ class GPT(nn.Module):
     def forward(self, idx, targets=None, kv_cache=None, loss_reduction='mean'):
         B, T = idx.size()
 
-        # Grab the rotary embeddings for the current sequence length (they are of shape (1, seq_len, 1, head_dim))
+        # Grab the rotary embeddings for the current sequence length (they are of shape (1, seq_len, 1, head_dim/2))
         assert T <= self.cos.size(1), f"Sequence length grew beyond the rotary embeddings cache: {T} > {self.cos.size(1)}"
         assert idx.device == self.cos.device, f"Rotary embeddings and idx are on different devices: {idx.device} != {self.cos.device}"
         assert self.cos.dtype == torch.bfloat16, "Rotary embeddings must be in bfloat16"

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"
@@ -67,9 +67,15 @@ cpu = [
 gpu = [  
     "torch>=2.8.0",  
 ]  
-  
-[tool.uv]  
-conflicts = [  
+
+[tool.uv]
+cache-keys = [
+    { file = "pyproject.toml" },
+    { file = "rustbpe/src/**/*.rs" },
+    { file = "rustbpe/Cargo.toml" },
+    { file = "rustbpe/Cargo.lock" }
+]
+conflicts = [
     [  
         { extra = "cpu" },  
         { extra = "gpu" },