418 lines
16 KiB
C++
418 lines
16 KiB
C++
#include <yams/daemon/resource/onnx_colbert_session.h>
|
|
#include <yams/daemon/resource/onnx_text_processing.h>
|
|
#include <yams/vector/tokenizer.h>
|
|
|
|
#include "onnx_gpu_provider.h"
|
|
#include <yams/daemon/resource/gpu_info.h>
|
|
|
|
#include <nlohmann/json.hpp>
|
|
#include <spdlog/spdlog.h>
|
|
|
|
#include "ort_cxx_api_wrapper.h"
|
|
#include "ort_runtime_loader.h"
|
|
|
|
#include <algorithm>
|
|
#include <cmath>
|
|
#include <cstdlib>
|
|
#include <filesystem>
|
|
#include <fstream>
|
|
#include <limits>
|
|
#include <mutex>
|
|
#include <sstream>
|
|
#include <string_view>
|
|
#include <unordered_map>
|
|
|
|
namespace fs = std::filesystem;
|
|
|
|
static std::recursive_mutex* g_colbert_onnx_mutex = new std::recursive_mutex();
|
|
static Ort::Env* g_colbert_onnx_env = nullptr;
|
|
static bool g_colbert_onnx_env_initialized = false;
|
|
|
|
static Ort::Env& get_colbert_ort_env() {
|
|
std::lock_guard<std::recursive_mutex> lock(*g_colbert_onnx_mutex);
|
|
if (!g_colbert_onnx_env_initialized) {
|
|
OrtThreadingOptions* threading_options = nullptr;
|
|
(void)Ort::GetApi().CreateThreadingOptions(&threading_options);
|
|
if (threading_options) {
|
|
(void)Ort::GetApi().SetGlobalIntraOpNumThreads(threading_options, 1);
|
|
(void)Ort::GetApi().SetGlobalInterOpNumThreads(threading_options, 1);
|
|
(void)Ort::GetApi().SetGlobalSpinControl(threading_options, 0);
|
|
g_colbert_onnx_env =
|
|
new Ort::Env(threading_options, ORT_LOGGING_LEVEL_WARNING, "YamsColbert");
|
|
Ort::GetApi().ReleaseThreadingOptions(threading_options);
|
|
} else {
|
|
g_colbert_onnx_env = new Ort::Env(ORT_LOGGING_LEVEL_WARNING, "YamsColbert");
|
|
}
|
|
g_colbert_onnx_env_initialized = true;
|
|
}
|
|
return *g_colbert_onnx_env;
|
|
}
|
|
|
|
namespace yams::daemon {
|
|
|
|
class OnnxColbertSession::Impl {
|
|
public:
|
|
Impl(const std::string& modelPath, const std::string& modelName, const ColbertConfig& config)
|
|
: modelPath_(modelPath), modelName_(modelName), config_(config), preprocessor_({}) {
|
|
if (std::getenv("YAMS_USE_MOCK_PROVIDER") || std::getenv("YAMS_SKIP_MODEL_LOADING") ||
|
|
std::getenv("YAMS_TEST_MODE")) {
|
|
testMode_ = true;
|
|
return;
|
|
}
|
|
|
|
const auto& runtimeInfo = yams::onnx_util::OrtRuntimeLoader::instance().ensureLoaded();
|
|
if (!runtimeInfo.available) {
|
|
throw std::runtime_error(runtimeInfo.errorMessage.empty()
|
|
? std::string("ONNX Runtime unavailable")
|
|
: runtimeInfo.errorMessage);
|
|
}
|
|
|
|
env_ = &get_colbert_ort_env();
|
|
sessionOptions_ = std::make_unique<Ort::SessionOptions>();
|
|
|
|
int intra = config_.num_threads > 0 ? config_.num_threads : 4;
|
|
if (const char* s = std::getenv("YAMS_ONNX_INTRA_OP_THREADS")) {
|
|
try {
|
|
int v = std::stoi(s);
|
|
if (v > 0 && v <= 64)
|
|
intra = v;
|
|
} catch (...) {
|
|
}
|
|
}
|
|
sessionOptions_->SetIntraOpNumThreads(intra);
|
|
sessionOptions_->SetInterOpNumThreads(1);
|
|
sessionOptions_->AddConfigEntry("session.intra_op.allow_spinning", "0");
|
|
sessionOptions_->AddConfigEntry("session.inter_op.allow_spinning", "0");
|
|
sessionOptions_->SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_BASIC);
|
|
sessionOptions_->EnableMemPattern();
|
|
sessionOptions_->EnableCpuMemArena();
|
|
|
|
// Attach GPU provider only if GPU is detected
|
|
const auto& gpuInfo = resource::detectGpu();
|
|
if (gpuInfo.detected) {
|
|
std::string ep = onnx_util::appendGpuProvider(*sessionOptions_);
|
|
spdlog::info("[ColBERT] GPU detected ({}), using {} execution provider", gpuInfo.name,
|
|
ep);
|
|
} else {
|
|
spdlog::debug("[ColBERT] No GPU detected, using CPU execution provider");
|
|
}
|
|
}
|
|
|
|
Result<void> loadModel() {
|
|
if (testMode_) {
|
|
maxSequenceLength_ = config_.max_sequence_length;
|
|
tokenDim_ = config_.token_dim;
|
|
isLoaded_ = true;
|
|
return Result<void>();
|
|
}
|
|
|
|
try {
|
|
auto options = sessionOptions_->Clone();
|
|
options.SetIntraOpNumThreads(1);
|
|
options.SetInterOpNumThreads(1);
|
|
session_ = std::make_unique<Ort::Session>(*env_, fs::path(modelPath_).c_str(), options);
|
|
|
|
Ort::AllocatorWithDefaultOptions allocator;
|
|
size_t numInputs = session_->GetInputCount();
|
|
size_t numOutputs = session_->GetOutputCount();
|
|
for (size_t i = 0; i < numInputs; ++i) {
|
|
auto name = session_->GetInputNameAllocated(i, allocator);
|
|
inputNames_.push_back(name.get());
|
|
}
|
|
for (size_t i = 0; i < numOutputs; ++i) {
|
|
auto name = session_->GetOutputNameAllocated(i, allocator);
|
|
outputNames_.push_back(name.get());
|
|
}
|
|
|
|
parseModelConfig();
|
|
isLoaded_ = true;
|
|
return Result<void>();
|
|
} catch (const Ort::Exception& e) {
|
|
return Error{ErrorCode::InternalError,
|
|
std::string("Failed to load ColBERT model: ") + e.what()};
|
|
} catch (const std::exception& e) {
|
|
return Error{ErrorCode::InternalError,
|
|
std::string("Failed to load ColBERT model: ") + e.what()};
|
|
}
|
|
}
|
|
|
|
Result<ColbertEmbeddings> encode(const std::string& text, bool isQuery) {
|
|
if (testMode_) {
|
|
return mockEmbeddings(text, isQuery);
|
|
}
|
|
if (!isLoaded_) {
|
|
if (auto r = loadModel(); !r)
|
|
return r.error();
|
|
}
|
|
|
|
auto tokenIds = tokenizeWithMarkers(text, isQuery);
|
|
auto trimmedTokens = preprocessor_.truncateTokens(tokenIds, maxSequenceLength_);
|
|
auto paddedTokens = preprocessor_.padTokens(trimmedTokens, maxSequenceLength_);
|
|
auto attentionMask = preprocessor_.generateAttentionMask(paddedTokens);
|
|
|
|
std::vector<int64_t> inputIds(paddedTokens.begin(), paddedTokens.end());
|
|
std::vector<int64_t> attention(attentionMask.begin(), attentionMask.end());
|
|
std::vector<int64_t> shape = {1, static_cast<int64_t>(maxSequenceLength_)};
|
|
|
|
auto memoryInfo = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
|
|
std::vector<Ort::Value> inputs;
|
|
inputs.push_back(Ort::Value::CreateTensor<int64_t>(memoryInfo, inputIds.data(),
|
|
inputIds.size(), shape.data(), 2));
|
|
inputs.push_back(Ort::Value::CreateTensor<int64_t>(memoryInfo, attention.data(),
|
|
attention.size(), shape.data(), 2));
|
|
|
|
std::vector<const char*> inNames;
|
|
for (auto& n : inputNames_)
|
|
inNames.push_back(n.c_str());
|
|
std::vector<const char*> outNames;
|
|
for (auto& n : outputNames_)
|
|
outNames.push_back(n.c_str());
|
|
|
|
auto outputs = session_->Run(Ort::RunOptions{nullptr}, inNames.data(), inputs.data(),
|
|
inputs.size(), outNames.data(), outNames.size());
|
|
if (outputs.empty()) {
|
|
return Error{ErrorCode::InternalError, "ColBERT returned no outputs"};
|
|
}
|
|
|
|
auto tensorInfo = outputs[0].GetTensorTypeAndShapeInfo();
|
|
auto shapeOut = tensorInfo.GetShape();
|
|
if (shapeOut.size() < 3) {
|
|
return Error{ErrorCode::InternalError, "ColBERT output has unexpected rank"};
|
|
}
|
|
size_t seqLen = static_cast<size_t>(shapeOut[1]);
|
|
size_t dim = static_cast<size_t>(shapeOut[2]);
|
|
|
|
const float* data = outputs[0].GetTensorData<float>();
|
|
ColbertEmbeddings result;
|
|
result.token_ids = std::move(trimmedTokens);
|
|
result.token_embeddings.resize(seqLen);
|
|
for (size_t i = 0; i < seqLen; ++i) {
|
|
result.token_embeddings[i].assign(data + i * dim, data + (i + 1) * dim);
|
|
}
|
|
|
|
if (!config_.skiplist.empty()) {
|
|
applySkiplist(result.token_embeddings, result.token_ids, config_.skiplist);
|
|
}
|
|
|
|
return result;
|
|
}
|
|
|
|
Result<std::vector<ColbertEmbeddings>> encodeBatch(const std::vector<std::string>& documents) {
|
|
std::vector<ColbertEmbeddings> out;
|
|
out.reserve(documents.size());
|
|
for (const auto& doc : documents) {
|
|
auto r = encode(doc, false);
|
|
if (!r)
|
|
return r.error();
|
|
out.push_back(std::move(r.value()));
|
|
}
|
|
return out;
|
|
}
|
|
|
|
std::vector<float> maxPoolEmbeddings(const ColbertEmbeddings& embeddings, size_t dim) const {
|
|
std::vector<float> pooled(dim, -std::numeric_limits<float>::infinity());
|
|
if (embeddings.token_embeddings.empty()) {
|
|
std::fill(pooled.begin(), pooled.end(), 0.0f);
|
|
return pooled;
|
|
}
|
|
for (const auto& token : embeddings.token_embeddings) {
|
|
for (size_t i = 0; i < dim && i < token.size(); ++i) {
|
|
pooled[i] = std::max(pooled[i], token[i]);
|
|
}
|
|
}
|
|
for (auto& v : pooled) {
|
|
if (!std::isfinite(v))
|
|
v = 0.0f;
|
|
}
|
|
return pooled;
|
|
}
|
|
|
|
void normalizeEmbedding(std::vector<float>& vec) const {
|
|
double norm = 0.0;
|
|
for (float v : vec) {
|
|
norm += static_cast<double>(v) * static_cast<double>(v);
|
|
}
|
|
norm = std::sqrt(norm);
|
|
if (norm > 1e-8) {
|
|
for (auto& v : vec) {
|
|
v = static_cast<float>(v / norm);
|
|
}
|
|
} else {
|
|
std::fill(vec.begin(), vec.end(), 0.0f);
|
|
}
|
|
}
|
|
|
|
float computeMaxSim(const ColbertEmbeddings& query, const ColbertEmbeddings& document) const {
|
|
float score = 0.0f;
|
|
for (const auto& q : query.token_embeddings) {
|
|
float maxSim = -std::numeric_limits<float>::infinity();
|
|
for (const auto& d : document.token_embeddings) {
|
|
float sim = dot(q, d);
|
|
if (sim > maxSim)
|
|
maxSim = sim;
|
|
}
|
|
score += maxSim;
|
|
}
|
|
return score;
|
|
}
|
|
|
|
bool isValid() const { return testMode_ ? isLoaded_ : (isLoaded_ && session_ != nullptr); }
|
|
|
|
size_t maxSequenceLength() const { return maxSequenceLength_; }
|
|
size_t tokenDim() const { return tokenDim_; }
|
|
|
|
private:
|
|
std::vector<int32_t> tokenizeWithMarkers(const std::string& text, bool isQuery) {
|
|
std::vector<int32_t> tokens;
|
|
tokens.push_back(isQuery ? config_.query_marker_id : config_.doc_marker_id);
|
|
auto inner = tokenizer_.encode(text);
|
|
tokens.insert(tokens.end(), inner.begin(), inner.end());
|
|
return tokens;
|
|
}
|
|
|
|
void applySkiplist(std::vector<std::vector<float>>& embeddings,
|
|
const std::vector<int32_t>& tokenIds, const std::vector<int32_t>& skiplist) {
|
|
if (embeddings.empty() || tokenIds.empty())
|
|
return;
|
|
for (size_t i = 0; i < tokenIds.size() && i < embeddings.size(); ++i) {
|
|
if (std::find(skiplist.begin(), skiplist.end(), tokenIds[i]) != skiplist.end()) {
|
|
std::fill(embeddings[i].begin(), embeddings[i].end(), 0.0f);
|
|
}
|
|
}
|
|
}
|
|
|
|
float dot(const std::vector<float>& a, const std::vector<float>& b) const {
|
|
float sum = 0.0f;
|
|
size_t n = std::min(a.size(), b.size());
|
|
for (size_t i = 0; i < n; ++i) {
|
|
sum += a[i] * b[i];
|
|
}
|
|
return sum;
|
|
}
|
|
|
|
ColbertEmbeddings mockEmbeddings(const std::string& text, bool isQuery) {
|
|
ColbertEmbeddings mock;
|
|
mock.token_ids.push_back(isQuery ? config_.query_marker_id : config_.doc_marker_id);
|
|
mock.token_ids.push_back(static_cast<int32_t>(text.size() % 32000));
|
|
mock.token_embeddings.assign(mock.token_ids.size(), std::vector<float>(tokenDim_, 0.0f));
|
|
if (!mock.token_embeddings.empty()) {
|
|
mock.token_embeddings[0][0] = 1.0f;
|
|
}
|
|
return mock;
|
|
}
|
|
|
|
void parseModelConfig() {
|
|
maxSequenceLength_ = config_.max_sequence_length;
|
|
tokenDim_ = config_.token_dim;
|
|
|
|
try {
|
|
fs::path mp(modelPath_);
|
|
fs::path cfg = mp.parent_path() / "config.json";
|
|
if (!fs::exists(cfg)) {
|
|
cfg = mp.parent_path().parent_path() / "config.json";
|
|
}
|
|
if (fs::exists(cfg)) {
|
|
std::ifstream in(cfg);
|
|
if (in.good()) {
|
|
nlohmann::json j;
|
|
in >> j;
|
|
if (j.contains("max_position_embeddings") &&
|
|
j["max_position_embeddings"].is_number_integer()) {
|
|
maxSequenceLength_ =
|
|
static_cast<size_t>(j["max_position_embeddings"].get<int>());
|
|
}
|
|
}
|
|
}
|
|
|
|
fs::path tok = mp.parent_path() / "tokenizer.json";
|
|
if (!fs::exists(tok)) {
|
|
tok = mp.parent_path().parent_path() / "tokenizer.json";
|
|
}
|
|
if (!tokenizer_.load(tok.string())) {
|
|
spdlog::warn("[ColBERT] tokenizer.json not loaded from {}",
|
|
tok.empty() ? "<empty>" : tok.string());
|
|
}
|
|
} catch (...) {
|
|
}
|
|
}
|
|
|
|
std::string modelPath_;
|
|
std::string modelName_;
|
|
ColbertConfig config_;
|
|
OnnxTextPreprocessor preprocessor_;
|
|
yams::vector::HuggingFaceTokenizer tokenizer_{};
|
|
|
|
Ort::Env* env_ = nullptr;
|
|
std::unique_ptr<Ort::SessionOptions> sessionOptions_;
|
|
std::unique_ptr<Ort::Session> session_;
|
|
std::vector<std::string> inputNames_;
|
|
std::vector<std::string> outputNames_;
|
|
|
|
size_t maxSequenceLength_ = 512;
|
|
size_t tokenDim_ = 48;
|
|
bool isLoaded_ = false;
|
|
bool testMode_ = false;
|
|
};
|
|
|
|
OnnxColbertSession::OnnxColbertSession(const std::string& modelPath, const std::string& modelName,
|
|
const ColbertConfig& config)
|
|
: modelPath_(modelPath), modelName_(modelName), config_(config),
|
|
pImpl_(std::make_unique<Impl>(modelPath, modelName, config)) {
|
|
if (auto r = pImpl_->loadModel(); !r) {
|
|
spdlog::error("[ColBERT] Eager load failed for '{}': {}", modelName, r.error().message);
|
|
} else {
|
|
maxSequenceLength_ = pImpl_->maxSequenceLength();
|
|
tokenDim_ = pImpl_->tokenDim();
|
|
}
|
|
}
|
|
|
|
OnnxColbertSession::~OnnxColbertSession() = default;
|
|
|
|
Result<ColbertEmbeddings> OnnxColbertSession::encodeQuery(const std::string& query) {
|
|
return pImpl_->encode(query, true);
|
|
}
|
|
|
|
Result<ColbertEmbeddings> OnnxColbertSession::encodeDocument(const std::string& document) {
|
|
return pImpl_->encode(document, false);
|
|
}
|
|
|
|
Result<std::vector<ColbertEmbeddings>>
|
|
OnnxColbertSession::encodeDocuments(const std::vector<std::string>& documents) {
|
|
return pImpl_->encodeBatch(documents);
|
|
}
|
|
|
|
Result<std::vector<float>>
|
|
OnnxColbertSession::encodeDocumentEmbedding(const std::string& document) {
|
|
auto encoded = pImpl_->encode(document, false);
|
|
if (!encoded)
|
|
return encoded.error();
|
|
auto pooled = pImpl_->maxPoolEmbeddings(encoded.value(), tokenDim_);
|
|
pImpl_->normalizeEmbedding(pooled);
|
|
return pooled;
|
|
}
|
|
|
|
Result<std::vector<std::vector<float>>>
|
|
OnnxColbertSession::encodeDocumentEmbeddings(const std::vector<std::string>& documents) {
|
|
auto batch = pImpl_->encodeBatch(documents);
|
|
if (!batch)
|
|
return batch.error();
|
|
std::vector<std::vector<float>> out;
|
|
out.reserve(batch.value().size());
|
|
for (const auto& doc : batch.value()) {
|
|
auto pooled = pImpl_->maxPoolEmbeddings(doc, tokenDim_);
|
|
pImpl_->normalizeEmbedding(pooled);
|
|
out.push_back(std::move(pooled));
|
|
}
|
|
return out;
|
|
}
|
|
|
|
float OnnxColbertSession::computeMaxSim(const ColbertEmbeddings& query,
|
|
const ColbertEmbeddings& document) const {
|
|
return pImpl_->computeMaxSim(query, document);
|
|
}
|
|
|
|
bool OnnxColbertSession::isValid() const {
|
|
return pImpl_->isValid();
|
|
}
|
|
|
|
} // namespace yams::daemon
|