File size: 2,053 Bytes
430b17b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | 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))
|