162 lines
5.3 KiB
C++
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
|