For max tokens using max new tokens instead of max_length (#1220)

santiatpml · santiadavani · web-flow · commit 1323a843cc1f · 2023-12-06T14:19:08.000-08:00
Co-authored-by: Santi Adavani &lt;santis@gmail.com&gt;
diff --git a/pgml-sdks/pgml/src/open_source_ai.rs b/pgml-sdks/pgml/src/open_source_ai.rs
@@ -236,7 +236,7 @@ mistralai/Mistral-7B-v0.1
         let md5_digest = md5::compute(to_hash.as_bytes());
         let fingerprint = uuid::Uuid::from_slice(&md5_digest.0)?;
 
-        let mut args = serde_json::json!({ "max_length": max_tokens, "temperature": temperature, "do_sample": true, "num_return_sequences": n });
+        let mut args = serde_json::json!({ "max_new_tokens": max_tokens, "temperature": temperature, "do_sample": true, "num_return_sequences": n });
         if let Some(t) = chat_template
             .or_else(|| try_get_model_chat_template(&model_name).map(|s| s.to_string()))
         {
@@ -322,7 +322,7 @@ mistralai/Mistral-7B-v0.1
         let md5_digest = md5::compute(to_hash.as_bytes());
         let fingerprint = uuid::Uuid::from_slice(&md5_digest.0)?;
 
-        let mut args = serde_json::json!({ "max_length": max_tokens, "temperature": temperature, "do_sample": true, "num_return_sequences": n });
+        let mut args = serde_json::json!({ "max_new_tokens": max_tokens, "temperature": temperature, "do_sample": true, "num_return_sequences": n });
         if let Some(t) = chat_template
             .or_else(|| try_get_model_chat_template(&model_name).map(|s| s.to_string()))
         {
