| from airllm import AutoModel | |
| import mlx.core as mx | |
| MAX_LENGTH = 128 | |
| model = AutoModel.from_pretrained("microsoft/Phi-3-mini-4k-instruct") | |
| model.tokenizer.pad_token = model.tokenizer.eos_token | |
| input_text = ["Write a C++ code for a simple calculator"] | |
| input_tokens = model.tokenizer( | |
| input_text, return_tensors="pt", truncation=True, max_length=MAX_LENGTH | |
| ) | |
| # convert torch tensor -> mlx array | |
| input_ids = mx.array(input_tokens["input_ids"].tolist()) | |
| generation_output = model.generate(input_ids, max_new_tokens=200) | |
| print(generation_output) | |