feat: increase seq_len and auto-update of dim on UI during compute

app.py

@@ -14,6 +14,7 @@ from src.extract import (
     MODEL_CHOICES,
     expand_kv_heads,
     extract_from_model,
+    get_model_head_dim,
     random_qk,
     select_head,
 )
@@ -90,6 +91,22 @@ def _safe_slider_max(n: int) -> int:
     return max(int(n), 1)
 
 
+def update_dimension(source: str, model_name: str):
+    if source.startswith("Random"):
+        return gr.update(minimum=4, maximum=128, value=32, step=2, interactive=True)
+    try:
+        model_dim = get_model_head_dim(model_name)
+        return gr.update(
+            minimum=4,
+            maximum=max(128, model_dim),
+            value=model_dim,
+            step=2,
+            interactive=False,
+        )
+    except Exception:
+        return gr.update()
+
+
 def compute(
     source: str,
     sentence: str,
@@ -249,7 +266,7 @@ def toggle_source(source: str):
     is_random = source.startswith("Random")
     return (
         gr.update(visible=is_random),
-        gr.update(visible=is_random),
+        gr.update(visible=True, interactive=is_random),
         gr.update(visible=is_random),
         gr.update(visible=not is_random),
         gr.update(visible=not is_random),
@@ -287,8 +304,8 @@ with gr.Blocks(title="RoPE Explorer") as demo:
                 )
             with gr.Row():
                 seq_len = gr.Slider(2, MAX_SEQ_LEN, value=16, step=1, label="Sequence length")
-                dim = gr.Slider(4, 128, value=32, step=2, label="Dimension (even)")
-                seed = gr.Number(value=0, label="Seed", precision=0)
+                dim = gr.Slider(4, 128, value=32, step=2, label="Dimension (even; per attention head)")
+                seed = gr.Number(value=42, label="Seed", precision=0)
             base = gr.Number(
                 value=10000,
                 label="RoPE base (overridden by config.rope_theta for real models)",
@@ -335,6 +352,16 @@ with gr.Blocks(title="RoPE Explorer") as demo:
         inputs=[source],
         outputs=[seq_len, dim, seed, sentence, model_name],
     )
+    source.change(
+        update_dimension,
+        inputs=[source, model_name],
+        outputs=[dim],
+    )
+    model_name.change(
+        update_dimension,
+        inputs=[source, model_name],
+        outputs=[dim],
+    )
 
     bulk_inputs = [state, which, head, mod_2pi]
     bulk_outputs = [bulk_main, bulk_norm, bulk_theta, bulk_freq]

src/extract.py

@@ -11,20 +11,24 @@ import numpy as np
 
 from src.rope import apply_rope
 
-MAX_SEQ_LEN = 64
+MAX_SEQ_LEN = 2048
 
 # Ungated Llama-like checkpoints (q_proj / k_proj + rotary). No HF token required.
 MODEL_CHOICES = [
+    "HuggingFaceM4/tiny-random-LlamaForCausalLM",
     "HuggingFaceTB/SmolLM2-135M",
     "Qwen/Qwen2.5-0.5B-Instruct",
     "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
 ]
-DEFAULT_MODEL = MODEL_CHOICES[0]
+DEFAULT_MODEL = MODEL_CHOICES[1]
 
 _cache: dict[str, Any] = {"name": None, "model": None, "tokenizer": None}
 
 
-def random_matrix(seq_len: int, dim: int, seed: int = 0) -> np.ndarray:
+_config_cache: dict[str, Any] = {}
+
+
+def random_matrix(seq_len: int, dim: int, seed: int = 42) -> np.ndarray:
     if dim % 2 != 0:
         raise ValueError(f"dim must be even for RoPE, got {dim}")
     seq_len = int(np.clip(seq_len, 1, MAX_SEQ_LEN))
@@ -35,7 +39,7 @@ def random_matrix(seq_len: int, dim: int, seed: int = 0) -> np.ndarray:
 def random_qk(
     seq_len: int,
     dim: int,
-    seed: int = 0,
+    seed: int = 42,
     base: float = 10000.0,
 ) -> dict[str, Any]:
     q = random_matrix(seq_len, dim, seed=seed)
@@ -104,6 +108,22 @@ def _require_hf():
     return torch, AutoModelForCausalLM, AutoTokenizer
 
 
+def get_model_head_dim(model_name: str) -> int:
+    """Return the Q/K dimension per attention head without loading model weights."""
+    if model_name not in _config_cache:
+        try:
+            from transformers import AutoConfig
+        except ImportError as exc:
+            raise RuntimeError("Real-model mode needs `transformers`.") from exc
+        _config_cache[model_name] = AutoConfig.from_pretrained(model_name)
+
+    config = _config_cache[model_name]
+    configured_head_dim = getattr(config, "head_dim", None)
+    if configured_head_dim is not None:
+        return int(configured_head_dim)
+    return int(config.hidden_size) // int(config.num_attention_heads)
+
+
 def get_model(model_name: str):
     """Load tokenizer + causal LM on CPU; cache the last selection."""
     torch, AutoModelForCausalLM, AutoTokenizer = _require_hf()