File size: 1,911 Bytes
e3e33cf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | """
Minimal example: load Quartz Micro Preview Base's weights and generate text.
This is a custom PyTorch architecture (DeepSeek-style fine-grained MoE, see
model.py / config.py in this repo) -- not a registered transformers
architecture, so there's no AutoModel support. Load it directly:
pip install torch safetensors tokenizers
python load_model.py
The tokenizer is a from-scratch 32K-vocab byte-level BPE (tokenizer/
vocab.json + merges.txt) trained on this model's own pretraining corpus --
it is NOT compatible with any other model's tokenizer, and no other
tokenizer is compatible with these weights.
"""
import os
import torch
from safetensors.torch import load_file
from tokenizers import ByteLevelBPETokenizer
from config import MoEConfig
from model import MoELanguageModel
HERE = os.path.dirname(__file__)
def load():
tok = ByteLevelBPETokenizer(
os.path.join(HERE, "tokenizer", "vocab.json"),
os.path.join(HERE, "tokenizer", "merges.txt"),
)
cfg = MoEConfig()
model = MoELanguageModel(cfg)
state_dict = load_file(os.path.join(HERE, "model.safetensors"))
model.load_state_dict(state_dict)
model.eval()
return model, tok
@torch.no_grad()
def generate(model, tok, prompt: str, max_new_tokens: int = 80, temperature: float = 0.8):
BOS_ID, EOS_ID = 1, 2
ids = [BOS_ID] + tok.encode(prompt).ids
x = torch.tensor([ids], dtype=torch.long)
for _ in range(max_new_tokens):
logits, _ = model(x)
next_logits = logits[0, -1] / temperature
probs = torch.softmax(next_logits, dim=-1)
next_id = torch.multinomial(probs, 1).item()
if next_id == EOS_ID:
break
x = torch.cat([x, torch.tensor([[next_id]])], dim=1)
return tok.decode(x[0, 1:].tolist())
if __name__ == "__main__":
model, tok = load()
print(generate(model, tok, "The history of the Roman Empire"))
|