replay-rec 0.19.0rc0__tar.gz → 0.20.0rc0__tar.gz

{replay_rec-0.19.0rc0 → replay_rec-0.20.0rc0}/PKG-INFO

@@ -1,53 +1,38 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: replay-rec
-Version: 0.19.0rc0
+Version: 0.20.0rc0
 Summary: RecSys Library
-Home-page: https://sb-ai-lab.github.io/RePlay/
-License: Apache-2.0
+License-Expression: Apache-2.0
+License-File: LICENSE
+License-File: NOTICE
 Author: AI Lab
-Requires-Python: >=3.8.1,<3.12
+Requires-Python: >=3.9, <3.13
+Classifier: Operating System :: Unix
 Classifier: Development Status :: 4 - Beta
 Classifier: Environment :: Console
 Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Science/Research
-Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Natural Language :: English
-Classifier: Operating System :: Unix
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
-Provides-Extra: all
-Provides-Extra: spark
-Provides-Extra: torch
-Provides-Extra: torch-openvino
-Requires-Dist: d3rlpy (>=2.0.4,<3.0.0)
-Requires-Dist: fixed-install-nmslib (==2.1.2)
-Requires-Dist: gym (>=0.26.0,<0.27.0)
-Requires-Dist: hnswlib (>=0.7.0,<0.8.0)
-Requires-Dist: implicit (>=0.7.0,<0.8.0)
-Requires-Dist: lightautoml (>=0.3.1,<0.4.0)
-Requires-Dist: lightfm (==1.17)
-Requires-Dist: lightning (>=2.0.2,<=2.4.0) ; extra == "torch" or extra == "torch-openvino" or extra == "all"
-Requires-Dist: llvmlite (>=0.32.1)
-Requires-Dist: numba (>=0.50)
-Requires-Dist: numpy (>=1.20.0)
-Requires-Dist: onnx (>=1.16.2,<1.17.0) ; extra == "torch-openvino" or extra == "all"
-Requires-Dist: openvino (>=2024.3.0,<2024.4.0) ; extra == "torch-openvino" or extra == "all"
-Requires-Dist: optuna (>=3.2.0,<3.3.0)
-Requires-Dist: pandas (>=1.3.5,<=2.2.2)
-Requires-Dist: polars (>=1.0.0,<1.1.0)
-Requires-Dist: psutil (>=6.0.0,<6.1.0)
-Requires-Dist: pyarrow (>=12.0.1)
-Requires-Dist: pyspark (>=3.0,<3.5) ; (python_full_version >= "3.8.1" and python_version < "3.11") and (extra == "spark" or extra == "all")
-Requires-Dist: pyspark (>=3.4,<3.5) ; (python_version >= "3.11" and python_version < "3.12") and (extra == "spark" or extra == "all")
-Requires-Dist: pytorch-ranger (>=0.1.1,<0.2.0) ; extra == "torch" or extra == "torch-openvino" or extra == "all"
-Requires-Dist: sb-obp (>=0.5.8,<0.6.0)
-Requires-Dist: scikit-learn (>=1.0.2,<2.0.0)
-Requires-Dist: scipy (>=1.8.1,<2.0.0)
-Requires-Dist: torch (>=1.8,<3.0.0) ; (python_version >= "3.9") and (extra == "torch" or extra == "torch-openvino" or extra == "all")
-Requires-Dist: torch (>=1.8,<=2.4.1) ; (python_version >= "3.8" and python_version < "3.9") and (extra == "torch" or extra == "torch-openvino" or extra == "all")
+Requires-Dist: d3rlpy (>=2.8.1,<2.9)
+Requires-Dist: implicit (>=0.7.2,<0.8)
+Requires-Dist: lightautoml (>=0.4.1,<0.5)
+Requires-Dist: lightning (>=2.0.2,<=2.4.0)
+Requires-Dist: numba (>=0.50,<1)
+Requires-Dist: numpy (>=1.20.0,<2)
+Requires-Dist: pandas (>=1.3.5,<2.4.0)
+Requires-Dist: polars (<2.0)
+Requires-Dist: psutil (<=7.0.0)
+Requires-Dist: pyarrow (<22.0)
+Requires-Dist: pyspark (>=3.0,<3.5)
+Requires-Dist: pytorch-optimizer (>=3.8.0,<4)
+Requires-Dist: sb-obp (>=0.5.10,<0.6)
+Requires-Dist: scikit-learn (>=1.6.1,<1.7.0)
+Requires-Dist: scipy (>=1.13.1,<1.14)
+Requires-Dist: setuptools
+Requires-Dist: torch (>=1.8,<3.0.0)
+Requires-Dist: tqdm (>=4.67,<5)
+Project-URL: Homepage, https://sb-ai-lab.github.io/RePlay/
 Project-URL: Repository, https://github.com/sb-ai-lab/RePlay
 Description-Content-Type: text/markdown
 
