Files
wehub-resource-sync 770d92cb1f
Lint / lint (push) Waiting to run
Windows CI / Windows (push) Waiting to run
Build Docs / Deploy Docs (push) Waiting to run
chore: import upstream snapshot with attribution
2026-07-13 13:23:58 +08:00

162 lines
5.3 KiB
C++

#ifndef MLC_LLM_JSON_FFI_CONV_TEMPLATE_H
#define MLC_LLM_JSON_FFI_CONV_TEMPLATE_H
#include <tvm/ffi/extra/json.h>
#include <iostream>
#include <map>
#include <optional>
#include <string>
#include <typeinfo>
#include <variant>
#include <vector>
#include "../serve/data.h"
#include "../support/result.h"
#include "openai_api_protocol.h"
using namespace mlc::llm::serve;
namespace mlc {
namespace llm {
namespace json_ffi {
/****************** Model vision config ******************/
/*! \brief Defines the Vision config of the model (if present) */
class ModelVisionConfig {
public:
int hidden_size;
int image_size;
int intermediate_size;
int num_attention_heads;
int num_hidden_layers;
int patch_size;
int projection_dim;
int vocab_size;
std::string dtype;
int num_channels;
double layer_norm_eps;
static ModelVisionConfig FromJSON(const tvm::ffi::json::Object& json_obj);
};
/****************** Model config ******************/
/*! \brief Defines the config of the model.
Populated from "model_config" field in mlc-chat-config.json */
class ModelConfig {
public:
int vocab_size;
int context_window_size;
int sliding_window_size;
int prefill_chunk_size;
int tensor_parallel_shards;
int pipeline_parallel_stages;
int max_batch_size;
std::optional<ModelVisionConfig> vision_config = std::nullopt;
static ModelConfig FromJSON(const tvm::ffi::json::Object& json_obj);
};
/****************** Conversation template ******************/
enum class MessagePlaceholders { SYSTEM, USER, ASSISTANT, TOOL, FUNCTION };
MessagePlaceholders MessagePlaceholderFromString(const std::string& role);
/**
* @brief A struct that specifies the convention template of conversation
* and contains the conversation history.
*/
struct Conversation {
// Optional name of the template.
std::optional<std::string> name = std::nullopt;
// The system prompt template, it optionally contains the system
// message placeholder, and the placeholder will be replaced with
// the system message below.
std::string system_template;
// The content of the system prompt (without the template format).
std::string system_message;
// The system token ids to be prepended at the beginning of tokenized
// generated prompt.
std::optional<std::vector<int>> system_prefix_token_ids = std::nullopt;
// Whether or not to append user role and separator after the system message.
// This is mainly for [INST] [/INST] style prompt format
bool add_role_after_system_message = true;
// The conversation roles
std::unordered_map<std::string, std::string> roles;
// The roles prompt template, it optionally contains the defaults
// message placeholders and will be replaced by actual content
std::unordered_map<std::string, std::string> role_templates;
// The conversation history messages.
// Each message is a pair of strings, denoting "(role, content)".
// The content can be None.
std::vector<ChatCompletionMessage> messages;
// The separators between messages when concatenating into a single prompt.
// List size should be either 1 or 2.
// - When size is 1, the separator will be used between adjacent messages.
// - When size is 2, seps[0] is used after user message, and
// seps[1] is used after assistant message.
std::vector<std::string> seps;
// The separator between the role and the content in a message.
std::string role_content_sep;
// The separator between the role and empty contents.
std::string role_empty_sep;
// The stop criteria
std::vector<std::string> stop_str;
std::vector<int> stop_token_ids;
// When true, strip `<think>...</think>` blocks (and any trailing whitespace)
// from historical assistant messages before rendering the prompt, mirroring
// Qwen3's official HF chat template. Only assistant messages that appear
// before the last user message are affected.
bool strip_reasoning_in_history = false;
Conversation();
/*!
* \brief Get the system text(with the prompt template) given the system prompt message
* \param system_msg The system prompt message.
* \return The created system text.
*/
std::string GetSystemText(const std::string& system_msg) const;
/*!
* \brief replace the content from role by the correct role text in template
* \param role The input role
* \param content The input content from the role
* \param fn_call_str The function calling string if any.
* \return The created text.
*/
std::string GetRoleText(const std::string& role, const std::string& content,
const std::optional<std::string>& fn_call_str) const;
/*! \brief Create a Conversation instance from the given JSON object. */
static Result<Conversation> FromJSON(const tvm::ffi::json::Object& json);
/*! \brief Parse and create a Conversation instance from the given JSON string. */
static Result<Conversation> FromJSON(const std::string& json_str);
};
/*! \brief Create the list of prompts from the messages based on the conversation template. */
Result<std::vector<Data>> CreatePrompt(const Conversation& conv,
const ChatCompletionRequest& request,
const ModelConfig& config, DLDevice device);
} // namespace json_ffi
} // namespace llm
} // namespace mlc
#endif // MLC_LLM_JSON_FFI_CONV_TEMPLATE_H