| from airllm import AutoModel | |
| import mlx.core as mx | |
| MAX_LENGTH = 128 | |
| # load llama model | |
| model = AutoModel.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct") | |
| model.tokenizer.pad_token = model.tokenizer.eos_token | |
| input_text = ["Write a C++ program for a simple calculator"] | |
| input_tokens = model.tokenizer( | |
| input_text, return_tensors="pt", truncation=True, max_length=MAX_LENGTH | |
| ) | |
| # convert to MLX tensor | |
| input_ids = mx.array(input_tokens["input_ids"].tolist()) | |
| output = model.generate(input_ids, max_new_tokens=300) | |
| print(output) | |