[Feat] add support for bm25 and tfidf #2567

jperez999 · 2025-02-03T21:41:05Z

This PR will add support for tfidf and BM25 preprocessing of sparse matrix. This PR supports encoding values in raft device COO or CSR sparse matrices. It breaks up the statistical recording (fit) phase and the transformation phase. This allows for batch fitting data and then transforming a target input. This class also allows for exporting/loading, so you can fit in one place and transform in a separate environment. This builds on #2353

jperez999 · 2025-02-03T21:45:48Z

cpp/include/raft/sparse/matrix/preprocessing.cuh

+    saveFile << fullIdLen << " ";
+    // serialize_mdspan<IndexType>(handle, oss, featIdCount_md.view());
+    for (int i = 0; i < vocabSize; i++) {
+      saveFile << featIdCount[i] << " ";


could not use the serialize_mdspan, it does not work well when reading from a file that has other information on it.

would probably be better to use a map here to save only indexes/values that are not zero. This would save time and space. Thoughts?

could not use the serialize_mdspan, it does not work well when reading from a file that has other information on it.

How did serialize_mdspan not work well? We have code like https://github.com/rapidsai/cuvs/blob/49298b22956fdf4d3966825ae9aa41e1aa94975b/cpp/src/neighbors/detail/dataset_serialize.hpp#L79-L83 that serializes both scalar values and mdspan values.

One of the advantages of serialize_mdspan is that its compatible with numpy, also we're using for most other serialization and would be nice to be consistent

fixed this, thanks for the tip

benfred · 2025-02-13T19:34:05Z

cpp/include/raft/sparse/matrix/preprocessing.cuh

+    saveFile << fullIdLen << " ";
+    // serialize_mdspan<IndexType>(handle, oss, featIdCount_md.view());
+    for (int i = 0; i < vocabSize; i++) {
+      saveFile << featIdCount[i] << " ";


could not use the serialize_mdspan, it does not work well when reading from a file that has other information on it.

How did serialize_mdspan not work well? We have code like https://github.com/rapidsai/cuvs/blob/49298b22956fdf4d3966825ae9aa41e1aa94975b/cpp/src/neighbors/detail/dataset_serialize.hpp#L79-L83 that serializes both scalar values and mdspan values.

One of the advantages of serialize_mdspan is that its compatible with numpy, also we're using for most other serialization and would be nice to be consistent

cpp/include/raft/sparse/matrix/preprocessing.cuh

benfred · 2025-02-13T19:36:00Z

cpp/include/raft/sparse/matrix/preprocessing.cuh

+  raft::sparse::op::coo_sort(
+    nnz, nnz, nnz, d_cols.data_handle(), d_rows.data_handle(), d_vals.data_handle(), stream);
+  IndexType* counts;
+  cudaMallocManaged(&counts, nnz * sizeof(IndexType));


Can we use RMM managed_memory_resource instead here?

benfred · 2025-02-13T19:51:40Z

cpp/include/raft/sparse/matrix/detail/preprocessing.cuh

+ *   The array holding the feature(column) occurrence counts for all fitted inputs.
+ * @param[in] counts
+ *   The array representing value changes in rows input.
+ * @param[in] out_values


nitpick:

Suggested change

* @param[in] out_values

* @param[out] out_values

benfred · 2025-02-13T20:15:42Z

cpp/include/raft/sparse/matrix/preprocessing.cuh

+template <typename ValueType = float, typename IndexType = int>
+class SparseEncoder {
+ private:
+  int* featIdCount;


Should this be IndexType instead of int?

Suggested change

int* featIdCount;

IndexType* featIdCount;

benfred · 2025-02-13T21:30:09Z

cpp/include/raft/sparse/matrix/preprocessing.cuh

+
+ * */
+template <typename ValueType, typename IndexType>
+void SparseEncoder<ValueType, IndexType>::fit(raft::resources& handle,


nitpick: handle should be const

Suggested change

void SparseEncoder<ValueType, IndexType>::fit(raft::resources& handle,

void SparseEncoder<ValueType, IndexType>::fit(raft::resources const& handle,

benfred · 2025-02-14T00:06:32Z

cpp/include/raft/sparse/matrix/preprocessing.cuh

+void SparseEncoder<ValueType, IndexType>::_fit_feats(IndexType* cols,
+                                                     IndexType* counts,
+                                                     IndexType nnz,
+                                                     IndexType* results)


This function should accept a raft handle - since its launching cuda kernels, it needs to use the cuda stream when creating these kernels

fixed, but my kernel don't use the handle. I wonder if I am missing something else?

benfred · 2025-02-14T00:08:01Z

cpp/include/raft/sparse/matrix/preprocessing.cuh

+  raft::sparse::matrix::detail::_scan<<<blockSize, num_blocks>>>(cols, nnz, counts);
+  raft::sparse::matrix::detail::_fit_compute_occurs<<<blockSize, num_blocks>>>(
+    cols, nnz, counts, results, numFeats);


I think we can probably simplify this -

This seems like this is doing something like np.bincount - counting up the document frequencies of each term so that we can compute the idf for tf-idf/bm25.

I think we can avoid having the _scan/_fit_compute_occurs kernels here entirely (and also remove the requirement for the coo_sort that is performed previously) by using cub. For example - other code inside of raft is using cub and HistogramEven to peform the bincount operation (like

raft/cpp/include/raft/stats/detail/silhouette_score.cuh

Lines 98 to 146 in 842afd7

/**

* @brief function to calculate the bincounts of number of samples in every label

* @tparam DataT: type of the data samples

* @tparam LabelT: type of the labels

* @param labels: the pointer to the array containing labels for every data sample (1 x nRows)

* @param binCountArray: pointer to the 1D array that contains the count of samples per cluster (1 x

* nLabels)

* @param nRows: number of data samples

* @param nUniqueLabels: number of Labels

* @param workspace: device buffer containing workspace memory

* @param stream: the cuda stream where to launch this kernel

*/

template <typename DataT, typename LabelT>

void countLabels(const LabelT* labels,

DataT* binCountArray,

int nRows,

int nUniqueLabels,

rmm::device_uvector<char>& workspace,

cudaStream_t stream)

{

int num_levels = nUniqueLabels + 1;

LabelT lower_level = 0;

LabelT upper_level = nUniqueLabels;

size_t temp_storage_bytes = 0;

rmm::device_uvector<int> countArray(nUniqueLabels, stream);

RAFT_CUDA_TRY(cub::DeviceHistogram::HistogramEven(nullptr,

temp_storage_bytes,

labels,

binCountArray,

num_levels,

lower_level,

upper_level,

nRows,

stream));

workspace.resize(temp_storage_bytes, stream);

RAFT_CUDA_TRY(cub::DeviceHistogram::HistogramEven(workspace.data(),

temp_storage_bytes,

labels,

binCountArray,

num_levels,

lower_level,

upper_level,

nRows,

stream));

}

etc)

…bm25-tfidf

…m25-tfidf

rhdong · 2025-02-25T18:34:01Z

cpp/tests/sparse/preprocess.cu

+TEST_P(SparsePreprocessBigTF, Result) { Run(true, false); }
+
+const std::vector<SparsePreprocessInputs<float, int>> sparse_preprocess_inputs = {
+  {


Since "std::vector" is used and CUDA kernels are involved in the main codes, it is strongly recommended to add test cases for multiple settings, especially the corner cases.Since "std::vector" is used and CUDA kernels are involved in the main codes, it is strongly recommended to add test cases for multiple settings, especially the corner cases.

More cases, better sleep 😜

rhdong · 2025-02-25T18:38:23Z

cpp/include/raft/sparse/matrix/preprocessing.cuh

+template <typename ValueType, typename IndexType>
+SparseEncoder<ValueType, IndexType>::SparseEncoder(int feats) : numFeats(feats)
+{
+  cudaMallocManaged(&featIdCount, feats * sizeof(int));


May need RAFT_CUDA_TRY

Suggested change

cudaMallocManaged(&featIdCount, feats * sizeof(int));

RAFT_CUDA_TRY(cudaMallocManaged(&featIdCount, feats * sizeof(int)));

rhdong · 2025-02-25T20:17:59Z

cpp/include/raft/sparse/matrix/preprocessing.cuh

+  auto d_rows = raft::make_device_vector<IndexType, int64_t>(handle, nnz);
+  auto d_cols = raft::make_device_vector<IndexType, int64_t>(handle, nnz);
+  auto d_vals = raft::make_device_vector<ValueType, int64_t>(handle, nnz);


If the d_xxx are temporary data, it would be better to use get_workspace_resource as memory pool for better performance, pls refer to https://github.com/rapidsai/raft/blob/branch-25.04/cpp/include/raft/cluster/detail/kmeans_balanced.cuh#L329-L330

rhdong · 2025-02-25T20:26:23Z

cpp/include/raft/sparse/matrix/preprocessing.cuh

+  cudaMemset(counts, 0, nnz * sizeof(IndexType));
+  _fit_feats(handle, d_cols.data_handle(), counts, nnz, featIdCount);
+  cudaFree(counts);
+  cudaDeviceSynchronize();


cudaFree will sync the device implicitly (before releasing a pointer), so there is no need to call cudaDeviceSynchronize separately. Furthermore, : I believe cudaMallocManaged is low performance, maybe you can try pinned host memory via cudaMallocHost when you're pretty sure the device memory is not sufficient.

rhdong · 2025-02-25T20:38:16Z

cpp/include/raft/sparse/matrix/preprocessing.cuh

+                          raft::device_uvector_policy,
+                          raft::PRESERVING> csr_in,
+  float* results,
+  bool bm25_on,


If we have more options on algorithms in the future extension, we may need an enum to avoid public API changing(after release)

rhdong · 2025-02-25T20:40:47Z

cpp/include/raft/sparse/matrix/preprocessing.cuh

+                             columns.data_handle(),
+                             values.data_handle(),
+                             stream);
+  raft::sparse::matrix::detail::_scan<<<num_blocks, blockSize>>>(rows.data_handle(), nnz, counts);


Since we have a stream, we'd better to apply it in the kernels

Suggested change

raft::sparse::matrix::detail::_scan<<<num_blocks, blockSize>>>(rows.data_handle(), nnz, counts);

raft::sparse::matrix::detail::_scan<<<num_blocks, blockSize, 0, stream>>>(rows.data_handle(), nnz, counts);

rhdong · 2025-02-25T20:40:58Z

cpp/include/raft/sparse/matrix/preprocessing.cuh

+                             values.data_handle(),
+                             stream);
+  raft::sparse::matrix::detail::_scan<<<num_blocks, blockSize>>>(rows.data_handle(), nnz, counts);
+  raft::sparse::matrix::detail::_transform<<<num_blocks, blockSize>>>(rows.data_handle(),


Suggested change

raft::sparse::matrix::detail::_transform<<<num_blocks, blockSize>>>(rows.data_handle(),

raft::sparse::matrix::detail::_transform<<<num_blocks, blockSize, 0, stream>>>(rows.data_handle(),

rhdong · 2025-02-25T20:44:04Z

cpp/include/raft/sparse/matrix/preprocessing.cuh

+                                                                      nnz,
+                                                                      numFeats,
+                                                                      bm25_on);
+  cudaFree(counts);


We should take care to use these explicitly and implicitly device-sync APIs, which will break the asynchronized parallel pipeline in the end-user environment and hurt performance. It would be good practice to try to use a stream and return the right of sync on a stream(cudaStreamSynchronize(stream)) to callers(end-users) as possible as we can.

rhdong · 2025-02-25T20:49:18Z

cpp/include/raft/sparse/matrix/preprocessing.cuh

+  auto featIdCount_h = raft::make_host_vector<IndexType, int64_t>(num_feats);
+  raft::copy(featIdCount_h.data_handle(), vals.data(), vals.size(), stream);
+  loadFile.close();
+  return new SparseEncoder<ValueType, IndexType>(


Considering using smart pointer like std::shared_ptr

rhdong · 2025-02-25T20:52:50Z

cpp/tests/util/preprocess_utils.cu

+namespace raft::util {
+
+template <typename T>
+void print_vals(raft::resources& handle, const raft::device_vector_view<T, size_t>& out)


It seems to be unused code, which should be removed.

rhdong · 2025-02-25T20:55:36Z

cpp/tests/sparse/preprocess.cu

+                                               5,
+                                               params.n_rows,
+                                               params.n_cols,
+                                               67584);


Fixed random number seeds will cause the generated test data to always be the same, which may not cover certain special input situations and thus miss potential bugs.

rhdong

well done

add support for bm25 and tfidf

afb74f2

jperez999 requested review from a team as code owners February 3, 2025 21:41

jperez999 self-assigned this Feb 3, 2025

github-actions bot added cpp CMake labels Feb 3, 2025

jperez999 added 3 - Ready for Review and removed cpp CMake 3 - Ready for Review labels Feb 3, 2025

jperez999 requested a review from cjnolet February 3, 2025 21:42

jperez999 added enhancement New feature or request feature request New feature or request non-breaking Non-breaking change labels Feb 3, 2025

jperez999 commented Feb 3, 2025

View reviewed changes

remove unnecessary commented function

d2cd397

github-actions bot added cpp CMake labels Feb 3, 2025

jperez999 and others added 7 commits February 4, 2025 13:47

add function comments and change wording to feats from vocab

43b0491

add handle descriptor for comment on function definition

54c6c30

add missing comments for functions

44c68f4

add missing param comment

60a94a1

Merge branch 'branch-25.04' into bm25-tfidf

9dee4ec

Merge branch 'branch-25.04' into bm25-tfidf

a6e2f00

Merge branch 'branch-25.04' into bm25-tfidf

4175a4d

benfred reviewed Feb 14, 2025

View reviewed changes

cjnolet and others added 2 commits February 19, 2025 15:40

Merge branch 'branch-25.04' into bm25-tfidf

faef8ab

Merge branch 'branch-25.04' of https://github.com/rapidsai/raft into …

507d46e

…bm25-tfidf

rhdong changed the title ~~add support for bm25 and tfidf~~ [Feat] add support for bm25 and tfidf Feb 24, 2025

jperez999 added 2 commits February 25, 2025 13:16

made changes to fix tests and add const for all handle args in functions

233b508

Merge branch 'bm25-tfidf' of https://github.com/jperez999/raft into b…

69815df

…m25-tfidf

rhdong reviewed Feb 25, 2025

View reviewed changes

fix serialization and deserialization

950e564

rhdong reviewed Feb 25, 2025

View reviewed changes

rhdong requested changes Feb 25, 2025

View reviewed changes

replace loop with memset

4cec160

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Feat] add support for bm25 and tfidf #2567

[Feat] add support for bm25 and tfidf #2567

jperez999 commented Feb 3, 2025 •

edited

Loading

jperez999 Feb 3, 2025

jperez999 Feb 4, 2025

benfred Feb 13, 2025

jperez999 Feb 25, 2025

benfred Feb 13, 2025

benfred Feb 13, 2025

benfred Feb 13, 2025

jperez999 Feb 25, 2025

benfred Feb 13, 2025

jperez999 Feb 25, 2025

benfred Feb 13, 2025

jperez999 Feb 25, 2025

benfred Feb 14, 2025

jperez999 Feb 25, 2025 •

edited

Loading

benfred Feb 14, 2025

rhdong Feb 25, 2025

rhdong Feb 25, 2025

rhdong Feb 25, 2025

rhdong Feb 25, 2025 •

edited

Loading

rhdong Feb 25, 2025 •

edited

Loading

rhdong Feb 25, 2025 •

edited

Loading

rhdong Feb 25, 2025

rhdong Feb 25, 2025

rhdong Feb 25, 2025 •

edited

Loading

rhdong Feb 25, 2025

rhdong Feb 25, 2025

rhdong Feb 25, 2025

rhdong left a comment

	void SparseEncoder<ValueType, IndexType>::fit(raft::resources& handle,
	void SparseEncoder<ValueType, IndexType>::fit(raft::resources const& handle,

	/**
	* @brief function to calculate the bincounts of number of samples in every label
	* @tparam DataT: type of the data samples
	* @tparam LabelT: type of the labels
	* @param labels: the pointer to the array containing labels for every data sample (1 x nRows)
	* @param binCountArray: pointer to the 1D array that contains the count of samples per cluster (1 x
	* nLabels)
	* @param nRows: number of data samples
	* @param nUniqueLabels: number of Labels
	* @param workspace: device buffer containing workspace memory
	* @param stream: the cuda stream where to launch this kernel
	*/
	template <typename DataT, typename LabelT>
	void countLabels(const LabelT* labels,
	DataT* binCountArray,
	int nRows,
	int nUniqueLabels,
	rmm::device_uvector<char>& workspace,
	cudaStream_t stream)
	{
	int num_levels = nUniqueLabels + 1;
	LabelT lower_level = 0;
	LabelT upper_level = nUniqueLabels;
	size_t temp_storage_bytes = 0;

	rmm::device_uvector<int> countArray(nUniqueLabels, stream);

	RAFT_CUDA_TRY(cub::DeviceHistogram::HistogramEven(nullptr,
	temp_storage_bytes,
	labels,
	binCountArray,
	num_levels,
	lower_level,
	upper_level,
	nRows,
	stream));

	workspace.resize(temp_storage_bytes, stream);

	RAFT_CUDA_TRY(cub::DeviceHistogram::HistogramEven(workspace.data(),
	temp_storage_bytes,
	labels,
	binCountArray,
	num_levels,
	lower_level,
	upper_level,
	nRows,
	stream));
	}

	cudaMallocManaged(&featIdCount, feats * sizeof(int));
	RAFT_CUDA_TRY(cudaMallocManaged(&featIdCount, feats * sizeof(int)));

	raft::sparse::matrix::detail::_scan<<<num_blocks, blockSize>>>(rows.data_handle(), nnz, counts);
	raft::sparse::matrix::detail::_scan<<<num_blocks, blockSize, 0, stream>>>(rows.data_handle(), nnz, counts);

	raft::sparse::matrix::detail::_transform<<<num_blocks, blockSize>>>(rows.data_handle(),
	raft::sparse::matrix::detail::_transform<<<num_blocks, blockSize, 0, stream>>>(rows.data_handle(),

[Feat] add support for bm25 and tfidf #2567

Are you sure you want to change the base?

[Feat] add support for bm25 and tfidf #2567

Conversation

jperez999 commented Feb 3, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jperez999 Feb 25, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rhdong Feb 25, 2025 • edited Loading

Choose a reason for hiding this comment

rhdong Feb 25, 2025 • edited Loading

Choose a reason for hiding this comment

rhdong Feb 25, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rhdong Feb 25, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rhdong left a comment

Choose a reason for hiding this comment

jperez999 commented Feb 3, 2025 •

edited

Loading

jperez999 Feb 25, 2025 •

edited

Loading

rhdong Feb 25, 2025 •

edited

Loading

rhdong Feb 25, 2025 •

edited

Loading

rhdong Feb 25, 2025 •

edited

Loading

rhdong Feb 25, 2025 •

edited

Loading