wehub-resource-sync
83 lines
2.5 KiB
C++
83 lines
2.5 KiB
C++
/*!
|
|
* Copyright (c) 2023-2025 by Contributors
|
|
* \file serve/request.cc
|
|
*/
|
|
|
|
#include "request.h"
|
|
|
|
#include <tvm/ffi/function.h>
|
|
#include <tvm/ffi/reflection/registry.h>
|
|
|
|
#include "data.h"
|
|
|
|
namespace mlc {
|
|
namespace llm {
|
|
namespace serve {
|
|
|
|
/****************** Request ******************/
|
|
|
|
TVM_FFI_STATIC_INIT_BLOCK() { RequestNode::RegisterReflection(); }
|
|
|
|
Request::Request(String id, Array<Data> inputs, GenerationConfig generation_cfg) {
|
|
if (generation_cfg->debug_config.special_request == SpecialRequestKind::kNone) {
|
|
TVM_FFI_ICHECK(!inputs.empty()) << "No input data is given.";
|
|
}
|
|
// Compute the total input length, or fall back to "-1" which means
|
|
// unknown due to the existence of untokenized data.
|
|
int prompt_tokens = 0;
|
|
for (Data input : inputs) {
|
|
if (const auto* token_data = input.as<TokenDataNode>()) {
|
|
prompt_tokens += token_data->token_ids.size();
|
|
} else if (const auto* image_data = input.as<ImageDataNode>()) {
|
|
prompt_tokens += image_data->GetLength();
|
|
} else {
|
|
prompt_tokens = -1;
|
|
break;
|
|
}
|
|
}
|
|
|
|
ObjectPtr<RequestNode> n = tvm::ffi::make_object<RequestNode>();
|
|
n->id = std::move(id);
|
|
n->inputs = std::move(inputs);
|
|
n->prompt_tokens = prompt_tokens;
|
|
n->generation_cfg = std::move(generation_cfg);
|
|
data_ = std::move(n);
|
|
}
|
|
|
|
Request Request::FromUntokenized(const Request& request, const Tokenizer& tokenizer) {
|
|
bool has_untokenized_input = false;
|
|
Array<Data> inputs;
|
|
inputs.reserve(request->inputs.size());
|
|
// Tokenize all text inputs.
|
|
for (Data input : request->inputs) {
|
|
if (const auto* text_data = input.as<TextDataNode>()) {
|
|
has_untokenized_input = true;
|
|
std::vector<int> token_ids = tokenizer->Encode(text_data->text);
|
|
inputs.push_back(TokenData(token_ids));
|
|
} else {
|
|
inputs.push_back(input);
|
|
}
|
|
}
|
|
|
|
// If there is no untokenized input, we don't need to create a new request.
|
|
if (!has_untokenized_input) {
|
|
TVM_FFI_ICHECK_NE(request->prompt_tokens, -1);
|
|
return request;
|
|
} else {
|
|
return Request(request->id, std::move(inputs), request->generation_cfg);
|
|
}
|
|
}
|
|
|
|
TVM_FFI_STATIC_INIT_BLOCK() {
|
|
namespace refl = tvm::ffi::reflection;
|
|
refl::GlobalDef()
|
|
.def("mlc.serve.RequestGetInputs", [](Request request) { return request->inputs; })
|
|
.def("mlc.serve.RequestGetGenerationConfigJSON", [](Request request) {
|
|
return tvm::ffi::json::Stringify(request->generation_cfg->AsJSON());
|
|
});
|
|
}
|
|
|
|
} // namespace serve
|
|
} // namespace llm
|
|
} // namespace mlc
|