Upload inference.py

inference.py

@@ -0,0 +1,63 @@
+import torch
+import torch.nn.functional as F
+from hrt import ModelConfig, HierarchicalRadialTransformerV7
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# 1. Initialize Configuration matching training
+cfg = ModelConfig(
+    d_model=768,
+    d_ff=3072,
+    n_outer_latents=512,
+    n_outer_cycles=6,
+    n_inner_cycles=8,
+    n_center_latents=16,
+    routing_k=64,
+    n_outer_heads=12,
+    n_inner_heads=12,
+    n_latent_heads=12,
+    vocab_size=257,
+    max_seq_len=131072,
+    use_qk_norm=True,
+    use_rezero=True,
+    use_compaction=True,
+    use_internalization=True,
+    use_jfb=True,
+    use_q_cache=True,
+)
+
+# 2. Load model & weights
+model = HierarchicalRadialTransformerV7(cfg).to(device)
+weights = torch.load("hrt_v7_148m_weights.pt", map_location=device)
+model.load_state_dict(weights["model"] if "model" in weights else weights)
+model.eval()
+
+# 3. Autoregressive Byte-level Generation
+def generate(prompt: str, max_new_bytes: int = 120, temp: float = 0.5, top_k: int = 5):
+    prompt_bytes = list(prompt.encode("utf-8"))
+    prompt_ids = torch.tensor([prompt_bytes], dtype=torch.long, device=device)
+    
+    with torch.no_grad():
+        prompt_emb = model.tok_emb(prompt_ids)
+        logits, cache = model._init_generation_cache(prompt_emb)
+        out_bytes = list(prompt_bytes)
+
+        for _ in range(max_new_bytes):
+            l = logits / max(temp, 1e-5)
+            if top_k > 0:
+                v, _ = torch.topk(l, min(top_k, l.size(-1)))
+                l[l < v[:, [-1]]] = float("-inf")
+            
+            nxt = torch.multinomial(F.softmax(l, dim=-1), num_samples=1)
+            nxt_id = nxt.item()
+            if nxt_id == 256:  # EOS
+                break
+                
+            out_bytes.append(nxt_id)
+            nxt_emb = model.tok_emb(nxt)
+            logits = model.step_generation(nxt_emb, cache)
+            
+    return bytes(out_bytes).decode("utf-8", errors="replace")
+
+# Test completion
+print(generate("def", max_new_bytes=100))