asimov_mlx_module/
lib.rs

1// This is free and unencumbered software released into the public domain.
2
3#![no_std]
4#![forbid(unsafe_code)]
5
6#[cfg(feature = "std")]
7extern crate std;
8
9use asimov_module::{prelude::*, tracing};
10use core::error::Error;
11
12extern crate alloc;
13use alloc::vec;
14
15#[derive(Clone, Debug, bon::Builder)]
16#[builder(on(String, into))]
17pub struct Options {
18    #[builder(default = "mlx-community/Llama-3.2-3B-Instruct-4bit")]
19    pub model: String,
20}
21
22#[cfg(feature = "std")]
23pub fn generate(input: impl AsRef<str>, options: &Options) -> Result<Vec<String>, Box<dyn Error>> {
24    use std::process::Stdio;
25
26    let mut cmd = std::process::Command::new("mlx_lm.generate");
27
28    cmd.env("NO_COLOR", "1"); // See: https://no-color.org
29    cmd.stdin(Stdio::null());
30    cmd.stdout(Stdio::piped());
31    cmd.stderr(Stdio::piped());
32    cmd.args(["--model", &options.model]);
33    cmd.args(["--prompt", input.as_ref()]);
34    cmd.args(["--verbose", "False"]);
35
36    let output = cmd
37        .spawn()
38        .inspect_err(|err| tracing::error!(?err, "unable to run mlx"))?
39        .wait_with_output()
40        .inspect_err(|err| tracing::error!(?err, "mlx execution failed"))?;
41
42    if !output.stderr.is_empty() {
43        String::from_utf8(output.stderr)
44            .inspect(|stderr| tracing::debug!(stderr))
45            .ok();
46    }
47
48    let stdout = String::from_utf8(output.stdout)?;
49
50    Ok(vec![stdout])
51}