update

benchmark/data_frame_benchmark.py

@@ -3,7 +3,7 @@
 import os
 import os.path as osp
 import time
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Optional
 
 import numpy as np
 import optuna
@@ -303,10 +303,10 @@ def test(
 
 
 def train_and_eval_with_cfg(
-    model_cfg: Dict[str, Any],
-    train_cfg: Dict[str, Any],
+    model_cfg: dict[str, Any],
+    train_cfg: dict[str, Any],
     trial: Optional[optuna.trial.Trial] = None,
-) -> Tuple[float, float]:
+) -> tuple[float, float]:
     # Use model_cfg to set up training procedure
     if args.model_type == 'FTTransformerBucket':
         # Use LinearBucketEncoder instead

benchmark/encoder/encoder_benchmark.py

@@ -1,7 +1,6 @@
 import time
 from argparse import ArgumentParser
 from contextlib import nullcontext
-from typing import Dict
 
 import torch
 from line_profiler import profile
@@ -115,7 +114,7 @@
 }
 
 
-def make_stype_encoder_dict() -> Dict[stype, StypeEncoder]:
+def make_stype_encoder_dict() -> dict[stype, StypeEncoder]:
     stype_encoder_dict = {}
     for stype_str, encoder_str in args.stype_kv:
         encoder_kwargs = encoder_str2encoder_cls_kwargs[encoder_str]

examples/llm_embedding.py

@@ -1,7 +1,6 @@
 import argparse
 import os
 import os.path as osp
-from typing import List
 
 import torch
 import torch.nn.functional as F
@@ -64,10 +63,10 @@ def __init__(self, model: str = "text-embedding-ada-002"):
         self.client = OpenAI(api_key=api_key)
         self.model = model
 
-    def __call__(self, sentences: List[str]) -> Tensor:
+    def __call__(self, sentences: list[str]) -> Tensor:
         from openai import Embedding
 
