No public description

PiperOrigin-RevId: 550954023

mediapipe/model_maker/python/text/text_classifier/dataset.py

@@ -15,11 +15,15 @@
 
 import csv
 import dataclasses
+import hashlib
+import os
 import random
+import tempfile
+from typing import List, Optional, Sequence
 
-from typing import Optional, Sequence
 import tensorflow as tf
 
+from mediapipe.model_maker.python.core.data import cache_files as cache_files_lib
 from mediapipe.model_maker.python.core.data import classification_dataset
 
 
@@ -46,21 +50,49 @@ class CSVParameters:
 class Dataset(classification_dataset.ClassificationDataset):
   """Dataset library for text classifier."""
 
+  def __init__(
+      self,
+      dataset: tf.data.Dataset,
+      label_names: List[str],
+      tfrecord_cache_files: Optional[cache_files_lib.TFRecordCacheFiles] = None,
+      size: Optional[int] = None,
+  ):
+    super().__init__(dataset, label_names, size)
+    if not tfrecord_cache_files:
+      tfrecord_cache_files = cache_files_lib.TFRecordCacheFiles(
+          cache_prefix_filename="tfrecord", num_shards=1
+      )
+    self.tfrecord_cache_files = tfrecord_cache_files
+
   @classmethod
-  def from_csv(cls,
-               filename: str,
-               csv_params: CSVParameters,
-               shuffle: bool = True) -> "Dataset":
+  def from_csv(
+      cls,
+      filename: str,
+      csv_params: CSVParameters,
+      shuffle: bool = True,
+      cache_dir: Optional[str] = None,
+      num_shards: int = 1,
+  ) -> "Dataset":
     """Loads text with labels from a CSV file.
 
     Args:
       filename: Name of the CSV file.
       csv_params: Parameters used for reading the CSV file.
       shuffle: If True, randomly shuffle the data.
+      cache_dir: Optional parameter to specify where to store the preprocessed
+        dataset. Only used for BERT models.
+      num_shards: Optional parameter for num shards of the preprocessed dataset.
+        Note that using more than 1 shard will reorder the dataset. Only used
+        for BERT models.
 
     Returns:
       Dataset containing (text, label) pairs and other related info.
     """
+    if cache_dir is None:
+      cache_dir = tempfile.mkdtemp()
+    # calculate hash for cache based off of files
+    hasher = hashlib.md5()
+    hasher.update(os.path.basename(filename).encode("utf-8"))
     with tf.io.gfile.GFile(filename, "r") as f:
       reader = csv.DictReader(
           f,
@@ -69,6 +101,9 @@ class Dataset(classification_dataset.ClassificationDataset):
           quotechar=csv_params.quotechar)
 
     lines = list(reader)
+    for line in lines:
+      hasher.update(str(line).encode("utf-8"))
+
     if shuffle:
       random.shuffle(lines)
 
@@ -81,9 +116,18 @@ class Dataset(classification_dataset.ClassificationDataset):
         index_by_label[line[csv_params.label_column]] for line in lines
     ]
     label_index_ds = tf.data.Dataset.from_tensor_slices(
-        tf.cast(label_indices, tf.int64))
+        tf.cast(label_indices, tf.int64)
+    )
     text_label_ds = tf.data.Dataset.zip((text_ds, label_index_ds))
 
+    hasher.update(str(num_shards).encode("utf-8"))
+    cache_prefix_filename = hasher.hexdigest()
+    tfrecord_cache_files = cache_files_lib.TFRecordCacheFiles(
+        cache_prefix_filename, cache_dir, num_shards
+    )
     return Dataset(
-        dataset=text_label_ds, label_names=label_names, size=len(texts)
+        dataset=text_label_ds,
+        label_names=label_names,
+        tfrecord_cache_files=tfrecord_cache_files,
+        size=len(texts),
     )