mlx-lm-utils 0.0.1

// Args:
//     conversation (Union[list[dict[str, str]], list[list[dict[str, str]]]]): A list of dicts
//         with "role" and "content" keys, representing the chat history so far.
//     tools (`list[Union[Dict, Callable]]`, *optional*):
//         A list of tools (callable functions) that will be accessible to the model. If the template does not
//         support function calling, this argument will have no effect. Each tool should be passed as a JSON Schema,
//         giving the name, description and argument types for the tool. See our
//         [chat templating guide](https://huggingface.co/docs/transformers/main/en/chat_templating#automated-function-conversion-for-tool-use)
//         for more information.
//     documents (`list[dict[str, str]]`, *optional*):
//         A list of dicts representing documents that will be accessible to the model if it is performing RAG
//         (retrieval-augmented generation). If the template does not support RAG, this argument will have no
//         effect. We recommend that each document should be a dict containing "title" and "text" keys. Please
//         see the RAG section of the [chat templating guide](https://huggingface.co/docs/transformers/main/en/chat_templating#arguments-for-RAG)
//         for examples of passing documents with chat templates.
//     chat_template (`str`, *optional*):
//         A Jinja template to use for this conversion. It is usually not necessary to pass anything to this
//         argument, as the model's template will be used by default.
//     add_generation_prompt (bool, *optional*):
//         If this is set, a prompt with the token(s) that indicate
//         the start of an assistant message will be appended to the formatted output. This is useful when you want to generate a response from the model.
//         Note that this argument will be passed to the chat template, and so it must be supported in the
//         template for this argument to have any effect.
//     continue_final_message (bool, *optional*):
//         If this is set, the chat will be formatted so that the final
//         message in the chat is open-ended, without any EOS tokens. The model will continue this message
//         rather than starting a new one. This allows you to "prefill" part of
//         the model's response for it. Cannot be used at the same time as `add_generation_prompt`.
//     tokenize (`bool`, defaults to `True`):
//         Whether to tokenize the output. If `False`, the output will be a string.
//     padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`):
//             Select a strategy to pad the returned sequences (according to the model's padding side and padding
//             index) among:

//         - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
//             sequence if provided).
//         - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
//             acceptable input length for the model if that argument is not provided.
//         - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different
//             lengths).
//     truncation (`bool`, defaults to `False`):
//         Whether to truncate sequences at the maximum length. Has no effect if tokenize is `False`.
//     max_length (`int`, *optional*):
//         Maximum length (in tokens) to use for padding or truncation. Has no effect if tokenize is `False`. If
//         not specified, the tokenizer's `max_length` attribute will be used as a default.
//     return_tensors (`str` or [`~utils.TensorType`], *optional*):
//         If set, will return tensors of a particular framework. Has no effect if tokenize is `False`. Acceptable
//         values are:
//         - `'tf'`: Return TensorFlow `tf.Tensor` objects.
//         - `'pt'`: Return PyTorch `torch.Tensor` objects.
//         - `'np'`: Return NumPy `np.ndarray` objects.
//         - `'jax'`: Return JAX `jnp.ndarray` objects.
//     return_dict (`bool`, defaults to `False`):
//         Whether to return a dictionary with named outputs. Has no effect if tokenize is `False`.
//     tokenizer_kwargs (`dict[str: Any]`, *optional*): Additional kwargs to pass to the tokenizer.
//     return_assistant_tokens_mask (`bool`, defaults to `False`):
//         Whether to return a mask of the assistant generated tokens. For tokens generated by the assistant,
//         the mask will contain 1. For user and system tokens, the mask will contain 0.
//         This functionality is only available for chat templates that support it via the `{% generation %}` keyword.
//     **kwargs: Additional kwargs to pass to the template renderer. Will be accessible by the chat template.

// Returns:
//     `Union[list[int], Dict]`: A list of token ids representing the tokenized chat so far, including control tokens. This
//     output is ready to pass to the model, either directly or via methods like `generate()`. If `return_dict` is
//     set, will return a dict of tokenizer outputs instead.
// """

use std::{
    collections::HashMap,
    fs::read_to_string,
    ops::{Deref, DerefMut},
    path::Path,
    str::FromStr,
};

use minijinja::{context, Environment, Template};
use serde::Serialize;
use tokenizers::Encoding;

use crate::error::Error;

/// Wrapper around [`tokenizers::Tokenizer`] and [`minijinja::Environment`]
/// providing more utilities.
pub struct Tokenizer {
    inner: tokenizers::Tokenizer,
    env: Environment<'static>,
}

