MLX.zig

MLX.zig offers an efficient, lightweight interface between the Zig programming language and Apple's MLX framework, specifically designed for machine learning on Apple Silicon devices. This repository serves as a powerful tool for developers looking to leverage Zig to create and execute machine learning models without reliance on additional build tools.

Key Features

• Pure Zig Build System: Builds without the need for CMake or other external tools, streamlining the development process.
• Comprehensive Dependency Management: Automatically resolves all necessary C/C++ dependencies such as MLX, MLX-C, and PCRE2 through Zig's robust build system.
• Pre-built Language Model Example: Contains a working example of a transformer-based language model, complete with tokenizer and text generation capabilities.
• Efficient Tokenization: Utilizes PCRE2 for high-speed text processing, ensuring reliable tokenization.
• Direct MLX Access: Provides direct bindings to the MLX C API, ensuring maximum performance when executing machine learning tasks.

Installation

This project requires an Apple Silicon Mac and Zig version 0.13.0. The complete setup involves cloning the repository and downloading the required model weights for deployment.

Example Usage

// Load tokenizer
var tokenizer = try Tokenizer.init(allocator, null);
defer tokenizer.deinit();

// Load transformer
var transformer = try Transformer.init(allocator, null);
defer transformer.deinit();

// Encode input string to token IDs (chat format)
const input_ids = try tokenizer.encodeChat(allocator, "You are a helpful assistant.", user_input);
defer allocator.free(input_ids);

// Generate new tokens
const output_ids = try transformer.generate(input_ids, num_tokens_to_generate);
defer allocator.free(output_ids);

In this example, tokens are encoded and processed using MLX’s tensor operations, allowing for the generation of new text based on user input.

Project Architecture

MLX.zig integrates seamlessly with Apple's ML framework through several core components:

1. Zig Build System: Compiles MLX, MLX-C, and PCRE2 from source, eliminating external dependencies.
2. Transformer Implementation: Features a Llama-style language model featuring attention mechanisms and caching for efficiency.
3. Tokenizer: Implements regex-based processing using PCRE2, capable of handling complex text patterns and special token scenarios.

Overall, MLX.zig simplifies the process of building machine learning applications on Apple Silicon, providing an effective alternative for developers seeking to utilize the performance benefits of Zig.