File size: 552 Bytes
283ba4a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 | 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)
|