MEGA

A from-scratch Mixture-of-Experts language model — reference implementation

Load and run

pip install "git+https://github.com/Frankenstein-Labs/mega-moe.git"

from mega.hub.serialization import from_pretrained
from mega.tokenizer.tokenizer import MegaTokenizer
import torch

model = from_pretrained("Frankenstein-Labs/mega-small")
tok = MegaTokenizer.from_pretrained("Frankenstein-Labs/mega-small")

ids = [tok.bos_token_id, *tok.encode("17 + 16 =")]
out = model(torch.tensor([ids]))
print("17 + 16 =", tok.decode([int(out.logits[0, -1].argmax())]))

Verified results