@@ -216,7 +201,6 @@ pip install replay-rec==XX.YY.ZZrc0
 In addition to the core package, several extras are also provided, including:
 - `[spark]`: Install PySpark functionality
 - `[torch]`: Install PyTorch and Lightning functionality
-- `[all]`: `[spark]` `[torch]`
 
 Example:
 ```bash
@@ -227,9 +211,41 @@ pip install replay-rec[spark]
 pip install replay-rec[spark]==XX.YY.ZZrc0
 ```
 
+Additionally, `replay-rec[torch]` may be installed with CPU-only version of `torch` by providing its respective index URL during installation:
+```bash
+# Install package with the CPU version of torch
+pip install replay-rec[torch] --extra-index-url https://download.pytorch.org/whl/cpu
+```
+
+
 To build RePlay from sources please use the [instruction](CONTRIBUTING.md#installing-from-the-source).
 
 
+### Optional features
+RePlay includes a set of optional features which require users to install optional dependencies manually. These features include:
+
+1) Hyperpearameter search via Optuna:
+```bash
+pip install optuna
+```
+
+2) Model compilation via OpenVINO:
+```bash
+pip install openvino onnx
+```
+
+3) Vector database and hierarchical search support:
+```bash
+pip install hnswlib fixed-install-nmslib
+```
+
+4) (Experimental) LightFM model support:
+```bash
+pip install ligfhtfm
+```
+> **_NOTE_** : LightFM is not officially supported for Python 3.12 due to discontinued maintenance of the library. If you wish to install it locally, you'll have to use a patched fork of LightFM, such as the [one used internally](https://github.com/daviddavo/lightfm).
+
+
 <a name="examples"></a>
 ## 📑  Resources
 

{replay_rec-0.19.0rc0 → replay_rec-0.20.0rc0}/README.md

@@ -163,7 +163,6 @@ pip install replay-rec==XX.YY.ZZrc0
 In addition to the core package, several extras are also provided, including:
 - `[spark]`: Install PySpark functionality
 - `[torch]`: Install PyTorch and Lightning functionality
-- `[all]`: `[spark]` `[torch]`
 
 Example:
 ```bash
@@ -174,9 +173,41 @@ pip install replay-rec[spark]
 pip install replay-rec[spark]==XX.YY.ZZrc0
 ```
 
+Additionally, `replay-rec[torch]` may be installed with CPU-only version of `torch` by providing its respective index URL during installation:
+```bash
+# Install package with the CPU version of torch
+pip install replay-rec[torch] --extra-index-url https://download.pytorch.org/whl/cpu
+```
+
+
 To build RePlay from sources please use the [instruction](CONTRIBUTING.md#installing-from-the-source).
 
 
+### Optional features
+RePlay includes a set of optional features which require users to install optional dependencies manually. These features include:
+
+1) Hyperpearameter search via Optuna:
+```bash
+pip install optuna
+```
+
+2) Model compilation via OpenVINO:
+```bash
+pip install openvino onnx
+```
+
+3) Vector database and hierarchical search support:
+```bash
+pip install hnswlib fixed-install-nmslib
+```
+
+4) (Experimental) LightFM model support:
+```bash
+pip install ligfhtfm
+```
+> **_NOTE_** : LightFM is not officially supported for Python 3.12 due to discontinued maintenance of the library. If you wish to install it locally, you'll have to use a patched fork of LightFM, such as the [one used internally](https://github.com/daviddavo/lightfm).
+
+
 <a name="examples"></a>
 ## 📑  Resources
 

{replay_rec-0.19.0rc0 → replay_rec-0.20.0rc0}/pyproject.toml

