tf-models-nightly 2.18.0.dev20240816__py2.py3-none-any.whl → 2.18.0.dev20240818__py2.py3-none-any.whl

official/recommendation/ranking/configs/config.py

@@ -138,6 +138,8 @@ class ModelConfig(hyperparams.Config):
   max_ids_per_chip_per_sample: int | None = None
   max_ids_per_table: Union[int, List[int]] | None = None
   max_unique_ids_per_table: Union[int, List[int]] | None = None
+  allow_id_dropping: bool = False
+  initialize_tables_on_host: bool = False
 
 
 @dataclasses.dataclass

official/recommendation/ranking/data/data_pipeline_multi_hot.py

@@ -45,15 +45,13 @@ class CriteoTsvReaderMultiHot:
                num_dense_features: int,
                vocab_sizes: List[int],
                multi_hot_sizes: List[int],
-               use_synthetic_data: bool = False,
-               use_cached_data: bool = False):
+               use_synthetic_data: bool = False):
     self._file_pattern = file_pattern
     self._params = params
     self._num_dense_features = num_dense_features
     self._vocab_sizes = vocab_sizes
     self._use_synthetic_data = use_synthetic_data
     self._multi_hot_sizes = multi_hot_sizes
-    self._use_cached_data = use_cached_data
 
   def __call__(self, ctx: tf.distribute.InputContext) -> tf.data.Dataset:
     params = self._params
@@ -146,7 +144,7 @@ class CriteoTsvReaderMultiHot:
         num_parallel_calls=tf.data.experimental.AUTOTUNE)
 
     dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
-    if self._use_cached_data:
+    if self._params.use_cached_data:
       dataset = dataset.take(1).cache().repeat()
 
     return dataset

official/recommendation/ranking/task.py

@@ -39,6 +39,8 @@ def _get_tpu_embedding_feature_config(
     max_ids_per_chip_per_sample: Optional[int] = None,
     max_ids_per_table: Optional[Union[int, List[int]]] = None,
     max_unique_ids_per_table: Optional[Union[int, List[int]]] = None,
+    allow_id_dropping: bool = False,
+    initialize_tables_on_host: bool = False,
 ) -> Tuple[
     Dict[str, tf.tpu.experimental.embedding.FeatureConfig],
     Optional[tf.tpu.experimental.embedding.SparseCoreEmbeddingConfig],
@@ -57,6 +59,10 @@ def _get_tpu_embedding_feature_config(
       sample.
     max_ids_per_table: Maximum number of embedding ids per table.
     max_unique_ids_per_table: Maximum number of unique embedding ids per table.
+    allow_id_dropping: bool to allow id dropping.
+    initialize_tables_on_host: bool : if the embedding table size is more than 
+      what HBM can handle, this flag will help initialize the full embedding
+      tables on host and then copy shards to HBM.
 
   Returns:
     A dictionary of feature_name, FeatureConfig pairs.
@@ -140,7 +146,8 @@ def _get_tpu_embedding_feature_config(
         max_ids_per_chip_per_sample=max_ids_per_chip_per_sample,
         max_ids_per_table=max_ids_per_table_dict,
         max_unique_ids_per_table=max_unique_ids_per_table_dict,
-        allow_id_dropping=False,
+        allow_id_dropping=allow_id_dropping,
+        initialize_tables_on_host=initialize_tables_on_host,
     )
 
   return feature_config, sparsecore_config
@@ -248,6 +255,8 @@ class RankingTask(base_task.Task):
             max_ids_per_chip_per_sample=self.task_config.model.max_ids_per_chip_per_sample,
             max_ids_per_table=self.task_config.model.max_ids_per_table,
             max_unique_ids_per_table=self.task_config.model.max_unique_ids_per_table,
+            allow_id_dropping=self.task_config.model.allow_id_dropping,
+            initialize_tables_on_host=self.task_config.model.initialize_tables_on_host,
         )
     )
 

