from airllm import AutoModel
import mlx.core as mx

MAX_LENGTH = 128

model = AutoModel.from_pretrained("mistralai/Mistral-7B-v0.1")

model.tokenizer.pad_token = model.tokenizer.eos_token

input_text = ["write a C++ code of a simple calculator"]

input_tokens = model.tokenizer(
    input_text, return_tensors="pt", truncation=True, max_length=MAX_LENGTH
)

input_ids = mx.array(input_tokens["input_ids"].tolist())

generation_output = model.generate(input_ids, max_new_tokens=300)

print(generation_output)