Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fixing some docstrings #7465

Merged
merged 1 commit into from
Apr 4, 2024
Merged

fixing some docstrings #7465

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 35 additions & 25 deletions haystack/utils/device.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,10 @@

class DeviceType(Enum):
"""
Represents device types supported by Haystack. This also includes
devices that are not directly used by models - for example, the disk
device is exclusively used in device maps for frameworks that support
offloading model weights to disk.
Represents device types supported by Haystack.

This also includes devices that are not directly used by models - for example, the disk device is exclusively used
in device maps for frameworks that support offloading model weights to disk.
"""

CPU = "cpu"
Expand Down Expand Up @@ -126,16 +126,24 @@ def mps() -> "Device":

@staticmethod
def from_str(string: str) -> "Device":
"""
Create a generic device from a string.

:returns:
The device.

"""
device_type_str, device_id = _split_device_string(string)
return Device(DeviceType.from_str(device_type_str), device_id)


@dataclass
class DeviceMap:
"""
A generic mapping from strings to devices. The semantics of the
strings are dependent on target framework. Primarily used to deploy
HuggingFace models to multiple devices.
A generic mapping from strings to devices.

The semantics of the strings are dependent on target framework. Primarily used to deploy HuggingFace models to
multiple devices.

:param mapping:
Dictionary mapping strings to devices.
Expand Down Expand Up @@ -226,8 +234,9 @@ def from_hf(hf_device_map: Dict[str, Union[int, str, "torch.device"]]) -> "Devic
@dataclass(frozen=True)
class ComponentDevice:
"""
A representation of a device for a component. This can be either
a single device or a device map.
A representation of a device for a component.

This can be either a single device or a device map.
"""

_single_device: Optional[Device] = field(default=None)
Expand All @@ -237,6 +246,7 @@ class ComponentDevice:
def from_str(cls, device_str: str) -> "ComponentDevice":
"""
Create a component device representation from a device string.

The device string can only represent a single device.

:param device_str:
Expand All @@ -251,6 +261,7 @@ def from_str(cls, device_str: str) -> "ComponentDevice":
def from_single(cls, device: Device) -> "ComponentDevice":
"""
Create a component device representation from a single device.

Disks cannot be used as single devices.

:param device:
Expand Down Expand Up @@ -287,6 +298,7 @@ def _validate(self):
def to_torch(self) -> "torch.device":
"""
Convert the component device representation to PyTorch format.

Device maps are not supported.

:returns:
Expand All @@ -304,6 +316,7 @@ def to_torch(self) -> "torch.device":
def to_torch_str(self) -> str:
"""
Convert the component device representation to PyTorch string format.

Device maps are not supported.

:returns:
Expand All @@ -320,6 +333,7 @@ def to_torch_str(self) -> str:
def to_spacy(self) -> int:
"""
Convert the component device representation to spaCy format.

Device maps are not supported.

:returns:
Expand Down Expand Up @@ -361,9 +375,9 @@ def convert_device(device: Device, *, gpu_id_only: bool = False) -> Union[int, s

def update_hf_kwargs(self, hf_kwargs: Dict[str, Any], *, overwrite: bool) -> Dict[str, Any]:
"""
Convert the component device representation to HuggingFace format
and add them as canonical keyword arguments to the keyword arguments
dictionary.
Convert the component device representation to HuggingFace format.

Add them as canonical keyword arguments to the keyword arguments dictionary.

:param hf_kwargs:
The HuggingFace keyword arguments dictionary.
Expand All @@ -385,8 +399,7 @@ def update_hf_kwargs(self, hf_kwargs: Dict[str, Any], *, overwrite: bool) -> Dic
@property
def has_multiple_devices(self) -> bool:
"""
Whether this component device representation contains multiple
devices.
Whether this component device representation contains multiple devices.
"""
self._validate()

Expand All @@ -395,8 +408,7 @@ def has_multiple_devices(self) -> bool:
@property
def first_device(self) -> Optional["ComponentDevice"]:
"""
Return either the single device or the first device in the
device map, if any.
Return either the single device or the first device in the device map, if any.

:returns:
The first device.
Expand All @@ -413,8 +425,7 @@ def first_device(self) -> Optional["ComponentDevice"]:
@staticmethod
def resolve_device(device: Optional["ComponentDevice"] = None) -> "ComponentDevice":
"""
Select a device for a component. If a device is specified,
it's used. Otherwise, the default device is used.
Select a device for a component. If a device is specified, it's used. Otherwise, the default device is used.

:param device:
The provided device, if any.
Expand All @@ -433,8 +444,7 @@ def resolve_device(device: Optional["ComponentDevice"] = None) -> "ComponentDevi

def to_dict(self) -> Dict[str, Any]:
"""
Convert the component device representation to a JSON-serializable
dictionary.
Convert the component device representation to a JSON-serializable dictionary.

:returns:
The dictionary representation.
Expand All @@ -450,8 +460,7 @@ def to_dict(self) -> Dict[str, Any]:
@classmethod
def from_dict(cls, dict: Dict[str, Any]) -> "ComponentDevice":
"""
Create a component device representation from a JSON-serialized
dictionary.
Create a component device representation from a JSON-serialized dictionary.

:param dict:
The serialized representation.
Expand All @@ -468,9 +477,10 @@ def from_dict(cls, dict: Dict[str, Any]) -> "ComponentDevice":

def _get_default_device() -> Device:
"""
Return the default device for Haystack. Precedence:
GPU > MPS > CPU. If PyTorch is not installed, only CPU
is available.
Return the default device for Haystack.

Precedence:
GPU > MPS > CPU. If PyTorch is not installed, only CPU is available.

:returns:
The default device.
Expand Down
Loading