Upload inspect_needle.py with huggingface_hub

inspect_needle.py

@@ -0,0 +1,488 @@
+"""
+inspect_needle.py — Introspect the Cactus/Needle Flax model, tokenizer, and prompt format.
+
+Outputs three markdown files:
+  ../notes/needle-internals.md
+  ../notes/tokenizer-internals.md
+  ../notes/prompt-format.md
+"""
+
+import json
+import os
+import sys
+import textwrap
+import inspect
+import traceback
+
+# Add the needle package to path
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'external', 'needle'))
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+# ── 1. Import needle symbols ──────────────────────────────────────────────────
+from needle.model.architecture import (
+    SimpleAttentionNetwork, TransformerConfig,
+    MultiHeadAttention, FeedForward, EncoderBlock, DecoderBlock,
+    Encoder, Decoder, ZCRMSNorm, make_causal_mask, make_padding_mask,
+)
+from needle.model.run import generate, _build_encoder_input, load_checkpoint
+from needle.dataset.tokenizer import (
+    NeedleTokenizer, get_tokenizer,
+    PAD_ID, EOS_ID, BOS_ID, UNK_ID, TOOL_CALL_ID, TOOLS_ID,
+    DEFAULT_MAX_ENC_LEN, DEFAULT_MAX_DEC_LEN, DEFAULT_MAX_GEN_LEN,
+    TOKENIZER_PREFIX,
+)
+
+print("=== Needle/Cactus Inspector ===\n")
+
+# ── 2. Instantiate model with default config ──────────────────────────────────
+config = TransformerConfig()
+print(f"TransformerConfig defaults:")
+for field in config.__dataclass_fields__:
+    print(f"  {field} = {getattr(config, field)}")
+
+model = SimpleAttentionNetwork(config)
+
+# Use init_all to initialize all parameters (both text + contrastive paths)
+dummy_src = jnp.zeros((1, 16), dtype=jnp.int32)
+dummy_tgt = jnp.zeros((1, 16), dtype=jnp.int32)
+variables = model.init(jax.random.PRNGKey(0), dummy_src, dummy_tgt, method='init_all')
+params = variables['params']
+
+print("\n=== Param Tree (paths + shapes) ===")
+
+def print_param_tree(tree, prefix='', indent=0):
+    """Recursively print param tree with shapes."""
+    lines = []
+    if isinstance(tree, dict):
+        for k, v in sorted(tree.items()):
+            if isinstance(v, dict):
+                lines.append('  ' * indent + f'{k}:')
+                lines.extend(print_param_tree(v, prefix=prefix+k+'.', indent=indent+1))
+            else:
+                shape = tuple(v.shape) if hasattr(v, 'shape') else repr(v)
+                dtype = str(v.dtype) if hasattr(v, 'dtype') else '?'
+                lines.append('  ' * indent + f'{k}: shape={shape}, dtype={dtype}')
+    return lines
+
+param_lines = print_param_tree(params)
+for line in param_lines:
+    print(line)
+
+# Build shape dict for JSON serialization
+def build_shape_dict(tree):
+    if isinstance(tree, dict):
+        return {k: build_shape_dict(v) for k, v in sorted(tree.items())}
+    else:
+        return {'shape': list(tree.shape), 'dtype': str(tree.dtype)}
+
+shape_dict = build_shape_dict(params)
+
+# Count parameters
+total_params = sum(x.size for x in jax.tree.leaves(params))
+print(f"\nTotal parameters: {total_params:,}")
+
+# ── 3. Tokenizer info ─────────────────────────────────────────────────────────
+print("\n=== Tokenizer ===")
+tok = get_tokenizer()
+print(f"Class: {type(tok).__name__}")
+print(f"vocab_size: {tok.vocab_size}")
+print(f"pad_token_id: {tok.pad_token_id}")
+print(f"eos_token_id: {tok.eos_token_id}")
+print(f"bos_token_id: {tok.bos_token_id}")
+print(f"tool_call_token_id: {tok.tool_call_token_id}")
+print(f"tools_token_id: {tok.tools_token_id}")
+print(f"attrs: {[a for a in dir(tok) if not a.startswith('_')]}")
+
+# Test encode/decode
+sample_text = "set a 5 min timer"
+encoded = tok.encode(sample_text)
+decoded = tok.decode(encoded)
+print(f"\nExample encode('{sample_text}'):")
+print(f"  -> {encoded}")
+print(f"  -> decode -> '{decoded}'")
+
+# Get vocab size from SentencePiece directly
+sp = tok.sp
+print(f"\nSentencePiece GetPieceSize(): {sp.GetPieceSize()}")
+print(f"Model type: BPE (confirmed in train_tokenizer source)")
+print(f"byte_fallback: True (confirmed in train_tokenizer source)")
+print(f"normalization_rule_name: identity (no Unicode normalization)")
+
+# ── 4. Prompt format capture ──────────────────────────────────────────────────
+print("\n=== Prompt Format Capture ===")
+
+query = "set a 5 min timer"
+tools = [{"name": "set_timer", "description": "Set a timer.", "parameters": {"time_human": {"type": "string", "description": "duration", "required": True}}}]
+tools_json = json.dumps(tools)
+
+captured_inputs = []
+original_encode = tok.encode
+
+def capture_encode(text, *a, **kw):
+    captured_inputs.append(('encode', text))
+    return original_encode(text, *a, **kw)
+
+tok.encode = capture_encode
+
+# Call _build_encoder_input directly
+enc_tokens = _build_encoder_input(tok, query, tools_json)
+print(f"Captured encode calls: {captured_inputs}")
+print(f"Encoder token IDs: {enc_tokens}")
+print(f"Encoder token count: {len(enc_tokens)}")
+
+# Decode the encoder input to see the full prompt
+decoded_enc = tok.decode(enc_tokens)
+print(f"Decoded encoder input: {repr(decoded_enc)}")
+
+# Also decode each segment
+tok.encode = original_encode  # restore
+q_toks = tok.encode(query)
+t_toks = tok.encode(tools_json)
+print(f"\nQuery tokens: {q_toks} -> '{tok.decode(q_toks)}'")
+print(f"tools_token_id (separator): {TOOLS_ID}")
+print(f"Tools tokens: {t_toks[:20]}... -> '{tok.decode(t_toks[:20])}'")
+
+# ── 5. Write notes files ──────────────────────────────────────────────────────
+
+NOTES_DIR = os.path.join(os.path.dirname(__file__), '..', 'notes')
+os.makedirs(NOTES_DIR, exist_ok=True)
+
+# ── 5a. needle-internals.md ───────────────────────────────────────────────────
+
+def fmt_shape_dict(d, indent=0):
+    """Format shape dict as markdown nested list."""
+    lines = []
+    prefix = '  ' * indent
+    for k, v in sorted(d.items()):
+        if isinstance(v, dict) and 'shape' not in v:
+            lines.append(f"{prefix}- **{k}**")
+            lines.extend(fmt_shape_dict(v, indent+1))
+        else:
+            shape = v.get('shape', '?')
+            dtype = v.get('dtype', '?')
+            lines.append(f"{prefix}  - `{k}`: shape={shape}, dtype={dtype}")
+    return lines
+
+needle_internals = f"""# Needle (Cactus) Flax Model Internals
+
+## Model: `SimpleAttentionNetwork`
+
+**Architecture:** Encoder–decoder transformer with shared embeddings, tied output projection, RoPE, bfloat16.
+
+### `TransformerConfig` Defaults
+
+| Field | Value | Notes |
+|---|---|---|
+| vocab_size | {config.vocab_size} | SentencePiece BPE vocab |
+| d_model | {config.d_model} | Hidden dimension |
+| num_heads | {config.num_heads} | Attention heads |
+| num_kv_heads | {config.num_kv_heads} | GQA key/value heads |
+| num_encoder_layers | {config.num_encoder_layers} | Encoder depth |
+| num_decoder_layers | {config.num_decoder_layers} | Decoder depth |
+| d_ff | {config.d_ff} | FFN hidden dim (unused by default) |
+| max_seq_len | {config.max_seq_len} | |
+| pad_token_id | {config.pad_token_id} | |
+| rope_theta | {config.rope_theta} | RoPE base frequency |
+| dtype | {config.dtype} | |
+| activation | {config.activation} | dual-ReLU gate (drelu) |
+| num_memory_slots | {config.num_memory_slots} | |
+| dropout_rate | {config.dropout_rate} | |
+| contrastive_dim | {config.contrastive_dim} | Contrastive projection output dim |
+| no_feedforward | {config.no_feedforward} | **True → FFN blocks are SKIPPED** |
+
+**Total parameters:** {total_params:,}
+
+**Important:** `no_feedforward=True` by default — `FeedForward` / `gate_proj` / `up_proj` / `down_proj` are NOT present in the saved param tree. The `ffn_gate` scalar also absent. This is a pure attention-only transformer.
+
+---
+
+## Layer-by-Layer Architecture
+
+### Encoder (`Encoder` module)
+
+Uses `nn.scan` over `num_encoder_layers={config.num_encoder_layers}` layers. Each `EncoderBlock` contains:
+
+1. **`attn_gate`** — scalar sigmoid gate `()` on attention residual
+2. **Pre-norm**: `ZCRMSNorm` (scale initialized to 0, applied as `(1+γ) * x / RMS(x)`)
+3. **`self_attn`** (`MultiHeadAttention`):
+   - `q_proj`: Dense `(d_model, d_model)` = `({config.d_model}, {config.d_model})`
+   - `k_proj`: Dense `(d_model, num_kv_heads * head_dim)` = `({config.d_model}, {config.num_kv_heads * (config.d_model // config.num_heads)})`
+   - `v_proj`: Dense `(d_model, num_kv_heads * head_dim)` = `({config.d_model}, {config.num_kv_heads * (config.d_model // config.num_heads)})`
+   - `out_proj`: Dense `(d_model, d_model)` = `({config.d_model}, {config.d_model})`
+   - `q_norm`, `k_norm`: `ZCRMSNorm` on Q and K (pre-RoPE)
+4. Residual: `x = x_in + gate * attn_out`
+5. (FFN block **skipped** because `no_feedforward=True`)
+6. **`final_norm`**: `ZCRMSNorm` after all layers (on encoder output)
+
+### Decoder (`Decoder` module)
+
+Uses `nn.scan` over `num_decoder_layers={config.num_decoder_layers}` layers. Each `DecoderBlock` contains:
+
+1. **`self_attn_gate`** — scalar on causal self-attention residual
+2. Pre-norm + **`self_attn`** (causal self-attention, same structure as encoder)
+3. Residual: `x = x_in + self_gate * self_attn_out`
+4. **`cross_attn_gate`** — scalar on cross-attention residual
+5. Pre-norm + **`cross_attn`** (cross-attention: Q from decoder, K/V from encoder)
+6. Residual: `x = x_in + cross_gate * cross_attn_out`
+7. (FFN block **skipped**)
+8. Final `ZCRMSNorm` after all layers
+
+### Top-level `SimpleAttentionNetwork` params
+
+- **`embedding`**: `nn.Embed` kernel `(vocab_size, d_model)` = `({config.vocab_size}, {config.d_model})`
+  Also used as tied output: `logits = hidden @ embedding.T`
+- **`log_temp`**: scalar `()` — contrastive temperature
+- **`contrastive_hidden`**: Dense `(d_model, d_model//4)` = `({config.d_model}, {config.d_model//4})`
+- **`contrastive_proj`**: Dense `(d_model//4, contrastive_dim)` = `({config.d_model//4}, {config.contrastive_dim})`, no bias
+
+---
+
+## `nn.scan` Parameter Layout
+
+Because Flax `nn.scan` with `variable_axes={{"params": 0}}` is used, each scanned layer stack is stored as a **single** parameter array with an extra leading dimension of size `num_layers`. So for encoder: `layers.attn_gate` has shape `(num_encoder_layers, ...)` = `({config.num_encoder_layers}, ...)`.
+
+**This is critical for Task 2C:** When mapping Flax → PyTorch weights, you must index `params['encoder']['layers']['...'][layer_idx]` to get per-layer weights.
+
+---
+
+## Full Flax `params` Tree (shapes)
+
+```json
+{json.dumps(shape_dict, indent=2)}
+```
+
+---
+
+## `SimpleAttentionNetwork.__call__` Source
+
+```python
+{inspect.getsource(SimpleAttentionNetwork.__call__)}
+```
+
+## `EncoderBlock.__call__` Source
+
+```python
+{inspect.getsource(EncoderBlock.__call__)}
+```
+
+## `DecoderBlock.__call__` Source
+
+```python
+{inspect.getsource(DecoderBlock.__call__)}
+```
+
+## `MultiHeadAttention.__call__` Source
+
+```python
+{inspect.getsource(MultiHeadAttention.__call__)}
+```
+"""
+
+with open(os.path.join(NOTES_DIR, 'needle-internals.md'), 'w') as f:
+    f.write(needle_internals)
+print(f"\nWrote needle-internals.md ({len(needle_internals.splitlines())} lines)")
+
+# ── 5b. tokenizer-internals.md ────────────────────────────────────────────────
+
+tokenizer_internals = f"""# Needle Tokenizer Internals
+
+## Tokenizer Class
+
+**Class:** `NeedleTokenizer` (defined in `external/needle/needle/dataset/tokenizer.py`)
+
+**Backend:** SentencePiece BPE — the model file is downloaded from HuggingFace on first use.
+
+**HuggingFace repo:** `Cactus-Compute/needle-tokenizer` (dataset repo)
+
+**Local model file path:** `{TOKENIZER_PREFIX}.model`
+(relative to the external/needle repo root: `needle/tokenizer/needle.model`)
+
+---
+
+## Tokenizer Properties
+
+| Property | Value | Source |
+|---|---|---|
+| `vocab_size` | {tok.vocab_size} | `sp.GetPieceSize()` |
+| `pad_token_id` | {PAD_ID} | hardcoded = 0 |
+| `eos_token_id` | {EOS_ID} | hardcoded = 1 |
+| `bos_token_id` | {BOS_ID} | hardcoded = 2 |
+| `unk_id` | {UNK_ID} | hardcoded = 3 |
+| `tool_call_token_id` | {TOOL_CALL_ID} | hardcoded = 4, symbol `<tool_call>` |
+| `tools_token_id` | {TOOLS_ID} | hardcoded = 5, symbol `<tools>` |
+
+---
+
+## Encode Algorithm
+
+**Algorithm:** SentencePiece **BPE** (Byte-Pair Encoding)
+
+**Configured with:**
+- `model_type="bpe"`
+- `byte_fallback=True` — unknown bytes represented as `<0xNN>` hex tokens
+- `normalization_rule_name="identity"` — **no Unicode normalization applied**
+- `user_defined_symbols=["<tool_call>", "<tools>"]` — fixed IDs 4 and 5
+
+**Encode call in NeedleTokenizer:**
+```python
+def encode(self, text):
+    return self.sp.Encode(text, out_type=int)
+```
+Returns a Python `list[int]`. Does NOT add BOS/EOS automatically.
+
+**Decode call:**
+```python
+def decode(self, ids):
+    return self.sp.Decode(list(ids))
+```
+
+---
+
+## Available Attributes on `NeedleTokenizer`
+
+```
+{[a for a in dir(tok) if not a.startswith('_')]}
+```
+
+The underlying SentencePiece processor is accessible as `tok.sp`.
+
+---
+
+## SentencePiece Model File
+
+The `.model` file is a serialized SentencePiece protobuf. For Task 5 (port tokenizer to JS/ONNX/etc.), you need either:
+1. The `.model` file directly (loadable by the `sentencepiece` library)
+2. Export the vocabulary + merge rules via `tok.sp.GetPieceSize()`, `tok.sp.IdToPiece(i)`, etc.
+
+To export the full vocabulary:
+```python
+vocab = [(i, tok.sp.IdToPiece(i), tok.sp.GetScore(i)) for i in range(tok.vocab_size)]
+```
+
+---
+
+## Example Encoding
+
+```
+encode("{sample_text}")
+  -> {encoded}
+  -> decode -> "{decoded}"
+```
+"""
+
+with open(os.path.join(NOTES_DIR, 'tokenizer-internals.md'), 'w') as f:
+    f.write(tokenizer_internals)
+print(f"Wrote tokenizer-internals.md ({len(tokenizer_internals.splitlines())} lines)")
+
+# ── 5c. prompt-format.md ─────────────────────────────────────────────────────
+
+# Get the decoded full prompt
+tok.encode = capture_encode
+enc_tokens2 = _build_encoder_input(tok, query, tools_json)
+tok.encode = original_encode  # restore
+
+# Reconstruct what each segment contains
+q_toks_raw = tok.encode(query)
+t_toks_raw = tok.encode(tools_json)
+
+MAX_ENC = DEFAULT_MAX_ENC_LEN  # 1024
+max_query = MAX_ENC - 2
+q_toks_trunc = q_toks_raw[:max_query]
+remaining = MAX_ENC - len(q_toks_trunc) - 1
+t_toks_trunc = t_toks_raw[:remaining]
+
+final_tokens = q_toks_trunc + [TOOLS_ID] + t_toks_trunc
+decoded_full = tok.decode(final_tokens)
+
+# Decoder starts with EOS token (1), then model predicts <tool_call> (4) first
+prompt_format = f"""# Needle Prompt Format
+
+## Overview
+
+Needle uses an **encoder–decoder** architecture. The query + tools are encoded by the encoder; the decoder generates the tool-call JSON.
+
+## Encoder Input Format
+
+```
+[query_tokens..., <tools>(id={TOOLS_ID}), tools_tokens...]
+```
+
+**Built by `_build_encoder_input(tokenizer, query, tools, max_enc_len={DEFAULT_MAX_ENC_LEN})`:**
+
+```python
+{inspect.getsource(_build_encoder_input)}
+```
+
+### Truncation Rules
+
+1. Query is truncated to `max_enc_len - 2 = {MAX_ENC - 2}` tokens
+2. Tools are truncated to `max_enc_len - len(q_toks) - 1` tokens (fills remaining space)
+3. The `<tools>` separator token (id={TOOLS_ID}) is always inserted between query and tools
+
+### Example: `query="set a 5 min timer"`, `tools=[{{"name": "set_timer", ...}}]`
+
+**Encode calls made by `_build_encoder_input`:**
+1. `tokenizer.encode("{query}")` → {q_toks_raw} ({len(q_toks_raw)} tokens)
+2. `tokenizer.encode('{tools_json[:80]}...')` → {t_toks_raw[:20]}... ({len(t_toks_raw)} tokens total)
+
+**Final encoder token sequence (len={len(final_tokens)}):**
+```
+{final_tokens}
+```
+
+**Decoded encoder input string:**
+```
+{repr(decoded_full)}
+```
+
+**Structural breakdown:**
+- Query tokens ({len(q_toks_trunc)} tokens): `{q_toks_trunc}`
+- Separator token (1 token): `[{TOOLS_ID}]` → `<tools>`
+- Tools tokens ({len(t_toks_trunc)} tokens): `{t_toks_trunc}`
+
+---
+
+## Decoder Input Format
+
+The decoder is initialized with a single **EOS token** (id={EOS_ID}) as prefix:
+
+```python
+dec_buffer = jnp.full((1, max_gen_len), pad_id, dtype=jnp.int32)
+dec_buffer = dec_buffer.at[0, 0].set(eos_id)   # prefix = [EOS]
+```
+
+At each step, the model predicts the next token autoregressively.
+
+**Expected output sequence:** `<tool_call>(id={TOOL_CALL_ID}) + answer_json_tokens + EOS(id={EOS_ID})`
+
+The `<tool_call>` prefix is stripped from the decoded output string in `generate()`.
+
+---
+
+## Tool Name Normalization
+
+Before building encoder input, tool names are converted to `snake_case` via `normalize_tools()`:
+- `camelCase` → `camel_case`
+- `PascalCase` → `pascal_case`
+- `dot.notation` → `dot_notation`
+- Non-alphanumeric chars → underscores
+
+After decoding, `restore_tool_names()` maps back to original names.
+
+---
+
+## Full `generate()` Signature
+
+```python
+{inspect.getsource(generate)}
+```
+"""
+
+with open(os.path.join(NOTES_DIR, 'prompt-format.md'), 'w') as f:
+    f.write(prompt_format)
+print(f"Wrote prompt-format.md ({len(prompt_format.splitlines())} lines)")
+
+print("\n=== All notes written successfully ===")