-        items: List[Embedding] = self.client.embeddings.create(
+        items: list[Embedding] = self.client.embeddings.create(
             input=sentences, model=self.model).data
         assert len(items) == len(sentences)
         embeddings = [
@@ -87,7 +86,7 @@ def __init__(self, model: str = "embed-english-v3.0"):
         self.model = model
         self.co = cohere.Client(api_key)
 
-    def __call__(self, sentences: List[str]) -> Tensor:
+    def __call__(self, sentences: list[str]) -> Tensor:
         from cohere import EmbedResponse
 
         response: EmbedResponse = self.co.embed(model=self.model,
@@ -106,13 +105,13 @@ def __init__(self, model: str = "voyage-01"):
         # Please run `pip install voyageai` to install the package
         self.model = model
 
-    def __call__(self, sentences: List[str]) -> Tensor:
+    def __call__(self, sentences: list[str]) -> Tensor:
         import voyageai  # noqa
 
         voyageai.api_key = api_key
         from voyageai import get_embeddings
 
-        items: List[List[float]] = get_embeddings(sentences, model=self.model)
+        items: list[list[float]] = get_embeddings(sentences, model=self.model)
         assert len(items) == len(sentences)
         embeddings = [torch.FloatTensor(item).view(1, -1) for item in items]
         return torch.cat(embeddings, dim=0)

examples/tutorial.py

@@ -5,7 +5,7 @@
 import argparse
 import math
 import os.path as osp
-from typing import Any, Dict, List
+from typing import Any
 
 import torch
 import torch.nn.functional as F
@@ -157,8 +157,8 @@ def __init__(
         out_channels: int,
         num_layers: int,
         # kwargs for feature encoder
-        col_stats: Dict[str, Dict[StatType, Any]],
-        col_names_dict: Dict[torch_frame.stype, List[str]],
+        col_stats: dict[str, dict[StatType, Any]],
+        col_names_dict: dict[torch_frame.stype, list[str]],
     ):
         super().__init__()
         # Specify what feature encoder to use for each stype.

test/data/test_multi_embedding_tensor.py

@@ -1,5 +1,5 @@
 import random
-from typing import List, Optional, Tuple, Union
+from typing import Optional, Union
 
 import pytest
 import torch
@@ -9,7 +9,7 @@
 
 
 def assert_equal(
-    tensor_list: List[torch.Tensor],
+    tensor_list: list[torch.Tensor],
     met: MultiEmbeddingTensor,
 ) -> None:
     assert len(tensor_list) == met.num_cols
@@ -22,9 +22,9 @@ def assert_equal(
 
 
 def row_select(
-    tensor_list: List[torch.Tensor],
-    index: Union[List[int], slice],
-) -> List[torch.Tensor]:
+    tensor_list: list[torch.Tensor],
+    index: Union[list[int], slice],
+) -> list[torch.Tensor]:
     """Selects rows from a list of column tensors.
 
     Args:
@@ -45,7 +45,7 @@ def get_fake_multi_embedding_tensor(
     num_cols: int,
     embedding_dim: Optional[int] = None,
     device: Optional[torch.device] = None,
-) -> Tuple[MultiEmbeddingTensor, List[torch.Tensor]]:
+) -> tuple[MultiEmbeddingTensor, list[torch.Tensor]]:
     tensor_list = []
     for _ in range(num_cols):
         embedding_dim = embedding_dim or random.randint(1, 5)

test/data/test_multi_nested_tensor.py

@@ -1,5 +1,5 @@
 import random
-from typing import List, Union
+from typing import Union
 
 import pytest
 import torch
@@ -9,7 +9,7 @@
 from torch_frame.testing import withCUDA
 
 
-def assert_equal(tensor_mat: List[List[Tensor]],
+def assert_equal(tensor_mat: list[list[Tensor]],
                  multi_nested_tensor: MultiNestedTensor):
     assert len(tensor_mat) == multi_nested_tensor.shape[0]
     assert len(tensor_mat[0]) == multi_nested_tensor.shape[1]
@@ -20,9 +20,9 @@ def assert_equal(tensor_mat: List[List[Tensor]],
 
 
 def column_select(
-    tensor_mat: List[List[Tensor]],
-    index: Union[List[int], slice],
-) -> List[List[Tensor]]:
+    tensor_mat: list[list[Tensor]],
+    index: Union[list[int], slice],
+) -> list[list[Tensor]]:
     new_tensor_mat = []
     for tensor_vec in tensor_mat:
         if isinstance(index, slice):

test/nn/test_simple_basecls.py

@@ -1,5 +1,3 @@
-from typing import List, Tuple
-
 import torch
 from torch import Tensor
 
@@ -15,7 +13,7 @@ def __init__(
             self,
             out_channels: int,
             num_numerical: int,
-            num_categories: List[int],
+            num_categories: list[int],
         ):
             super().__init__()
 
@@ -31,7 +29,7 @@ def __init__(
                 for num_category in num_categories
             ])
 
-        def forward(self, tf: TensorFrame) -> Tuple[Tensor, List[str]]:
+        def forward(self, tf: TensorFrame) -> tuple[Tensor, list[str]]:
             xs = []
             for i, lin in enumerate(self.lins):
                 xs.append(lin(tf.feat_dict[torch_frame.numerical][:, i:i + 1]))

torch_frame/data/mapper.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Iterable, Mapping
+from collections.abc import Iterable, Mapping
+from typing import Any, Callable
 
 import numpy as np
 import pandas as pd

torch_frame/data/multi_embedding_tensor.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Sequence
+from collections.abc import Sequence
 
 import torch
 from torch import Tensor

torch_frame/data/multi_nested_tensor.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
-from typing import Sequence, cast
+from collections.abc import Sequence
+from typing import cast
 
 import torch
 from torch import Tensor

torch_frame/data/multi_tensor.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 import copy
-from typing import Any, Callable, Sequence, TypeVar
+from collections.abc import Sequence
+from typing import Any, Callable, TypeVar
 
 import torch
 from torch import Tensor

torch_frame/nn/encoder/encoder.py

@@ -1,5 +1,4 @@
 from abc import ABC, abstractmethod
-from typing import List, Tuple
 
 from torch import Tensor
 from torch.nn import Module
@@ -16,7 +15,7 @@ class FeatureEncoder(Module, ABC):
     value handling.
     """
     @abstractmethod
-    def forward(self, tf: TensorFrame) -> Tuple[Tensor, List[str]]:
+    def forward(self, tf: TensorFrame) -> tuple[Tensor, list[str]]:
         r"""Encode :class:`TensorFrame` object into a tuple
         :obj:`(x, col_names)`.
 

torch_frame/typing.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
+from collections.abc import Mapping
 from enum import Enum
-from typing import Dict, List, Mapping, Union
+from typing import Union
 
 import pandas as pd
 import torch
@@ -109,14 +110,14 @@ def is_timestamp_strategy(self) -> bool:
 Series = pd.Series
 DataFrame = pd.DataFrame
 
-IndexSelectType = Union[int, List[int], range, slice, Tensor]
-ColumnSelectType = Union[str, List[str]]
+IndexSelectType = Union[int, list[int], range, slice, Tensor]
+ColumnSelectType = Union[str, list[str]]
 TextTokenizationMapping = Mapping[str, Tensor]
-TextTokenizationOutputs = Union[List[TextTokenizationMapping],
+TextTokenizationOutputs = Union[list[TextTokenizationMapping],
                                 TextTokenizationMapping]
 TensorData = Union[
     Tensor,
     MultiNestedTensor,
     MultiEmbeddingTensor,
-    Dict[str, MultiNestedTensor],
+    dict[str, MultiNestedTensor],
 ]