impl FromStr for Tokenizer {
    type Err = tokenizers::Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        tokenizers::Tokenizer::from_str(s).map(Self::from_tokenizer)
    }
}

impl Tokenizer {
    pub fn from_tokenizer(tokenizer: tokenizers::Tokenizer) -> Self {
        let mut env = Environment::new();
        env.set_unknown_method_callback(minijinja_contrib::pycompat::unknown_method_callback);
        Self {
            inner: tokenizer,
            env,
        }
    }

    pub fn from_file(file: impl AsRef<Path>) -> tokenizers::Result<Self> {
        tokenizers::Tokenizer::from_file(file).map(Self::from_tokenizer)
    }

    pub fn from_bytes(bytes: impl AsRef<[u8]>) -> tokenizers::Result<Self> {
        tokenizers::Tokenizer::from_bytes(bytes).map(Self::from_tokenizer)
    }

    pub fn apply_chat_template<'a, I, R, T>(
        &'a mut self,
        model_template: String,
        args: ApplyChatTemplateArgs<'a, I, R, T>,
    ) -> Result<Vec<String>, Error>
    where
        I: IntoIterator<Item = Chat<'a, R, T>>,
        R: Serialize + 'a,
        T: Serialize + ToString + 'a,
    {
        apply_chat_template(&mut self.env, model_template, args)
    }

    pub fn apply_chat_template_and_encode<'a, I, R, T>(
        &mut self,
        model_template: String,
        args: ApplyChatTemplateArgs<'a, I, R, T>,
    ) -> Result<Vec<Encoding>, Error>
    where
        I: IntoIterator<Item = Chat<'a, R, T>>,
        R: Serialize + 'a,
        T: Serialize + ToString + 'a,
    {
        let Self { inner, env } = self;

        let rendered_chats = apply_chat_template(env, model_template, args)?;
        inner
            .encode_batch(rendered_chats, false)
            .map_err(Into::into)
    }
}

impl Deref for Tokenizer {
    type Target = tokenizers::Tokenizer;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl DerefMut for Tokenizer {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

#[derive(Debug, Clone, Copy, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum Role {
    User,
    Assistant,
}

#[derive(Debug, Clone, Serialize)]
pub enum Content {
    String(String),
    Map(HashMap<String, String>),
}

#[derive(Debug, Clone, Serialize)]
pub struct Conversation<R, T> {
    pub role: R,
    pub content: T,
}

#[derive(Debug, Clone, Serialize)]
#[serde(untagged)]
pub enum Chat<'a, R, T> {
    Borrowed(&'a [Conversation<R, T>]),
    Owned(Vec<Conversation<R, T>>),
}

impl<R, T> Deref for Chat<'_, R, T> {
    type Target = [Conversation<R, T>];

    fn deref(&self) -> &Self::Target {
        match self {
            Chat::Borrowed(conversations) => conversations,
            Chat::Owned(conversations) => conversations,
        }
    }
}

impl<R, T> From<Vec<Conversation<R, T>>> for Chat<'_, R, T> {
    fn from(value: Vec<Conversation<R, T>>) -> Self {
        Chat::Owned(value)
    }
}

impl<'a, R, T> From<&'a [Conversation<R, T>]> for Chat<'a, R, T> {
    fn from(value: &'a [Conversation<R, T>]) -> Self {
        Chat::Borrowed(value)
    }
}

#[derive(Debug, Clone, Serialize)]
pub struct Document {
    pub title: String,
    pub text: String,
}

pub enum Padding {
    Longest,
    MaxLength,
}

pub enum Truncation {
    MaxLength(usize),
}

#[derive(Default)]
pub struct ApplyChatTemplateArgs<'a, I, R = Role, T = String>
where
    I: IntoIterator<Item = Chat<'a, R, T>>,
    R: Serialize + 'a,
    T: Serialize + ToString + 'a,
{
    // pub conversations: &'a [Conversation<R, T>],
    pub conversations: I,
    // pub tools: Option<Box<dyn FnOnce()>>, // TODO
    pub documents: Option<&'a [Document]>,
    pub model_id: &'a str,
    pub chat_template_id: Option<&'a str>,
    pub add_generation_prompt: Option<bool>,
    pub continue_final_message: Option<bool>,
}

pub fn load_model_chat_template_from_str(content: &str) -> std::io::Result<Option<String>> {
    serde_json::from_str::<serde_json::Value>(content)
        .map(|value| {
            value
                .get("chat_template")
                .and_then(|value| value.as_str())
                .map(ToString::to_string)
        })
        .map_err(Into::into)
}

