src.model.model module

Abstract base class and utilities for model management.

This module provides a base Model class for managing different model types and utilities for handling GPU memory optimization during inference.

class src.model.model.Model(model_id: str, device: str, model_swapping: bool, **load_kwargs)

Bases: ABC

Abstract base class for managing machine learning models.

Provides common functionality for loading, device management, and model swapping between CPU and GPU for efficient memory usage.

property model_id

Get the model identifier.

Returns:

The model ID used for initialization.

src.model.model.gpu_inference_wrapper(inference_func: Callable[[Any], Any])

Decorator to manage GPU memory by swapping models between CPU and GPU.

Sends the model to GPU before inference and returns it to CPU after, allowing efficient memory management. Only activates if model_swapping is enabled in the Model instance.

Parameters:

inference_func – The inference function to wrap.

Returns:

Wrapped function that manages device placement.