Numpy parquet faster

embedding_reader/parquet_numpy_reader.py

@@ -94,20 +94,19 @@ def __call__(self, batch_size, start=0, end=None, max_piece_size=None, parallel_
         cols = PIECES_BASE_COLUMNS + metadata_columns
         Piece = namedtuple("Count", cols)
 
-        def read_piece(piece):
+        def read_piece(t):
+            (piece, table) = t
             try:
                 start = piece.piece_start
                 end = piece.piece_end
                 path = piece.filename
                 metadata_path = piece.metadata_path
                 header_offset = piece.header_offset
 
-                with self.metadata_fs.open(metadata_path, "rb") as f:
-                    length = end - start
-                    table = pq.read_table(f, use_threads=False)
-                    id_columns = self.metadata_column_names
-                    table_slice = table.slice(start, length)
-                    ids = table_slice.select(id_columns).to_pandas()
+                length = end - start
+                id_columns = self.metadata_column_names
+                table_slice = table.slice(start, length)
+                ids = table_slice.select(id_columns).to_pandas()
 
                 with self.fs.open(path, "rb") as f:
                     length = end - start
@@ -128,13 +127,22 @@ def read_piece(piece):
         semaphore = Semaphore(parallel_pieces)
 
         stopped = False
+        # from path to table and file
+        open_parquet_files = {}
 
-        def piece_generator(pieces):
+        def piece_generator(pieces, open_parquet_files):
+            current_parquet_file = None
             for piece in (Piece(*parts) for parts in zip(*[pieces[col] for col in cols])):
                 if stopped:
                     break
                 semaphore.acquire()
-                yield piece
+                if piece.metadata_path not in open_parquet_files:
+                    file = self.metadata_fs.open(piece.metadata_path, "rb")
+                    table = pq.read_table(file, use_threads=True)
+                    open_parquet_files[piece.metadata_path] = {"file": file, "table": table}
+                if current_parquet_file != piece.metadata_path:
+                    current_parquet_file = piece.metadata_path
+                yield (piece, open_parquet_files[piece.metadata_path]["table"])
 
         batch = None
         batch_meta = None
@@ -143,7 +151,7 @@ def piece_generator(pieces):
         if show_progress:
             pbar = tqdm(total=len(pieces))
         with ThreadPool(parallel_pieces) as p:
-            for err, (data, meta, piece) in p.imap(read_piece, piece_generator(pieces)):
+            for err, (data, meta, piece) in p.imap(read_piece, piece_generator(pieces, open_parquet_files)):
                 if err is not None:
                     stopped = True
                     semaphore.release()
@@ -166,6 +174,8 @@ def piece_generator(pieces):
                         batch = None
                         batch_meta = None
                         batch_offset = 0
+                        open_parquet_files[piece.metadata_path]["file"].close()
+                        del open_parquet_files[piece.metadata_path]
 
                     if show_progress:
                         pbar.update(1)

examples/inference_example.py

@@ -0,0 +1,76 @@
+from embedding_reader import EmbeddingReader
+import fire
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+import numpy as np
+import fsspec
+import math
+import pandas as pd
+
+def load_safety_model():
+    """load the safety model"""
+    import autokeras as ak  # pylint: disable=import-outside-toplevel
+    from tensorflow.keras.models import load_model  # pylint: disable=import-outside-toplevel
+    from os.path import expanduser  # pylint: disable=import-outside-toplevel
+
+    home = expanduser("~")
+
+    cache_folder = home + "/.cache/clip_retrieval"
+    model_dir = cache_folder + "/clip_autokeras_binary_nsfw"
+    if not os.path.exists(model_dir):
+        os.makedirs(cache_folder, exist_ok=True)
+
+        from urllib.request import urlretrieve  # pylint: disable=import-outside-toplevel
+
+        path_to_zip_file = cache_folder + "/clip_autokeras_binary_nsfw.zip"
+        url_model = (
+            "https://raw.githubusercontent.com/LAION-AI/CLIP-based-NSFW-Detector/main/clip_autokeras_binary_nsfw.zip"
+        )
+        urlretrieve(url_model, path_to_zip_file)
+        import zipfile  # pylint: disable=import-outside-toplevel
+
+        with zipfile.ZipFile(path_to_zip_file, "r") as zip_ref:
+            zip_ref.extractall(cache_folder)
+
+    loaded_model = load_model(model_dir, custom_objects=ak.CUSTOM_OBJECTS)
+    loaded_model.predict(np.random.rand(10 ** 3, 768).astype("float32"), batch_size=10 ** 3)
+
+    return loaded_model
+
+import mmh3
+def compute_hash(url, text):
+  if url is None:
+    url = ''
+
+  if text is None:
+    text = ''
+
+  total = (url + text).encode("utf-8")
+  return mmh3.hash64(total)[0]
+
+def main(embedding_folder, metadata_folder, output_folder, batch_size=10**6, end=None):
+    """main function"""
+    reader = EmbeddingReader(embedding_folder, metadata_folder=metadata_folder, file_format="parquet_npy", meta_columns=["url", "caption"])
+    fs, relative_output_path = fsspec.core.url_to_fs(output_folder)
+    fs.mkdirs(relative_output_path, exist_ok=True)
+
+    model = load_safety_model()
+
+    total = reader.count
+    batch_count = math.ceil(total // batch_size)
+    padding = int(math.log10(batch_count)) + 1
+
+    for i, (embeddings, ids) in enumerate(reader(batch_size=batch_size, start=0, end=end, parallel_pieces=10, max_piece_size=10**4)):
+        predictions = model.predict(embeddings, batch_size=embeddings.shape[0])
+        batch = np.hstack(predictions)
+        padded_id = str(i).zfill(padding)
+        output_file_path = os.path.join(relative_output_path, padded_id + ".parquet")
+        df = pd.DataFrame(batch, columns=["prediction"])
+        df["hash"] = [compute_hash(x, y) for x, y in zip(ids['url'], ids['caption'])]
+        df["url"] = ids['url']
+        with fs.open(output_file_path, "wb") as f:
+            df.to_parquet(f)
+
+
+if __name__ == '__main__':
+    fire.Fire(main)