Add device parameter for TransformerModel in models.py

[ScientistIzaak]

Added "device" parameter to TransformerModel in model.py for Improved GPU/CPU Compatibility

Description

This PR introduces the device parameter to the TransformersModel class, allowing users to explicitly specify the device ("cpu", "cuda", etc.) for model loading and inference. This change improves flexibility and control over model deployment environments, especially when running on systems with mixed hardware setups.

Changes Made

1. Device Detection:

   • Added logic to automatically detect and set the device (cuda if available, otherwise cpu).
   • Modified the model instantiation and data tensor operations to respect the detected device.
   • Added a to(device) call to move the model and tensors to the appropriate device.

2. Backward Compatibility:

   • No breaking changes introduced. If the device is not specified, the script defaults to the best available option.

3. Enhance Logging:

   • Added logging to indicate when the default device is used.

Example Usage

Before

model = TransformersModel(model_id="my-model")

After

model = TransformersModel(model_id="my-model", device="cuda")

This ensures the model explicitly runs on the GPU.

Motivation

Explicit hardware selection is crucial in scenarios involving mixed hardware environments. Many users work in systems where both CPU and GPU resources are available, and having the ability to specify the device simplifies configuration and improves usability.

By adding the device parameter, users no longer need to modify internal logic or rely on automatic device detection, making the library more adaptable to diverse deployment environments.

Testing

• Device Compatibility: Tested both GPU ("cuda") and CPU ("cpu").
• Fallback Behavior: Verified correct fallback to default device when no device parameter is provided or when the specified device is unavailable.
• Backward Compatibility: Ensured that existing code without a device parameter continues to work as expected.

Checklist

• Added device parameter to TransformersModel.
• Updated docstrings to include the new parameter description.
• Verified backward compatibility with existing implementations.
• Enhanced logging to clarify default behavior and errors.

Impact

This update is fully backward-compatible. Users who do not specify the device parameter will experience no change in behavior, as the default device ("cuda" if available, otherwise "cpu") is automatically selected.

[aymeric-roucher]

This is a really nice addition, thank you @ScientistIzaak ! Let's merge this.