MoR-TC-v1

#inference model
import json
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from transformers import GPT2Tokenizer
from safetensors.torch import load_file
from huggingface_hub import snapshot_download
import sys

class Config:
    def __init__(self, **kwargs):
        self.vocab_size = 50257
        self.d_model = 1024
        self.n_head = 16
        self.d_k = self.d_model // self.n_head
        self.d_ff = 4096
        self.max_depth = 4
        self.num_recursive_layers = 6
        self.balancing_weight = 0.01
        self.temperature = 1.0
        self.seq_len = 512
        self.batch_size = 16
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        for key, value in kwargs.items():
            setattr(self, key, value)
        if hasattr(self, 'd_model') and hasattr(self, 'n_head'):
            self.d_k = self.d_model // self.n_head

class RecursiveLayer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.w_q = nn.Linear(config.d_model, config.d_model)
        self.w_k = nn.Linear(config.d_model, config.d_model)
        self.w_v = nn.Linear(config.d_model, config.d_model)
        self.attn_out = nn.Linear(config.d_model, config.d_model)
        self.ffn = nn.Sequential(
            nn.Linear(config.d_model, config.d_ff),
            nn.GELU(),
            nn.Linear(config.d_ff, config.d_model)
        )
        self.norm1 = nn.LayerNorm(config.d_model)
        self.norm2 = nn.LayerNorm(config.d_model)

    def forward(self, h, K_cache, V_cache, active_mask, causal_mask):
        batch_size, seq_len, _ = h.shape
        q = self.w_q(h).view(batch_size, seq_len, self.config.n_head, self.config.d_k)
        q = q.permute(0, 2, 1, 3)
        active_mask_expanded = active_mask.unsqueeze(-1)
        attn_scores = torch.matmul(q, K_cache.transpose(-2, -1)) / math.sqrt(self.config.d_k)
        attn_scores = attn_scores.masked_fill(causal_mask == 0, float('-inf'))
        attn_probs = F.softmax(attn_scores, dim=-1)
        attn_out = torch.matmul(attn_probs, V_cache)
        attn_out = attn_out.permute(0, 2, 1, 3).contiguous()
        attn_out = attn_out.view(batch_size, seq_len, self.config.d_model)
        attn_out = self.attn_out(attn_out)
        attn_out = attn_out * active_mask_expanded
        h = h + attn_out
        h = self.norm1(h)
        ffn_out = self.ffn(h) * active_mask_expanded
        h = h + ffn_out
        h = self.norm2(h)
        return h

class Router(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.linear = nn.Sequential(
            nn.Linear(config.d_model, config.d_model // 2),
            nn.GELU(),
            nn.Linear(config.d_model // 2, config.max_depth)
        )
        self.temperature = config.temperature

    def forward(self, h, train=True):
        logits = self.linear(h)
        if train:
            probs = F.gumbel_softmax(logits, tau=self.temperature, dim=-1)
            return probs, F.softmax(logits, dim=-1)
        else:
            probs = F.softmax(logits, dim=-1)
            return probs, probs

class MixtureRecursions(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.embed = nn.Embedding(config.vocab_size, config.d_model)
        self.pos_embed = nn.Embedding(config.seq_len, config.d_model)
        self.first_layer = nn.Sequential(
            nn.Linear(config.d_model, config.d_model),
            nn.GELU(),
            nn.LayerNorm(config.d_model)
        )
        self.recursive_layers = nn.ModuleList([
            RecursiveLayer(config) for _ in range(config.num_recursive_layers)
        ])
        self.router = Router(config)
        self.final_norm = nn.LayerNorm(config.d_model)
        self.head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, x, targets=None):
        config = self.embed.weight.device  # Gets device
        device = x.device
        batch_size, seq_len = x.shape
        pos_ids = torch.arange(0, seq_len, dtype=torch.long, device=device)
        pos_emb = self.pos_embed(pos_ids)
        tok_emb = self.embed(x)
        h = tok_emb + pos_emb
        h = self.first_layer(h)
        initial_h = h.clone()
        router_probs, router_soft = self.router(h)
        assigned_depths = router_probs.argmax(dim=-1) + 1
        K_cache, V_cache = [], []
        for layer in self.recursive_layers:
            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)
            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)
            K_cache.append(K.permute(0, 2, 1, 3))
            V_cache.append(V.permute(0, 2, 1, 3))
        causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=device)).view(1, 1, seq_len, seq_len)
        for depth in range(1, self.recursive_layers[0].config.max_depth + 1):
            active_mask = (assigned_depths >= depth)
            layer_idx = (depth - 1) % self.recursive_layers[0].config.num_recursive_layers
            h = self.recursive_layers[layer_idx](
                h, 
                K_cache[layer_idx], 
                V_cache[layer_idx], 
                active_mask,
                causal_mask
            )
        h = self.final_norm(h)
        logits = self.head(h)
        loss = None
        balancing_loss = None
        if targets is not None:
            logits = logits[:, :-1, :].contiguous()
            targets = targets[:, 1:].contiguous()
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
            router_decision = router_probs.sum(dim=[0, 1])
            router_decision = router_decision / (batch_size * seq_len)
            balancing_loss = torch.var(router_decision) * self.recursive_layers[0].config.balancing_weight
            return logits, loss, balancing_loss
        return logits, loss, balancing_loss