{tf_models_nightly-2.18.0.dev20240816.dist-info → tf_models_nightly-2.18.0.dev20240818.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tf-models-nightly
-Version: 2.18.0.dev20240816
+Version: 2.18.0.dev20240818
 Summary: TensorFlow Official Models
 Home-page: https://github.com/tensorflow/models
 Author: Google Inc.

{tf_models_nightly-2.18.0.dev20240816.dist-info → tf_models_nightly-2.18.0.dev20240818.dist-info}/RECORD

@@ -877,16 +877,16 @@ official/recommendation/popen_helper.py,sha256=TMWMwsW1DF15YCJ0RG9pE3wsL6njLu8Ed
 official/recommendation/stat_utils.py,sha256=BjWRO2jzmAJyqeRUw0hMHhQqlyDNrdNvQJKAmbDJ4Rc,3076
 official/recommendation/ranking/__init__.py,sha256=7oiypy0N82PDw9aSdcJBLVoGTd_oRSUOdvuJhMv4leQ,609
 official/recommendation/ranking/common.py,sha256=fUpBln1auJWwvX5BJIYJYxRQDzXPqCzbS6NYAcn0KaQ,3998
-official/recommendation/ranking/task.py,sha256=NLoU4ZrCzrxfqg3h-xk_S8kVN1ZedZ-PKjdULmwusQU,13908
+official/recommendation/ranking/task.py,sha256=Xw4VjbrdONt2pNZf_1e35A2swb5_fcYrqZqmEZ1mmAg,14480
 official/recommendation/ranking/task_test.py,sha256=vPN_5oq1tWF3r7GuRTuAppKSF_mP3qoGFzGRoX09ylw,2891
 official/recommendation/ranking/train.py,sha256=_7zC2SVsOOOBN--If0XlWw4gtdaEeTT04PbCQRSzieo,6604
 official/recommendation/ranking/train_test.py,sha256=n96JN-gI5wj4cbZNJCwQ9XLViKt0ZAg9f9dtaTpAdV8,8756
 official/recommendation/ranking/configs/__init__.py,sha256=7oiypy0N82PDw9aSdcJBLVoGTd_oRSUOdvuJhMv4leQ,609
-official/recommendation/ranking/configs/config.py,sha256=_VICZEsRaLS4gbKeVv6T4RDnktqU9ughgx0s-GfgMgA,14876
+official/recommendation/ranking/configs/config.py,sha256=M0DgZvJ8CNtl6wUMutepQbn9joBtQtNQ-iEggbUxgec,14952
 official/recommendation/ranking/configs/config_test.py,sha256=4_W0YUwJ2q0D8HdPdIvpJcAsaySjb_pDMWj26G7T1Ws,1474
 official/recommendation/ranking/data/__init__.py,sha256=7oiypy0N82PDw9aSdcJBLVoGTd_oRSUOdvuJhMv4leQ,609
 official/recommendation/ranking/data/data_pipeline.py,sha256=jtNyeAEZQFcFJaatTJhxYPcokN56OtETbF2fFqzWk8k,7578
-official/recommendation/ranking/data/data_pipeline_multi_hot.py,sha256=b8JGO0UyFOEjWn43NRXhEAs2wVBzEJuewwnwOGqTluM,12280
+official/recommendation/ranking/data/data_pipeline_multi_hot.py,sha256=ghhxQh-L5OQU6cn0mfvo2S73Ppi50KfbU90Q2bn9dwo,12197
 official/recommendation/ranking/data/data_pipeline_multi_hot_test.py,sha256=-NKJ4-0wwSjveCY7sMP9lChDkN6lWH5VahYMHUa9z2A,2949
 official/recommendation/ranking/data/data_pipeline_test.py,sha256=34hKGxUBRjxyh9d2c7_eBhISpXI5F5PfgPfqBQoHM3Q,2573
 official/recommendation/uplift/__init__.py,sha256=_jZilTPWKu-MfMaz1IgBjEW6wqkK3FNZ1QAP4a8my3I,990
@@ -1212,9 +1212,9 @@ tensorflow_models/tensorflow_models_test.py,sha256=nc6A9K53OGqF25xN5St8EiWvdVbda
 tensorflow_models/nlp/__init__.py,sha256=4tA5Pf4qaFwT-fIFOpX7x7FHJpnyJT-5UgOeFYTyMlc,807
 tensorflow_models/uplift/__init__.py,sha256=mqfa55gweOdpKoaQyid4A_4u7xw__FcQeSIF0k_pYmI,999
 tensorflow_models/vision/__init__.py,sha256=zBorY_v5xva1uI-qxhZO3Qh-Dii-Suq6wEYh6hKHDfc,833
-tf_models_nightly-2.18.0.dev20240816.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
-tf_models_nightly-2.18.0.dev20240816.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
-tf_models_nightly-2.18.0.dev20240816.dist-info/METADATA,sha256=UsMkLoPCO2wYVldAx7uq9M9XlJ30HmXsgScGm24ZaCY,1432
-tf_models_nightly-2.18.0.dev20240816.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
-tf_models_nightly-2.18.0.dev20240816.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
-tf_models_nightly-2.18.0.dev20240816.dist-info/RECORD,,
+tf_models_nightly-2.18.0.dev20240818.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.18.0.dev20240818.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.18.0.dev20240818.dist-info/METADATA,sha256=HqIcfgOhVj0ZIvYNVNRLiZ8crrTfPwNKsa7O5n0dYDw,1432
+tf_models_nightly-2.18.0.dev20240818.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.18.0.dev20240818.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.18.0.dev20240818.dist-info/RECORD,,