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)