docstr: In this function, the chief worker (rank 0):
  - initializes a Chroma client and creates a Chroma collection. The collection is persisted in the
    directory "/tmp/chroma" of the container. The "/tmp" directory in the container is a bind-mount of the
    "/tmp" directory on the agent machine (see distributed.yaml file).
    - Note that our set-up is for demonstration purpose only. Production use cases should use a
      shared file system directory bind-mounted to all agent machines and experiment containers.
  - reads in and insert embedding files generated from all workers to the Chroma collection
docstr_tok: ["In","this","function",",","the","chief","worker","(","rank","0",")",":","-","initializes","a","Chroma","client","and","creates","a","Chroma","collection",".","The","collection","is","persisted","in","the","directory","``","\/tmp\/chroma","''","of","the","container",".","The","``","\/tmp","''","directory","in","the","container","is","a","bind-mount","of","the","``","\/tmp","''","directory","on","the","agent","machine","(","see","distributed.yaml","file",")",".","-","Note","that","our","set-up","is","for","demonstration","purpose","only",".","Production","use","cases","should","use","a","shared","file","system","directory","bind-mounted","to","all","agent","machines","and","experiment","containers",".","-","reads","in","and","insert","embedding","files","generated","from","all","workers","to","the","Chroma","collection"]
code: def on_finish(self):
      """
      In this function, the chief worker (rank 0):
      - initializes a Chroma client and creates a Chroma collection. The collection is persisted in the
        directory "/tmp/chroma" of the container. The "/tmp" directory in the container is a bind-mount of the
        "/tmp" directory on the agent machine (see distributed.yaml file).
        - Note that our set-up is for demonstration purpose only. Production use cases should use a
          shared file system directory bind-mounted to all agent machines and experiment containers.
      - reads in and insert embedding files generated from all workers to the Chroma collection
      """
      if self.rank == 0:
          chroma_dir = "/tmp/chroma"
          os.makedirs(chroma_dir, exist_ok=True)
          chroma_client = chromadb.PersistentClient(chroma_dir)
          collection = chroma_client.get_or_create_collection(
              name="scidocs_embedding"
          )
  
          embeddings = []
          documents = []
          ids = []
  
          for file in os.listdir(self.output_dir):
              file_path = pathlib.Path(self.output_dir, file)
              batches = torch.load(file_path, map_location="cuda:0")
              for batch in batches:
                  for embedding in batch["embeddings"]:
                      embedding = embedding.tolist()
                      embeddings.append(embedding)
                  ids += batch["id"]
                  documents += batch["text"]
  
          collection.upsert(
              embeddings=embeddings, documents=documents, ids=ids
          )
          logging.info(
              f"Embedding contains {collection.count()} entries"
          )
  
          # Clean-up temporary embedding files
          shutil.rmtree(self.output_dir)
  
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 4 for calling the EmbeddingProcessor obj's underlying member method code with example usage: obj.on_finish() without return types