docstr: In this function, each worker persists the in-memory embeddings to the file system of the agent machine.
     - Note that our set-up is for demonstration purpose only. Production use cases should save to a
       shared file system directory bind-mounted to all agent machines and experiment containers.
  File names contain rank and batch index information to avoid duplication between:
  - files created by different workers
  - files created by the same worker for different batches of input data
docstr_tok: ["In","this","function",",","each","worker","persists","the","in-memory","embeddings","to","the","file","system","of","the","agent","machine",".","-","Note","that","our","set-up","is","for","demonstration","purpose","only",".","Production","use","cases","should","save","to","a","shared","file","system","directory","bind-mounted","to","all","agent","machines","and","experiment","containers",".","File","names","contain","rank","and","batch","index","information","to","avoid","duplication","between",":","-","files","created","by","different","workers","-","files","created","by","the","same","worker","for","different","batches","of","input","data"]
code: def on_checkpoint_start(self):
      """
      In this function, each worker persists the in-memory embeddings to the file system of the agent machine.
         - Note that our set-up is for demonstration purpose only. Production use cases should save to a
           shared file system directory bind-mounted to all agent machines and experiment containers.
      File names contain rank and batch index information to avoid duplication between:
      - files created by different workers
      - files created by the same worker for different batches of input data
      """
      if len(self.output) == 0:
          return
      file_name = f"bert_embedding_worker_{self.rank}_end_batch_{self.last_index}"
      file_path = pathlib.Path(self.output_dir, file_name)
      torch.save(self.output, file_path)
      self.output = []
  
file_path: determined/examples/features/torch_batch_process_embeddings/bert_embedding_generation.py
imports: import logging
  import os
  import pathlib
  import shutil
  import chromadb
  import torch
  from datasets import load_dataset
  from transformers import BertModel, BertTokenizer
  from determined.pytorch import experimental
label_desc: Use image node_id 3 for calling the EmbeddingProcessor obj's underlying member method code with example usage: obj.on_checkpoint_start() without return types