pub fn load_model_chat_template_from_file(
    file: impl AsRef<Path>,
) -> std::io::Result<Option<String>> {
    let content = read_to_string(file)?;
    load_model_chat_template_from_str(&content)
}

// chat_template = self.get_chat_template(chat_template, tools)

// if isinstance(conversation, (list, tuple)) and (
//     isinstance(conversation[0], (list, tuple)) or hasattr(conversation[0], "messages")
// ):
//     conversations = conversation
//     is_batched = True
// else:
//     conversations = [conversation]
//     is_batched = False

// if continue_final_message:
//     if add_generation_prompt:
//         raise ValueError(
//             "continue_final_message and add_generation_prompt are not compatible. Use continue_final_message when you want the model to continue the final message, and add_generation_prompt when you want to add a header that will prompt it to start a new assistant message instead."
//         )
//     if return_assistant_tokens_mask:
//         raise ValueError("continue_final_message is not compatible with return_assistant_tokens_mask.")

// template_kwargs = {**self.special_tokens_map, **kwargs}  # kwargs overwrite special tokens if both are present
// rendered_chat, generation_indices = render_jinja_template(
//     conversations=conversations,
//     tools=tools,
//     documents=documents,
//     chat_template=chat_template,
//     return_assistant_tokens_mask=return_assistant_tokens_mask,
//     continue_final_message=continue_final_message,
//     add_generation_prompt=add_generation_prompt,
//     **template_kwargs,
// )

// if not is_batched:
//     rendered_chat = rendered_chat[0]

// if tokenize:
//     out = self(
//         rendered_chat,
//         padding=padding,
//         truncation=truncation,
//         max_length=max_length,
//         add_special_tokens=False,
//         return_tensors=return_tensors,
//         **tokenizer_kwargs,
//     )
//     if return_dict:
//         if return_assistant_tokens_mask:
//             assistant_masks = []
//             if is_batched or return_tensors:
//                 input_ids = out["input_ids"]
//             else:
//                 input_ids = [out["input_ids"]]
//             for i in range(len(input_ids)):
//                 current_mask = [0] * len(input_ids[i])
//                 for assistant_start_char, assistant_end_char in generation_indices[i]:
//                     start_token = out.char_to_token(i, assistant_start_char)
//                     end_token = out.char_to_token(i, assistant_end_char - 1)
//                     if start_token is None:
//                         # start_token is out of bounds maybe due to truncation.
//                         break
//                     for token_id in range(start_token, end_token + 1 if end_token else len(input_ids[i])):
//                         current_mask[token_id] = 1
//                 assistant_masks.append(current_mask)

//             if not is_batched and not return_tensors:
//                 assistant_masks = assistant_masks[0]

//             out["assistant_masks"] = assistant_masks

//             if return_tensors:
//                 out.convert_to_tensors(tensor_type=return_tensors)

//         return out
//     else:
//         return out["input_ids"]
// else:
//     return rendered_chat

// def render_jinja_template(
//     conversations: list[list[dict[str, str]]],
//     tools: Optional[list[Union[dict, Callable]]] = None,
//     documents: Optional[list[dict[str, str]]] = None,
//     chat_template: Optional[str] = None,
//     return_assistant_tokens_mask: Optional[bool] = False,
//     continue_final_message: Optional[bool] = False,
//     add_generation_prompt: Optional[bool] = False,
//     **kwargs,
// ) -> str:
//     if return_assistant_tokens_mask and not re.search(r"\{\%-?\s*generation\s*-?\%\}", chat_template):
//         logger.warning_once(
//             "return_assistant_tokens_mask==True but chat template does not contain `{% generation %}` keyword."
//         )

//     # Compilation function uses a cache to avoid recompiling the same template
//     compiled_template = _compile_jinja_template(chat_template)

//     # We accept either JSON schemas or functions for tools. If we get functions, we convert them to schemas
//     if tools is not None:
//         tool_schemas = []
//         for tool in tools:
//             if isinstance(tool, dict):
//                 tool_schemas.append(tool)
//             elif isfunction(tool):
//                 tool_schemas.append(get_json_schema(tool))
//             else:
//                 raise ValueError(
//                     "Tools should either be a JSON schema, or a callable function with type hints "
//                     "and a docstring suitable for auto-conversion to a schema."
//                 )
//     else:
//         tool_schemas = None

//     if documents is not None:
//         for document in documents:
//             if not isinstance(document, dict):
//                 raise TypeError("Documents should be a list of dicts with 'title' and 'text' keys!")

