| 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) | |