Skip to content

Commit

Permalink
Try to show utils documentation
Browse files Browse the repository at this point in the history
  • Loading branch information
frostedoyster committed Dec 4, 2023
1 parent 9c69733 commit ded5c0d
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 2 deletions.
6 changes: 6 additions & 0 deletions docs/src/dev-docs/utils-api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -5,3 +5,9 @@ This is the API for the ``utils`` module of ``metatensor-models``.

.. automodule:: metatensor_models.utils.data
:members:

.. automodule:: metatensor_models.utils.data.dataset
:members:

.. automodule:: metatensor_models.utils.data.readers
:members:
11 changes: 9 additions & 2 deletions src/metatensor_models/utils/data/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,11 @@ class Dataset(torch.utils.data.Dataset):
def __init__(
self, structures: List[rascaline.torch.System], targets: Dict[str, TensorMap]
):
"""Creates a dataset from a list of `rascaline.torch.System`s and
a list of dictionaries of `TensorMap`s."""
"""
Creates a dataset from a list of `rascaline.torch.System` objects
and a dictionary of targets where the keys are strings and the
values are `TensorMap` objects.
"""

for tensor_map in targets.values():
n_structures = (
Expand All @@ -35,8 +38,10 @@ def __len__(self):
def __getitem__(self, index):
"""
Generates one sample of data.
Args:
index: The index of the item in the dataset.
Returns:
A tuple containing the structure and targets for the given index.
"""
Expand All @@ -59,9 +64,11 @@ def __getitem__(self, index):
def collate_fn(batch):
"""
Creates a batch from a list of samples.
Args:
batch: A list of samples, where each sample is a tuple containing a
structure and targets.
Returns:
A tuple containing the structures and targets for the batch.
"""
Expand Down

0 comments on commit ded5c0d

Please sign in to comment.