84 lines
2.3 KiB
C++
84 lines
2.3 KiB
C++
// sherpa-onnx/csrc/fast-clustering.cc
|
|
//
|
|
// Copyright (c) 2024 Xiaomi Corporation
|
|
|
|
#include "sherpa-onnx/csrc/fast-clustering.h"
|
|
|
|
#include <vector>
|
|
|
|
#include "Eigen/Dense"
|
|
#include "fastcluster-all-in-one.h" // NOLINT
|
|
|
|
namespace sherpa_onnx {
|
|
|
|
class FastClustering::Impl {
|
|
public:
|
|
explicit Impl(const FastClusteringConfig &config) : config_(config) {}
|
|
|
|
std::vector<int32_t> Cluster(float *features, int32_t num_rows,
|
|
int32_t num_cols) {
|
|
if (num_rows <= 0) {
|
|
return {};
|
|
}
|
|
|
|
if (num_rows == 1) {
|
|
return {0};
|
|
}
|
|
|
|
Eigen::Map<
|
|
Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>
|
|
m(features, num_rows, num_cols);
|
|
m.rowwise().normalize();
|
|
|
|
std::vector<double> distance((num_rows * (num_rows - 1)) / 2);
|
|
|
|
int32_t k = 0;
|
|
for (int32_t i = 0; i != num_rows; ++i) {
|
|
auto v = m.row(i);
|
|
for (int32_t j = i + 1; j != num_rows; ++j) {
|
|
double cosine_similarity = v.dot(m.row(j));
|
|
double consine_dissimilarity = 1 - cosine_similarity;
|
|
|
|
if (consine_dissimilarity < 0) {
|
|
consine_dissimilarity = 0;
|
|
}
|
|
|
|
distance[k] = consine_dissimilarity;
|
|
++k;
|
|
}
|
|
}
|
|
|
|
std::vector<int32_t> merge(2 * (num_rows - 1));
|
|
std::vector<double> height(num_rows - 1);
|
|
|
|
fastclustercpp::hclust_fast(num_rows, distance.data(),
|
|
fastclustercpp::HCLUST_METHOD_SINGLE,
|
|
merge.data(), height.data());
|
|
|
|
std::vector<int32_t> labels(num_rows);
|
|
if (config_.num_clusters > 0) {
|
|
fastclustercpp::cutree_k(num_rows, merge.data(), config_.num_clusters,
|
|
labels.data());
|
|
} else {
|
|
fastclustercpp::cutree_cdist(num_rows, merge.data(), height.data(),
|
|
config_.threshold, labels.data());
|
|
}
|
|
|
|
return labels;
|
|
}
|
|
|
|
private:
|
|
FastClusteringConfig config_;
|
|
};
|
|
|
|
FastClustering::FastClustering(const FastClusteringConfig &config)
|
|
: impl_(std::make_unique<Impl>(config)) {}
|
|
|
|
FastClustering::~FastClustering() = default;
|
|
|
|
std::vector<int32_t> FastClustering::Cluster(float *features, int32_t num_rows,
|
|
int32_t num_cols) {
|
|
return impl_->Cluster(features, num_rows, num_cols);
|
|
}
|
|
} // namespace sherpa_onnx
|