Question/Bug: Trying to use sentence-transformers (v3.0.0) with Google Colab's TPU

[harisarang]

I am trying to use sentence transformers for model retrieval fine-tuning use case. But I am not able to run it on Google Colab's TPUv2. My kernel is restarting if I am importing sentence_transformers itself.

Here is the snippet of the code

!pip install torch~=2.3.0 torch_xla[tpu]~=2.3.0 -f https://storage.googleapis.com/libtpu-releases/index.html
!pip install sentence-transformers
!pip install datasets
!pip install transformers[torch]
!pip install accelerate -U

from datasets import load_dataset, Dataset
from sentence_transformers import (
    SentenceTransformer,
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
    SentenceTransformerModelCardData,
)
from sentence_transformers.losses import CosineSimilarityLoss, ContrastiveLoss, MultipleNegativesRankingLoss
from sentence_transformers.training_args import BatchSamplers
from sentence_transformers.evaluation import InformationRetrievalEvaluator

The import itself makes the kernel to restart,
ERROR: Unknown command line flag 'xla_latency_hiding_scheduler_rerun'

Am I doing something wrong here ?

Here, is the link to the notebook: Colab Notebook

P.S. v3.0.0 makes training a lot more easier a good feature/refactor of fine tuning from previous versions. Kudos to @tomaarsen and other contributors.

Thanks in advance

[tomaarsen]

Hello!

This looks quite similar to this issue: huggingface/transformers#28609 (comment)
Could you try pip uninstall tensorflow before running the Sentence Transformers related imports?

• Tom Aarsen

[harisarang]

Thanks a lot @tomaarsen. This worked

[tomaarsen]

Glad to hear it!