RAM explosion when using mine hard negatives

[JINO-ROHIT]

Im trying to run this example snippet, but after the batches get processed, the RAM explodes exponentially and crashes.

I tested using my local setup(32 GB RAM and rtx 4060) and colab(16 GB and T4) with cuda enabled

from sentence_transformers.util import mine_hard_negatives
from sentence_transformers import SentenceTransformer
from datasets import load_dataset
# Load a Sentence Transformer model
model = SentenceTransformer("all-MiniLM-L6-v2")
# Load a dataset to mine hard negatives from
dataset = load_dataset("sentence-transformers/natural-questions", split="train")

dataset = mine_hard_negatives(
     dataset=dataset,
     model=model,
     range_min=10,
     range_max=50,
     max_score=0.8,
     margin=0.1,
     num_negatives=5,
     sampling_strategy="random",
     batch_size=128,
     use_faiss=False,
)

[JINO-ROHIT]

i suspect it happens during the dataset creation when we have large datasets/mined examples and loading them all into memory

[JINO-ROHIT]

hi @tomaarsen did you by any chance get to check this?