//     rendered = []
//     all_generation_indices = []
//     for chat in conversations:
//         if hasattr(chat, "messages"):
//             # Indicates it's a Conversation object
//             chat = chat.messages
//         if return_assistant_tokens_mask:
//             rendered_chat, generation_indices = _render_with_assistant_indices(
//                 compiled_template=compiled_template,
//                 messages=chat,
//                 tools=tool_schemas,
//                 documents=documents,
//                 add_generation_prompt=add_generation_prompt,
//                 **kwargs,
//             )
//             all_generation_indices.append(generation_indices)
//         else:
//             rendered_chat = compiled_template.render(
//                 messages=chat,
//                 tools=tool_schemas,
//                 documents=documents,
//                 add_generation_prompt=add_generation_prompt,
//                 **kwargs,
//             )
//         if continue_final_message:
//             final_message = chat[-1]["content"]
//             if isinstance(final_message, (list, tuple)):
//                 for content_block in reversed(final_message):
//                     if "text" in content_block:
//                         # Pick the last text block in the message (the first one we hit while iterating in reverse)
//                         final_message = content_block["text"]
//                         break
//                 else:
//                     raise ValueError(
//                         "continue_final_message is set but we could not find any text to continuein the final message!"
//                     )
//             if final_message.strip() not in rendered_chat:
//                 raise ValueError(
//                     "continue_final_message is set but the final message does not appear in the chat after "
//                     "applying the chat template! This can happen if the chat template deletes portions of "
//                     "the final message. Please verify the chat template and final message in your chat to "
//                     "ensure they are compatible."
//                 )
//             final_msg_loc = rendered_chat.rindex(final_message.strip())
//             if rendered_chat[final_msg_loc : final_msg_loc + len(final_message.lstrip())] == final_message:
//                 # The template preserves spacing or the message doesn't have trailing spacing, so things are simple
//                 rendered_chat = rendered_chat[: final_msg_loc + len(final_message.lstrip())]
//             else:
//                 # The message has trailing spacing that was trimmed, so we must be more cautious
//                 rendered_chat = rendered_chat[: final_msg_loc + len(final_message.strip())]
//         rendered.append(rendered_chat)

//     return rendered, all_generation_indices

pub fn apply_chat_template<'a, I, R, T>(
    env: &mut Environment<'static>,
    model_template: String,
    args: ApplyChatTemplateArgs<'a, I, R, T>,
) -> Result<Vec<String>, Error>
where
    I: IntoIterator<Item = Chat<'a, R, T>>,
    R: Serialize + 'a,
    T: Serialize + ToString + 'a,
{
    let ApplyChatTemplateArgs {
        conversations,
        // tools,
        documents,
        model_id,
        chat_template_id,
        add_generation_prompt,
        continue_final_message,
    } = args;

    let add_generation_prompt = add_generation_prompt.unwrap_or(false);
    let continue_final_message = continue_final_message.unwrap_or(false);

    let template = match chat_template_id {
        Some(chat_template_id) => env.get_template(chat_template_id)?,
        None => match env.get_template(model_id) {
            Ok(template) => template,
            Err(_) => {
                env.add_template_owned(model_id.to_owned(), model_template)?;
                env.get_template(model_id)
                    .expect("Newly added template must be present")
            }
        },
    };

    // TODO: handle tool

    // TODO: allow return_generation_indices

    render_jinja_tempalte(
        template,
        conversations,
        documents,
        Some(add_generation_prompt),
        Some(continue_final_message),
    )
}

// TODO: render with assistant indices
fn render_jinja_tempalte<'a, R, T>(
    template: Template,
    conversations: impl IntoIterator<Item = Chat<'a, R, T>>,
    documents: Option<&'a [Document]>,
    add_generation_prompt: Option<bool>,
    continue_final_message: Option<bool>,
) -> Result<Vec<String>, Error>
where
    R: Serialize + 'a,
    T: Serialize + ToString + 'a,
{
    let add_generation_prompt = add_generation_prompt.unwrap_or(false);
    let continue_final_message = continue_final_message.unwrap_or(false);

    // TODO: what does checking for "messages" key do in the python code?
    let mut rendered = Vec::new();
    for chat in conversations {
        let mut rendered_chat = template.render(context! {
            messages => chat,
            documents => documents,
            add_generation_prompt => add_generation_prompt,
        })?;

        if continue_final_message {
            let Some(final_message) = chat.last().map(|chat| &chat.content) else {
                continue;
            };

            let final_message_str = final_message.to_string();

            if !rendered_chat.contains(final_message_str.trim()) {
                return Err(Error::FinalMsgNotInChat);
            }

            let final_msg_loc = rendered_chat.rfind(&final_message_str.trim()).unwrap();
            let final_msg_len = final_message_str.trim_start().len();
            rendered_chat = if rendered_chat[final_msg_loc..final_msg_loc + final_msg_len]
                == final_message_str
            {
                // The template preserves spacing or the message doesn't have trailing spacing, so things are simple
                rendered_chat[..final_msg_loc + final_msg_len].to_string()
            } else {
                // The message has trailing spacing that was trimmed, so we must be more cautious
                rendered_chat[..final_msg_loc + final_message_str.trim().len()].to_string()
            };
        }
        rendered.push(rendered_chat);
    }

    Ok(rendered)
}

