generate.py

"""
Unified Autoregressive Generation & MoE Routing CLI for Sparse-AST / BWM
"""

import os
import argparse
import torch
import torch.nn.functional as F
from model import SparseAST, TopKSparseASTEnsemble

def generate_backbone(model, prompt, max_new_tokens=40, temperature=0.6, top_k=40):
    model.eval()
    device = next(model.parameters()).device
    tokens = list(prompt.encode("utf-8", "ignore"))
    
    with torch.no_grad():
        for _ in range(max_new_tokens):
            x = torch.tensor([tokens[-model.seq_len:]], dtype=torch.long, device=device)
            logits = model(x)[:, -1, :]
            
            if temperature <= 0.05:
                next_tok = int(logits.argmax(dim=-1).item())
            else:
                scaled = logits / temperature
                if top_k > 0:
                    v, _ = torch.topk(scaled, min(top_k, scaled.size(-1)))
                    scaled[scaled < v[:, [-1]]] = -float("inf")
                probs = F.softmax(scaled, dim=-1)
                next_tok = int(torch.multinomial(probs, num_samples=1).item())
            tokens.append(next_tok)
            
    return bytes([t for t in tokens if t < 256]).decode("utf-8", errors="replace")

def main():
    parser = argparse.ArgumentParser(description="Generate Blender Python code using Sparse-AST / BWM")
    parser.add_argument("--model", type=str, default="100M-32", help="Model variant: 3M-32, 10M-32, 100M-32, 200M-32, moe")
    parser.add_argument("--prompt", type=str, default="import bpy\n# Create primitive cube\nbpy.ops.mesh.", help="Prompt text")
    parser.add_argument("--tokens", type=int, default=40, help="Max tokens to generate")
    parser.add_argument("--temperature", type=float, default=0.6, help="Sampling temperature")
    args = parser.parse_args()
    
    repo_dir = os.path.dirname(os.path.abspath(__file__))
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    print("=" * 80)
    print(f"   SPARSE-AST / BWM CODE GENERATION: Model [{args.model}] on [{device}]")
    print("=" * 80)
    
    if args.model.lower() in ["moe", "topk-moe"]:
        ensemble = TopKSparseASTEnsemble.from_pretrained(repo_dir, subfolder="TopK-MoE", device=device)
        x = torch.tensor([list(args.prompt.encode("utf-8", "ignore"))], dtype=torch.long, device=device)
        _, topk_indices, gate_weights = ensemble(x, k=2)
        print(f"\nPrompt: {repr(args.prompt)}")
        print("Top-2 Experts selected for next token:")
        for idx, w in zip(topk_indices[0, -1].tolist(), gate_weights[0, -1].tolist()):
            exp_name = ensemble.expert_names[idx] if idx < len(ensemble.expert_names) else f"Expert {idx}"
            print(f"  {exp_name}: Weight {w:.4f}")
    else:
        sub = args.model if args.model.endswith("-32") or args.model.endswith("-512") else f"{args.model}-32"
        model = SparseAST.from_pretrained(repo_dir, subfolder=sub, device=device)
        print(f"\nPrompt:\n{args.prompt}")
        completion = generate_backbone(model, args.prompt, max_new_tokens=args.tokens, temperature=args.temperature)
        print(f"\nGenerated Completion:\n{completion}")

if __name__ == "__main__":
    main()