@@ -1,35 +1,28 @@
 [build-system]
 requires = [
-    "poetry-core>=1.0.0",
+    "poetry-core>=2.0.0",
     "poetry-dynamic-versioning>=1.0.0,<2.0.0",
+    "setuptools",
 ]
 build-backend = "poetry_dynamic_versioning.backend"
 
-[tool.black]
-line-length = 120
-target-versions = ["py38", "py39", "py310", "py311"]
-
-[tool.poetry]
+[project]
 name = "replay-rec"
-packages = [{include = "replay"}]
 license = "Apache-2.0"
 description = "RecSys Library"
 authors = [
-    "AI Lab",
-    "Alexey Vasilev",
-    "Anna Volodkevich",
-    "Alexey Grishanov",
-    "Yan-Martin Tamm",
-    "Boris Shminke",
-    "Alexander Sidorenko",
-    "Roza Aysina",
+    {name = "AI Lab"},
+    {name = "Alexey Vasilev"},
+    {name = "Anna Volodkevich"},
+    {name = "Alexey Grishanov"},
+    {name = "Yan-Martin Tamm"},
+    {name = "Boris Shminke"},
+    {name = "Alexander Sidorenko"},
+    {name = "Roza Aysina"},
 ]
 readme = "README.md"
-homepage = "https://sb-ai-lab.github.io/RePlay/"
-repository = "https://github.com/sb-ai-lab/RePlay"
 classifiers = [
     "Operating System :: Unix",
-    "Intended Audience :: Science/Research",
     "Development Status :: 4 - Beta",
     "Environment :: Console",
     "Intended Audience :: Developers",
@@ -37,51 +30,46 @@ classifiers = [
     "Natural Language :: English",
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
 ]
-exclude = [
-    "replay/conftest.py",
+requires-python = ">=3.9, <3.13"
+dependencies = [
+    "setuptools",
+    "numpy (>=1.20.0,<2)",
+    "pandas (>=1.3.5,<2.4.0)",
+    "polars (<2.0)",
+    "scipy (>=1.13.1,<1.14)",
+    "scikit-learn (>=1.6.1,<1.7.0)",
+    "pyarrow (<22.0)",
+    "tqdm (>=4.67,<5)",
+    "torch (>=1.8,<3.0.0)",
+    "lightning (>=2.0.2,<=2.4.0)",
+    "pytorch-optimizer (>=3.8.0,<4)",
+    "lightautoml (>=0.4.1,<0.5)",
+    "numba (>=0.50,<1)",
+    "sb-obp (>=0.5.10,<0.6)",
+    "d3rlpy (>=2.8.1,<2.9)",
+    "implicit (>=0.7.2,<0.8)",
+    "pyspark (>=3.0,<3.5)",
+    "psutil (<=7.0.0)",
 ]
-version = "0.19.0.preview"
+dynamic = ["dependencies"]
+version = "0.20.0.preview"
 
-[tool.poetry.dependencies]
-python = ">=3.8.1, <3.12"
-numpy = ">=1.20.0"
-pandas = ">=1.3.5, <=2.2.2"
-polars = "~1.0.0"
-optuna = "~3.2.0"
-scipy = "^1.8.1"
-psutil = "~6.0.0"
-scikit-learn = "^1.0.2"
-pyarrow = ">=12.0.1"
-openvino = {version = "~2024.3.0", optional = true}
-onnx = {version = "~1.16.2", optional = true}
-fixed-install-nmslib = "2.1.2"
-hnswlib = "^0.7.0"
-pyspark = [
-    {version = ">=3.4,<3.5", python = ">=3.11,<3.12"},
-    {version = ">=3.0,<3.5", python = ">=3.8.1,<3.11"},
-]
-torch = [
-    {version = ">=1.8, <3.0.0", python = ">=3.9", optional = true},
-    {version = ">=1.8, <=2.4.1", python = ">=3.8,<3.9", optional = true},
-]
-lightning = ">=2.0.2, <=2.4.0"
-pytorch-ranger = "^0.1.1"
-lightfm = "1.17"
-lightautoml = "~0.3.1"
-numba = ">=0.50"
-llvmlite = ">=0.32.1"
-sb-obp = "^0.5.8"
-d3rlpy = "^2.0.4"
-implicit = "~0.7.0"
-gym = "^0.26.0"
+[project.urls]
+homepage = "https://sb-ai-lab.github.io/RePlay/"
+repository = "https://github.com/sb-ai-lab/RePlay"
 
-[tool.poetry.extras]
-spark = ["pyspark"]
-torch = ["torch", "pytorch-ranger", "lightning"]
-torch-openvino = ["torch", "pytorch-ranger", "lightning", "openvino", "onnx"]
-all = ["pyspark", "torch", "pytorch-ranger", "lightning", "openvino", "onnx"]
+[tool.black]
+line-length = 120
+target-version = ["py39", "py310", "py311", "py312"]
+
+[tool.poetry]
+packages = [{include = "replay"}]
+exclude = [
+    "replay/conftest.py",
+]
 
 [tool.poetry.group.dev.dependencies]
+coverage-conditional-plugin = "^0.9.0"
 jupyter = "~1.0.0"
 jupyterlab = "^3.6.0"
 pytest = ">=7.1.0"
@@ -102,31 +90,29 @@ filelock = "~3.14.0"
 
 [tool.poetry-dynamic-versioning]
 enable = false
-format-jinja = """0.19.0{{ env['PACKAGE_SUFFIX'] }}"""
+format-jinja = """0.20.0{{ env['PACKAGE_SUFFIX'] }}"""
 vcs = "git"
 
 [tool.ruff]
 exclude = [".git", ".venv", "__pycache__", "env", "venv", "docs", "projects", "examples"]
-extend-select = ["C90", "T10", "T20", "UP004"]
 line-length = 120
+
+[tool.ruff.lint]
 select = ["ARG", "C4", "E", "EM", "ERA", "F", "FLY", "I", "INP", "ISC", "N", "PERF", "PGH", "PIE", "PYI", "Q", "RUF", "SIM", "TID", "W"]
+extend-select = ["C90", "T10", "T20", "UP004"]
+ignore = ["SIM115"]
+mccabe = {max-complexity = 13}
+isort = {combine-as-imports = true, force-wrap-aliases = true}
 
-[tool.ruff.flake8-quotes]
+[tool.ruff.lint.flake8-quotes]
 docstring-quotes = "double"
 inline-quotes = "double"
 multiline-quotes = "double"
 
-[tool.ruff.flake8-unused-arguments]
+[tool.ruff.lint.flake8-unused-arguments]
 ignore-variadic-names = false
 
-[tool.ruff.isort]
-combine-as-imports = true
-force-wrap-aliases = true
-
-[tool.ruff.mccabe]
-max-complexity = 13
-
-[tool.ruff.per-file-ignores]
+[tool.ruff.lint.per-file-ignores]
 "*/" = ["PERF203", "RUF001", "RUF002", "RUF012", "E402"]
 "__init__.py" = ["F401"]
 "replay/utils/model_handler.py" = ["F403", "F405"]

replay_rec-0.20.0rc0/replay/__init__.py

@@ -0,0 +1,7 @@
+"""RecSys library"""
+
+# NOTE: This ensures distutils monkey-patching is performed before any
+# functionality removed in Python 3.12 is used in downstream packages (like lightfm)
+import setuptools as _
+
+__version__ = "0.20.0.preview"

{replay_rec-0.19.0rc0 → replay_rec-0.20.0rc0}/replay/data/dataset.py

@@ -5,8 +5,9 @@
 from __future__ import annotations
 
 import json
+from collections.abc import Iterable, Sequence
 from pathlib import Path
-from typing import Callable, Dict, Iterable, List, Optional, Sequence, Union
+from typing import Callable, Optional, Union
 
 import numpy as np
 from pandas import read_parquet as pd_read_parquet
@@ -315,7 +316,7 @@ class Dataset:
         :returns: Loaded Dataset.
         """
         base_path = Path(path).with_suffix(".replay").resolve()
-        with open(base_path / "init_args.json", "r") as file:
+        with open(base_path / "init_args.json") as file:
             dataset_dict = json.loads(file.read())
 
         if dataframe_type not in ["pandas", "spark", "polars", None]:
@@ -436,14 +437,14 @@ class Dataset:
         )
 
     def _get_feature_source_map(self):
-        self._feature_source_map: Dict[FeatureSource, DataFrameLike] = {
+        self._feature_source_map: dict[FeatureSource, DataFrameLike] = {
             FeatureSource.INTERACTIONS: self.interactions,
             FeatureSource.QUERY_FEATURES: self.query_features,
             FeatureSource.ITEM_FEATURES: self.item_features,
         }
 
     def _get_ids_source_map(self):
-        self._ids_feature_map: Dict[FeatureHint, DataFrameLike] = {
+        self._ids_feature_map: dict[FeatureHint, DataFrameLike] = {
             FeatureHint.QUERY_ID: self.query_features if self.query_features is not None else self.interactions,
             FeatureHint.ITEM_ID: self.item_features if self.item_features is not None else self.interactions,
         }
@@ -499,10 +500,10 @@ class Dataset:
             )
         return FeatureSchema(features_list=features_list + filled_features)
 
-    def _fill_unlabeled_features_sources(self, feature_schema: FeatureSchema) -> List[FeatureInfo]:
+    def _fill_unlabeled_features_sources(self, feature_schema: FeatureSchema) -> list[FeatureInfo]:
         features_list = list(feature_schema.all_features)
 
-        source_mapping: Dict[str, FeatureSource] = {}
+        source_mapping: dict[str, FeatureSource] = {}
         for source in FeatureSource:
             dataframe = self._feature_source_map[source]
             if dataframe is not None:
@@ -524,7 +525,7 @@ class Dataset:
         self._set_cardinality(features_list=features_list)
         return features_list
 
-    def _get_unlabeled_columns(self, source: FeatureSource, feature_schema: FeatureSchema) -> List[FeatureInfo]:
+    def _get_unlabeled_columns(self, source: FeatureSource, feature_schema: FeatureSchema) -> list[FeatureInfo]:
         set_source_dataframe_columns = set(self._feature_source_map[source].columns)
         set_labeled_dataframe_columns = set(feature_schema.columns)
         unlabeled_columns = set_source_dataframe_columns - set_labeled_dataframe_columns
@@ -534,13 +535,13 @@ class Dataset:
         ]
         return unlabeled_features_list
 
-    def _fill_unlabeled_features(self, source: FeatureSource, feature_schema: FeatureSchema) -> List[FeatureInfo]:
+    def _fill_unlabeled_features(self, source: FeatureSource, feature_schema: FeatureSchema) -> list[FeatureInfo]:
         unlabeled_columns = self._get_unlabeled_columns(source=source, feature_schema=feature_schema)
         self._set_features_source(feature_list=unlabeled_columns, source=source)
         self._set_cardinality(features_list=unlabeled_columns)
         return unlabeled_columns
 
-    def _set_features_source(self, feature_list: List[FeatureInfo], source: FeatureSource) -> None:
+    def _set_features_source(self, feature_list: list[FeatureInfo], source: FeatureSource) -> None:
         for feature in feature_list:
             feature._set_feature_source(source)
 
@@ -610,9 +611,9 @@ class Dataset:
             if self.is_pandas:
                 try:
                     data[column] = data[column].astype(int)
-                except Exception:
+                except Exception as exc:
                     msg = f"IDs in {source.name}.{column} are not encoded. They are not int."
-                    raise ValueError(msg)
+                    raise ValueError(msg) from exc
 
         if self.is_pandas:
             is_int = np.issubdtype(dict(data.dtypes)[column], int)
@@ -775,10 +776,10 @@ def check_dataframes_types_equal(dataframe: DataFrameLike, other: DataFrameLike)
 
     :returns: True if dataframes have same type.
     """
-    if isinstance(dataframe, PandasDataFrame) and isinstance(other, PandasDataFrame):
-        return True
-    if isinstance(dataframe, SparkDataFrame) and isinstance(other, SparkDataFrame):
-        return True
-    if isinstance(dataframe, PolarsDataFrame) and isinstance(other, PolarsDataFrame):
-        return True
-    return False
+    return any(
+        [
+            isinstance(dataframe, PandasDataFrame) and isinstance(other, PandasDataFrame),
+            isinstance(dataframe, SparkDataFrame) and isinstance(other, SparkDataFrame),
+            isinstance(dataframe, PolarsDataFrame) and isinstance(other, PolarsDataFrame),
+        ]
+    )

{replay_rec-0.19.0rc0 → replay_rec-0.20.0rc0}/replay/data/dataset_utils/dataset_label_encoder.py

@@ -6,7 +6,8 @@ Contains classes for encoding categorical data
 """
 
 import warnings
-from typing import Dict, Iterable, Iterator, Optional, Sequence, Set, Union
+from collections.abc import Iterable, Iterator, Sequence
+from typing import Optional, Union
 
 from replay.data import Dataset, FeatureHint, FeatureSchema, FeatureSource, FeatureType
 from replay.preprocessing import LabelEncoder, LabelEncodingRule, SequenceEncodingRule
@@ -45,9 +46,9 @@ class DatasetLabelEncoder:
         """
         self._handle_unknown_rule = handle_unknown_rule
         self._default_value_rule = default_value_rule
-        self._encoding_rules: Dict[str, LabelEncodingRule] = {}
+        self._encoding_rules: dict[str, LabelEncodingRule] = {}
 
-        self._features_columns: Dict[Union[FeatureHint, FeatureSource], Sequence[str]] = {}
+        self._features_columns: dict[Union[FeatureHint, FeatureSource], Sequence[str]] = {}
 
     def fit(self, dataset: Dataset) -> "DatasetLabelEncoder":
         """
@@ -161,7 +162,7 @@ class DatasetLabelEncoder:
         """
         self._check_if_initialized()
 
-        columns_set: Set[str]
+        columns_set: set[str]
         columns_set = {columns} if isinstance(columns, str) else {*columns}
 
         def get_encoding_rules() -> Iterator[LabelEncodingRule]:

{replay_rec-0.19.0rc0 → replay_rec-0.20.0rc0}/replay/data/nn/schema.py

@@ -1,17 +1,8 @@
+from collections import OrderedDict
+from collections.abc import ItemsView, Iterable, Iterator, KeysView, Mapping, Sequence, ValuesView
 from typing import (
-    Dict,
-    ItemsView,
-    Iterable,
-    Iterator,
-    KeysView,
-    List,
-    Mapping,
     Optional,
-    OrderedDict,
-    Sequence,
-    Set,
     Union,
-    ValuesView,
 )
 
 import torch
@@ -20,7 +11,7 @@ from replay.data import FeatureHint, FeatureSource, FeatureType
 
 # Alias
 TensorMap = Mapping[str, torch.Tensor]
-MutableTensorMap = Dict[str, torch.Tensor]
+MutableTensorMap = dict[str, torch.Tensor]
 
 
 class TensorFeatureSource:
@@ -79,7 +70,7 @@ class TensorFeatureInfo:
         feature_type: FeatureType,
         is_seq: bool = False,
         feature_hint: Optional[FeatureHint] = None,
-        feature_sources: Optional[List[TensorFeatureSource]] = None,
+        feature_sources: Optional[list[TensorFeatureSource]] = None,
         cardinality: Optional[int] = None,
         padding_value: int = 0,
         embedding_dim: Optional[int] = None,
@@ -154,13 +145,13 @@ class TensorFeatureInfo:
         self._feature_hint = hint
 
     @property
-    def feature_sources(self) -> Optional[List[TensorFeatureSource]]:
+    def feature_sources(self) -> Optional[list[TensorFeatureSource]]:
         """
         :returns: List of sources feature came from.
         """
         return self._feature_sources
 
-    def _set_feature_sources(self, sources: List[TensorFeatureSource]) -> None:
+    def _set_feature_sources(self, sources: list[TensorFeatureSource]) -> None:
         self._feature_sources = sources
 
     @property
@@ -276,7 +267,7 @@ class TensorSchema(Mapping[str, TensorFeatureInfo]):
 
         :returns: New tensor schema of given features.
         """
-        features: Set[TensorFeatureInfo] = set()
+        features: set[TensorFeatureInfo] = set()
         for feature_name in features_to_keep:
             features.add(self._tensor_schema[feature_name])
         return TensorSchema(list(features))
@@ -432,7 +423,7 @@ class TensorSchema(Mapping[str, TensorFeatureInfo]):
             return None
         return rating_features.item().name
 
-    def _get_object_args(self) -> Dict:
+    def _get_object_args(self) -> dict:
         """
         Returns list of features represented as dictionaries.
         """
@@ -456,7 +447,7 @@ class TensorSchema(Mapping[str, TensorFeatureInfo]):
         return features
 
     @classmethod
-    def _create_object_by_args(cls, args: Dict) -> "TensorSchema":
+    def _create_object_by_args(cls, args: dict) -> "TensorSchema":
         features_list = []
         for feature_data in args:
             feature_data["feature_sources"] = (