#[cfg(test)]
mod tests {
    use minijinja::Environment;
    use std::path::Path;

    use crate::tokenizer::{
        apply_chat_template, load_model_chat_template_from_file, ApplyChatTemplateArgs,
        Conversation, Role,
    };

    // let model_id = "mlx-community/Qwen3-4B-bf16".to_string();
    const CACHED_TEST_MODEL_DIR: &str = "../cache/Qwen3-4B-bf16";

    // TODO: how to test this in CI? the model files might be too large
    #[test]
    fn test_load_chat_template_from_file() {
        let file = Path::new(CACHED_TEST_MODEL_DIR).join("tokenizer_config.json");
        let chat_template = load_model_chat_template_from_file(file).unwrap().unwrap();
        assert!(!chat_template.is_empty());
    }

    #[test]
    fn test_apply_chat_template() {
        let file = Path::new(CACHED_TEST_MODEL_DIR).join("tokenizer_config.json");
        let model_chat_template = load_model_chat_template_from_file(file).unwrap().unwrap();
        assert!(!model_chat_template.is_empty());

        let model_id = "mlx-community/Qwen3-4B-bf16".to_string();
        let conversations = vec![Conversation {
            role: Role::User,
            content: "hello",
        }];
        let args = ApplyChatTemplateArgs {
            conversations: [conversations.into()],
            documents: None,
            model_id: &model_id,
            chat_template_id: None,
            add_generation_prompt: None,
            continue_final_message: None,
        };

        let mut env = Environment::new();
        env.set_unknown_method_callback(minijinja_contrib::pycompat::unknown_method_callback);

        let rendered_chat = apply_chat_template(&mut env, model_chat_template, args).unwrap();
        println!("{:?}", rendered_chat);
    }

    #[test]
    fn test_tokenizer_apply_chat_template() {
        let tokenizer_file = Path::new(CACHED_TEST_MODEL_DIR).join("tokenizer.json");
        let tokenizer_config_file = Path::new(CACHED_TEST_MODEL_DIR).join("tokenizer_config.json");

        let model_id = "mlx-community/Qwen3-4B-bf16".to_string();

        let conversations = vec![Conversation {
            role: Role::User,
            content: "hello",
        }];

        let mut tokenizer = super::Tokenizer::from_file(tokenizer_file).unwrap();

        let model_chat_template = load_model_chat_template_from_file(tokenizer_config_file)
            .unwrap()
            .unwrap();
        assert!(!model_chat_template.is_empty());

        let args = ApplyChatTemplateArgs {
            conversations: [conversations.into()],
            documents: None,
            model_id: &model_id,
            chat_template_id: None,
            add_generation_prompt: None,
            continue_final_message: None,
        };

        let rendered_chat = tokenizer
            .apply_chat_template(model_chat_template, args)
            .unwrap();
        println!("{:?}", rendered_chat);
    }

    #[test]
    fn test_tokenizer_apply_chat_template_and_encode() {
        let tokenizer_file = Path::new(CACHED_TEST_MODEL_DIR).join("tokenizer.json");
        let tokenizer_config_file = Path::new(CACHED_TEST_MODEL_DIR).join("tokenizer_config.json");

        let model_id = "mlx-community/Qwen3-4B-bf16".to_string();

        let conversations = vec![Conversation {
            role: Role::User,
            content: "hello",
        }];
        let mut tokenizer = super::Tokenizer::from_file(tokenizer_file).unwrap();

        let model_chat_template = load_model_chat_template_from_file(tokenizer_config_file)
            .unwrap()
            .unwrap();
        assert!(!model_chat_template.is_empty());

        let args = ApplyChatTemplateArgs {
            conversations: [conversations.into()],
            documents: None,
            model_id: &model_id,
            chat_template_id: None,
            add_generation_prompt: None,
            continue_final_message: None,
        };

        let encodings = tokenizer
            .apply_chat_template_and_encode(model_chat_template, args)
            .unwrap();
        println!("{:?}", encodings.iter().map(|e| e.get_ids()).flatten());
    }
}