nemo_curator.stages.text.embedders.utils

View as Markdown

Module Contents

Functions

NameDescription
create_list_series_from_1d_or_2d_arCreate a cudf list series from 2d arrays.

API

nemo_curator.stages.text.embedders.utils.create_list_series_from_1d_or_2d_ar(
ar: typing.Any,
index: cudf.Index
) -> cudf.Series

Create a cudf list series from 2d arrays. This code comes from https://github.com/rapidsai/crossfit/blob/76f74d0d927cf76313a3960d7dd5575d1dff2f06/crossfit/backend/cudf/series.py#L20-L32

Parameters:

ar
cp.ndarray

any object that can be converted to a cupy array (cupy, numpy, torch, etc.)

index
cudf.Index

index of the the dataframe to be returned

Returns: cudf.Series

cudf.Series: cudf series with the index respected