main function

void main()

Implementation

void main() {
  print("--- Stable Tensor-Engine MuZero-GPT: Training & Generation ---");

  final Map<String, int> stoi = {
    "hello": 0,
    "world": 1,
    "the": 2,
    "quick": 3,
    "brown": 4,
    "fox": 5,
    ".": 6,
    "<start>": 7,
    "jumps": 8,
    "over": 9,
    "lazy": 10,
    "dog": 11,
  };
  final Map<int, String> itos = stoi.map((k, v) => MapEntry(v, k));

  final model = TransformerDecoder(
    vocabSize: 12,
    embedSize: 64, // Increased slightly for better latent separation
    encoderEmbedSize: 64,
    blockSize: 16,
    numLayers: 2,
    numHeads: 4,
  );

  final agent = MuZeroGreedyAgent(model, 64);

  // The full sequence to learn
  final List<int> trainingData = [
    stoi["<start>"]!,
    stoi["the"]!,
    stoi["quick"]!,
    stoi["brown"]!,
    stoi["fox"]!,
    stoi["jumps"]!,
    stoi["over"]!,
    stoi["the"]!,
    stoi["lazy"]!,
    stoi["dog"]!,
    stoi["."]!,
  ];

  // --- TRAIN ---
  // We use 600 epochs to give the alternating modes enough time to co-evolve
  trainMuZero(agent, trainingData, 1000);

  print("\n--- Generation After Training ---");
  // We start with "<start>" and let the Dynamics head imagine the rest
  generateMuZeroPure(agent, [stoi["<start>"]!], 12, itos);
}