# --- Download and load everything as before ---
repo_id = "liminerity/MoR-TC-v1"
model_dir = snapshot_download(repo_id=repo_id)
tokenizer = GPT2Tokenizer.from_pretrained(model_dir)
with open(f"{model_dir}/config.json", 'r') as f:
    hf_config = json.load(f)

config_map = {
    'vocab_size': 'vocab_size',
    'dim': 'd_model',
    'num_layers': 'num_recursive_layers',
    'num_heads': 'n_head',
    'max_recursion': 'max_depth',
    'max_position_embeddings': 'seq_len',
    'balancing_weight': 'balancing_weight',
    'temperature': 'temperature'
}
mapped_config = {config_map[k]: v for k, v in hf_config.items() if k in config_map}
mapped_config['d_ff'] = hf_config['ffn_expansion'] * mapped_config['d_model']
config = Config(**mapped_config)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MixtureRecursions(config).to(device)

weights = load_file(f"{model_dir}/model.safetensors", device=str(device))
model.load_state_dict(weights)
model.eval()

# --- Autoregressive Generation Loop ---
def autoregressive_generate(
    model, tokenizer, input_text, max_new_tokens=500, temperature=1.0, line_width=71):

    model.eval()
    device = next(model.parameters()).device
    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
    current_ids = input_ids

    # Start buffer with the prompt
    buffer = tokenizer.decode(current_ids[0])

    # Don't print until a whole line is available
    print_end = False

    for _ in range(max_new_tokens):
        if current_ids.shape[1] >= config.seq_len:
            current_ids = current_ids[:, -config.seq_len:]

        with torch.no_grad():
            logits, _, _ = model(current_ids)

        next_token_logits = logits[0, -1, :] / temperature
        probs = torch.softmax(next_token_logits, dim=-1)
        next_token_id = torch.multinomial(probs, num_samples=1).item()

        # Append new token
        current_ids = torch.cat(
            [current_ids, torch.tensor([[next_token_id]], device=device)], dim=1
        )

        # Append to buffer
        buffer += tokenizer.decode([next_token_id])

        # Print as many complete lines as possible
        while len(buffer) >= line_width:
            print(buffer[:line_width])
            buffer = buffer[line_width:]
            print_end = True

    # Print remaining buffer (incomplete final line)
    if buffer:
        print(buffer)

# Test streaming generation of 500 tokens, 71 chars per line
input_text = "The future of AI is"
autoregressive_generate(model, tokenizer, input_text, max_new_tokens=500, temperature=config.temperature, line_width=71)