load chat_template from file

Author: minghuaw

mlx-lm-utils/Cargo.toml

@@ -10,7 +10,8 @@ license.workspace = true
 documentation.workspace = true
 
 [dependencies]
-minijinja = "2"
+minijinja = { version = "2", features = ["loader"] }
 serde = { version = "1", features = ["derive"] }
+serde_json = "1"
 thiserror = "2"
 tokenizers = "0.21"

mlx-lm-utils/src/tokenizer.rs

@@ -65,7 +65,7 @@
 //     set, will return a dict of tokenizer outputs instead.
 // """
 
-use std::{borrow::Cow, collections::HashMap};
+use std::{borrow::Cow, collections::HashMap, fs::read_to_string, path::Path};
 
 use minijinja::{context, Environment, Template};
 use serde::{Deserialize, Serialize};
@@ -74,7 +74,7 @@ use crate::error::Error;
 
 #[derive(Serialize)]
 #[serde(untagged)]
-pub enum Content<T: Serialize=()> {
+pub enum Content<T: Serialize = ()> {
     String(String),
     Map(HashMap<String, String>),
     Typed(T),
@@ -115,6 +115,21 @@ pub struct TokenizeOptions {
     pub return_assistant_tokens_mask: Option<bool>,
 }
 
+pub fn load_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_chat_template_from_file(file: impl AsRef<Path>) -> std::io::Result<Option<String>> {
+    let content = read_to_string(file)?;
+    load_chat_template_from_str(&content)
+}
+
 // chat_template = self.get_chat_template(chat_template, tools)
 
 // if isinstance(conversation, (list, tuple)) and (
@@ -192,7 +207,6 @@ pub struct TokenizeOptions {
 // else:
 //     return rendered_chat
 
-
 // def render_jinja_template(
 //     conversations: list[list[dict[str, str]]],
 //     tools: Optional[list[Union[dict, Callable]]] = None,
@@ -286,7 +300,6 @@ pub struct TokenizeOptions {
 
 //     return rendered, all_generation_indices
 
-
 pub fn apply_chat_template<'a>(
     env: &'a mut Environment<'a>,
     model_template: &'a str,
@@ -311,19 +324,20 @@ pub fn apply_chat_template<'a>(
             Ok(template) => template,
             Err(_) => {
                 env.add_template(model_id, model_template)?;
-                env.get_template(model_id).expect("Newly added template must be present")
-            },
+                env.get_template(model_id)
+                    .expect("Newly added template must be present")
+            }
         },
     };
 
     // TODO: what about list of list of conversations
-    
+
     // TODO: handle tool
 
     // TODO: handle documents``
 
     // TODO: allow return_generation_indices
-    
+
     let rendered_chat = template.render(context! {
         messages => conversations,
         documents => documents,

mlx-lm/src/utils/mod.rs

@@ -58,7 +58,7 @@ fn index_out_of_bound_exception() -> Exception {
 
 #[allow(non_snake_case)]
 pub(crate) fn quantized_scaled_dot_product_attention(
-    queries: &Array,
+    queries: Array,
     mut q_keys: QuantizedKeys,
     mut q_values: QuantizedValues,
     scale: f32,
@@ -97,8 +97,8 @@ pub(crate) fn quantized_scaled_dot_product_attention(
         // TODO: handle str type mask
 
         if mask.dtype() == Dtype::Bool {
-            // scores = mlx_rs::ops::r#where(mask, scores, b)
-            todo!("need to add finfo.min equivalent to Dtype")
+            let finfo_min = scores.dtype().finfo_min()?;
+            scores = mlx_rs::ops::r#where(mask, scores, Array::from_f64(finfo_min))?;
         } else {
             scores += mask;
         }
@@ -113,8 +113,6 @@ pub(crate) fn quantized_scaled_dot_product_attention(
     Ok(out)
 }
 
-// type QuantizedKeyValue = (Array, Array, Array);
-
 pub struct QuantizedKeys {
     pub keys: Array,
     pub scales: Array,
@@ -162,7 +160,7 @@ impl From<QuantizedValues> for MaybeQuantizedValues {
 }
 
 pub(crate) fn scaled_dot_product_attention<C>(
-    queries: &Array,
+    queries: Array,
     keys: impl Into<MaybeQuantizedKeys>,
     values: impl Into<MaybeQuantizedValues>,
     cache: Option<C>,

mlx-rs/src/dtype.rs

@@ -1,6 +1,9 @@
+use half::{bf16, f16};
 use mlx_internal_macros::generate_test_cases;
 use strum::EnumIter;
 
+use crate::error::{Exception, InexactDtypeError};
+
 generate_test_cases! {
     /// Array element type
     #[derive(
@@ -83,6 +86,30 @@ impl Dtype {
     pub fn from_promoting_types(a: Dtype, b: Dtype) -> Self {
         a.promote_with(b)
     }
+
+    /// Minimum value of the float point types. Returns `Err(_)` if the type is not
+    /// float point
+    pub fn finfo_min(&self) -> Result<f64, InexactDtypeError> {
+        match self {
+            Dtype::Float16 => Ok(f16::MIN.to_f64_const()),
+            Dtype::Float32 => Ok(f32::MIN as f64),
+            Dtype::Complex64 => Ok(f32::MIN as f64),
+            Dtype::Bfloat16 => Ok(bf16::MIN.to_f64_const()),
+            _ => Err(InexactDtypeError(*self))
+        }
+    }
+
+    /// Maximum value of the float point types. Returns `Err(_)` if the type is not
+    /// float point
+    pub fn finfo_max(&self) -> Result<f64, InexactDtypeError> {
+        match self {
+            Dtype::Float16 => Ok(f16::MAX.to_f64_const()),
+            Dtype::Float32 => Ok(f32::MAX as f64),
+            Dtype::Complex64 => Ok(f32::MAX as f64),
+            Dtype::Bfloat16 => Ok(bf16::MAX.to_f64_const()),
+            _ => Err(InexactDtypeError(*self))
+        }
+    }
 }
 
 pub(crate) trait TypePromotion {

mlx-rs/src/error.rs

@@ -378,6 +378,18 @@ impl From<MultiHeadAttentionBuildError> for TransformerBulidError {
     }
 }
 
+/// The dtype is not a float-point type
+#[derive(Debug, Error)]
+#[error("[finfo] dtype {:?} is not inexact", .0)]
+pub struct InexactDtypeError(pub Dtype);
+
+impl From<InexactDtypeError> for Exception {
+    #[track_caller]
+    fn from(value: InexactDtypeError) -> Self {
+        Exception::custom(value.to_string())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::array;