Upload 4 files

README.md

@@ -3,6 +3,8 @@ license: mit
 language:
 - en
 ---
+<img src="himoe_visual.png">
+
 # HiMoE — Hierarchical Mixture of Experts
 
 > *A Matryoshka-inspired two-level routing architecture for efficient large-scale language modelling.*

train.py

@@ -47,7 +47,7 @@ class HiMoEConfig:
     num_experts:  int   = 8    # Level-2 choices per MoE
     # Training
     batch_size:   int   = 32
-    max_iters:    int   = 3000
+    max_iters:    int   = 750 # for testing, increase to 3000 for actual training
     eval_interval:int   = 50
     eval_iters:   int   = 20
     lr:           float = 3e-4
@@ -377,7 +377,7 @@ def load_model(model_dir: str, device: str) -> tuple:
     cfg.model_dir = model_dir
     vocab_size = meta["vocab_size"]
     stoi       = meta["stoi"]
-    itos       = meta["itos"]
+    itos       = {int(k): v for k, v in meta["itos"].items()}
     step       = meta["step"]
 
     model = HiMoEModel(cfg, vocab_size).to(device)
@@ -553,6 +553,23 @@ def train(cfg: HiMoEConfig, resume: bool = False):
                   f"lr {lr_now:.2e} | "
                   f"ETA {eta/60:.1f}m")
             save_model(model, cfg, vocab_size, stoi, itos, step)
+            
+            # Generate sample and save routing log periodically for visualization
+            model.eval()
+            with torch.no_grad():
+                # Workaround for MPS generation hangs: move to CPU for sampling
+                original_device = next(model.parameters()).device
+                model.to("cpu")
+                context = torch.zeros((1, 1), dtype=torch.long, device="cpu")
+                gen_ids, r_log = model.generate(context, max_new_tokens=400, temperature=0.8, top_k=40)
+                smp = "".join(itos[i] for i in gen_ids[0].tolist())
+                with open(os.path.join(cfg.model_dir, "sample.txt"), "w") as f:
+                    f.write(smp)
+                with open(os.path.join(cfg.model_dir, "routing_log.json"), "w") as f:
+                    json.dump(r_log, f, indent=2)
+                model.to(original_device)
+            model.train()
+
 
         # forward + backward
         x, y              = get_batch(train_data, cfg.block_size,
@@ -587,7 +604,7 @@ def train(cfg: HiMoEConfig, resume: bool = False):
     with open(os.path.join(cfg.model_dir, "sample.txt"), "w") as f:
         f.write(sample)
     with open(os.path.join(cfg.model_dir, "routing_log.json"), "w") as f:
-        json.dump(routing_log[:50], f, indent=2)   # first 50 tokens
+        json.dump(routing_log, f, indent=2)   # save full log for visualization
 
     print(f"\n[himoe] Sample + routing log saved to '{cfg.model_dir}/'")
 

visualizer.py

@@ -0,0 +1,122 @@
+import os
+import json
+import torch
+from PIL import Image, ImageDraw, ImageFont
+import numpy as np
+
+def visualize_routing():
+    model_dir = "model"
+    sample_file = os.path.join(model_dir, "sample.txt")
+    routing_file = os.path.join(model_dir, "routing_log.json")
+    output_file = "himoe_visual.png"
+
+    if not os.path.exists(sample_file) or not os.path.exists(routing_file):
+        print(f"Error: Required files missing in {model_dir}")
+        return
+
+    with open(sample_file, "r") as f:
+        text = f.read()
+    
+    with open(routing_file, "r") as f:
+        routing_log = json.load(f)
+
+    # Use Layer 0 for visualization by default
+    layer_idx = 0
+    
+    chars = list(text)
+    if len(chars) > len(routing_log):
+        # Skip the context character
+        chars = chars[1:]
+    
+    n = min(len(chars), len(routing_log))
+    chars = chars[:n]
+    routing_log = routing_log[:n]
+    # --- Setup Visuals ---
+    char_w, char_h = 24, 36 # Larger for zoom
+    cols = 60 
+    
+    # We need to calculate rows based on text AND newlines
+    current_col = 0
+    total_rows = 1
+    for char in chars:
+        if char == "\n":
+            current_col = 0
+            total_rows += 1
+        else:
+            current_col += 1
+            if current_col >= cols:
+                current_col = 0
+                total_rows += 1
+                
+    margin = 50
+    legend_w = 300
+    img_w = cols * char_w + margin * 2 + legend_w
+    img_h = max(total_rows * char_h + margin * 3, 1000)
+    
+    img = Image.new("RGB", (img_w, img_h), (20, 20, 25))
+    draw = ImageDraw.Draw(img)
+    
+    try:
+        font = ImageFont.truetype("/System/Library/Fonts/Supplemental/Courier New.ttf", 22)
+    except:
+        font = ImageFont.load_default()
+
+    # --- Color Mapping ---
+    def get_color(moe_id, exp_id):
+        h = (moe_id * 60) % 360
+        l = 30 + (exp_id * 7) # 30% to 79%
+        import colorsys
+        r, g, b = colorsys.hls_to_rgb(h/360, l/100, 0.7)
+        return (int(r*255), int(g*255), int(b*255))
+
+    moe_colors = [get_color(i, 4) for i in range(6)]
+
+    # --- Draw Text ---
+    curr_r, curr_c = 0, 0
+    for i in range(n):
+        char = chars[i]
+        
+        # Handle newline or wrap
+        if char == "\n" or curr_c >= cols:
+            curr_r += 1
+            curr_c = 0
+            if char == "\n": continue # Skip drawing the newline char itself
+            
+        x = margin + curr_c * char_w
+        y = margin + curr_r * char_h
+        
+        moe_id = routing_log[i]["moe"][layer_idx][0]
+        exp_id = routing_log[i]["exp"][layer_idx][0]
+        
+        bg_color = get_color(moe_id, exp_id)
+        draw.rectangle([x, y, x + char_w - 1, y + char_h - 1], fill=bg_color)
+        
+        if not char.isspace():
+            text_color = (255, 255, 255) if bg_color[0]*0.299 + bg_color[1]*0.587 + bg_color[2]*0.114 < 128 else (0, 0, 0)
+            draw.text((x + 4, y + 4), char, fill=text_color, font=font)
+        
+        curr_c += 1
+
+    # --- Draw Legend ---
+    lx = margin + cols * char_w + 40
+    ly = margin
+    draw.text((lx, ly), "HiMoE Routing Legend", fill=(255, 255, 255), font=font)
+    ly += 40
+    
+    for mi in range(6):
+        draw.text((lx, ly), f"MoE Block {mi+1}", fill=moe_colors[mi], font=font)
+        ly += 25
+        # Show a few expert shades
+        for ei in [0, 3, 7]:
+            ex = lx + 20
+            c = get_color(mi, ei)
+            draw.rectangle([ex, ly, ex + 15, ly + 15], fill=c)
+            draw.text((ex + 25, ly - 2), f"Exp {ei+1}", fill=(200, 200, 200), font=font)
+            ly += 20
+        ly += 10
+
+    img.save(output_file)
+    print(f"Visualization saved to {output_file}")
+
+if __name__ == "__main__":
+    visualize_routing()