bug fix

alibaba · Dec 4, 2024 · dcce8da · dcce8da
1 parent 2cc79c7
commit dcce8da
Show file tree

Hide file tree

Showing 283 changed files with 87 additions and 35 deletions.
diff --git a/.github/workflows/codestyle_ci.yml b/.github/workflows/codestyle_ci.yml
diff --git a/.github/workflows/pytyping_ci.yml b/.github/workflows/pytyping_ci.yml
diff --git a/.github/workflows/unittest_ci.yml b/.github/workflows/unittest_ci.yml
diff --git a/.github/workflows/unittest_cpu_ci.yml b/.github/workflows/unittest_cpu_ci.yml
diff --git a/.gitignore b/.gitignore
@@ -39,4 +39,4 @@ protoc*
 docs/source/intro.md
 docs/source/proto.html
 
-.vscode/
+.vscode/
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
diff --git a/.pyre_configuration b/.pyre_configuration
diff --git a/.readthedocs.yaml b/.readthedocs.yaml
diff --git a/.ruff.toml b/.ruff.toml
diff --git a/LICENSE b/LICENSE
diff --git a/MANIFEST.in b/MANIFEST.in
diff --git a/README.md b/README.md
diff --git a/data/.license_header.txt b/data/.license_header.txt
diff --git a/data/test/spiece.model b/data/test/spiece.model
diff --git a/data/test/tokenizer.json b/data/test/tokenizer.json
diff --git a/docker/Dockerfile b/docker/Dockerfile
diff --git a/docker/pip.conf b/docker/pip.conf
diff --git a/docs/Makefile b/docs/Makefile
diff --git a/docs/images/intro.png b/docs/images/intro.png
diff --git a/docs/images/models/dbmtl.png b/docs/images/models/dbmtl.png
diff --git a/docs/images/models/dbmtl_mmoe.png b/docs/images/models/dbmtl_mmoe.png
diff --git a/docs/images/models/deepfm.png b/docs/images/models/deepfm.png
diff --git a/docs/images/models/din.png b/docs/images/models/din.png
diff --git a/docs/images/models/dssm_neg_sampler.png b/docs/images/models/dssm_neg_sampler.png
diff --git a/docs/images/models/feature_groups_din.png b/docs/images/models/feature_groups_din.png
diff --git a/docs/images/models/mmoe.png b/docs/images/models/mmoe.png
diff --git a/docs/images/models/multi_tower.png b/docs/images/models/multi_tower.png
diff --git a/docs/images/models/ple.png b/docs/images/models/ple.png
diff --git a/docs/images/models/tdm.png b/docs/images/models/tdm.png
diff --git a/docs/images/qrcode/dinggroup1.png b/docs/images/qrcode/dinggroup1.png
diff --git a/docs/images/qrcode/dinggroup2.png b/docs/images/qrcode/dinggroup2.png
diff --git a/docs/make.bat b/docs/make.bat
diff --git a/docs/source/conf.py b/docs/source/conf.py
diff --git a/docs/source/develop.md b/docs/source/develop.md
diff --git a/docs/source/faq.md b/docs/source/faq.md
diff --git a/docs/source/feature/data.md b/docs/source/feature/data.md
diff --git a/docs/source/feature/feature.md b/docs/source/feature/feature.md
diff --git a/docs/source/index.rst b/docs/source/index.rst
diff --git a/docs/source/models/dbmtl.md b/docs/source/models/dbmtl.md
diff --git a/docs/source/models/deepfm.md b/docs/source/models/deepfm.md
diff --git a/docs/source/models/din.md b/docs/source/models/din.md
diff --git a/docs/source/models/dssm.md b/docs/source/models/dssm.md
diff --git a/docs/source/models/feature_group.md b/docs/source/models/feature_group.md
diff --git a/docs/source/models/loss.md b/docs/source/models/loss.md
diff --git a/docs/source/models/mmoe.md b/docs/source/models/mmoe.md
diff --git a/docs/source/models/multi_target.rst b/docs/source/models/multi_target.rst
diff --git a/docs/source/models/multi_tower.md b/docs/source/models/multi_tower.md
diff --git a/docs/source/models/ple.md b/docs/source/models/ple.md
diff --git a/docs/source/models/rank.rst b/docs/source/models/rank.rst
diff --git a/docs/source/models/recall.rst b/docs/source/models/recall.rst
diff --git a/docs/source/models/tdm.md b/docs/source/models/tdm.md
diff --git a/docs/source/models/user_define.md b/docs/source/models/user_define.md
diff --git a/docs/source/quick_start.rst b/docs/source/quick_start.rst
diff --git a/docs/source/quick_start/dlc_tutorial.md b/docs/source/quick_start/dlc_tutorial.md
diff --git a/docs/source/quick_start/local_tutorial.md b/docs/source/quick_start/local_tutorial.md
diff --git a/docs/source/quick_start/local_tutorial_tdm.md b/docs/source/quick_start/local_tutorial_tdm.md
diff --git a/docs/source/quick_start/local_tutorial_u2i_vec.md b/docs/source/quick_start/local_tutorial_u2i_vec.md
diff --git a/docs/source/reference.md b/docs/source/reference.md
diff --git a/docs/source/usage/convert_easyrec_config_to_tzrec_config.md b/docs/source/usage/convert_easyrec_config_to_tzrec_config.md
diff --git a/docs/source/usage/eval.md b/docs/source/usage/eval.md
diff --git a/docs/source/usage/export.md b/docs/source/usage/export.md
diff --git a/docs/source/usage/feature_selection.md b/docs/source/usage/feature_selection.md
diff --git a/docs/source/usage/predict.md b/docs/source/usage/predict.md
diff --git a/docs/source/usage/serving.md b/docs/source/usage/serving.md
diff --git a/docs/source/usage/train.md b/docs/source/usage/train.md
diff --git a/examples/dbmtl_taobao.config b/examples/dbmtl_taobao.config
diff --git a/examples/dbmtl_taobao_jrc.config b/examples/dbmtl_taobao_jrc.config
diff --git a/examples/dbmtl_taobao_seq.config b/examples/dbmtl_taobao_seq.config
diff --git a/examples/deepfm_criteo.config b/examples/deepfm_criteo.config
diff --git a/examples/dssm_taobao.config b/examples/dssm_taobao.config
diff --git a/examples/dssm_v2_taobao.config b/examples/dssm_v2_taobao.config
diff --git a/examples/mmoe_taobao.config b/examples/mmoe_taobao.config
diff --git a/examples/multi_tower_din_taobao.config b/examples/multi_tower_din_taobao.config
diff --git a/examples/multi_tower_taobao.config b/examples/multi_tower_taobao.config
diff --git a/examples/ple_taobao.config b/examples/ple_taobao.config
diff --git a/examples/tdm_taobao.config b/examples/tdm_taobao.config
diff --git a/requirements.txt b/requirements.txt
diff --git a/requirements/docs.txt b/requirements/docs.txt
diff --git a/requirements/runtime.txt b/requirements/runtime.txt
diff --git a/requirements/test.txt b/requirements/test.txt
diff --git a/scripts/build_docker.sh b/scripts/build_docker.sh
diff --git a/scripts/build_wheel.sh b/scripts/build_wheel.sh
diff --git a/scripts/ci_test.sh b/scripts/ci_test.sh
diff --git a/scripts/doc/build_doc_pre_work.sh b/scripts/doc/build_doc_pre_work.sh
diff --git a/scripts/doc/build_docs.sh b/scripts/doc/build_docs.sh
diff --git a/scripts/gen_proto.sh b/scripts/gen_proto.sh
diff --git a/scripts/pyre_check.py b/scripts/pyre_check.py
diff --git a/setup.cfg b/setup.cfg
diff --git a/setup.py b/setup.py
diff --git a/tzrec/__init__.py b/tzrec/__init__.py
diff --git a/tzrec/acc/__init__.py b/tzrec/acc/__init__.py
diff --git a/tzrec/acc/_aten_lowering_pass.py b/tzrec/acc/_aten_lowering_pass.py
diff --git a/tzrec/acc/_decompositions.py b/tzrec/acc/_decompositions.py
diff --git a/tzrec/acc/trt_utils.py b/tzrec/acc/trt_utils.py
diff --git a/tzrec/acc/utils.py b/tzrec/acc/utils.py
diff --git a/tzrec/benchmark/__init__.py b/tzrec/benchmark/__init__.py
diff --git a/tzrec/benchmark/benchmark.py b/tzrec/benchmark/benchmark.py
diff --git a/tzrec/benchmark/configs/base_eval_metric.json b/tzrec/benchmark/configs/base_eval_metric.json
diff --git a/tzrec/benchmark/configs/criteo/deepfm.config b/tzrec/benchmark/configs/criteo/deepfm.config
diff --git a/tzrec/benchmark/configs/taobao/dbmtl.config b/tzrec/benchmark/configs/taobao/dbmtl.config
diff --git a/tzrec/benchmark/configs/taobao/dbmtl_has_sequence.config b/tzrec/benchmark/configs/taobao/dbmtl_has_sequence.config
diff --git a/tzrec/benchmark/configs/taobao/dbmtl_jrc.config b/tzrec/benchmark/configs/taobao/dbmtl_jrc.config
diff --git a/tzrec/benchmark/configs/taobao/mmoe.config b/tzrec/benchmark/configs/taobao/mmoe.config
diff --git a/tzrec/benchmark/configs/taobao/mmoe_has_sequence.config b/tzrec/benchmark/configs/taobao/mmoe_has_sequence.config
diff --git a/tzrec/benchmark/configs/taobao/ple.config b/tzrec/benchmark/configs/taobao/ple.config
diff --git a/tzrec/benchmark/configs/taobao/ple_has_sequence.config b/tzrec/benchmark/configs/taobao/ple_has_sequence.config
diff --git a/tzrec/benchmark/configs/taobao_ccp/dbmtl.config b/tzrec/benchmark/configs/taobao_ccp/dbmtl.config
diff --git a/tzrec/benchmark/configs/taobao_ccp/mmoe.config b/tzrec/benchmark/configs/taobao_ccp/mmoe.config
diff --git a/tzrec/benchmark/configs/taobao_ccp/ple.config b/tzrec/benchmark/configs/taobao_ccp/ple.config
diff --git a/tzrec/constant.py b/tzrec/constant.py
diff --git a/tzrec/datasets/__init__.py b/tzrec/datasets/__init__.py
diff --git a/tzrec/datasets/csv_dataset.py b/tzrec/datasets/csv_dataset.py
diff --git a/tzrec/datasets/csv_dataset_test.py b/tzrec/datasets/csv_dataset_test.py
diff --git a/tzrec/datasets/data_parser.py b/tzrec/datasets/data_parser.py
@@ -155,7 +155,7 @@ def parse(self, input_data: Dict[str, pa.Array]) -> Dict[str, torch.Tensor]:
 
         for label_name in self._labels:
             output_data[label_name] = _to_tensor(input_data[label_name].to_numpy())
