improve: stop using deprecated functionality

Author: samedii

datastream/dataset.py

@@ -1,26 +1,28 @@
 from __future__ import annotations
-from pydantic import BaseModel
+
+import inspect
+import random
+import string
+import textwrap
+from functools import lru_cache
+from pathlib import Path
 from typing import (
-    Tuple,
     Callable,
-    Union,
-    List,
-    TypeVar,
-    Generic,
     Dict,
-    Optional,
+    Generic,
     Iterable,
+    List,
+    Optional,
+    Tuple,
+    TypeVar,
+    Union,
 )
-from pathlib import Path
-from functools import lru_cache
-import string
-import random
-import textwrap
-import inspect
+
 import numpy as np
 import pandas as pd
-from datastream import tools
+from pydantic import BaseModel, ConfigDict
 
+from datastream import tools
 
 T = TypeVar("T")
 R = TypeVar("R")
@@ -53,9 +55,10 @@ class Dataset(BaseModel, Generic[T]):
     length: int
     get_item: Callable[[pd.DataFrame, int], T]
 
-    class Config:
-        arbitrary_types_allowed = True
-        allow_mutation = False
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        frozen=True,
+    )
 
     @staticmethod
     def from_subscriptable(subscriptable) -> Dataset:
@@ -96,7 +99,7 @@ def from_dataframe(dataframe: pd.DataFrame) -> Dataset[pd.Series]:
 
     @staticmethod
     def from_paths(paths: Iterable[str, Path], pattern: str) -> Dataset[pd.Series]:
-        """
+        r"""
         Create ``Dataset`` from paths using regex pattern that extracts information
         from the path itself.
         :func:`Dataset.__getitem__` will return a row from the dataframe and
@@ -154,7 +157,7 @@ def __eq__(self: Dataset[T], other: Dataset[R]) -> bool:
         return True
 
     def replace(self, **kwargs):
-        new_dict = self.dict()
+        new_dict = self.model_dump()
         new_dict.update(**kwargs)
         return type(self)(**new_dict)
 

datastream/datastream.py

@@ -4,7 +4,7 @@
 
 import numpy as np
 import torch
-from pydantic import BaseModel, PositiveInt
+from pydantic import BaseModel, ConfigDict, PositiveInt
 
 from datastream import Dataset
 from datastream.samplers import (
@@ -40,9 +40,10 @@ class Datastream(BaseModel, Generic[T]):
     dataset: Dataset
     sampler: Optional[torch.utils.data.Sampler]
 
-    class Config:
-        arbitrary_types_allowed = True
-        allow_mutation = False
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        frozen=True,
+    )
 
     def __init__(self, dataset: Dataset[T], sampler: torch.utils.data.Sampler = None):
         if len(dataset) == 0:

datastream/samplers/merge_sampler.py

@@ -5,7 +5,8 @@
 from typing import Callable, Iterable, Tuple
 
 import torch
-from pydantic import BaseModel
+import torch.utils.data
+from pydantic import BaseModel, ConfigDict
 
 from datastream import Dataset
 from datastream.tools import repeat_map_chain
@@ -19,9 +20,10 @@ class MergeSampler(BaseModel, torch.utils.data.Sampler):
     from_mapping: Callable[[int], Tuple[int, int]]
     merged_samplers: Iterable
 
-    class Config:
-        arbitrary_types_allowed = True
-        allow_mutation = False
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        frozen=True,
+    )
 
     def __init__(self, samplers, datasets, ns):
         BaseModel.__init__(

datastream/samplers/multi_sampler.py

@@ -1,11 +1,14 @@
 from __future__ import annotations
-from pydantic import BaseModel
-from typing import Tuple, Iterable
+
 from itertools import chain, islice
+from typing import Iterable, Tuple
+
 import torch
-from datastream.tools import repeat_map_chain
-from datastream.samplers import StandardSampler
+from pydantic import BaseModel, ConfigDict
+
 from datastream import Dataset
+from datastream.samplers import StandardSampler
+from datastream.tools import repeat_map_chain
 
 
 # TODO: write custom sampler that avoid replacement between samplers
@@ -15,9 +18,10 @@ class MultiSampler(BaseModel, torch.utils.data.Sampler):
     length: int
     merged_samplers: Iterable
 
-    class Config:
-        arbitrary_types_allowed = True
-        allow_mutation = False
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        frozen=True,
+    )
 
     def __init__(self, samplers, dataset):
         BaseModel.__init__(

datastream/samplers/repeat_sampler.py

@@ -1,7 +1,10 @@
 from __future__ import annotations
-from pydantic import BaseModel
+
 from typing import Iterable
+
 import torch
+import torch.utils.data
+from pydantic import BaseModel, ConfigDict
 
 
 class RepeatSampler(BaseModel, torch.utils.data.Sampler):
@@ -10,8 +13,7 @@ class RepeatSampler(BaseModel, torch.utils.data.Sampler):
     epoch_bound: bool = False
     queue: Iterable
 
-    class Config:
-        arbitrary_types_allowed = True
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
     def __init__(self, sampler, length, epoch_bound=False):
         """

datastream/samplers/sequential_sampler.py

@@ -1,14 +1,17 @@
 from __future__ import annotations
-from pydantic import BaseModel
+
 import torch
+import torch.utils.data
+from pydantic import BaseModel, ConfigDict
 
 
 class SequentialSampler(BaseModel, torch.utils.data.Sampler):
     sampler: torch.utils.data.SequentialSampler
 
-    class Config:
-        arbitrary_types_allowed = True
-        allow_mutation = False
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        frozen=True,
+    )
 
     def __init__(self, length):
         BaseModel.__init__(

datastream/samplers/standard_sampler.py

@@ -1,16 +1,19 @@
 from __future__ import annotations
-from pydantic import BaseModel
+
 import torch
+import torch.utils.data
+from pydantic import BaseModel, ConfigDict
 
 
 class StandardSampler(BaseModel, torch.utils.data.Sampler):
     proportion: float
     replacement: bool
     sampler: torch.utils.data.WeightedRandomSampler
 
-    class Config:
-        arbitrary_types_allowed = True
-        allow_mutation = False
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        frozen=True,
+    )
 
     def __init__(self, length, proportion=1.0, replacement=False):
         BaseModel.__init__(

datastream/samplers/zip_sampler.py

@@ -1,11 +1,15 @@
 from __future__ import annotations
-from pydantic import BaseModel
-from typing import Tuple, Callable, Iterable
+
 from functools import partial
 from itertools import islice
+from typing import Callable, Iterable, Tuple
+
 import torch
-from datastream.tools import starcompose, repeat_map_chain
+import torch.utils.data
+from pydantic import BaseModel, ConfigDict
+
 from datastream import Dataset
+from datastream.tools import repeat_map_chain, starcompose
 
 
 class ZipSampler(BaseModel, torch.utils.data.Sampler):
@@ -15,9 +19,10 @@ class ZipSampler(BaseModel, torch.utils.data.Sampler):
     from_mapping: Callable[[int], Tuple[int, ...]]
     zipped_samplers: Iterable
 
-    class Config:
-        arbitrary_types_allowed = True
-        allow_mutation = False
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        frozen=True,
+    )
 
     def __init__(self, samplers, datasets):
         BaseModel.__init__(

datastream/tools/verify_split.py

@@ -1,9 +1,10 @@
 import json
 from pathlib import Path
-from pydantic import validate_arguments
 
+from pydantic import validate_call
 
-@validate_arguments
+
+@validate_call
 def verify_split(old_path: Path, new_path: Path):
     """
     Verify that no keys from an old split are present in a different new split.