Merge pull request #306 from InfiniTensor/issue/305

issue/305 - feat: Add support for mistral model type

Author: wooway777

csrc/models/mistral/mistral_for_causal_lm.cpp

@@ -0,0 +1,37 @@
+#include "mistral_for_causal_lm.hpp"
+#include "../models_registry.hpp"
+
+namespace infinilm::models::mistral {
+
+std::shared_ptr<infinilm::config::ModelConfig> create_mistral_model_config(std::shared_ptr<infinilm::config::ModelConfig> model_config) {
+    const std::string &model_type = model_config->get<std::string>("model_type");
+    if ("mistral" != model_type) {
+        throw std::runtime_error(
+            "infinilm::models::mistral::create_mistral_model_config: model_type is not mistral");
+    }
+
+    nlohmann::json &config_json = model_config->get_config_json();
+
+    if (!config_json.contains("head_dim")) {
+        size_t head_dim = model_config->get<size_t>("hidden_size")
+            / model_config->get<size_t>("num_attention_heads");
+        config_json["head_dim"] = head_dim;
+    }
+
+    if (!config_json.contains("attention_bias")) {
+        config_json["attention_bias"] = false;
+    }
+
+    return model_config;
+}
+
+} // namespace infinilm::models::mistral
+
+namespace {
+
+INFINILM_REGISTER_CAUSAL_LM_MODEL(
+    mistral,
+    infinilm::models::mistral::MistralForCausalLM,
+    infinilm::models::mistral::create_mistral_model_config);
+
+} // namespace

csrc/models/mistral/mistral_for_causal_lm.hpp

@@ -0,0 +1,24 @@
+#pragma once
+
+#include "../../layers/common_modules.hpp"
+#include <memory>
+
+namespace infinilm::models::mistral {
+
+using MistralMLP = infinilm::layers::MLP;
+
+using MistralAttention = infinilm::layers::attention::Attention;
+
+using MistralDecoderLayer = infinilm::layers::causal_lm_templates::TextDecoderLayer<MistralAttention, MistralMLP>;
+
+using MistralModel = infinilm::layers::causal_lm_templates::TextModel<MistralDecoderLayer>;
+
+using MistralForCausalLM = infinilm::layers::causal_lm_templates::TextCausalLM<MistralModel>;
+
+} // namespace infinilm::models::mistral
+
+namespace infinilm::models::mistral {
+
+std::shared_ptr<infinilm::config::ModelConfig> create_mistral_model_config(std::shared_ptr<infinilm::config::ModelConfig> model_config);
+
+} // namespace infinilm::models::mistral

examples/jiuge.py

@@ -221,14 +221,17 @@ def test(
     # prompt = "山东最高的山是？"
     if isinstance(prompts, str):
         prompts = [prompts]
-    input_contents = [
-        tokenizer.apply_chat_template(
-            conversation=[{"role": "user", "content": prompt}],
-            add_generation_prompt=True,
-            tokenize=False,
-        )
-        for prompt in prompts
-    ]
+    if hasattr(tokenizer, 'chat_template') and tokenizer.chat_template is not None:
+        input_contents = [
+            tokenizer.apply_chat_template(
+                conversation=[{"role": "user", "content": prompt}],
+                add_generation_prompt=True,
+                tokenize=False,
+            )
+            for prompt in prompts
+        ]
+    else:
+        input_contents = prompts
 
     # input_ids_list = tokenizer.batch_encode_plus(input_contents)[
     #     "input_ids"

python/infinilm/auto_config.py

@@ -33,5 +33,7 @@ def from_pretrained(model_path):
             return LlamaConfig(**config_dict)
         elif config_dict["model_type"] in ["qwen3_next" , "minicpm_sala" , "qwen3_vl" , "qwen3_moe"]:
             return LlamaConfig(**config_dict)
+        elif config_dict["model_type"] == "mistral":
+            return LlamaConfig(**config_dict)
 
         raise ValueError(f"Unsupported model type `{config_dict['model_type']}`.")