| import json |
|
|
| import mlx.core as mx |
| import mlx.nn as nn |
|
|
| from src.tokenizer import BPETokenizer |
| from src.transformer import MLXGPT as GPT |
|
|
|
|
| def build_model(config_path: str = "config.json") -> GPT: |
| with open(config_path) as f: |
| config = json.load(f) |
|
|
| return GPT( |
| vocab_size=config["vocab_size"], |
| embedding_dim=config["embedding_dim"], |
| num_layers=config["num_layers"], |
| num_heads=config["num_heads"], |
| num_query_groups=config["num_query_groups"], |
| hidden_dim=config["hidden_dim"], |
| max_seq_len=config["max_seq_len"], |
| dropout=0.0, |
| ) |
|
|
|
|
| def main(): |
| device = mx.default_device() |
| print(f"Running on {device}") |
|
|
| model = build_model() |
| model.load_weights("model.safetensors") |
| model.eval() |
| mx.eval(model.parameters()) |
|
|
| tokenizer = BPETokenizer.load("tokenizer") |
| prompt = "The future of machine learning" |
| prompt_ids = mx.array([tokenizer.encode(prompt)], dtype=mx.int32) |
|
|
| generated = model.generate( |
| prompt_ids, |
| max_new_tokens=40, |
| temperature=0.8, |
| top_k=50, |
| ) |
| mx.eval(generated) |
| print(tokenizer.decode(generated[0].tolist())) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