-        
+
         for weight in self._sample_weights:
             output_data[weight] = _to_tensor(input_data[weight].to_numpy())
 
@@ -326,7 +326,7 @@ def to_batch(
         labels = {}
         for label_name in self._labels:
             labels[label_name] = input_data[label_name]
-        
+
         sample_weights = {}
         for weight in self._sample_weights:
             sample_weights[weight] = input_data[weight]

diff --git a/tzrec/datasets/data_parser_test.py b/tzrec/datasets/data_parser_test.py
diff --git a/tzrec/datasets/dataset.py b/tzrec/datasets/dataset.py
diff --git a/tzrec/datasets/dataset_test.py b/tzrec/datasets/dataset_test.py
diff --git a/tzrec/datasets/odps_dataset.py b/tzrec/datasets/odps_dataset.py
diff --git a/tzrec/datasets/odps_dataset_test.py b/tzrec/datasets/odps_dataset_test.py
diff --git a/tzrec/datasets/odps_dataset_v1.py b/tzrec/datasets/odps_dataset_v1.py
diff --git a/tzrec/datasets/odps_dataset_v1_test.py b/tzrec/datasets/odps_dataset_v1_test.py
diff --git a/tzrec/datasets/parquet_dataset.py b/tzrec/datasets/parquet_dataset.py
diff --git a/tzrec/datasets/parquet_dataset_test.py b/tzrec/datasets/parquet_dataset_test.py
diff --git a/tzrec/datasets/sampler.py b/tzrec/datasets/sampler.py
diff --git a/tzrec/datasets/sampler_test.py b/tzrec/datasets/sampler_test.py
diff --git a/tzrec/datasets/utils.py b/tzrec/datasets/utils.py
@@ -136,7 +136,7 @@ def to(self, device: torch.device, non_blocking: bool = False) -> "Batch":
             sample_weights={
                 k: v.to(device=device, non_blocking=non_blocking)
                 for k, v in self.sample_weights.items()
-            }
+            },
         )
 
     def record_stream(self, stream: torch.Stream) -> None:

