nemo_deploy.utils#

Module Contents#

Functions#

typedict2tensor

Converts a type dictionary class into a tuple of PyTriton Tensor objects.

nemo_checkpoint_version

Determines the version of a NeMo checkpoint from its file structure.

str_list2numpy

Converts a list of strings to a numpy array of UTF-8 encoded bytes.

str_ndarray2list

Converts a numpy array of UTF-8 encoded bytes back to a list of strings.

ndarray2img

Converts a numpy array of images to a list of PIL Image objects.

cast_output

Casts input data to a numpy array with the required dtype.

broadcast_list

Broadcasts a list of text data to all processes.

Data#

NEMO2

NEMO1

API#

nemo_deploy.utils.NEMO2 = 'NEMO 2.0'#
nemo_deploy.utils.NEMO1 = 'NEMO 1.0'#
nemo_deploy.utils.typedict2tensor(
typedict_class,
overwrite_kwargs: Optional[Dict[str, Any]] = None,
defaults: Optional[Dict[str, Any]] = None,
)#

Converts a type dictionary class into a tuple of PyTriton Tensor objects.

This function takes a class with type hints and converts each typed field into a PyTriton Tensor specification, handling nested list types and mapping Python types to numpy dtypes.

Parameters:

  • typedict_class – A class with type hints that will be converted to Tensor specs

  • overwrite_kwargs – Optional dictionary of kwargs to override default Tensor parameters

  • defaults – Optional dictionary of default values (unused)

Returns:

A tuple of PyTriton Tensor objects, one for each typed field in the input class

Return type:

tuple

Raises:

Exception – If an unsupported type is encountered during type mapping

nemo_deploy.utils.nemo_checkpoint_version(path: str) → str#

Determines the version of a NeMo checkpoint from its file structure.

Examines the provided checkpoint path to determine if it follows the NeMo 2.0 or NeMo 1.0 format based on the presence of ‘context’ and ‘weights’ directories.

Parameters:

path (str) – Path to the NeMo checkpoint file or directory

Returns:

Version string - either NEMO2 or NEMO1 constant indicating the checkpoint version

Return type:

str

nemo_deploy.utils.str_list2numpy(str_list: List[str]) → numpy.ndarray#

Converts a list of strings to a numpy array of UTF-8 encoded bytes.

Takes a list of strings and converts it to a numpy array with an additional dimension, then encodes the strings as UTF-8 bytes.

Parameters:

str_list (List[str]) – List of strings to convert

Returns:

Numpy array of UTF-8 encoded bytes with shape (N, 1) where N is the length of the input list

Return type:

np.ndarray

nemo_deploy.utils.str_ndarray2list(str_ndarray: numpy.ndarray) → List[str]#

Converts a numpy array of UTF-8 encoded bytes back to a list of strings.

Takes a numpy array of UTF-8 encoded bytes and decodes it back to strings, removing any extra dimensions, and returns the result as a Python list.

Parameters:

str_ndarray (np.ndarray) – Numpy array of UTF-8 encoded bytes, typically with shape (N, 1) where N is the length of the resulting list

Returns:

List of decoded strings

Return type:

List[str]

nemo_deploy.utils.ndarray2img(
img_ndarray: numpy.ndarray,
) → List[PIL.Image.Image]#

Converts a numpy array of images to a list of PIL Image objects.

Takes a numpy array containing one or more images and converts each image to a PIL Image object using Image.fromarray().

Parameters:

img_ndarray (np.ndarray) – Numpy array of images, where each image is a 2D or 3D array representing pixel values

Returns:

List of PIL Image objects created from the input array

Return type:

List[Image.Image]

nemo_deploy.utils.cast_output(data, required_dtype)#

Casts input data to a numpy array with the required dtype.

Takes input data that may be a torch.Tensor, numpy array, or other sequence type and converts it to a numpy array with the specified dtype. For string dtypes, the data is encoded as UTF-8 bytes. The output array is ensured to have at least 2 dimensions.

Parameters:

  • data – Input data to cast. Can be a torch.Tensor, numpy array, or sequence type that can be converted to a numpy array.

  • required_dtype – The desired numpy dtype for the output array.

Returns:

A numpy array containing the input data cast to the required dtype, with at least 2 dimensions.

Return type:

np.ndarray

nemo_deploy.utils.broadcast_list(data, src=0, group=None)#

Broadcasts a list of text data to all processes.

Parameters:

  • data (list) – List of strings to broadcast.

  • src (int, optional) – Source rank. Defaults to 0.

  • group (ProcessGroup, optional) – The process group to work on. If None, the default process group will be used.