Views
No views yet
1#inference model
2import json
3import torch
4import torch.nn as nn
5import torch.nn.functional as F
6import math
7from transformers import GPT2Tokenizer
8from safetensors.torch import load_file
9from huggingface_hub import snapshot_download
10import sys
11
12class Config:
13 def __init__(self, **kwargs):
14 self.vocab_size = 50257
15 self.d_model = 1024
16 self.n_head = 16
17 self.d_k = self.d_model // self.n_head
18 self.d_ff = 4096
19 self.max_depth = 4
20 self.num_recursive_layers = 6
21 self.balancing_weight = 0.01
22 self.temperature = 1.0
23 self.seq_len = 512
24 self.batch_size = 16
25 self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
26 for key, value in kwargs.items():
27 setattr(self, key, value)
28 if hasattr(self, 'd_model') and hasattr(self, 'n_head'):
29 self.d_k = self.d_model // self.n_head
30
31class RecursiveLayer(nn.Module):
32 def __init__(self, config):
33 super().__init__()
34 self.config = config
35 self.w_q = nn.Linear(config.d_model, config.d_model)
36 self.w_k = nn.Linear(config.d_model, config.d_model)
37 self.w_v = nn.Linear(config.d_model, config.d_model)
38 self.attn_out = nn.Linear(config.d_model, config.d_model)
39 self.ffn = nn.Sequential(
40 nn.Linear(config.d_model, config.d_ff),
41 nn.GELU(),
42 nn.Linear(config.d_ff, config.d_model)
43 )
44 self.norm1 = nn.LayerNorm(config.d_model)
45 self.norm2 = nn.LayerNorm(config.d_model)
46
47 def forward(self, h, K_cache, V_cache, active_mask, causal_mask):
48 batch_size, seq_len, _ = h.shape
49 q = self.w_q(h).view(batch_size, seq_len, self.config.n_head, self.config.d_k)
50 q = q.permute(0, 2, 1, 3)
51 active_mask_expanded = active_mask.unsqueeze(-1)
52 attn_scores = torch.matmul(q, K_cache.transpose(-2, -1)) / math.sqrt(self.config.d_k)
53 attn_scores = attn_scores.masked_fill(causal_mask == 0, float('-inf'))
54 attn_probs = F.softmax(attn_scores, dim=-1)
55 attn_out = torch.matmul(attn_probs, V_cache)
56 attn_out = attn_out.permute(0, 2, 1, 3).contiguous()
57 attn_out = attn_out.view(batch_size, seq_len, self.config.d_model)
58 attn_out = self.attn_out(attn_out)
59 attn_out = attn_out * active_mask_expanded
60 h = h + attn_out
61 h = self.norm1(h)
62 ffn_out = self.ffn(h) * active_mask_expanded
63 h = h + ffn_out
64 h = self.norm2(h)
65 return h
66
67class Router(nn.Module):
68 def __init__(self, config):
69 super().__init__()
70 self.linear = nn.Sequential(
71 nn.Linear(config.d_model, config.d_model // 2),
72 nn.GELU(),
73 nn.Linear(config.d_model // 2, config.max_depth)
74 )
75 self.temperature = config.temperature
76
77 def forward(self, h, train=True):
78 logits = self.linear(h)
79 if train:
80 probs = F.gumbel_softmax(logits, tau=self.temperature, dim=-1)
81 return probs, F.softmax(logits, dim=-1)
82 else:
83 probs = F.softmax(logits, dim=-1)
84 return probs, probs
85
86class MixtureRecursions(nn.Module):
87 def __init__(self, config):
88 super().__init__()
89 self.embed = nn.Embedding(config.vocab_size, config.d_model)
90 self.pos_embed = nn.Embedding(config.seq_len, config.d_model)
91 self.first_layer = nn.Sequential(
92 nn.Linear(config.d_model, config.d_model),
93 nn.GELU(),
94 nn.LayerNorm(config.d_model)
95 )
96 self.recursive_layers = nn.ModuleList([
97 RecursiveLayer(config) for _ in range(config.num_recursive_layers)
98 ])
99 self.router = Router(config)
100 self.final_norm = nn.LayerNorm(config.d_model)
101 self.head = nn.Linear(config.d_model, config.vocab_size, bias=False)
102 self.apply(self._init_weights)
103
104 def _init_weights(self, module):
105 if isinstance(module, nn.Linear):
106 nn.init.normal_(module.weight, mean=0.0, std=0.02)
107 if module.bias is not None:
108 nn.init.zeros_(module.bias)
109 elif isinstance(module, nn.Embedding):
110 nn.init.normal_(module.weight, mean=0.0, std=0.02)
111
112 def forward(self, x, targets=None):
113 config = self.embed.weight.device # Gets device
114 device = x.device
115 batch_size, seq_len = x.shape
116 pos_ids = torch.arange(0, seq_len, dtype=torch.long, device=device)
117 pos_emb = self.pos_embed(pos_ids)
118 tok_emb = self.embed(x)
119 h = tok_emb + pos_emb
120 h = self.first_layer(h)
121 initial_h = h.clone()
122 router_probs, router_soft = self.router(h)
123 assigned_depths = router_probs.argmax(dim=-1) + 1
124 K_cache, V_cache = [], []
125 for layer in self.recursive_layers:
126 K = layer.w_k(initial_h).view(batch_size, seq_len, self.embed.embedding_dim // self.router.linear[0].in_features * self.recursive_layers[0].config.n_head, self.recursive_layers[0].config.d_k)
127 V = layer.w_v(initial_h).view(batch_size, seq_len, self.embed.embedding_dim // self.router.linear[0].in_features * self.recursive_layers[0].config.n_head, self.recursive_layers[0].config.d_k)
128 K_cache.append(K.permute(0, 2, 1, 3))
129 V_cache.append(V.permute(0, 2, 1, 3))
130 causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=device)).view(1, 1, seq_len, seq_len)
131 for depth in range(1, self.recursive_layers[0].config.max_depth + 1):
132 active_mask = (assigned_depths >= depth)
133 layer_idx = (depth - 1) % self.recursive_layers[0].config.num_recursive_layers
134 h = self.recursive_layers[layer_idx](
135 h,
136 K_cache[layer_idx],
137 V_cache[layer_idx],
138 active_mask,
139 causal_mask
140 )
141 h = self.final_norm(h)
142 logits = self.head(h)
143 loss = None
144 balancing_loss = None
145 if targets is not None:
146 logits = logits[:, :-1, :].contiguous()
147 targets = targets[:, 1:].contiguous()
148 loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
149 router_decision = router_probs.sum(dim=[0, 1])
150 router_decision = router_decision / (batch_size * seq_len)
151 balancing_loss = torch.var(router_decision) * self.recursive_layers[0].config.balancing_weight
152 return logits, loss, balancing_loss
153 return logits, loss, balancing_loss
154
155# --- Download and load everything as before ---
156repo_id = "liminerity/MoR-TC-v1"
157model_dir = snapshot_download(repo_id=repo_id)
158tokenizer = GPT2Tokenizer.from_pretrained(model_dir)
159with open(f"{model_dir}/config.json", 'r') as f:
160 hf_config = json.load(f)
161
162config_map = {
163 'vocab_size': 'vocab_size',
164 'dim': 'd_model',
165 'num_layers': 'num_recursive_layers',
166 'num_heads': 'n_head',
167 'max_recursion': 'max_depth',
168 'max_position_embeddings': 'seq_len',
169 'balancing_weight': 'balancing_weight',
170 'temperature': 'temperature'
171}
172mapped_config = {config_map[k]: v for k, v in hf_config.items() if k in config_map}
173mapped_config['d_ff'] = hf_config['ffn_expansion'] * mapped_config['d_model']
174config = Config(**mapped_config)
175
176device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
177model = MixtureRecursions(config).to(device)
178
179weights = load_file(f"{model_dir}/model.safetensors", device=str(device))
180model.load_state_dict(weights)
181model.eval()
182
183# --- Autoregressive Generation Loop ---
184def autoregressive_generate(
185 model, tokenizer, input_text, max_new_tokens=500, temperature=1.0, line_width=71):
186
187 model.eval()
188 device = next(model.parameters()).device
189 input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
190 current_ids = input_ids
191
192 # Start buffer with the prompt
193 buffer = tokenizer.decode(current_ids[0])
194
195 # Don't print until a whole line is available
196 print_end = False
197
198 for _ in range(max_new_tokens):
199 if current_ids.shape[1] >= config.seq_len:
200 current_ids = current_ids[:, -config.seq_len:]
201
202 with torch.no_grad():
203 logits, _, _ = model(current_ids)
204
205 next_token_logits = logits[0, -1, :] / temperature
206 probs = torch.softmax(next_token_logits, dim=-1)
207 next_token_id = torch.multinomial(probs, num_samples=1).item()
208
209 # Append new token
210 current_ids = torch.cat(
211 [current_ids, torch.tensor([[next_token_id]], device=device)], dim=1
212 )
213
214 # Append to buffer
215 buffer += tokenizer.decode([next_token_id])
216
217 # Print as many complete lines as possible
218 while len(buffer) >= line_width:
219 print(buffer[:line_width])
220 buffer = buffer[line_width:]
221 print_end = True
222
223 # Print remaining buffer (incomplete final line)
224 if buffer:
225 print(buffer)
226
227# Test streaming generation of 500 tokens, 71 chars per line
228input_text = "The future of AI is"
229autoregressive_generate(model, tokenizer, input_text, max_new_tokens=500, temperature=config.temperature, line_width=71)