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)