diff --git a/tzrec/eval.py b/tzrec/eval.py
diff --git a/tzrec/export.py b/tzrec/export.py
diff --git a/tzrec/features/__init__.py b/tzrec/features/__init__.py
diff --git a/tzrec/features/combo_feature.py b/tzrec/features/combo_feature.py
diff --git a/tzrec/features/combo_feature_test.py b/tzrec/features/combo_feature_test.py
diff --git a/tzrec/features/expr_feature.py b/tzrec/features/expr_feature.py
diff --git a/tzrec/features/expr_feature_test.py b/tzrec/features/expr_feature_test.py
diff --git a/tzrec/features/feature.py b/tzrec/features/feature.py
diff --git a/tzrec/features/feature_test.py b/tzrec/features/feature_test.py
diff --git a/tzrec/features/id_feature.py b/tzrec/features/id_feature.py
diff --git a/tzrec/features/id_feature_test.py b/tzrec/features/id_feature_test.py
diff --git a/tzrec/features/lookup_feature.py b/tzrec/features/lookup_feature.py
diff --git a/tzrec/features/lookup_feature_test.py b/tzrec/features/lookup_feature_test.py
diff --git a/tzrec/features/match_feature.py b/tzrec/features/match_feature.py
diff --git a/tzrec/features/match_feature_test.py b/tzrec/features/match_feature_test.py
diff --git a/tzrec/features/overlap_feature.py b/tzrec/features/overlap_feature.py
diff --git a/tzrec/features/overlap_feature_test.py b/tzrec/features/overlap_feature_test.py
diff --git a/tzrec/features/raw_feature.py b/tzrec/features/raw_feature.py
diff --git a/tzrec/features/raw_feature_test.py b/tzrec/features/raw_feature_test.py
diff --git a/tzrec/features/sequence_feature.py b/tzrec/features/sequence_feature.py
diff --git a/tzrec/features/sequence_feature_test.py b/tzrec/features/sequence_feature_test.py
diff --git a/tzrec/features/tokenize_feature.py b/tzrec/features/tokenize_feature.py
diff --git a/tzrec/features/tokenize_feature_test.py b/tzrec/features/tokenize_feature_test.py
diff --git a/tzrec/loss/__init__.py b/tzrec/loss/__init__.py
diff --git a/tzrec/loss/jrc_loss.py b/tzrec/loss/jrc_loss.py
diff --git a/tzrec/loss/jrc_loss_test.py b/tzrec/loss/jrc_loss_test.py
diff --git a/tzrec/main.py b/tzrec/main.py
@@ -218,14 +218,18 @@ def _get_dataloader(
 
 
 def _create_model(
-    model_config: ModelConfig, features: List[BaseFeature], labels: List[str], sample_weights: List[str] = []
+    model_config: ModelConfig,
+    features: List[BaseFeature],
+    labels: List[str],
+    sample_weights: Optional[List[str]] = None,
 ) -> BaseModel:
     """Build model.
 
     Args:
         model_config (ModelConfig): easyrec model config.
         features (list): list of features.
         labels (list): list of label names.
+        sample_weights (list): list of sample weight names
 
     Return:
         model: a EasyRec Model.
@@ -538,7 +542,7 @@ def train_and_evaluate(
         pipeline_config.model_config,
         features,
         list(data_config.label_fields),
-        list(data_config.sample_weight_fields)
+        list(data_config.sample_weight_fields),
     )
     model = TrainWrapper(model)
 

diff --git a/tzrec/metrics/__init__.py b/tzrec/metrics/__init__.py
diff --git a/tzrec/metrics/grouped_auc.py b/tzrec/metrics/grouped_auc.py
diff --git a/tzrec/metrics/grouped_auc_test.py b/tzrec/metrics/grouped_auc_test.py
diff --git a/tzrec/metrics/recall_at_k.py b/tzrec/metrics/recall_at_k.py
diff --git a/tzrec/metrics/recall_at_k_test.py b/tzrec/metrics/recall_at_k_test.py
diff --git a/tzrec/models/__init__.py b/tzrec/models/__init__.py
diff --git a/tzrec/models/dbmtl.py b/tzrec/models/dbmtl.py
@@ -9,7 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 import torch
 from torch import nn
@@ -35,7 +35,11 @@ class DBMTL(MultiTaskRank):
     """
 
     def __init__(
-        self, model_config: ModelConfig, features: List[BaseFeature], labels: List[str], sample_weights: List[str] = []
+        self,
+        model_config: ModelConfig,
+        features: List[BaseFeature],
+        labels: List[str],
+        sample_weights: Optional[List[str]] = None,
     ) -> None:
         super().__init__(model_config, features, labels, sample_weights)
         assert model_config.WhichOneof("model") == "dbmtl", (

diff --git a/tzrec/models/dbmtl_test.py b/tzrec/models/dbmtl_test.py
diff --git a/tzrec/models/deepfm.py b/tzrec/models/deepfm.py
@@ -9,7 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 import torch
 from torch import nn
@@ -34,7 +34,11 @@ class DeepFM(RankModel):
     """
 
     def __init__(
-        self, model_config: ModelConfig, features: List[BaseFeature], labels: List[str], sample_weights: List[str] = []
+        self,
+        model_config: ModelConfig,
+        features: List[BaseFeature],
+        labels: List[str],
+        sample_weights: Optional[List[str]] = None,
     ) -> None:
         super().__init__(model_config, features, labels, sample_weights)
         self.init_input()

diff --git a/tzrec/models/deepfm_test.py b/tzrec/models/deepfm_test.py
diff --git a/tzrec/models/dssm.py b/tzrec/models/dssm.py
@@ -99,7 +99,7 @@ def __init__(
         model_config: model_pb2.ModelConfig,
         features: List[BaseFeature],
         labels: List[str],
-        sample_weights: List[str] = []
+        sample_weights: Optional[List[str]] = None,
     ) -> None:
         super().__init__(model_config, features, labels, sample_weights)
         name_to_feature_group = {x.group_name: x for x in model_config.feature_groups}

diff --git a/tzrec/models/dssm_test.py b/tzrec/models/dssm_test.py
diff --git a/tzrec/models/dssm_v2.py b/tzrec/models/dssm_v2.py
@@ -10,7 +10,7 @@
 # limitations under the License.
 
 from collections import OrderedDict
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 import torch
 import torch.nn.functional as F
@@ -87,7 +87,7 @@ def __init__(
         model_config: model_pb2.ModelConfig,
         features: List[BaseFeature],
         labels: List[str],
-        sample_weights: List[str] = []
+        sample_weights: Optional[List[str]] = None,
     ) -> None:
         super().__init__(model_config, features, labels, sample_weights)
         name_to_feature_group = {x.group_name: x for x in model_config.feature_groups}

diff --git a/tzrec/models/dssm_v2_test.py b/tzrec/models/dssm_v2_test.py
diff --git a/tzrec/models/match_model.py b/tzrec/models/match_model.py
@@ -159,7 +159,11 @@ class MatchModel(BaseModel):
     """
 
     def __init__(
-        self, model_config: ModelConfig, features: List[BaseFeature], labels: List[str], sample_weights: List[str] = []
+        self,
+        model_config: ModelConfig,
+        features: List[BaseFeature],
+        labels: List[str],
+        sample_weights: Optional[List[str]] = None,
     ) -> None:
         super().__init__(model_config, features, labels, sample_weights)
         self._num_class = model_config.num_class
@@ -190,7 +194,7 @@ def _init_loss_impl(self, loss_cfg: LossConfig, suffix: str = "") -> None:
         assert (
             loss_type == "softmax_cross_entropy"
         ), "match model only support softmax_cross_entropy loss now."
-        self._loss_modules[loss_name] = nn.CrossEntropyLoss(reduction='none')
+        self._loss_modules[loss_name] = nn.CrossEntropyLoss(reduction="none")
 
     def init_loss(self) -> None:
         """Initialize loss modules."""
@@ -210,7 +214,9 @@ def _loss_impl(
     ) -> Dict[str, torch.Tensor]:
         losses = {}
         label = batch.labels[label_name]
-        sample_weight = batch.sample_weights[self._sample_weight] if self._sample_weight else 1.0
+        sample_weight = (
+            batch.sample_weights[self._sample_weight] if self._sample_weight else 1.0
+        )
 
         loss_type = loss_cfg.WhichOneof("loss")
         loss_name = loss_type + suffix
@@ -223,7 +229,9 @@ def _loss_impl(
             label = _arange_int_label(pred)
         else:
             label = _zero_int_label(pred)
-        losses[loss_name] = torch.mean(self._loss_modules[loss_name](pred, label) * sample_weight)
+        losses[loss_name] = torch.mean(
+            self._loss_modules[loss_name](pred, label) * sample_weight
+        )
         return losses
 
     def loss(

diff --git a/tzrec/models/match_model_test.py b/tzrec/models/match_model_test.py
diff --git a/tzrec/models/mmoe.py b/tzrec/models/mmoe.py
@@ -9,7 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 import torch
 from torch import nn
@@ -34,7 +34,11 @@ class MMoE(MultiTaskRank):
     """
 
     def __init__(
-        self, model_config: ModelConfig, features: List[BaseFeature], labels: List[str], sample_weights: List[str] = []
+        self,
+        model_config: ModelConfig,
+        features: List[BaseFeature],
+        labels: List[str],
+        sample_weights: Optional[List[str]] = None,
     ) -> None:
         super().__init__(model_config, features, labels, sample_weights)
 

diff --git a/tzrec/models/mmoe_test.py b/tzrec/models/mmoe_test.py
diff --git a/tzrec/models/model.py b/tzrec/models/model.py
@@ -44,7 +44,11 @@ class BaseModel(nn.Module, metaclass=_meta_cls):
     """
 
     def __init__(
-        self, model_config: ModelConfig, features: List[BaseFeature], labels: List[str], sample_weights: List[str] = []
+        self,
+        model_config: ModelConfig,
+        features: List[BaseFeature],
+        labels: List[str],
+        sample_weights: Optional[List[str]] = None,
     ) -> None:
         super().__init__()
         self._base_model_config = model_config

diff --git a/tzrec/models/multi_task_rank.py b/tzrec/models/multi_task_rank.py
@@ -30,7 +30,11 @@ class MultiTaskRank(RankModel):
     """
 
     def __init__(
-        self, model_config: ModelConfig, features: List[BaseFeature], labels: List[str], sample_weights: List[str] = []
+        self,
+        model_config: ModelConfig,
+        features: List[BaseFeature],
+        labels: List[str],
+        sample_weights: Optional[List[str]] = None,
     ) -> None:
         super().__init__(model_config, features, labels, sample_weights)
         self._task_tower_cfgs = list(self._model_config.task_towers)

diff --git a/tzrec/models/multi_task_rank_test.py b/tzrec/models/multi_task_rank_test.py
diff --git a/tzrec/models/multi_tower.py b/tzrec/models/multi_tower.py
@@ -9,7 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 import torch
 from torch import nn
@@ -33,7 +33,11 @@ class MultiTower(RankModel):
     """
 
     def __init__(
-        self, model_config: ModelConfig, features: List[BaseFeature], labels: List[str], sample_weights: List[str] = []
+        self,
+        model_config: ModelConfig,
+        features: List[BaseFeature],
+        labels: List[str],
+        sample_weights: Optional[List[str]] = None,
     ) -> None:
         super().__init__(model_config, features, labels, sample_weights)
 

diff --git a/tzrec/models/multi_tower_din.py b/tzrec/models/multi_tower_din.py
@@ -9,7 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 import torch
 from torch import nn
@@ -34,7 +34,11 @@ class MultiTowerDIN(RankModel):
     """
 
     def __init__(
-        self, model_config: ModelConfig, features: List[BaseFeature], labels: List[str], sample_weights: List[str] = []
+        self,
+        model_config: ModelConfig,
+        features: List[BaseFeature],
+        labels: List[str],
+        sample_weights: Optional[List[str]] = None,
     ) -> None:
         super().__init__(model_config, features, labels, sample_weights)
 

diff --git a/tzrec/models/multi_tower_din_test.py b/tzrec/models/multi_tower_din_test.py
diff --git a/tzrec/models/multi_tower_din_trt.py b/tzrec/models/multi_tower_din_trt.py
@@ -10,7 +10,7 @@
 # limitations under the License.
 
 # Copyright (c) Alibaba, Inc. and its affiliates.
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 import torch
 from torch import nn
@@ -54,7 +54,7 @@ def __init__(
         model_config: ModelConfig,
         features: List[BaseFeature],
         labels: List[str],
-        sample_weights: List[str] = []
+        sample_weights: Optional[List[str]] = None,
     ) -> None:
         super().__init__(model_config, features, labels, sample_weights)
 
@@ -129,7 +129,11 @@ class MultiTowerDINTRT(RankModel):
     """
 
     def __init__(
-        self, model_config: ModelConfig, features: List[BaseFeature], labels: List[str], sample_weights: List[str] = []
+        self,
+        model_config: ModelConfig,
+        features: List[BaseFeature],
+        labels: List[str],
+        sample_weights: Optional[List[str]] = None,
     ) -> None:
         super().__init__(model_config, features, labels, sample_weights)
         self.embedding_group = EmbeddingGroup(

diff --git a/tzrec/models/multi_tower_test.py b/tzrec/models/multi_tower_test.py
diff --git a/tzrec/models/ple.py b/tzrec/models/ple.py
@@ -9,7 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 import torch
 from torch import nn
@@ -35,7 +35,11 @@ class PLE(MultiTaskRank):
     """
 
     def __init__(
-        self, model_config: ModelConfig, features: List[BaseFeature], labels: List[str], sample_weights: List[str] = []
+        self,
+        model_config: ModelConfig,
+        features: List[BaseFeature],
+        labels: List[str],
+        sample_weights: Optional[List[str]] = None,
     ) -> None:
         super().__init__(model_config, features, labels, sample_weights)
         assert model_config.WhichOneof("model") == "ple", (

diff --git a/tzrec/models/ple_test.py b/tzrec/models/ple_test.py
diff --git a/tzrec/models/rank_model.py b/tzrec/models/rank_model.py
@@ -9,13 +9,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 import torch
 import torchmetrics
 from torch import nn
 
-from tzrec.datasets.utils import BASE_DATA_GROUP, Batch, Optional
+from tzrec.datasets.utils import BASE_DATA_GROUP, Batch
 from tzrec.features.feature import BaseFeature
 from tzrec.loss.jrc_loss import JRCLoss
 from tzrec.metrics.grouped_auc import GroupedAUC
@@ -50,7 +50,7 @@ def __init__(
         model_config: model_pb2.ModelConfig,
         features: List[BaseFeature],
         labels: List[str],
-        sample_weights: List[str] = []
+        sample_weights: Optional[List[str]] = None,
     ) -> None:
         super().__init__(model_config, features, labels, sample_weights)
         self._num_class = model_config.num_class

diff --git a/tzrec/models/rank_model_test.py b/tzrec/models/rank_model_test.py
diff --git a/tzrec/models/tdm.py b/tzrec/models/tdm.py
@@ -9,7 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 import torch
 from torch import nn
@@ -36,7 +36,11 @@ class TDM(RankModel):
     """
 
     def __init__(
-        self, model_config: ModelConfig, features: List[BaseFeature], labels: List[str], sample_weights: List[str] = []
+        self,
+        model_config: ModelConfig,
+        features: List[BaseFeature],
+        labels: List[str],
+        sample_weights: Optional[List[str]] = None,
     ) -> None:
         super().__init__(model_config, features, labels, sample_weights)
         self.embedding_group = EmbeddingGroup(

diff --git a/tzrec/models/tdm_test.py b/tzrec/models/tdm_test.py
diff --git a/tzrec/modules/__init__.py b/tzrec/modules/__init__.py
diff --git a/tzrec/modules/embedding.py b/tzrec/modules/embedding.py
diff --git a/tzrec/modules/embedding_test.py b/tzrec/modules/embedding_test.py
diff --git a/tzrec/modules/extraction_net.py b/tzrec/modules/extraction_net.py
diff --git a/tzrec/modules/extraction_net_test.py b/tzrec/modules/extraction_net_test.py
diff --git a/tzrec/modules/fm.py b/tzrec/modules/fm.py
diff --git a/tzrec/modules/fm_test.py b/tzrec/modules/fm_test.py
diff --git a/tzrec/modules/interaction.py b/tzrec/modules/interaction.py
diff --git a/tzrec/modules/interaction_test.py b/tzrec/modules/interaction_test.py
diff --git a/tzrec/modules/mlp.py b/tzrec/modules/mlp.py
diff --git a/tzrec/modules/mlp_test.py b/tzrec/modules/mlp_test.py
diff --git a/tzrec/modules/mmoe.py b/tzrec/modules/mmoe.py
diff --git a/tzrec/modules/mmoe_test.py b/tzrec/modules/mmoe_test.py
diff --git a/tzrec/modules/sequence.py b/tzrec/modules/sequence.py
diff --git a/tzrec/modules/sequence_test.py b/tzrec/modules/sequence_test.py
diff --git a/tzrec/modules/task_tower.py b/tzrec/modules/task_tower.py
diff --git a/tzrec/modules/task_tower_test.py b/tzrec/modules/task_tower_test.py
diff --git a/tzrec/modules/variational_dropout.py b/tzrec/modules/variational_dropout.py
diff --git a/tzrec/modules/variational_dropout_test.py b/tzrec/modules/variational_dropout_test.py
diff --git a/tzrec/optim/__init__.py b/tzrec/optim/__init__.py
diff --git a/tzrec/optim/lr_scheduler.py b/tzrec/optim/lr_scheduler.py
diff --git a/tzrec/optim/lr_scheduler_test.py b/tzrec/optim/lr_scheduler_test.py
diff --git a/tzrec/optim/optimizer_builder.py b/tzrec/optim/optimizer_builder.py
diff --git a/tzrec/predict.py b/tzrec/predict.py
diff --git a/tzrec/protos/__init__.py b/tzrec/protos/__init__.py
diff --git a/tzrec/protos/data.proto b/tzrec/protos/data.proto
diff --git a/tzrec/protos/eval.proto b/tzrec/protos/eval.proto
diff --git a/tzrec/protos/export.proto b/tzrec/protos/export.proto
diff --git a/tzrec/protos/feature.proto b/tzrec/protos/feature.proto
diff --git a/tzrec/protos/loss.proto b/tzrec/protos/loss.proto
diff --git a/tzrec/protos/metric.proto b/tzrec/protos/metric.proto
diff --git a/tzrec/protos/model.proto b/tzrec/protos/model.proto
diff --git a/tzrec/protos/models/__init__.py b/tzrec/protos/models/__init__.py
diff --git a/tzrec/protos/models/match_model.proto b/tzrec/protos/models/match_model.proto
diff --git a/tzrec/protos/models/multi_task_rank.proto b/tzrec/protos/models/multi_task_rank.proto
diff --git a/tzrec/protos/models/rank_model.proto b/tzrec/protos/models/rank_model.proto
diff --git a/tzrec/protos/module.proto b/tzrec/protos/module.proto
diff --git a/tzrec/protos/optimizer.proto b/tzrec/protos/optimizer.proto
diff --git a/tzrec/protos/pipeline.proto b/tzrec/protos/pipeline.proto
diff --git a/tzrec/protos/sampler.proto b/tzrec/protos/sampler.proto
diff --git a/tzrec/protos/seq_encoder.proto b/tzrec/protos/seq_encoder.proto
diff --git a/tzrec/protos/tower.proto b/tzrec/protos/tower.proto
diff --git a/tzrec/protos/train.proto b/tzrec/protos/train.proto
diff --git a/tzrec/tests/__init__.py b/tzrec/tests/__init__.py
diff --git a/tzrec/tests/configs/dbmtl_has_sequence_mock.config b/tzrec/tests/configs/dbmtl_has_sequence_mock.config
diff --git a/tzrec/tests/configs/dbmtl_has_sequence_variational_dropout_mock.config b/tzrec/tests/configs/dbmtl_has_sequence_variational_dropout_mock.config
diff --git a/tzrec/tests/configs/dssm_fg_mock.config b/tzrec/tests/configs/dssm_fg_mock.config
diff --git a/tzrec/tests/configs/dssm_mock.config b/tzrec/tests/configs/dssm_mock.config
diff --git a/tzrec/tests/configs/dssm_v2_fg_mock.config b/tzrec/tests/configs/dssm_v2_fg_mock.config
diff --git a/tzrec/tests/configs/dssm_variational_dropout_mock.config b/tzrec/tests/configs/dssm_variational_dropout_mock.config
diff --git a/tzrec/tests/configs/multi_tower_din_fg_mock.config b/tzrec/tests/configs/multi_tower_din_fg_mock.config
diff --git a/tzrec/tests/configs/multi_tower_din_mock.config b/tzrec/tests/configs/multi_tower_din_mock.config
diff --git a/tzrec/tests/configs/multi_tower_din_trt_fg_mock.config b/tzrec/tests/configs/multi_tower_din_trt_fg_mock.config
diff --git a/tzrec/tests/configs/tdm_fg_mock.config b/tzrec/tests/configs/tdm_fg_mock.config
diff --git a/tzrec/tests/run.py b/tzrec/tests/run.py
diff --git a/tzrec/tests/train_eval_export_test.py b/tzrec/tests/train_eval_export_test.py
diff --git a/tzrec/tests/utils.py b/tzrec/tests/utils.py
diff --git a/tzrec/tools/__init__.py b/tzrec/tools/__init__.py
diff --git a/tzrec/tools/add_feature_info_to_config.py b/tzrec/tools/add_feature_info_to_config.py
diff --git a/tzrec/tools/add_feature_info_to_config_test.py b/tzrec/tools/add_feature_info_to_config_test.py
diff --git a/tzrec/tools/convert_easyrec_config_to_tzrec_config.py b/tzrec/tools/convert_easyrec_config_to_tzrec_config.py
diff --git a/tzrec/tools/convert_easyrec_config_to_tzrec_config_test.py b/tzrec/tools/convert_easyrec_config_to_tzrec_config_test.py
diff --git a/tzrec/tools/create_faiss_index.py b/tzrec/tools/create_faiss_index.py
diff --git a/tzrec/tools/create_fg_json.py b/tzrec/tools/create_fg_json.py
diff --git a/tzrec/tools/create_online_infer_data.py b/tzrec/tools/create_online_infer_data.py
diff --git a/tzrec/tools/feature_selection.py b/tzrec/tools/feature_selection.py
diff --git a/tzrec/tools/hitrate.py b/tzrec/tools/hitrate.py
diff --git a/tzrec/tools/list_distcp_param.py b/tzrec/tools/list_distcp_param.py
diff --git a/tzrec/tools/tdm/__init__.py b/tzrec/tools/tdm/__init__.py
diff --git a/tzrec/tools/tdm/cluster_tree.py b/tzrec/tools/tdm/cluster_tree.py
diff --git a/tzrec/tools/tdm/gen_tree/__init__.py b/tzrec/tools/tdm/gen_tree/__init__.py
diff --git a/tzrec/tools/tdm/gen_tree/tree_builder.py b/tzrec/tools/tdm/gen_tree/tree_builder.py
diff --git a/tzrec/tools/tdm/gen_tree/tree_builder_test.py b/tzrec/tools/tdm/gen_tree/tree_builder_test.py
diff --git a/tzrec/tools/tdm/gen_tree/tree_cluster.py b/tzrec/tools/tdm/gen_tree/tree_cluster.py
diff --git a/tzrec/tools/tdm/gen_tree/tree_cluster_test.py b/tzrec/tools/tdm/gen_tree/tree_cluster_test.py
diff --git a/tzrec/tools/tdm/gen_tree/tree_generator.py b/tzrec/tools/tdm/gen_tree/tree_generator.py
diff --git a/tzrec/tools/tdm/gen_tree/tree_generator_test.py b/tzrec/tools/tdm/gen_tree/tree_generator_test.py
diff --git a/tzrec/tools/tdm/gen_tree/tree_search_util.py b/tzrec/tools/tdm/gen_tree/tree_search_util.py
diff --git a/tzrec/tools/tdm/gen_tree/tree_search_util_test.py b/tzrec/tools/tdm/gen_tree/tree_search_util_test.py
diff --git a/tzrec/tools/tdm/init_tree.py b/tzrec/tools/tdm/init_tree.py
diff --git a/tzrec/tools/tdm/retrieval.py b/tzrec/tools/tdm/retrieval.py
diff --git a/tzrec/train_eval.py b/tzrec/train_eval.py
diff --git a/tzrec/utils/__init__.py b/tzrec/utils/__init__.py
diff --git a/tzrec/utils/checkpoint_util.py b/tzrec/utils/checkpoint_util.py
diff --git a/tzrec/utils/checkpoint_util_test.py b/tzrec/utils/checkpoint_util_test.py
diff --git a/tzrec/utils/config_util.py b/tzrec/utils/config_util.py
diff --git a/tzrec/utils/config_util_test.py b/tzrec/utils/config_util_test.py
diff --git a/tzrec/utils/dist_util.py b/tzrec/utils/dist_util.py
diff --git a/tzrec/utils/faiss_util.py b/tzrec/utils/faiss_util.py
diff --git a/tzrec/utils/faiss_util_test.py b/tzrec/utils/faiss_util_test.py
diff --git a/tzrec/utils/load_class.py b/tzrec/utils/load_class.py
diff --git a/tzrec/utils/load_class_test.py b/tzrec/utils/load_class_test.py
diff --git a/tzrec/utils/logging_util.py b/tzrec/utils/logging_util.py
diff --git a/tzrec/utils/misc_util.py b/tzrec/utils/misc_util.py
diff --git a/tzrec/utils/plan_util.py b/tzrec/utils/plan_util.py
diff --git a/tzrec/utils/plan_util_test.py b/tzrec/utils/plan_util_test.py
diff --git a/tzrec/utils/test_util.py b/tzrec/utils/test_util.py
diff --git a/tzrec/version.py b/tzrec/version.py